Exécuter une charge de travail par lot avec Pathways

Ce guide utilise une charge de travail d'entraînement JAX qui s'exécute jusqu'à son terme dans le même cluster que les composants du contrôleur Pathways pour montrer comment exécuter une charge de travail par lot.

Vous pouvez exécuter des charges de travail par lot sur un cluster Google Kubernetes Engine (GKE) ou sur un cluster Cluster Director avec Slurm :

  • Sur GKE, l'achèvement de la charge de travail JAX met fin aux composants du cluster Pathways.
  • Sur Cluster Director, les composants du contrôleur Pathways s'exécutent en tant que job Slurm distinct. Ils continuent donc de s'exécuter une fois la charge de travail JAX terminée, jusqu'à ce que vous les arrêtiez.

Cette page contient des instructions pour les deux plates-formes. Dans chaque section, sélectionnez l'onglet GKE ou Cluster Director.

Avant de commencer

Avant de commencer

GKE

Remplissez les conditions préalables suivantes :

Cluster Director

Remplissez les conditions préalables suivantes :

Créer une image d'entraînement à l'aide de MaxText

MaxText est un projet de grand modèle de langage (LLM) Open Source développé par Google. Il est écrit en JAX et conçu pour s'exécuter efficacement sur les TPU et les GPU Google Cloud .

GKE

Pour créer une image Docker MaxText à l'aide de la dernière version de JAX stable à partir du dépôt GitHub Open Source, exécutez la commande suivante :

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

Remplacez les éléments suivants :

  • PROJECT_ID : ID de votre projetGoogle Cloud .
  • USER : votre ID utilisateurGoogle Cloud .

Cette commande transfère l'image Kubernetes MaxText vers gcr.io/PROJECT_ID/USER_runner. Vous pouvez utiliser cette image Docker pour exécuter l'entraînement sur des TPU à l'aide du backend Pathways.

Cluster Director

Vous n'avez pas besoin de créer une image d'entraînement : setup_env.sh a installé MaxText dans l'environnement Pathways dans /home/pathways, et run_maxtext.sh l'exécute.

Exécuter une charge de travail par lot

GKE

  1. Si le binaire gcluster ne se trouve pas dans votre variable d'environnement PATH, ajoutez le répertoire contenant gcluster à votre PATH :

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    Remplacez CLUSTER_TOOLKIT_PATH par le chemin absolu vers le répertoire contenant le binaire gcluster, par exemple $HOME/cluster-toolkit.

  2. Envoyez l'image Docker MaxText prédéfinie à l'aide de la commande gcluster job submit :

    gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.train /deps/src/maxtext/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

    Remplacez les éléments suivants :

    • BUCKET_NAME : bucket Cloud Storage pour les sorties.
    • WORKLOAD : nom unique permettant d'identifier votre charge de travail (22 caractères ou moins en raison des limites d'étiquettes DNS).
    • CLUSTER : nom de votre cluster GKE.
    • PROJECT_ID : ID de votre projetGoogle Cloud .
    • LOCATION : région ou zone de votre cluster GKE.
    • WORKLOAD_NODEPOOL_COUNT : nombre de pools de nœuds de tranche TPU.
    • COMPUTE_TYPE : type de machine TPU, par exemple ct6e-standard-4t. Pour en savoir plus sur les types de TPU compatibles avec chaque version de TPU, consultez Versions de TPU.
    • TOPOLOGY : topologie de l'emplacement du TPU, par exemple 2x4.
    • USER : votre ID utilisateurGoogle Cloud .
    • RUN_NAME : nom attribué par l'utilisateur pour identifier l'exécution du workflow.

    Pour en savoir plus sur les options d'envoi de jobs, consultez le guide d'envoi de jobs Cluster Toolkit.

  3. Suivez la progression de votre charge de travail à l'aide de la commande gcluster job logs :

    gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

    Le résultat ressemble à ce qui suit :

    completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

  4. Facultatif : Pour annuler la charge de travail avant qu'elle ne se termine, utilisez la commande gcluster job cancel :

    gcluster job cancel WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION

Cluster Director

Exécutez les commandes suivantes sur le nœud de connexion du cluster à partir de ~/pathways. Les scripts se recherchent mutuellement et écrivent leurs journaux dans ce répertoire.

  1. Démarrez le service Pathways :

    cd ~/pathways
SLICE_SHAPE=TOPOLOGY \
    MIN_SLICES=MIN_SLICES \
    MAX_SLICES=MAX_SLICES \
    GCS_BUCKET=gs://BUCKET_NAME \
    sbatch start_pathways.sh

    Remplacez les éléments suivants :

    • TOPOLOGY : topologie TPU de chaque tranche. Il doit correspondre à la valeur acceleratorTopology de votre ensemble de nœuds TPU (par exemple, 2x4).
    • MIN_SLICES : nombre minimal de tranches en cours d'exécution requis pour démarrer et poursuivre l'entraînement.
    • MAX_SLICES : nombre maximal de tranches à utiliser, normalement le nombre de tranches dans l'ensemble de nœuds.
    • BUCKET_NAME : bucket Cloud Storage pour les données de travail et les points de contrôle Pathways.

    Si votre partition TPU n'est pas la partition par défaut du cluster, définissez également PARTITION sur son nom.

    Le service est prêt lorsque son journal, pathways-service-JOB_ID.out (où JOB_ID est l'ID de job Slurm), affiche un résultat semblable à celui-ci :

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
 Running TPU Slices : 2 / 2 (min=2)

  2. Envoyez la charge de travail d'entraînement MaxText. Vous pouvez l'envoyer juste après le job de service : il attend jusqu'à une heure que le service soit prêt.

    sbatch run_maxtext.sh

    Le job MaxText attend que le service Pathways dispose d'au moins MIN_SLICES tranches en cours d'exécution avant de commencer l'entraînement. La charge de travail entraîne Llama 3.1 8B sur des données synthétiques. Les paramètres de tranche et le bucket proviennent du service. Ne les définissez donc pas sur ce job. Vous pouvez définir les variables facultatives suivantes :

    • STEPS : nombre d'étapes d'entraînement. La valeur par défaut est 30.
    • CHECKPOINT_PERIOD : nombre d'étapes entre les points de contrôle. La valeur par défaut est 5.
    • BASE_OUTPUT_DIR : chemin d'accès Cloud Storage où MaxText écrit les points de contrôle. La valeur par défaut est gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1 : entraîne une variante à 8 couches au lieu des 32 couches complètes, pour les tranches 2x2.

    Suivez la progression de votre charge de travail (remplacez JOB_ID par l'ID de job Slurm du job MaxText). MaxText écrit son journal dans l'erreur standard :

    tail -f maxtext-llama31-JOB_ID.err

    Chaque étape d'entraînement affiche une ligne semblable à celle-ci :

    completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238

    Vous pourrez envoyer à nouveau run_maxtext.sh ultérieurement sans redémarrer le service.

  3. Lorsque vous avez terminé, arrêtez la charge de travail, puis le service Pathways :

    scancel -u ${USER} -n maxtext-llama31
scancel -u ${USER} -n pathways-service

    L'annulation du service annule également les jobs de nœud de calcul Pathways. Les VM TPU sont des nœuds statiques. Elles continuent donc de s'exécuter et deviennent disponibles pour d'autres jobs Slurm.

Si vous rencontrez des problèmes, consultez Résoudre les problèmes liés à Pathways sur Cloud.

Étapes suivantes