Tipos y topologías de TPU

Configuraciones de TPU

Cuando creas una TPU, debes especificar una configuración de TPU. Puedes especificar una configuración de TPU en términos de TensorCores o chips de TPU. La especificación de una configuración de TPU en términos de chips solo es necesaria en situaciones avanzadas en las que entrenas un modelo grande y ajustas su rendimiento.

Configuración basada en TensorCore

Puedes usar una configuración basada en TensorCore con las TPU v2 - v4. Describes configuraciones basadas en TensorCore con una versión de TPU y la cantidad de TensorCore que deseas usar. Por ejemplo, v4-128 especifica una configuración con TPU v4 y 128 TensorCores. Debes especificar una configuración basada en TensorCore con la marca --accelerator-type del comando gcloud compute tpus tpu-vm create:

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

Configuración de TPU basada en chips

Puedes usar una configuración basada en chips de TPU con TPU v4 y versiones posteriores de TPU. Describes configuraciones basadas en chips con una versión de TPU y una topología (una disposición física de chips TPU). Especificas una versión de TPU con una string como “v4”. Debes especificar una topología con una tupla triple que describa cómo se organizan los chips en tres dimensiones. Especifica una configuración basada en chips con las marcas --version y --topology del comando create gcloud compute tpus tpu-vm. Por ejemplo, el siguiente comando especifica TPU v4 como la versión de TPU y 2x2x4 (16 chips, cada uno con 2 TensorCores) como la topología: 

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=2x2x4
    --version=tpu-vm-tf-2.12.0

Puedes encontrar información adicional sobre las diferentes versiones y topologías de Cloud TPU en la Arquitectura del sistema.

La representación de una topología v4 es AxBxC, en la que A<=B<=C y A, B, C son todos <= 4 o todos los múltiplos enteros de 4. Los valores A, B y C son los recuentos de chips en cada una de las tres dimensiones. Las topologías en las que 2A=B=C o 2A=2B=C también tienen variantes de topología optimizadas para la comunicación integral, por ejemplo, 4×4×8, 8×8×16 y 12×12×24.

Para obtener más información sobre las topologías de v4, consulta Configuraciones de TPU v4.

Variantes de topología

La topología predeterminada asociada con un recuento de chips determinado es la que más se asemeja a un cubo (consulta Formas de topología). Es probable que esta forma sea la mejor opción para el entrenamiento de AA paralelo de datos. Otras topologías pueden ser útiles para cargas de trabajo con varios tipos de paralelismo (por ejemplo, paralelismo de modelos y datos, o partición espacial de una simulación). Estas cargas de trabajo tienen un mejor rendimiento si la topología coincide con el paralelismo que se usa. Por ejemplo, si se coloca el paralelismo de modelos de 4 vías en la dimensión X y el paralelismo de datos de 256 partes en las dimensiones Y y Z coincide con una topología de 4 × 16 × 16.

Los modelos con varias dimensiones de paralelismo tienen un mejor rendimiento con sus dimensiones de paralelismo asignadas a las dimensiones de topología de TPU. Por lo general, son modelos de Data Language grandes (LLM) de datos y modelos paralelos. Por ejemplo, para una porción de pod de TPU v4 con topología de 8 × 16 × 16, sus dimensiones son 8, 16 y 16. Es más eficaz usar paralelismo de modelos de 8 o 16 vías (asignado a una de las dimensiones de topología de TPU física). Un paralelismo de modelos de 4 vías sería subóptimo con esta topología, ya que no está alineado con ninguna de las dimensiones de la topología de TPU, pero sería óptimo con una topología 4x16x32 en la misma cantidad de chips.

Topologías de versión 4 pequeñas

Cloud TPU es compatible con los siguientes fragmentos de TPU v4, de menos de 64 chips, un cubo 4x4x4. Puedes crear estas topologías pequeñas de v4 con su nombre basado en TensorCore (por ejemplo, v4-32) o su topología (por ejemplo, 2x2x4):

Nombre (basado en el recuento de TensorCore) Cantidad de chips Topología
v4‐8 4 2 × 2 × 1
v4‐16 8 2 × 2 × 2
v4‐32 16 2 × 2 × 4
v4‐64 32 2 × 4 × 4

Topologías de v4 de gran tamaño

Las secciones de TPU v4 están disponibles en incrementos de 64 chips, con formas que son múltiplos de 4 en las tres dimensiones. Las dimensiones también deben estar en orden ascendente. En la siguiente tabla, se muestran varios ejemplos. Algunas de estas topologías son topologías “personalizadas” que solo se pueden iniciar con la API de topología porque tienen la misma cantidad de chips que una topología con nombre más usada.

Nombre (basado en el recuento de TensorCore) Cantidad de chips Topología
v4‐128 64 4 × 4 × 4
v4-256 128 4 × 4 × 8
v4‐512 256 4 × 8 × 8
Consulta API de Topology 256 4 × 4 × 16
v4-1,024 512 8 × 8 × 8
v4‐1536 768 8 × 8 × 12
v4‐2048 1,024 8 × 8 × 16
Consulta API de Topology 1,024 4 × 16 × 16
v4-4,096 2,048 8 × 16 × 16
… … …

API de Topology

A fin de crear porciones de pod de Cloud TPU con topología personalizada, la API de TPU gcloud se puede usar de la siguiente manera:

     $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --subnetwork=tpusubnet \
        --type=v4 \
        --topology=4x4x16 \
        --version=runtime-version

Formas de la topología

En la siguiente imagen, se muestran tres formas de topología posibles y su configuración de marca topology asociada.

imagen

Compatibilidad con el tipo de TPU

Puedes cambiar el tipo de TPU a otro que tenga la misma cantidad de TensorCores o chips (por ejemplo, v3-128 y v4-128) y ejecutar la secuencia de comandos de entrenamiento sin cambios en el código. Sin embargo, si cambias a un tipo de TPU con una cantidad mayor o menor de chips o TensorCores, deberás realizar una optimización y un ajuste significativo. Para obtener más información, consulta Entrenamiento en pods de TPU.

¿Qué sigue?