Esegui un carico di lavoro batch con Pathways

Questa guida utilizza un workload di addestramento JAX che viene eseguito fino al completamento nello stesso cluster dei componenti del controller Pathways per dimostrare come eseguire un workload batch.

Puoi eseguire workload batch su un cluster Google Kubernetes Engine (GKE) o su un cluster Cluster Director con Slurm:

  • Su GKE, il completamento del workload JAX termina i componenti del cluster Pathways.
  • In Cluster Director, i componenti del controller Pathways vengono eseguiti come un job Slurm separato, quindi continuano a essere eseguiti dopo il completamento del workload JAX, finché non li interrompi.

Questa pagina contiene le istruzioni per entrambe le piattaforme. In ogni sezione, seleziona la scheda GKE o Cluster Director.

Prima di iniziare

Prima di iniziare

GKE

Completa i seguenti prerequisiti:

Cluster Director

Completa i seguenti prerequisiti:

Crea un'immagine di addestramento utilizzando MaxText

MaxText è un progetto di modello linguistico di grandi dimensioni (LLM) open source sviluppato da Google. È scritto in JAX e progettato per essere eseguito in modo efficiente su TPU e GPU Google Cloud .

GKE

Per creare un'immagine Docker MaxText utilizzando l'ultima versione di JAX stabile dal repository GitHub open source, esegui questo comando:

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

Sostituisci quanto segue:

  • PROJECT_ID: il tuo ID progettoGoogle Cloud .
  • USER: il tuo ID utenteGoogle Cloud .

Questo comando esegue il push dell'immagine Kubernetes MaxText in gcr.io/PROJECT_ID/USER_runner. Puoi utilizzare questa immagine Docker per eseguire l'addestramento sulle TPU utilizzando il backend Pathways.

Cluster Director

Non è necessario creare un'immagine di addestramento: setup_env.sh ha installato MaxText nell'ambiente Pathways in /home/pathways e run_maxtext.sh la esegue.

Esegui un workload batch

GKE

  1. Se il programma binario gcluster non si trova nella variabile di ambiente PATH, aggiungi la directory contenente gcluster alla variabile PATH:

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    Sostituisci CLUSTER_TOOLKIT_PATH con il percorso assoluto della directory contenente il file binario gcluster, ad esempio $HOME/cluster-toolkit.

  2. Invia l'immagine Docker MaxText predefinita utilizzando il comando 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"

    Sostituisci quanto segue:

    • BUCKET_NAME: il bucket Cloud Storage per gli output.
    • WORKLOAD: un nome univoco per identificare il workload (massimo 22 caratteri a causa dei limiti delle etichette DNS).
    • CLUSTER: il nome del tuo cluster GKE.
    • PROJECT_ID: il tuo ID progettoGoogle Cloud .
    • LOCATION: la regione o la zona del tuo cluster GKE.
    • WORKLOAD_NODEPOOL_COUNT: il numero di node pool di sezioni TPU.
    • COMPUTE_TYPE: il tipo di macchina TPU, ad esempio ct6e-standard-4t. Per ulteriori informazioni sui tipi di TPU supportati per ogni versione di TPU, consulta Versioni di TPU.
    • TOPOLOGY: la topologia di posizionamento della TPU, ad esempio 2x4.
    • USER: il tuo ID utenteGoogle Cloud .
    • RUN_NAME: un nome assegnato dall'utente per identificare l'esecuzione del workflow.

    Per saperne di più sulle opzioni di invio dei job, consulta la guida all'invio dei job di Cluster Toolkit.

  3. Segui l'avanzamento del workload utilizzando il comando gcluster job logs:

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

    L'output è simile al seguente:

    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. (Facoltativo) Per annullare il workload prima del completamento, utilizza il comando gcluster job cancel:

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

Cluster Director

Esegui questi comandi sul nodo di accesso del cluster da ~/pathways. Gli script si cercano a vicenda e scrivono i log in quella directory.

  1. Avvia il servizio Pathways:

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

    Sostituisci quanto segue:

    • TOPOLOGY: la topologia TPU di ogni sezione. Deve corrispondere al valore acceleratorTopology del tuo nodeset TPU, ad esempio 2x4.
    • MIN_SLICES: il numero minimo di sezioni in esecuzione necessarie per avviare e continuare l'addestramento.
    • MAX_SLICES: il numero massimo di sezioni da utilizzare, in genere il numero di sezioni nel nodeset.
    • BUCKET_NAME: il bucket Cloud Storage per i dati temporanei e i checkpoint di Pathways.

    Se la partizione TPU non è la partizione predefinita del cluster, imposta anche PARTITION sul suo nome.

    Il servizio è pronto quando il relativo log, pathways-service-JOB_ID.out (dove JOB_ID è l'ID job Slurm), mostra un output simile al seguente:

    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. Invia il workload di addestramento MaxText. Puoi inviarlo subito dopo il job del servizio: attende fino a un'ora che il servizio sia pronto.

    sbatch run_maxtext.sh

    Il job MaxText attende che il servizio Pathways abbia almeno MIN_SLICES slice in esecuzione prima di avviare l'addestramento. Il workload addestra Llama 3.1 8B su dati sintetici. Le impostazioni di suddivisione e il bucket provengono dal servizio, quindi non impostarli in questo job. Puoi impostare le seguenti variabili facoltative:

    • STEPS: il numero di passaggi di addestramento. Il valore predefinito è 30.
    • CHECKPOINT_PERIOD: il numero di passaggi tra i checkpoint. Il valore predefinito è 5.
    • BASE_OUTPUT_DIR: il percorso Cloud Storage in cui MaxText scrive i checkpoint. Il valore predefinito è gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1: addestra una variante a 8 strati anziché a 32 strati completi, per 2x2 sezioni.

    Segui l'avanzamento del workload (sostituisci JOB_ID con l'ID job Slurm del job MaxText). MaxText scrive il log nell'errore standard:

    tail -f maxtext-llama31-JOB_ID.err

    Ogni passaggio di addestramento stampa una riga simile alla seguente:

    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

    Puoi inviare nuovamente run_maxtext.sh in un secondo momento senza riavviare il servizio.

  3. Al termine, interrompi il workload e poi il servizio Pathways:

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

    L'annullamento del servizio comporta anche l'annullamento dei job dei worker Pathways. Le VM TPU sono nodi statici, quindi continuano a essere eseguite e diventano disponibili per altri job Slurm.

Se riscontri problemi, consulta Risolvi i problemi relativi a Pathways on Cloud.

Passaggi successivi