Executar uma carga de trabalho em lote com os programas

Para os fins deste documento, as cargas de trabalho em lote são definidas como cargas de trabalho do JAX que são executadas até a conclusão e implantadas no mesmo cluster do GKE que o cluster do Pathways, especificamente ao lado dos componentes do controlador do Pathways (servidor proxy IFRT e gerenciador de recursos do Pathways). A conclusão da carga de trabalho do JAX encerra os componentes do cluster do Pathways. Este guia usa uma carga de trabalho de treinamento do JAX para demonstrar isso.

Antes de começar

Você precisa ter:

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 altamente eficiente e escalonável, sendo executado com eficiência em Cloud TPUs e GPUs do Google Cloud.

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

Esse comando envia a imagem do MaxText no 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.

Executar uma carga de trabalho em lote com o Cluster Toolkit

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.

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"

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

Substitua:

  • WORKLOAD: um nome exclusivo para identificar sua carga de trabalho. Devido aos limites de rótulo de DNS, ele precisa ter até 22 caracteres.
  • CLUSTER: o nome 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, consulte Versões de TPU.
  • TOPOLOGY: a topologia de posicionamento da TPU (por exemplo, 2x4)
  • PROJECT_ID: o ID do projeto do Google Cloud .
  • LOCATION: a região ou zona do cluster do GKE
  • USER: seu ID do usuário do Google Cloud
  • BUCKET_NAME: o bucket do Cloud Storage para saídas
  • RUN_NAME: um nome atribuído pelo usuário para identificar a execução do fluxo de trabalho.

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

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

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

A seguir