Executar uma carga de trabalho em lote com os programas

Para fins deste documento, as cargas de trabalho em lote são definidas como cargas de trabalho 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 programa de aprendizado. 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 de maneira eficiente em 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/dependencies/scripts
gcloud config set project PROJECT_ID
bash ./docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
bash ./docker_upload_runner.sh CLOUD_IMAGE_NAME=USER_runner # This script needs bash version >= 4.2 to execute.

Esse comando envia a imagem do Kubernetes do MaxText para gcr.io/$PROJECT_ID/${USER}_runner. Você pode 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

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=ZONE \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m MaxText.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 22 caracteres ou menos.
  • 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 aceitos em 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
  • ZONE: a zona em que você planeja executar sua carga de trabalho.
  • USER: seu Google Cloud ID de usuário
  • 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 sua carga de trabalho usando o comando gcluster job logs:

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --main-only=false

completed step: 1, seconds: 0.484, TFLOP/s/device: 87.349, Tokens/s/device: 2117.382, total_weights: 2945, loss: 10.888
completed step: 2, seconds: 0.407, TFLOP/s/device: 103.699, Tokens/s/device: 2513.735, total_weights: 3253, loss: 9.697
completed step: 3, seconds: 0.248, TFLOP/s/device: 170.300, Tokens/s/device: 4128.167, total_weights: 3154, loss: 9.641
completed step: 4, seconds: 0.216, TFLOP/s/device: 195.122, Tokens/s/device: 4729.880, total_weights: 3119, loss: 9.547
completed step: 5, seconds: 0.272, TFLOP/s/device: 155.298, Tokens/s/device: 3764.512, total_weights: 2837, loss: 10.179
completed step: 6, seconds: 0.472, TFLOP/s/device: 89.489, Tokens/s/device: 2169.266, total_weights: 3069, loss: 9.776

Para cancelar a carga de trabalho antes da conclusão, use o comando gcluster job cancel:

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

A seguir