Entrenamiento de Cloud TPU v5p

Cloud TPU v5p es la Cloud TPU de quinta generación de Google Cloud y la sucesora de la TPU v4. La v5p está optimizada para el entrenamiento a gran escala y es una plataforma líder en el desarrollo de LLM fundamentales, modelos de difusión y la IA generativa. En un nivel alto, v5p proporciona hasta el doble del rendimiento que v4, a la vez que empaqueta 2 veces más TPU en un Pod (6,000 porciones más grandes en comparación con 3,000 en v4), lo que brinda un rendimiento de hasta 4 veces a nivel de Pod. También se ejecuta con una frecuencia de reloj más alta (1.75 GHz en comparación con 1.05 GHz), agrega SparseCore para incorporaciones a gran escala y triplica la capacidad de memoria de gran ancho de banda (HBM).

Conceptos de Cloud TPU v5p

Si eres nuevo en Cloud TPU, consulta la página principal de la documentación de TPU.

Los conceptos de Cloud TPU (por ejemplo, porciones, hosts y TensorCores) y la arquitectura del sistema de Cloud TPU para todas las versiones de Cloud TPU se describen en la página Arquitectura del sistema de Cloud TPU.

Cada versión de Cloud TPU requiere tipos de acelerador específicos para el entrenamiento o la inferencia. Estos tipos de acelerador se describen en Configuraciones de v5p.

Administra recursos de TPU

Todos los comandos que puedes usar para administrar las VM de TPU se describen en Administra las TPU o en la Guía del usuario de recursos en cola para administrar los recursos en cola.

Configuración del framework

En esta sección, se describe el proceso de configuración general para el entrenamiento de modelos mediante JAX o PyTorch con TPU v5p.

Configuración de JAX

Si tienes formas de porción de más de 4 chips, tendrás varias VM en una porción. En este caso, debes usar la marca --worker=all para ejecutar la instalación en todas las VM de TPU con un solo comando:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Puedes ejecutar el siguiente comando para verificar la cantidad de dispositivos (los resultados que se muestran aquí se produjeron con una porción v5p-32). Con este código, se verifica que todo esté instalado correctamente. Para ello, verifica que JAX vea los TensorCores de Cloud TPU y pueda ejecutar operaciones básicas:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

El resultado será similar al siguiente: 

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() muestra la cantidad total de chips en la porción determinada. jax.local_device_count() indica el recuento de chips a los que puede acceder una sola VM en esta porción.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

El resultado será similar al siguiente: 

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Usa --node=all para ejecutar el comando en todos los trabajadores de Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Prueba los instructivos de JAX en este documento para comenzar con el entrenamiento de v5p con JAX.

Configuración de PyTorch

El entorno de ejecución de PJRT es el único compatible con v5p, y PyTorch 2.1 y versiones posteriores usa PJRT como el entorno de ejecución predeterminado para todas las versiones de TPU. En esta sección, se describe cómo comenzar a usar PJRT en Pods v5p con PyTorch/XLA 2.2.0 para todos los trabajadores.

