Ejecuta una carga de trabajo por lotes con Pathways

A los efectos de este documento, las cargas de trabajo por lotes se definen como cargas de trabajo de JAX que se ejecutan hasta completarse y se implementan en el mismo clúster de GKE que el clúster de Pathways, específicamente junto con los componentes del controlador de Pathways (servidor proxy de IFRT y administrador de recursos de Pathways). Cuando se completa la carga de trabajo de JAX, se finalizan los componentes del clúster de Pathways. En esta guía, se usa una carga de trabajo de entrenamiento de JAX para demostrar esto.

Antes de comenzar

Asegúrate de tener lo siguiente:

Compila una imagen de entrenamiento con MaxText

MaxText es un proyecto de modelo de lenguaje grande (LLM) de código abierto desarrollado por Google. Está escrito en JAX y diseñado para ser altamente eficiente y escalable, y se ejecuta de manera eficiente en las TPU y GPU de Google Cloud.

Para compilar una imagen de Docker de MaxText con la versión estable más reciente de JAX del repositorio de GitHub de OSS, ejecuta el siguiente 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.

Este comando envía la imagen de Kubernetes de MaxText a gcr.io/$PROJECT_ID/${USER}_runner. Puedes usar esta imagen de Docker para ejecutar el entrenamiento en las TPU con el backend de Pathways.

Ejecuta una carga de trabajo por lotes con Cluster Toolkit

Envía la imagen de Docker de MaxText compilada previamente con el 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 obtener más información sobre las opciones de envío de trabajos, consulta la Guía de envío de trabajos de Cluster Toolkit.

Reemplaza lo siguiente:

  • WORKLOAD: Es un nombre único para identificar tu carga de trabajo. Debido a los límites de etiquetas de DNS, debe tener 22 caracteres o menos.
  • CLUSTER: Es el nombre de tu clúster de GKE.
  • WORKLOAD_NODEPOOL_COUNT: La cantidad de grupos de nodos de porción de TPU
  • COMPUTE_TYPE: Es el tipo de máquina de TPU (por ejemplo, ct6e-standard-4t). Para obtener más información sobre los tipos de TPU compatibles con cada versión de TPU, consulta Versiones de TPU.
  • TOPOLOGY: Es la topología de ubicación de la TPU (por ejemplo, 2x4).
  • PROJECT_ID: ID del proyecto de Google Cloud
  • ZONE: Es la zona en la que planeas ejecutar tu carga de trabajo.
  • USER: Tu Google Cloud ID de usuario
  • BUCKET_NAME: Es el bucket de Cloud Storage para las salidas.
  • RUN_NAME: Es un nombre asignado por el usuario para identificar la ejecución del flujo de trabajo.

Para seguir el progreso de tu carga de trabajo, usa el 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 la carga de trabajo antes de que se complete, usa el comando gcluster job cancel:

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

¿Qué sigue?