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 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 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 OSS, execute o seguinte comando:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext/dependencies/scripts
gcloud config set project PROJECT
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/${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 usando o XPK

Agora é possível enviar a imagem Docker do MaxText pré-criada usando o XPK com o mesmo comando usado anteriormente.

xpk workload create-pathways \
--workload=WORKLOAD \
--cluster=CLUSTER \
--num-slices=WORKLOAD_NODEPOOL_COUNT \
--tpu-type=TPU_TYPE \
--project=PROJECT \
--zone=ZONE \
--docker-image='gcr.io/PROJECT/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"

Substitua:

  • WORKLOAD: um nome exclusivo para identificar sua carga de trabalho.
  • CLUSTER: o nome do cluster do GKE
  • WORKLOAD_NODEPOOL_COUNT : o número máximo de vezes que o job pode ser reiniciado
  • TPU_TYPE: o tipo de TPU especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de TPU disponíveis em cada versão, consulte Versões de TPU.
  • PROJECT : 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
  • RUN_NAME : um nome atribuído pelo usuário para identificar a execução do fluxo de trabalho

O resultado será assim:

[XPK] Follow your Pathways workload and other resources here : https://console.cloud.google.com/logs/query;query=resource.type%3D"k8s_container"%0Aresource.labels.project_id%3D"<project-name>"%0Aresource.labels.location%3D"<your-zone>"%0Aresource.labels.cluster_name%3D"<your-cluster-name>"%0Aresource.labels.pod_name:"<your-pod-name>"%0Aseverity>%3DDEFAULT

Use o link na saída do comando XPK anterior para acompanhar o progresso da sua carga de trabalho. Para filtrar os registros do contêiner JAX, escolha jax-tpu no filtro "Nome do contêiner".

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

A carga de trabalho é concluída após o número especificado de etapas. Se você quiser encerrar antes do tempo, use o seguinte comando:

xpk workload delete --workload=WORKLOAD --cluster=CLUSTER --project=PROJECT --zone=ZONE

A seguir