Entrenamiento de Cloud TPU v5p
Cloud TPU v5p es la quinta generación de Cloud TPU de Google Cloud y sucesoras de la TPU v4. v5p está optimizado para entrenamiento a gran escala y ser una plataforma líder para el desarrollo de los LLM, los modelos de difusión y la IA generativa. En un nivel alto, v5p proporciona hasta el doble del rendimiento de v4 y, al mismo tiempo, incluye el doble de TPU en un pod (6,000 de la porción más grande en comparación con 3,000 en v4), lo que genera hasta 4 veces el rendimiento a nivel del pod. También se ejecuta con un frecuencia de reloj (1.75 Ghz frente a 1.05 Ghz), agrega SparseCore para incorporaciones a gran escala y triplica la capacidad de memoria de ancho de banda alto (HBM).
Conceptos de Cloud TPU v5p
Si eres nuevo en Cloud TPU, consulta la Página principal de documentación de TPU.
conceptos de Cloud TPU (por ejemplo, slices, hosts y TensorCores) y la arquitectura de sistema de Cloud TPU para todas las Cloud TPU se describen en el Arquitectura del sistema de Cloud TPU .
Cada versión de Cloud TPU requiere tipos de aceleradores específicos para 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 tus VMs de TPU se describen en Administrar TPU o Guía del usuario de recursos en cola para administrar recursos en cola.
Configuración del framework
En esta sección, se describe el proceso de configuración general para el entrenamiento de modelos con JAX o PyTorch con TPU v5p.
Configuración para JAX
Si tienes cortes de más de 4 chips, tendrás varias VMs.
en una porción. En este caso, debes usar la marca --worker=all
para ejecutar la instalación en todas las VMs 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). Este código comprueba que todo esté instalado correctamente verificando que JAX ve los tensores de Cloud TPU y puede 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 ejemplo:
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 de 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 ejemplo:
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 obtener comenzamos con el entrenamiento v5p con JAX.
Configuración para PyTorch
El entorno de ejecución de PJRT es el único entorno de ejecución compatible con v5p, y PyTorch 2.1 y versiones posteriores usan 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 para mostrar los dispositivos de TPU disponibles (los resultados que se muestran aquí se produjeron con un segmento 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 de este documento para obtener comenzaste con el entrenamiento de v5p con PyTorch.
Supervisar y generar perfiles
Cloud TPU v5p admite la supervisión y la generación de perfiles con los mismos métodos que las generaciones anteriores de Cloud TPU. Puedes Lee Crea perfiles de tu modelo con las herramientas de Cloud TPU para obtener más información sobre la generación de perfiles y la supervisión de las VMs de Cloud TPU para obtener más información sobre la supervisión.
Instructivos de entrenamiento
Esta sección se centra en los instructivos de entrenamiento para una sola porción.
Adaptar estos instructivos al entrenamiento
de varias porciones puede ser
se logra agregando la marca --node=all
a los comandos SSH.
Para conocer más detalles y 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 difusió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 de texto a imagen latente que genera y fotorrealistas a partir de cualquier entrada de texto. Para obtener más información, consulta los siguientes recursos:
Configurar
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 ID DEL PROYECTO Nombre del proyecto de Google Cloud ACCELERATOR_TYPE Consulta las versiones de TPU de tu versión de TPU. ZONA Consulta el documento de regiones y zonas de TPU para conocer las zonas compatibles. RUNTIME_VERSION Para v5p, usa v2-alpha-tpuv5 para RUNTIME_VERSION. SERVICE_ACCOUNT Esta es la dirección de la 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 recurso en cola. QUEUED_RESOURCE_ID El ID de texto asignado por el usuario de la solicitud de recurso en cola. Consulta el documento Recursos en cola para obtener información sobre los recursos en cola. QUOTA_TYPE Puede ser reserved
ospot
. Si ninguna de estas opciones se especifica, el La configuración predeterminada eson-demand
. Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.VALID_UNTIL_DURATION Es la duración de la validez de la solicitud. Consulta Recursos en cola para obtener información sobre las distintas duraciones válidas. -
gcloud 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 tu VM de TPU una vez que tu recurso en cola en el estado
ACTIVE
. Ejecuta el siguiente comando para verificar 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 tiene el estado
ACTIVE
, la salida será similar al siguiente ejemplo:state: ACTIVE
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'
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 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'
Entrenar el modelo
Entrena el modelo con un búfer preasignado 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 tu TPU y la solicitud de recursos en cola al final de la sesión o para quitar las solicitudes de recursos en cola que estén en el estado "FAILED". Para borrar un recurso en cola, borra las porciones y, luego, el recurso 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 porciones y la solicitud de recurso 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
La secuencia de comandos de entrenamiento de Stable Diffusion se ejecutó en las versiones 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 |
Capacidad de procesamiento (ejemplos/s) |
12.10 |
18/08 |
19.10 |
MaxText
En este instructivo, se muestra cómo entrenar el MaxText con un modelo sintético conjunto de datos en Cloud TPU.
MaxText es una herramienta de código abierto, escalable de alto rendimiento y arbitraria LLM probado y escrito en Python/JAX puro orientado a las Cloud TPU. MaxText 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 configura tu entorno de Cloud TPU.
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 ID DEL PROYECTO Nombre del proyecto de Google Cloud TPU_NAME Un nombre definido por el usuario para tu TPU. ACCELERATOR_TYPE Consulta la página Versiones de TPU para ver la versión de TPU que tienes. ZONA Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles. RUNTIME_VERSION Para v5p, usa v2-alpha-tpuv5 para la versión del entorno de ejecución. RUN_NAME Es el nombre de la ejecución del experimento que proporcionó el usuario. Configuración opcional recomendada para Multislice:
export NETWORK_NAME=your_network_name export FIREWALL_RULE_NAME=your_firewall_rule_name
Si ejecutas cargas de trabajo de Multislice y deseas un rendimiento óptimo de la red, 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. Si bien es opcional, este paso puede mejorar significativamente el rendimiento, en especial cuando se aumenta la cantidad de rebanadas en la red del centro de datos (DCN). Creación de notas una red requiere el permiso
compute.networks.create
en el proyecto. En los siguientes ejemplos, se muestra cómo crear una red dedicada y un 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}
Clona el repositorio de MaxText
git clone https://github.com/google/maxtext.git
Entrenar el modelo
En las siguientes secciones, se describen dos opciones para entrenar MaxText.
Opción 1
Si quieres que una secuencia de comandos administre todo el flujo de trabajo, desde el aprovisionamiento Cloud TPU y la instalación de dependencias para ejecutar el modelo y eliminando 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 a lo 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 un dispositivo
Cloud TPU, usa
la secuencia de comandos multihost_runner.py
para usar el recurso
Configura 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
Crea un recurso de TPU.
gcloud 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 con SSH una vez que
QueuedResource
está en el estadoACTIVE
: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 tu recurso en cola se encuentre en estado ACTIVO, la salida será similar al siguiente:
state: ACTIVE
Conéctate a la TPU con SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
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'
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 las comparativas
La secuencia de comandos de entrenamiento de MaxText se ejecutó de 32B a 1160B con precisión bf16. Los resultados de estas ejecuciones se muestran en la siguiente tabla.
Cant. de parámetros |
Tipo de acelerador |
TFLOP/chip/s |
Utilización de los 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 millones |
v5p-2048 |
2.49E+02 |
64,05% |
1,024 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 5p-512 y v5p-1024 con pesos int8 y bf16. La siguiente tabla muestra 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 |
Utilización de los flops del modelo (MFU) |
66,98% |
88,85% |
67,09% |
90.23% |
Instructivos de TensorFlow
Entrena ResNet en un solo host v5p
En este instructivo, se describe cómo entrenar ImageNet en una TPU v5p-8
con un conjunto de datos falso. Si quieres usar un conjunto de datos diferente, consulta
Prepara el conjunto de datos.
Configurar
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.17.0-pjrt export TPU_NAME=your-tpu-name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
En este instructivo, usa
v5p-8
comoACCELERATOR_TYPE
.-
gcloud 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 recurso en cola está en el estado
ACTIVE
Para verificar el estado de tu recurso en cola, usa el siguiente comando:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project ${PROJECT_ID} \ --zone ${ZONE}
Conéctate a la TPU con SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
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
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
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 un v5p para varios hosts
En este instructivo, se describe cómo entrenar ImageNet en v5p-16
o más grande con
un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.
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.17.0-pod-pjrt export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
puede serv5p-16
o mayor.-
gcloud 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 recurso en cola 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}
Conéctate a tu TPU (trabajador cero) con SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project ${PROJECT_ID} \ --zone ${ZONE}
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
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
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/XLA
Llama 2
En este instructivo, se explica cómo entrenar el modelo de Llama 2 7B en v5p con una bifurcación del repositorio HuggingFace en PyTorch/XLA con modelos Paralelización de gráficos de procesamiento de AA (GSPMD).
Configuración
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 la 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
Crea un recurso de TPU
gcloud 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 con SSH una vez que
QueuedResource
está en el estadoACTIVE
: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 tu recurso en cola se encuentra en estado ACTIVO, la salida será similar al siguiente ejemplo:
state: ACTIVE
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 '
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'
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"
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
establece la marca --spmd_dcn_parallelism
en el número de porciones.
La SPMD_USER_GUIDE proporciona una guía del usuario más detallada que explica todos los diferentes entornos variables y botones de activación de la secuencia de comandos de HF. Cabe destacar que el LIBTPU_INIT_ARGS se incorporará a PyTorch/XLA y de forma predeterminada en versiones futuras.
Limpia
Borra tus TPU y recursos en cola.
Resultados de comparativas
La capacidad de procesamiento de los tres tamaños de modelo de Llama 2 se incluye en los siguientes elementos: desde una tabla de particiones.
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 |
Forma de malla de fragmentación |
(4, 1) |
(64, 1) |
(16, 4) |
Utilización de los flops del modelo (MFU) |
56,67% |
55.80% |
51,85% |
Asistencia y comentarios
Todos los comentarios son bienvenidos. Para compartir comentarios o solicitar asistencia, completa el formulario de comentarios o asistencia de Cloud TPU.