Se usó la API de Cloud Translation para traducir esta página.

Entrena un modelo con la TPU7x (Ironwood)

En este documento, se describe cómo aprovisionar recursos de TPU7x y se proporciona un ejemplo de implementación de una carga de trabajo de entrenamiento con MaxText y XPK.

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

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

Aprovisiona TPU

Para aprovisionar y administrar la TPU7x, sigue los pasos que se indican a continuación:

GKE: Puedes usar GKE con el objetivo de aprovisionar y administrar TPU como un grupo de aceleradores para las cargas de trabajo de aprendizaje automático alojadas en contenedores. Usa Google Cloud CLI para crear manualmente tu instancia de clúster de GKE y personalizar o expandir con precisión los entornos de producción de GKE existentes. Para obtener más información, consulta Acerca de las TPU 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 AA aprovisionen TPU y ejecuten trabajos de entrenamiento sin necesidad de tener un conocimiento profundo 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 XPK en GitHub.
GKE y TPU Cluster Director: TPU Cluster Director está disponible a través de una reserva en modo All Capacity, que te brinda acceso completo a toda tu capacidad reservada (sin retenciones) y visibilidad completa de la topología del hardware de TPU, el estado de utilización y el estado. Para obtener más información, consulta la Descripción general del modo All Capacity.

Implementa 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 administración y la ejecución de cargas de trabajo de aprendizaje automático.

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

Antes de comenzar

Antes de comenzar, completa los siguientes pasos:

Asegúrate de tener un proyecto de Google Cloud con la facturación habilitada.
Obtén acceso a TPU7x. Para obtener más información, comunícate con tu equipo de cuentas.
Asegúrate de que la cuenta que usas con XPK tenga los roles que se indican en el repositorio de GitHub de XPK.

Instala XPK y las dependencias

Instala XPK. Sigue las instrucciones del repositorio de XPK en GitHub.
Instala Docker siguiendo las instrucciones que te proporcione el 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
```
Configura 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"
```
Reemplaza lo siguiente:
- YOUR_PROJECT_ID: Es el ID del proyecto de Google Cloud .
- YOUR_ZONE: Es la zona en la que se creará el clúster. En la versión preliminar, solo se admite us-central1-c.
- YOUR_CLUSTER_NAME el nombre del clúster nuevo.
- YOUR_ACCELERATOR_TYPE: La versión y la topología de la TPU. Por ejemplo, tpu7x-4x4x8 Para obtener una lista de las topologías compatibles, consulta Configuraciones compatibles.
- YOUR_RESERVATION_NAME: el nombre de tu reserva Para las reservas compartidas, usa projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: Es el nombre de tu bucket de Cloud Storage, que será el directorio de salida para el entrenamiento de modelos.

Si no tienes un bucket de Cloud Storage existente, 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

Crea un clúster de una sola NIC y un solo segmento

Sigue las instrucciones de la sección Configura la MTU para optimizar la configuración de tu red.
Propaga 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 usando 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}"
```
Establecer la marca --cluster-cpu-machine-type en n1-standard-8 (o un valor mayor) garantiza que el grupo de nodos predeterminado tenga suficiente CPU para los pods del sistema, por ejemplo, el webhook de JobSet, lo que evita errores. De forma predeterminada, XPK usa e2-standard-16. Algunas zonas solo admiten tipos de CPU específicos, por lo que es posible que debas cambiar entre los tipos n1, n2 y e2. De lo contrario, es posible que recibas errores de cuota.
Agrega 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"
```
Reemplaza lo siguiente:
- EXCLUSION_START_TIME: Es la hora de inicio seleccionada para la exclusión de mantenimiento en formato YYYY-MM-DDTHH:MM:SSZ.
- EXCLUSION_END_TIME: Es la hora de finalización seleccionada para la exclusión de mantenimiento en formato YYYY-MM-DDTHH:MM:SSZ.

Compila o sube la imagen de Docker de MaxText

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

Compila de forma local

Los siguientes comandos copian tu 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 de forma local. Puedes usar tu imagen local directamente en el comando de la carga de trabajo de xpk.

Subir imagen (opcional)

Después de compilar la imagen de Docker de forma local con 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 la ejecución correcta de este comando, deberías ver la imagen de MaxText en gcr.io con el nombre gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Define el comando de entrenamiento de MaxText

Prepara el comando para ejecutar tu secuencia de comandos de entrenamiento dentro del contenedor de Docker.

El modelo MaxText 1B es una configuración dentro del framework de MaxText diseñada para entrenar un modelo de lenguaje con aproximadamente 1,000 millones de parámetros. Usa este modelo para experimentar con escalas de chips pequeñas. 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"

Implementa la carga de trabajo de entrenamiento

Ejecuta el comando xpk workload create para implementar tu trabajo de entrenamiento. Debes especificar la marca --base-docker-image para usar la imagen base de 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 dentro del clúster. En este ejemplo, se agrega $(date +%H%M) al nombre de la carga de trabajo para garantizar la unicidad.