TPU v3

En este documento, se describen la arquitectura y las configuraciones compatibles de Cloud TPU v3.

Arquitectura del sistema

Cada chip de TPU v3 contiene dos TensorCores. Cada TensorCore tiene dos 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 un pod de TPU v3.

Especificaciones clave Valores de Pods de la versión 3
Procesamiento máximo por chip 123 teraflops (bf16)
Capacidad y ancho de banda de HBM2 32 GiB, 900 GBps
Potencia mínima, media o máxima medida 123/220/262 O
Tamaño del pod de TPU 1,024 chips
Topología de interconexión Torus 2D
Procesamiento máximo por Pod 126 petaflops (bf16)
Ancho de banda de Allreduce por Pod 340 TB/s
Ancho de banda de bisección por Pod 6.4 TB/s

En el siguiente diagrama, se ilustra un chip TPU v3.

Diagrama de un chip TPU v3

Los detalles de la arquitectura y las características de rendimiento de TPU v3 están disponibles en Una supercomputadora específica del dominio para entrenar redes neuronales profundas.

Beneficios de rendimiento de TPU v3 sobre v2

El aumento de los FLOPS por TensorCore y la capacidad de memoria en las configuraciones de TPU v3 pueden mejorar el rendimiento de los modelos de las siguientes maneras:

  • Las configuraciones de TPU v3 proporcionan importantes beneficios de rendimiento por TensorCore para los modelos que dependen del procesamiento. Es posible que los modelos que dependen de la memoria en las configuraciones de TPU v2 no logren la misma mejora de rendimiento si también dependen de la memoria en las configuraciones de TPU v3.

  • En los casos en que los datos no caben en la memoria en las configuraciones de TPU v2, TPU v3 puede proporcionar un rendimiento mejorado y una reprocesamiento reducida de los valores intermedios (rematerialización).

  • Las configuraciones de TPU v3 pueden ejecutar nuevos modelos con tamaños de lote que no cabían en las configuraciones de TPU v2. Por ejemplo, TPU v3 podría permitir modelos ResNet más profundos e imágenes más grandes con RetinaNet.

Los modelos que dependen casi totalmente de entradas ("entrada") en TPU v2 debido a que los pasos de entrenamiento están esperando entradas también podrían depender de las entradas con Cloud TPU v3. La guía sobre rendimiento de la canalización puede ayudarte a resolver problemas con las entradas.

Parámetros de configuración

Un pod de TPU v3 se compone de 1,024 chips interconectados con vínculos de alta velocidad. Para crear un dispositivo TPU v3 o una porción de pod, usa la marca --accelerator-type en el comando de creación de TPU (gcloud compute tpus tpu-vm). Para especificar el tipo de acelerador, especifica la versión de TPU y la cantidad de núcleos de TPU. Por ejemplo, para una sola TPU v3, usa --accelerator-type=v3-8. Para una porción de pod v3 con 128 TensorCores, usa --accelerator-type=v3-128.

El siguiente comando muestra cómo crear una porción de pod de TPU v3 con 128 tensores de núcleos:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v3-128 \
    --version=tpu-vm-tf-2.16.1-pjrt

En la siguiente tabla, se enumeran los tipos de TPU v3 compatibles:

Versión de TPU Finalización de compatibilidad
v3-8 (Fecha de finalización aún no establecida)
v3-32 (Fecha de finalización aún no establecida)
v3-128 (Fecha de finalización aún no establecida)
v3-256 (Fecha de finalización aún no establecida)
v3-512 (Fecha de finalización aún no establecida)
v3-1024 (Fecha de finalización aún no establecida)
v3-2048 (Fecha de finalización aún no establecida)

Para obtener más información sobre la administración de TPU, consulta Administra TPU. Para obtener más información sobre la arquitectura del sistema de Cloud TPU, consulta Arquitectura del sistema.