Ejecuta una carga de trabajo por lotes con Pathways

En esta guía, se usa una carga de trabajo de entrenamiento de JAX que se ejecuta hasta completarse en el mismo clúster que los componentes del controlador de Pathways para demostrar cómo ejecutar una carga de trabajo por lotes.

Puedes ejecutar cargas de trabajo por lotes en un clúster de Google Kubernetes Engine (GKE) o en un clúster de Cluster Director con Slurm:

  • En GKE, la finalización de la carga de trabajo de JAX finaliza los componentes del clúster de Pathways.
  • En Cluster Director, los componentes del controlador de Pathways se ejecutan como un trabajo de Slurm independiente, por lo que siguen ejecutándose después de que se completa la carga de trabajo de JAX, hasta que los detengas.

Esta página contiene instrucciones para ambas plataformas. En cada sección, selecciona la pestaña GKE o Cluster Director.

Antes de comenzar

Antes de comenzar

GKE

Completa los siguientes requisitos previos:

Cluster Director

Completa los siguientes requisitos previos:

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 ejecutarse de manera eficiente en TPU y GPU de Google Cloud .

GKE

Para compilar una imagen de Docker de MaxText con la versión más reciente estable de JAX del repositorio de GitHub de código abierto, 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

Reemplaza lo siguiente:

  • PROJECT_ID: Es el ID del proyecto deGoogle Cloud .
  • USER: Es tu ID de usuario deGoogle Cloud .

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.

Cluster Director

No es necesario que compiles una imagen de entrenamiento: setup_env.sh instaló MaxText en el entorno de Pathways en /home/pathways, y run_maxtext.sh lo ejecuta.

Ejecuta una carga de trabajo por lotes

GKE

  1. 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.

  2. 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"

    Reemplaza lo siguiente:

    • BUCKET_NAME: Es el bucket de Cloud Storage para los resultados.
    • WORKLOAD: Es un nombre único para identificar tu carga de trabajo (22 caracteres o menos debido a los límites de las etiquetas de DNS).
    • CLUSTER: Es el nombre del clúster de GKE.
    • PROJECT_ID: Es el ID del proyecto deGoogle Cloud .
    • LOCATION: Es la región o zona 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.
    • USER: Es tu ID de usuario deGoogle Cloud .
    • RUN_NAME: Es un nombre asignado por el usuario para identificar la ejecución del flujo de trabajo.

    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.

  3. 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

    El resultado es similar a lo siguiente:

    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

  4. Opcional: 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

Cluster Director

Ejecuta los siguientes comandos en el nodo de acceso del clúster desde ~/pathways. Las secuencias de comandos se buscan entre sí y escriben sus registros en ese directorio.

  1. Inicia el servicio de Pathways:

    cd ~/pathways
SLICE_SHAPE=TOPOLOGY \
    MIN_SLICES=MIN_SLICES \
    MAX_SLICES=MAX_SLICES \
    GCS_BUCKET=gs://BUCKET_NAME \
    sbatch start_pathways.sh

    Reemplaza lo siguiente:

    • TOPOLOGY: la topología de TPU de cada porción. Debe coincidir con el acceleratorTopology de tu conjunto de nodos de TPU, por ejemplo, 2x4.
    • MIN_SLICES: Es la cantidad mínima de porciones en ejecución que se requieren para iniciar y continuar el entrenamiento.
    • MAX_SLICES: Es la cantidad máxima de porciones que se usarán, normalmente la cantidad de porciones en el conjunto de nodos.
    • BUCKET_NAME: Es el bucket de Cloud Storage para los datos temporales y los puntos de control de Pathways.

    Si tu partición de TPU no es la predeterminada del clúster, también establece PARTITION en su nombre.

    El servicio está listo cuando su registro, pathways-service-JOB_ID.out (en el que JOB_ID es el ID de trabajo de Slurm), muestra un resultado similar al siguiente:

    Pathways Runtime Service is READY!
 Discovery file    : /home/USER/pathways/cluster.env
 IFRT Proxy Target : grpc://RM_IP:29000
 Running TPU Slices : 2 / 2 (min=2)

  2. Envía la carga de trabajo de entrenamiento de MaxText. Puedes enviarlo justo después del trabajo de servicio: espera hasta una hora para que el servicio esté listo.

    sbatch run_maxtext.sh

    El trabajo de MaxText espera hasta que el servicio de Pathways tenga al menos MIN_SLICES porciones en ejecución antes de comenzar el entrenamiento. La carga de trabajo entrena a Llama 3.1 8B con datos sintéticos. La configuración de porciones y el bucket provienen del servicio, por lo que no debes establecerlos en este trabajo. Puedes establecer las siguientes variables opcionales:

    • STEPS: Es la cantidad de pasos de entrenamiento. El valor predeterminado es 30.
    • CHECKPOINT_PERIOD: Es la cantidad de pasos entre los puntos de control. El valor predeterminado 5.
    • BASE_OUTPUT_DIR: Es la ruta de acceso de Cloud Storage en la que MaxText escribe los puntos de control. El valor predeterminado es gs://BUCKET_NAME/maxtext-runs.
    • REDUCED_LAYERS=1: Entrena una variante de 8 capas en lugar de las 32 capas completas para las porciones de 2x2.

    Sigue el progreso de tu carga de trabajo (reemplaza JOB_ID por el ID de trabajo de Slurm del trabajo de MaxText). MaxText escribe su registro en el error estándar:

    tail -f maxtext-llama31-JOB_ID.err

    Cada paso de entrenamiento imprime una línea similar a la siguiente:

    completed step: 24, seconds: 0.326, live slice count: 2, TFLOP/s/device: 88.199, Tokens/s/device: 6289.27, total_weights: 32768, loss: 6.238

    Puedes volver a enviar run_maxtext.sh más tarde sin reiniciar el servicio.

  3. Cuando termines, detén la carga de trabajo y, luego, el servicio de Pathways:

    scancel -u ${USER} -n maxtext-llama31
scancel -u ${USER} -n pathways-service

    Si cancelas el servicio, también se cancelarán los trabajos de los trabajadores de Pathways. Las VMs de TPU son nodos estáticos, por lo que siguen ejecutándose y están disponibles para otros trabajos de Slurm.

Si tienes problemas, consulta Soluciona problemas de Pathways en Cloud.

¿Qué sigue?