Instala dependencias

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Usa una secuencia de comandos de Python con PJRT para validar la instalación y mostrar los dispositivos de TPU disponibles (los resultados que se muestran aquí se produjeron con una porción v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Usa --node=all para ejecutar el comando en todos los trabajadores de Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Prueba los instructivos de PyTorch en este documento para comenzar con el entrenamiento de v5p con PyTorch.

Supervisar y generar perfiles

Cloud TPU v5p admite la supervisión y la creación de perfiles con los mismos métodos que las generaciones anteriores de Cloud TPU. Puedes leer Perfila tu modelo con herramientas de Cloud TPU para obtener más información sobre la generación de perfiles y sobre la supervisión de VMs de Cloud TPU para obtener más información sobre la supervisión.

Instructivos de entrenamiento

Esta sección se enfoca en instructivos de entrenamiento para una sola porción. Puedes adaptar estos instructivos al entrenamiento de Multislice si agregas la marca --node=all a los comandos SSH. Para obtener más detalles y las prácticas recomendadas, consulta la Introducción a Multislice.

Instructivos de JAX

Train Diffusion 2.1

En este instructivo, se muestra cómo entrenar el modelo de dispersión estable de HuggingFace con el conjunto de datos de Pokémon en Cloud TPU v5p.

El modelo de dispersión estable es un modelo latente de texto a imagen que genera imágenes fotorrealistas a partir de cualquier entrada de texto. Para obtener más información, consulta los siguientes recursos:

Configurar

  1. Crea variables de entorno: 

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

    Descripciones de las marcas de comandos

    Variable Descripción
    PROJECT_ID Nombre del proyecto de Google Cloud
    ACCELERATOR_TYPE Consulta la página de versiones de TPU para conocer tu versión de TPU.
    ZONA Consulta el documento sobre las regiones y zonas de TPU para conocer las zonas admitidas.
    RUNTIME_VERSION En el caso de la versión 5p, usa v2-alpha-tpuv5 para la RUNTIME_VERSION.
    SERVICE_ACCOUNT Esta es la dirección de tu cuenta de servicio, que puedes encontrar en la consola de Google Cloud -> IAM -> Cuentas de servicio. Por ejemplo: tpu-service-account@myprojectID.iam.gserviceaccount.com
    TPU_NAME El ID de texto asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recursos en cola.
    QUEUED_RESOURCE_ID El ID de texto asignado por el usuario de la solicitud de recursos en cola. Consulta el documento Recursos en cola para obtener información sobre los recursos en cola.
    QUOTA_TYPE Puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado de RESOURCE_TYPE es on-demand. Consulta la sección quotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.
    VALID_UNTIL_DURATION Es el tiempo durante el cual la solicitud es válida. Consulta Recursos en cola para obtener información sobre las diferentes duraciones válidas.

  2. Crea un recurso de TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_UNTIL_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Podrás establecer una conexión SSH a la VM de TPU una vez que el recurso en cola esté en el estado ACTIVE. Ejecuta el siguiente comando para comprobar el estado de tu recurso en cola:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
--project ${PROJECT_ID} --zone ${ZONE}

    Cuando el recurso en cola está en el estado ACTIVE, el resultado será similar al siguiente:

    state: ACTIVE

  3. Instala JAX y sus dependencias.

    # compatible with v5p: only jax version 0.4.19 and later \
# jax 0.4.19 requires py 3.10 \

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} --zone=${ZONE} --worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

  4. Descarga el repositorio de HuggingFace y los requisitos para instalarlos.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

  5. Entrenar el modelo

    Entrena el modelo con un búfer previamente asignado de 4 GB.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

Limpia

Borra la TPU y la solicitud de recursos en cola al final de la sesión o para quitar las solicitudes de recursos en cola que tengan el estado “FALLID”. Para borrar un recurso en cola, borra las porciones y, luego, la solicitud de recursos en cola en 2 pasos:

   gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID}
   --zone=${ZONE} --quiet

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE} --quiet

También puedes usar --force para borrar las partes y la solicitud de recursos en cola en un solo paso:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Resultados de comparativas

El script de entrenamiento de dispersión estable se ejecutaba en v5p-8, v5p-32 y v5p-128. En la siguiente tabla, se muestra la capacidad de procesamiento.

v5p-8

v5p-32

v5p-128

Paso de entrenamiento

150

150

150

Tamaño del lote global

32

64

64

Rendimiento (ejemplos/s)

12.10

18,08

19,10

MaxText

En este instructivo, se muestra cómo entrenar el modelo MaxText con un conjunto de datos sintético en Cloud TPU.

MaxText es un LLM de alto rendimiento, arbitrariamente escalable, de código abierto y bien probado, escrito en Python/JAX puro dirigido a Cloud TPU. MaxText les brinda a los investigadores y desarrolladores una herramienta accesible y adaptable para avanzar en las fronteras de la investigación y el desarrollo del procesamiento de lenguaje natural (PLN).

Antes de ejecutar este instructivo, debes configurar tu entorno de Cloud TPU.

  1. Configura variables de entorno

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

    Descripciones de las marcas de comandos

    Variable Descripción
    PROJECT_ID Nombre del proyecto de Google Cloud
    TPU_NAME Es un nombre definido por el usuario para la TPU.
    ACCELERATOR_TYPE Consulta la página de versiones de TPU para conocer tu versión de TPU.
    ZONA Consulta el documento sobre las regiones y zonas de TPU para conocer las zonas admitidas.
    RUNTIME_VERSION En el caso de v5p, usa v2-alpha-tpuv5 para la versión del entorno de ejecución.
    RUN_NAME Nombre de la ejecución del experimento proporcionado por el usuario.

    Se recomienda la configuración opcional para Multislice:

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    Si ejecutas cargas de trabajo Multislice y deseas un rendimiento de red óptimo, considera crear una red dedicada con una unidad de transmisión máxima (MTU) de 8,896 bytes y configurar las reglas de firewall adecuadas. Aunque es opcional, este paso puede mejorar significativamente el rendimiento, en especial cuando se escala verticalmente la cantidad de porciones en la red del centro de datos (DCN). Ten en cuenta que crear una red requiere el permiso compute.networks.create en el proyecto. En los siguientes ejemplos, se muestra cómo crear una red dedicada y una regla de firewall.

    Crea una red dedicada:

    gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT_ID} \
