TPU v6e

En este documento, se describen la arquitectura y las configuraciones admitidas de la Cloud TPU v6e (Trillium).

Trillium es el acelerador de IA de última generación de las Cloud TPU. En todas las plataformas técnicas, como la API y los registros, y en todo este documento, se hará referencia a Trillium como v6e.

Con una capacidad de 256 chips por Pod, la TPU v6e comparte muchas similitudes con la TPU v5e. Este sistema está optimizado para ser el producto de mayor valor para el entrenamiento, el ajuste y la publicación de transformadores, modelos de texto a imagen y redes neuronales convolucionales (CNN).

Arquitectura del sistema

Cada chip v6e contiene un TensorCore. Cada TensorCore tiene 2 unidades de multiplicación de matrices (MXU), una unidad vectorial y una unidad escalar. En la siguiente tabla, se muestran las especificaciones clave y sus valores para la TPU v6e en comparación con la TPU v5e.

Configuraciones admitidas

En la siguiente tabla, se muestran las formas de corte 2D que se admiten para la versión 6e:

Los segmentos con 8 chips (v6e-8) conectados a una sola VM están optimizados para la inferencia, lo que permite que los 8 chips se usen en una sola carga de trabajo de procesamiento. Puedes realizar inferencias con varios hosts usando Pathways on Cloud. Para obtener más información, consulta Cómo realizar inferencias en varios hosts con Pathways.

Para obtener información sobre la cantidad de VMs para cada topología, consulta Tipos de VM.

Tipos de VM

Cada VM de TPU v6e puede contener 1, 4 u 8 chips. Las porciones de 4 chips y más pequeñas tienen el mismo nodo de acceso a la memoria no uniforme (NUMA). Para obtener más información sobre los nodos NUMA, consulta Acceso a memoria no uniforme en Wikipedia.

Las porciones de v6e se crean con VMs de medio host, cada una con 4 chips de TPU. Existen dos excepciones a esta regla:

• v6e-1: Es una VM con un solo chip, diseñada principalmente para pruebas.
• v6e-8: Es una VM de host completo que se optimizó para un caso de uso de inferencia con los 8 chips conectados a una sola VM.

En la siguiente tabla, se muestra una comparación de los tipos de VM de TPU v6e:

Cómo especificar la configuración de v6e

Cuando asignas un segmento de TPU v6e con la API de TPU, especificas su tamaño y forma con el parámetro AcceleratorType.

Si usas GKE, usa la marca --machine-type para especificar un tipo de máquina que admita la TPU que deseas usar. Para obtener más información, consulta Planifica las TPU en GKE en la documentación de GKE.

Utilizar AcceleratorType

Cuando asignas recursos de TPU, usas AcceleratorType para especificar la cantidad de TensorCores en una porción. El valor que especifiques para AcceleratorType es una cadena con el formato v$VERSION-$TENSORCORE_COUNT. Por ejemplo, v6e-8 especifica una porción de TPU v6e con 8 TensorCores.

En el siguiente ejemplo, se muestra cómo crear un segmento de TPU v6e con 32 TensorCores usando AcceleratorType:

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=zone \
        --accelerator-type=v6e-32 \
        --version=v2-alpha-tpuv6e

¿Qué sigue?

• Ejecuta el entrenamiento y la inferencia con la TPU v6e