Ejecuta una carga de trabajo por lotes con Pathways

Para los fines 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 dentro del 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). La finalización de la carga de trabajo de JAX finaliza los componentes del clúster de Pathways. En esta guía, se usa una carga de trabajo de entrenamiento de JAX para demostrarlo.

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 escalable y de alto rendimiento, y se ejecuta de manera eficiente en las Cloud TPU y GPU de Google Cloud.

Para compilar una imagen de Docker de MaxText con la versión más reciente de JAX estable del repositorio de GitHub de OSS, ejecuta el siguiente comando:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

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 TPU con el backend de Pathways.

Ejecuta una carga de trabajo por lotes con Cluster Toolkit

Si el objeto binario gcluster no está en tu variable de entorno PATH, agrega el directorio que contiene gcluster a tu PATH:

export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

Reemplaza CLUSTER_TOOLKIT_PATH por la ruta de acceso absoluta al directorio que contiene el objeto binario gcluster, como $HOME/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=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.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: un nombre único para identificar tu carga de trabajo; debido a los límites de las etiquetas de DNS, este debe tener 22 caracteres o menos
  • CLUSTER: Es el nombre de tu clúster de GKE.
  • WORKLOAD_NODEPOOL_COUNT: Es 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: la topología de ubicación de la TPU (por ejemplo, 2x4)
  • PROJECT_ID: Es el ID del proyecto de Google Cloud .
  • LOCATION: Es la región o zona de tu clúster de GKE.
  • USERTu ID de usuario de Google Cloud
  • BUCKET_NAME: Es el bucket de Cloud Storage para los resultados.
  • RUN_NAME: Es un nombre asignado por el usuario para identificar la ejecución del flujo de trabajo.

Sigue el progreso de tu carga de trabajo con el comando gcluster job logs:

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

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=LOCATION

¿Qué sigue?