--subnet-mode=auto \
--bgp-routing-mode=regional

    Crea una regla de firewall:

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

  2. Clona el repositorio de MaxText

    git clone https://github.com/google/maxtext.git

  3. Entrenar el modelo

    En las siguientes secciones, se describen dos opciones para entrenar MaxText.

    Opción 1

    Si deseas una secuencia de comandos que administre todo el flujo de trabajo, desde el aprovisionamiento de Cloud TPU y la instalación de dependencias hasta la ejecución del modelo y la eliminación de recursos, puedes usar multihost_job.py.

    cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

    Después de iniciar la secuencia de comandos, deberías ver un mensaje similar al siguiente en el registro. Se hace referencia a la ubicación del registro en el mensaje de salida. Haz clic en el primer vínculo para acceder a los registros de todos los trabajadores una vez que se complete el aprovisionamiento de TPU.

    ------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID
Opción 2

Para ejecutar la secuencia de comandos de entrenamiento varias veces en una Cloud TPU aprovisionada, usa la secuencia de comandos multihost_runner.py a fin de usar el recurso.

  1. Configura las variables para crear una TPU.

    export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type

    --node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

  2. Crea un recurso de TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Podrás conectarte a las VMs de TPU mediante SSH una vez que tu QueuedResource esté en el estado ACTIVE:

    Usa el comando describe para consultar el estado del recurso en cola.

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVO, el resultado será similar al siguiente:

     state: ACTIVE

  3. Conéctate a la TPU con SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Instala dependencias

    export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path

    cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'

  5. Ejecuta el modelo con varias secuencias de comandos de configuración, como 32b.sh, 64b.sh. Si ejecutas la secuencia de comandos desde una VM de TPU, debes agregar la marca --INTERNAL_IP=true.

    python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

Limpia

Borra la TPU y los recursos en cola.

Resultados de comparativas

La secuencia de comandos de entrenamiento de MaxText se ejecutó de 32 B a 1,160 B con precisión bf16. Los resultados de estas ejecuciones se muestran en la siguiente tabla.

Cantidad de parámetros

Tipo de acelerador

TFLOP/chip/s

Uso de flops del modelo

(MFU)

32,000 millones

v5p-128

3.28E+02

71,47%

64,000 millones

v5p-128

3.23E+02

70,31%

128,000 millones

V5P-256

3.15E+02

68,68%

128,000 millones

v5p-512

3.15E+02

68,53%

256,000 millones

V5P-1024

3.16E+02

68,82%

512,000 MILLONES

V5P-1024

2.94E+02

63,99%

1,024 mil millones

V5P-2048

2.49E+02

64,05%

1,024 mil millones

V5P-4096

2.97E+02

64,80%

1,160 mil millones

V5P-7680

2.95E+02

64,27%

1,160 mil millones

v5p-12288

3.04E+02

66,23%

El modelo del parámetro 256B se probó en las versiones v5p-512 y v5p-1024 con ponderaciones int8 y bf16. En la siguiente tabla, se muestran los resultados de estas pruebas.

v5p-512

v5p-512

V5P-1024

V5P-1024

Tamaño del lote global

(tokens)

5.24E+05

5.24E+05

1.05E+06

1.05E+06

Precisión

bf16

int8

bf16

int8

TFLOP/chip/s

307

408

308

414

Uso de flops del modelo

(MFU)

66,98%

88,85%

67,09%

90,23%

Instructivos de TensorFlow

Entrena ResNet en una v5p de host único

En este instructivo, se describe cómo entrenar ImageNet en una TPU v5p-8 con un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

