Esta página se ha traducido con Cloud Translation API.

Entrenar un modelo con TPU7x (Ironwood)

En este documento se describe cómo aprovisionar recursos de TPU v7 y se ofrece un ejemplo de cómo desplegar una carga de trabajo de entrenamiento con MaxText y XPK.

TPU7x es la primera versión de la familia Ironwood,la TPU de séptima generación de Google Cloud. La generación Ironwood se ha diseñado para el entrenamiento y la inferencia de IA a gran escala. Para obtener más información, consulta TPU7x.

Para ver más ejemplos optimizados para TPU7x, consulta Recetas de entrenamiento para TPU Ironwood en GitHub.

Provisionar TPUs

Puedes aprovisionar y gestionar TPU v7 con los siguientes métodos:

GKE puedes usar GKE para aprovisionar y gestionar las TPUs como un conjunto de aceleradores para tus cargas de trabajo de aprendizaje automático en contenedores. Usa la CLI de Google Cloud para crear manualmente tu instancia de clúster de GKE si quieres personalizarla con precisión o ampliar los entornos de producción de GKE que ya tengas. Para obtener más información, consulta el artículo Acerca de las TPUs en GKE.
GKE y XPK: XPK es una herramienta de línea de comandos que simplifica la creación de clústeres y la ejecución de cargas de trabajo en GKE. Está diseñada para que los profesionales del aprendizaje automático puedan aprovisionar TPUs y ejecutar tareas de entrenamiento sin necesidad de tener conocimientos profundos de Kubernetes. Usa XPK para crear rápidamente clústeres de GKE y ejecutar cargas de trabajo para pruebas de concepto y pruebas. Para obtener más información, consulta el repositorio de GitHub de XPK.
GKE y TPU Cluster Director: TPU Cluster Director está disponible a través de una reserva en el modo Toda la capacidad, que te da acceso completo a toda la capacidad reservada (sin retenciones) y visibilidad total de la topología del hardware de TPU, el estado de utilización y el estado de salud. Para obtener más información, consulta el resumen del modo Capacidad total.

Desplegar una carga de trabajo de entrenamiento con MaxText y XPK

Usa el kit de procesamiento acelerado (XPK) para crear clústeres de GKE para pruebas de concepto y pruebas. XPK es una herramienta de línea de comandos diseñada para simplificar el aprovisionamiento, la gestión y la ejecución de cargas de trabajo de aprendizaje automático.

En las siguientes secciones se muestra cómo desplegar una carga de trabajo de entrenamiento con MaxText y XPK.

Antes de empezar

Antes de empezar, completa los siguientes pasos:

Asegúrate de que tienes un Google Cloud proyecto con la facturación habilitada.
Obtener acceso a TPU7x. Para obtener más información, póngase en contacto con el equipo de su cuenta.
Asegúrate de que la cuenta que usas con XPK tenga los roles que se indican en el repositorio de GitHub de XPK.

Instalar XPK y las dependencias

Instala XPK. Sigue las instrucciones del repositorio de GitHub de XPK.
Instala Docker siguiendo las instrucciones que te proporcione tu administrador o las instrucciones de instalación oficiales. Una vez instalado, ejecuta los siguientes comandos para configurar Docker y probar la instalación:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
Define las siguientes variables de entorno:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Haz los cambios siguientes:
- YOUR_PROJECT_ID: tu ID de proyecto Google Cloud .
- YOUR_ZONE: la zona en la que se creará el clúster. En la vista previa, solo se admite us-central1-c.
- YOUR_CLUSTER_NAME: el nombre del nuevo clúster.
- YOUR_ACCELERATOR_TYPE: la versión y la topología de la TPU. Por ejemplo, tpu7x-4x4x8. Para ver una lista de las topologías admitidas, consulta Configuraciones admitidas.
- YOUR_RESERVATION_NAME: el nombre de tu reserva. Para las reservas compartidas, usa projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: el nombre de tu segmento de Cloud Storage, que será el directorio de salida de la preparación del modelo.

Si no tienes ningún segmento de Cloud Storage, crea uno con el siguiente comando:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Crear un clúster de una sola NIC y una sola partición

Sigue las instrucciones de la sección Configurar MTU para optimizar la configuración de tu red.
Rellena la variable ${CLUSTER_ARGUMENTS}, que usarás en el comando xpk cluster create:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Crea tu clúster de GKE con grupos de nodos TPU7x mediante el comando xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Si asignas el valor n1-standard-8 (o un valor superior) a la marca --cluster-cpu-machine-type, te aseguras de que el grupo de nodos predeterminado tenga suficientes CPUs para los pods del sistema, como el webhook JobSet, y evitas errores. De forma predeterminada, XPK usa e2-standard-16. Algunas zonas solo admiten tipos de CPU específicos, por lo que es posible que tengas que cambiar entre los tipos n1, n2 y e2. De lo contrario, es posible que se produzcan errores de cuota.
Añade una exclusión de mantenimiento para evitar las actualizaciones del clúster:
```
gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"
```
Haz los cambios siguientes:
- EXCLUSION_START_TIME: la hora de inicio que has seleccionado para la exclusión de mantenimiento en formato YYYY-MM-DDTHH:MM:SSZ.
- EXCLUSION_END_TIME: la hora de finalización que has seleccionado para la exclusión de mantenimiento en formato YYYY-MM-DDTHH:MM:SSZ.

Compila o sube la imagen Docker de MaxText

Puedes compilar una imagen Docker de forma local con las secuencias de comandos que proporciona MaxText o usar una imagen precompilada.

Compilar localmente

Los siguientes comandos copian el directorio local en el contenedor:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Después de ejecutar los comandos correctamente, deberías ver una imagen llamada maxtext_base_image creada localmente. Puedes usar tu imagen local directamente en el comando de carga de trabajo xpk.

Subir imagen (opcional)

Después de compilar la imagen Docker de forma local siguiendo las instrucciones de la sección anterior, puedes subirla al registro con el siguiente comando:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Después de ejecutar este comando correctamente, deberías ver la imagen MaxText en gcr.io con el nombre gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Definir el comando de entrenamiento de MaxText

Prepara el comando para ejecutar la secuencia de comandos de entrenamiento en el contenedor Docker.

El modelo MaxText 1B es una configuración del framework MaxText diseñada para entrenar un modelo de lenguaje con aproximadamente 1000 millones de parámetros. Usa este modelo para experimentar con chips de tamaño pequeño. El rendimiento no está optimizado.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

Desplegar la carga de trabajo de entrenamiento

Ejecuta el comando xpk workload create para implementar tu tarea de entrenamiento. Debes especificar la marca --base-docker-image para usar la imagen base MaxText o la marca --docker-image y la imagen que quieras usar. Puedes incluir la marca --enable-debug-logs para habilitar el registro de depuración.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

Los nombres de las cargas de trabajo deben ser únicos en el clúster. En este ejemplo, se añade $(date +%H%M) al nombre de la carga de trabajo para que sea único.