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 do IFRT e gerenciador de recursos do Pathways). A conclusão da carga de trabalho do JAX encerra os componentes do cluster do programa. 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 do 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 do Docker para executar o treinamento em TPUs usando o back-end do Pathways.

Executar uma carga de trabalho em lote usando a API PathwaysJob

O manifesto a seguir implanta os componentes do Pathways e executa uma carga de trabalho do MaxText usando a API PathwaysJob. A carga de trabalho é encapsulada no contêiner main e exerce train.py.

Copie o YAML a seguir em um arquivo chamado pathways-job-batch-training.yaml e atualize os valores editáveis.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  workers:
  - type: TPU_MACHINE_TYPE
    topology: TOPOLOGY
    numSlices: WORKLOAD_NODEPOOL_COUNT
  pathwaysDir: gs://BUCKET_NAME
  controller:
    deploymentMode: default
    template:
      spec:
        containers:
        - name: main
          image: gcr.io/PROJECT/USER_runner
          command:
          - bash
          - -c
          - |
            python3 -m MaxText.train /deps/src/MaxText/configs/base.yml \
            base_output_directory=gs://BUCKET_NAME \
            run_name=RUN_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

Substitua:

  • USER : seu Google Cloud ID de usuário
  • MAX_RESTARTS : o número máximo de vezes que o job pode ser reiniciado
  • TPU_MACHINE_TYPE : o tipo de máquina da TPU
  • TOPOLOGY : a topologia da TPU v4 ou mais recente. Para mais informações sobre as versões de TPU e as topologias compatíveis, consulte Versões de TPU.
  • WORKLOAD_NODEPOOL_COUNT : o número de pools de nós usados por uma carga de trabalho do Pathways
  • BUCKET_NAME : um bucket do Cloud Storage para armazenar arquivos temporários
  • PROJECT : o ID do projeto do Google Cloud
  • RUN_NAME : um nome atribuído pelo usuário para identificar a execução do fluxo de trabalho

Implante o YAML PathwaysJob da seguinte maneira:

kubectl apply -f pathways-job-batch-training.yaml

Para conferir a instância PathwaysJob criada pelo comando anterior, use:

kubectl get pathwaysjob

O resultado será semelhante a este:

NAME             AGE
pathways-trial   9s

Para modificar um atributo da instância PathwaysJob, exclua a instância PathwaysJob, modifique o YAML e aplique-o para criar uma nova instância PathwaysJob.

Para acompanhar o progresso da sua carga de trabalho, acesse a Análise de registros do contêiner JAX e escolha main no filtro "Nome do contêiner".

Você vai ver registros como este, que indicam que o treinamento está progredindo. A carga de trabalho será concluída após 30 etapas.

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 excluir a instância PathwaysJob, use o seguinte comando:

kubectl delete -f pathways-job-batch-training.yaml

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

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

A seguir