Executar uma carga de trabalho em lote com os programas

Este guia usa uma carga de trabalho de treinamento do JAX que é executada até a conclusão no mesmo cluster que os componentes do controlador do Pathways para demonstrar como executar uma carga de trabalho em lote.

É possível executar cargas de trabalho em lote em um cluster do Google Kubernetes Engine (GKE) ou em um cluster do Cluster Director com o Slurm:

  • No GKE, a conclusão da carga de trabalho do JAX encerra os componentes do cluster do Pathways.
  • No Cluster Director, os componentes do controlador do Pathways são executados como um job separado do Slurm. Por isso, eles continuam em execução após a conclusão da carga de trabalho do JAX, até que você os interrompa.

Esta página tem instruções para as duas plataformas. Em cada seção, selecione a guia GKE ou Cluster Director.

Antes de começar

Antes de começar

GKE

Siga os seguintes pré-requisitos:

Cluster Director

Siga os seguintes pré-requisitos:

Criar uma imagem de treinamento usando o MaxText

O MaxText é um projeto de modelo de linguagem grande (LLM) de código aberto desenvolvido pelo Google. Ele é escrito em JAX e projetado para ser executado com eficiência em Google Cloud TPUs e GPUs.

GKE

Para criar uma imagem Docker do MaxText usando a versão estável mais recente do JAX do repositório de código aberto do GitHub, execute o seguinte 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

Substitua:

  • PROJECT_ID: o ID do projeto doGoogle Cloud .
  • USER: seu ID de usuário doGoogle Cloud .

Esse comando envia a imagem do MaxText do Kubernetes para gcr.io/PROJECT_ID/USER_runner. É possível usar essa imagem Docker para executar o treinamento em TPUs usando o back-end do Pathways.

Cluster Director

Não é necessário criar uma imagem de treinamento: o setup_env.sh instalou o MaxText no ambiente do Pathways em /home/pathways, e o run_maxtext.sh o executa.

Executar uma carga de trabalho em lote

GKE

  1. Se o binário gcluster não estiver na variável de ambiente PATH, adicione o diretório que contém gcluster ao PATH:

    export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

    Substitua CLUSTER_TOOLKIT_PATH pelo caminho absoluto para o diretório que contém o binário gcluster, como $HOME/cluster-toolkit.

  2. Envie a imagem Docker pré-criada do MaxText usando o 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"

    Substitua:

    • BUCKET_NAME: o bucket do Cloud Storage para saídas.
    • WORKLOAD: um nome exclusivo para identificar sua carga de trabalho (22 caracteres ou menos devido aos limites de rótulo de DNS).
    • CLUSTER: o nome do cluster do GKE.
    • PROJECT_ID: o ID do projeto doGoogle Cloud .
    • LOCATION: a região ou zona do cluster do GKE.
    • WORKLOAD_NODEPOOL_COUNT: o número de pools de nós de fração de TPU.
    • COMPUTE_TYPE: o tipo de máquina da TPU, por exemplo, ct6e-standard-4t. Para mais informações sobre os tipos de TPU compatíveis com cada versão de TPU, consulte Versões de TPU.
    • TOPOLOGY: a topologia de posicionamento da TPU, por exemplo, 2x4.
    • USER: seu ID de usuário doGoogle Cloud .
    • RUN_NAME: um nome atribuído pelo usuário para identificar a execução do fluxo de trabalho.

    Para mais informações sobre as opções de envio de jobs, consulte o Guia de envio de jobs do Cluster Toolkit.

  3. Acompanhe o progresso da carga de trabalho usando o comando gcluster job logs:

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

    O resultado será o seguinte:

    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. Opcional: para cancelar a carga de trabalho antes que ela seja concluída, use o comando gcluster job cancel:

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

Cluster Director

Execute os comandos a seguir no nó de login do cluster em ~/pathways. Os scripts procuram uns aos outros e gravam os registros nesse diretório.

  1. Inicie o serviço do Pathways:

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

    Substitua:

    • TOPOLOGY: a topologia de TPU de cada fração. Ele precisa corresponder ao acceleratorTopology do conjunto de nós da TPU, por exemplo, 2x4.
    • MIN_SLICES: o número mínimo de frações em execução necessárias para iniciar e continuar o treinamento.
    • MAX_SLICES: o número máximo de frações a serem usadas, normalmente o número de frações no conjunto de nós.
    • BUCKET_NAME: o bucket do Cloud Storage para dados temporários e checkpoints do Pathways.

    Se a partição de TPU não for a padrão do cluster, defina PARTITION como o nome dela.

    O serviço estará pronto quando o registro, pathways-service-JOB_ID.out (em que JOB_ID é o ID da tarefa do Slurm), mostrar uma saída semelhante a esta:

    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. Envie a carga de trabalho de treinamento do MaxText. É possível enviá-lo logo após o job de serviço: ele aguarda até uma hora para que o serviço esteja pronto.

    sbatch run_maxtext.sh

    O job do MaxText aguarda até que o serviço Pathways tenha pelo menos MIN_SLICES frações em execução antes de iniciar o treinamento. A carga de trabalho treina o Llama 3.1 8B em dados sintéticos. As configurações de fração e o bucket vêm do serviço, então não os defina neste job. É possível definir as seguintes variáveis opcionais:

    • STEPS: o número de etapas de treinamento. O padrão é 30.
    • CHECKPOINT_PERIOD: o número de etapas entre os checkpoints. O padrão é 5.
    • BASE_OUTPUT_DIR: o caminho do Cloud Storage em que o MaxText grava checkpoints. O padrão é gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1: treina uma variante de 8 camadas em vez das 32 camadas completas, para frações 2x2.

    Acompanhe o progresso da carga de trabalho (substitua JOB_ID pelo código da tarefa do Slurm do job do MaxText). O MaxText grava o registro no erro padrão:

    tail -f maxtext-llama31-JOB_ID.err

    Cada etapa de treinamento imprime uma linha semelhante a esta:

    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

    É possível enviar run_maxtext.sh novamente mais tarde sem reiniciar o serviço.

  3. Quando terminar, interrompa a carga de trabalho e o serviço do Pathways:

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

    O cancelamento do serviço também cancela os jobs de worker do Pathways. As VMs de TPU são nós estáticos, então continuam em execução e ficam disponíveis para outros jobs do Slurm.

Se você tiver problemas, consulte Resolver problemas do Pathways no Cloud.

A seguir