Executar uma carga de trabalho em lote com os programas

Visualizar

Esse produto ou recurso está sujeito aos "Termos de ofertas pré-GA" na seção "Termos gerais de serviço" dos Termos específicos de serviço. Os produtos e recursos pré-GA estão disponíveis "no estado em que se encontram" e podem ter suporte limitado. Para mais informações, consulte as descrições da fase de lançamento.

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 Pathways. Este guia usa uma carga de trabalho de treinamento do JAX para demonstrar isso.

Antes de começar

Você precisa ter:

• Criado um cluster do GKE.
• Instalado o XPK
• Instalado ferramentas do Kubernetes
• Ativado a API TPU
• API do Google Kubernetes Engine ativada

Criar uma imagem de treinamento usando o MaxText

MaxText é um projeto de modelo de linguagem grande (LLM, na sigla em inglês) 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 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 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 executa 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 : seu Google Cloud ID do projeto
• RUN_NAME : um nome atribuído pelo usuário para identificar a execução do fluxo de trabalho

É possível implantar 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.

É possível acompanhar o progresso da carga de trabalho navegando até a Análise de registros para o contêiner JAX. Para isso, escolha main no filtro "Nome do contêiner".

Você verá registros como o seguinte, que indica que o treinamento está em andamento. 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 do 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 do Cloud TPU que você quer criar. Para mais informações sobre os tipos de TPU compatíveis com cada versão de TPU, consulte Versões de TPU
• PROJECT : seu Google Cloud ID do projeto
• 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 carga de trabalho. É possível filtrar os registros do contêiner JAX escolhendo 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 encerrá-la prematuramente, use o seguinte comando:

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

A seguir

• Realizar a inferência com vários hosts usando o Pathways
• Executar uma carga de trabalho interativa com o Pathways
• Portar cargas de trabalho do JAX para o Pathways
• Treinamento resiliente com o Pathways
• Resolver problemas do Pathways na nuvem