Configurar

  1. Crea variables de entorno:

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    Para este instructivo, usa v5p-8 como ACCELERATOR_TYPE.

  2. Crea un recurso de TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --${QUOTA_TYPE}

    Podrás conectarte a la VM de TPU con SSH una vez que el recurso en cola esté en el estado ACTIVE. Para verificar el estado del recurso en cola, usa el siguiente comando:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Conéctate a la TPU con SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Configura algunas variables de entorno

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Cambia al directorio del repositorio de modelos y los requisitos de instalación.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Entrenar el modelo

  1. Ejecuta la secuencia de comandos de entrenamiento.

    python3 official/vision/train.py \
  --tpu=local \
  --experiment=resnet_imagenet \
  --mode=train_and_eval \
  --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
  --model_dir=${MODEL_DIR} \
  --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Limpia

Borra la TPU y los recursos en cola.

Entrena ResNet en una v5p de varios hosts

En este instructivo, se describe cómo entrenar ImageNet en v5p-16 o más grandes mediante un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

  1. Crea variables de entorno:

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name
export ZONE=us-east1-c
export ACCELERATOR_TYPE=v5p-16
export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pod-pjrt
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE puede ser v5p-16 o mayor.

  2. Crea un recurso de TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --node-id ${TPU_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --${QUOTA_TYPE}

    Podrás conectarte a la VM de TPU con SSH una vez que el recurso en cola esté en el estado ACTIVE.

    Usa el comando describe para consultar el estado de tu recurso en cola:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Conéctate a la TPU (worker cero) con SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  4. Configura algunas variables de entorno

    export TPU_NAME=your_tpu_name
export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0

  5. Cambia al directorio del repositorio de modelos y los requisitos de instalación.

    cd $MODELS_REPO && git checkout r2.15.0
pip install -r official/requirements.txt

Entrenar el modelo

  1. Ejecuta la secuencia de comandos de entrenamiento.

    python3 official/vision/train.py \
  --tpu=${TPU_NAME} \
  --experiment=resnet_imagenet \
  --mode=train_and_eval \
  --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
  --model_dir=${MODEL_DIR} \
  --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Limpia

Borra la TPU y los recursos en cola.

PyTorch y XLA

Llama 2

En este instructivo, se cubrirá cómo entrenar el modelo Llama 2 7B en v5p con una bifurcación del repositorio HuggingFace en PyTorch/XLA con paralelización general y escalable para gráficos de procesamiento de AA (GSPMD).

Configuración

  1. Crea variables para el ID del proyecto, el tipo de acelerador, la zona, la versión del entorno de ejecución y el nombre de TPU.

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_DURATION=1d

  2. Crea un recurso de TPU

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
--node-id ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--accelerator-type ${ACCELERATOR_TYPE} \
--runtime-version ${RUNTIME_VERSION} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Podrás conectarte a la VM de TPU mediante SSH una vez que QueuedResource esté en el estado ACTIVE:

    Usa el comando describe para consultar el estado del recurso en cola.

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}

    Cuando el recurso en cola esté en el estado ACTIVO, el resultado será similar al siguiente: 

     state: ACTIVE

  3. Instala Pytorch/XLA y las dependencias requeridas.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

  4. Descarga el repositorio de HuggingFace y los requisitos de instalación.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

  5. Descarga la configuración del modelo 7B.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

  6. Entrenar el modelo

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

Si ejecutas en un entorno de varias porciones, debes establecer la marca --spmd_dcn_parallelism en la cantidad de porciones.

SPMD_USER_GUIDE proporciona una guía del usuario más detallada en la que se explican todas las diferentes variables de entorno y los botones de activación y desactivación de la secuencia de comandos de HF. Para tener en cuenta, el LIBTPU_INIT_ARGS se incorporará a PyTorch/XLA y se activará de forma predeterminada en versiones futuras.

Limpia

Borra la TPU y los recursos en cola.

Resultados de comparativas

La capacidad de procesamiento de los tres tamaños de modelos de Llama 2 se incluye en la siguiente tabla.

v5p-8

v5p-128

v5p-128

Tamaño del modelo

7,000 millones

13,000 millones

70,000 millones

Tamaño del lote global

96

1024

128

Fragmentación de la forma de la malla

(4, 1)

(64, 1)

(16 y 4)

Uso de flops del modelo

(MFU)

56,67%

55,80%

51,85%

Asistencia y comentarios

Agradecemos todos los comentarios. Para compartir comentarios o solicitar asistencia, completa el formulario de comentarios o asistencia de Cloud TPU.