Configuraciones de TPU

Configuraciones de TPU v4

Un pod de TPU v4 consta de 4,096 chips interconectados con vínculos de alta velocidad reconfigurables. Las redes flexibles de TPU v4 te permiten conectar los chips en un Podslice del mismo tamaño de varias maneras. Cuando creas un podslice de TPU, debes especificar la versión de TPU y la cantidad de recursos TPU que necesitas. Cuando creas una porción de pod de TPU v4, puedes especificar su tipo y tamaño de las siguientes dos maneras:

  1. Usa AcceleratorType

    Esta es la opción recomendada cuando no especificas una topología.

  2. Usa AcceleratorConfig

    Recomendamos usar AcceleratorConfig cuando desees personalizar la topología física. Por lo general, esto es necesario para el ajuste de rendimiento con 256 chips de Podslices.

Para configurar las TPU v4 con AcceleratorType, usa la marca --accelerator-type cuando crees tu podslice de TPU. Configura --accelerator-type como una string que contenga la versión de TPU y la cantidad de TensorCores que deseas usar. Por ejemplo, para crear un Podslice v4 con 32 TensorCores, usarías --accelerator-type=v4-32.

Para configurar las TPU v4 con AcceleratorConfig, usa las marcas --version y --topology. Establece --version en la versión de TPU que deseas usar y --topology en la disposición física de los chips TPU en el Podslice. Por ejemplo, para crear un Podslice v4 con 16 chips, usarías --version=v4 y --topology=2x2x4.

Debes especificar una topología de TPU con una tupla de 3, 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 de todos a todos, por ejemplo, 4 × 4 × 8, 8 × 8 × 16 y 12 × 12 × 24. Estas se conocen como topologías de "tori torcido". Se puede especificar una configuración de tori torcido para cualquier topología que pueda admitirla. Consulta las configuraciones de variantes para obtener una lista de las topologías trenzadas admitidas.

Con el siguiente comando, se crea una porción de pod de TPU v4 con 128 chips TPU organizados en un arreglo de 4 x 4 x 8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=4x4x8
    --version=tpu-vm-tf-2.13.0

Puedes crear un Podslice de TPU v4 con la misma cantidad de chips/TensorCores mediante --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-512
    --version=tpu-vm-tf-2.13.0

El número después de la versión de TPU (v4) especifica la cantidad de TensorCores. Hay dos TensorCores en una TPU v4, por lo que la cantidad de chips TPU sería 512/2 = 256.

En las siguientes ilustraciones, se muestran algunas topologías comunes de TPU v4.

imagen

Los Podslices más grandes se pueden compilar a partir de uno o más “cubos” de 4×4×4.

Variantes de topología

Algunas topologías que contienen la misma cantidad de chips se pueden organizar de diferentes maneras. Por ejemplo, se puede configurar un Podslice de TPU con 512 chips (1, 024 TensorCores) mediante las siguientes topologías: 4x4x32, 4x8x16 u 8x8x8. Una porción de pod de TPU con 2,048 chips (4,096 TensorCores) ofrece aún más opciones de topología: 4 x 4 x 128, 4 x 8 x 64, 4 x 16 x 32 y 8 x 16 x 16. Un Podslice de TPU con 2,048 chips (4,096 TensorCores) ofrece aún más opciones de topología: 4 x 4 x 128, 4 x 8 x 64, 4 x 16 x 32 y 8 x 16 x 16.

La topología predeterminada asociada con un recuento de chips determinado es la más similar a un cubo (consulta Formas de topología). Es probable que esta forma sea la mejor opción para el entrenamiento de AA con datos paralelos. Otras topologías pueden ser útiles para cargas de trabajo con varios tipos de paralelismo (por ejemplo, paralelismo de modelo 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 usado. Por ejemplo, ubicar el paralelismo del modelo de 4 vías en la dimensión X y el paralelismo de datos de 256 vías en las dimensiones Y y Z coincide con una topología de 4 × 16 × 16.

Los modelos con múltiples dimensiones de paralelismo tienen un mejor rendimiento cuando las dimensiones de paralelismo se asignan a las dimensiones de topología de TPU. Por lo general, son modelos de lenguaje extenso (LLM) paralelos de datos y modelos. Por ejemplo, para una porción de pod de TPU v4 con una topología de 8 x 16 x 16, las dimensiones de la topología de TPU son 8, 16 y 16. Es más eficaz usar un paralelismo de modelos de 8 o 16 vías (asignado a una de las dimensiones de la 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 de 4 × 16 × 32 en la misma cantidad de chips.

Las configuraciones de TPU v4 constan de dos grupos: el que tiene topologías de menos de 64 chips (topologías pequeñas) y los que tienen topologías de más de 64 chips (topologías grandes).

Topologías v4 pequeñas

Cloud TPU es compatible con los siguientes pods TPU v4 más pequeños que 64 chips, un cubo 4x4x4. Puedes crear estas topologías v4 pequeñas mediante el nombre basado en TensorCore (por ejemplo, v4-32) o su topología (por ejemplo, 2x2x4):

Nombre (según la cantidad de TensorCore) Cantidad de chips Topología
Versión 4-8 4 2x2x1
Versión 4-16 8 2x2x2
Versión 4-32 16 2x2x4
Versión 4-64 32 2x4x4

Topologías v4 grandes

Los Pods 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 creciente. En la siguiente tabla, se muestran varios ejemplos. Algunas de estas topologías son topologías “personalizadas” que solo se pueden crear con las marcas --type y --topology porque hay más de una forma de organizar los chips.

Con el siguiente comando, se crea un podslice de TPU v4 con 512 chips TPU organizados en un arreglo de 8x8x8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=8x8x8
    --version=tpu-vm-tf-2.13.0

Puedes crear un Podslice de TPU v4 con la misma cantidad de chips/TensorCores mediante --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-1024
    --version=tpu-vm-tf-2.13.0
Nombre (según la cantidad de TensorCore) Cantidad de chips Topología
Versión 4-128 64 4 × 4 × 4
v4-256 128 4 × 4 × 8
v4-512 256 4 × 8 × 8
N/A: Se deben usar las marcas --type y --topology. 256 4 × 4 × 16
v4-1024 512 8 × 8 × 8
v4-1536 768 8 × 8 × 12
v4-2048 1,024 8 × 8 × 16
N/A: Se deben usar las marcas --type y --topology 1,024 4 × 16 × 16
v4-4096 2,048 8 × 16 × 16
… … …

Configuraciones de TPU v3

Un pod de TPU v3 consta de 1,024 chips interconectados con vínculos de alta velocidad. Para crear un dispositivo TPU v3 o un Podslice, usa la marca --accelerator-type para el comando 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. Para una sola TPU v3, usa --accelerator-type=v3-8. Para un Podslice v3 con 128 TensorCores, usa --accelerator-type=v3-128.

En el siguiente comando, se muestra cómo crear un Podslice de TPU v3 con 128 TensorCores:

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

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 cómo administrar TPU, consulta Administra TPU. Para obtener más información sobre las diferentes versiones de Cloud TPU, consulta Arquitectura del sistema.

Configuraciones de TPU v2

Un pod de TPU v2 consta de 512 chips interconectados con vínculos de alta velocidad reconfigurables. Si quieres crear un Podslice de TPU v2, usa la marca --accelerator-type para el comando 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. Para una sola TPU v2, usa --accelerator-type=v2-8. Para un Podslice v2 con 128 núcleos de Tensor, usa --accelerator-type=v2-128.

En el siguiente comando, se muestra cómo crear un Podslice de TPU v2 con 128 TensorCores:

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

Para obtener más información sobre cómo administrar TPU, consulta Administra TPU. Para obtener más información sobre las diferentes versiones de Cloud TPU, consulta Arquitectura del sistema.

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

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

Compatibilidad con tipos 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 ajustes y optimización significativos. Para obtener más información, consulta Entrenamiento en pods de TPU.

Versiones de software de TPU

La versión del software de TPU que debes usar depende de la arquitectura de TPU (VM de TPU o Nodo de TPU) y del framework de AA que usas (TensorFlow, PyTorch o JAX).

VM de TPU

Las versiones de software de TPU están disponibles para los frameworks de TensorFlow, PyTorch y JAX.

TensorFlow

Usa la versión de software de TPU que coincida con la versión de TensorFlow con la que se escribió tu modelo. Por ejemplo, si usas TensorFlow 2.13.0, utiliza la versión de software de TPU tpu-vm-tf-2.13.0. Estas son las versiones de software de VM de TPU de TensorFlow compatibles actuales para TPU:

  • tpu-vm-tf-2.13.0
  • tpu-vm-tf-2.12.1
  • tpu-vm-tf-2.12.0
  • tpu-vm-tf-2.11.1
  • tpu-vm-tf-2.11.0
  • tpu-vm-tf-2.10.1
  • tpu-vm-tf-2.10.0
  • tpu-vm-tf-2.9.3
  • tpu-vm-tf-2.9.1
  • tpu-vm-tf-2.8.4
  • tpu-vm-tf-2.8.3
  • tpu-vm-tf-2.8.0
  • tpu-vm-tf-2.7.4
  • tpu-vm-tf-2.7.3

Si usas un Podslice, agrega -pod a la versión de software de TPU que deseas usar. Por ejemplo, tpu-vm-tf-2.13.0-pod

VM de TPU con TPU v4

Si entrenas un modelo en TPU v4 con TensorFlow, usa una de las versiones de v4 que se muestran en la siguiente tabla.

Versión de TensorFlow Versión de software de TPU
2.13.0 tpu-vm-tf-2.13.0, tpu-vm-tf-2.13.0-pod
2.12.1 tpu-vm-tf-2.12.1, tpu-vm-tf-2.12.1-pod
2.12.0 tpu-vm-tf-2.12.0, tpu-vm-tf-2.12.0-pod
2.11.1 tpu-vm-tf-2.11.1, tpu-vm-tf-2.11.1-pod
2.11.0 tpu-vm-tf-2.11.0, tpu-vm-tf-2.11.0-pod
2.10.0 tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4
2.10.1 tpu-vm-tf-2.10.1, tpu-vm-tf-2.10.1-pod
2.9.3 tpu-vm-tf-2.9.3-v4, tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4, tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4, tpu-vm-tf-2.9.1-pod-v4

Para obtener más información sobre las versiones de parches de TensorFlow, consulta Versiones de parches de TensorFlow compatibles.

Las VMs de TPU se crean con TensorFlow y la biblioteca Libtpu correspondiente preinstalada. Si creas tu propia imagen de VM, especifica las siguientes versiones del software de TPU de TensorFlow y las versiones de Libtpu correspondientes:

Versión de TensorFlow Versión de libtpu.so
2.13.0 1.7.0
2.12.1 1.6.1
2.12.0 1.6.0
2.11.1 1.5.1
2.11.0 1.5.0
2.10.1 1.4.1
2.10.0 1.4.0
2.9.3 1.3.2
2.9.1 1.3.0
2.8.3 1.2.3
2.8.* 1.2.0
2.7.3 1.1.2

PyTorch

Usa la versión de software de TPU que coincida con la versión de PyTorch con la que se escribió tu modelo. Por ejemplo, si usas PyTorch 1.13 y TPU v2 o v3, usa la versión del software de TPU tpu-vm-pt-1.13. Si usas TPU v4, usa la versión del software de TPU tpu-vm-v4-pt-1.13. Se usa la misma versión de software de TPU para los pods de TPU (p.ej., v2-32, v3-128 y v4-32). Estas son las versiones del software de TPU compatibles actuales:

TPU v2/v3:
  • tpu-vm-pt-2.0 (pytorch-2.0)
  • tpu-vm-pt-1.13 (pytorch-1.13)
  • tpu-vm-pt-1.12 (pytorch-1.12)
  • tpu-vm-pt-1.11 (pytorch-1.11)
  • tpu-vm-pt-1.10 (pytorch-1.10)
  • Versión 2-alpha (pytorch-1.8.1)
TPU v4:
  • tpu-vm-v4-pt-2.0 (pytorch-2.0)
  • tpu-vm-v4-pt-1.13 (pytorch-1.13)

Cuando creas una VM de TPU, la versión más reciente de PyTorch está preinstalada en la VM de TPU. La versión correcta de libtpu.so se instala automáticamente cuando instalas PyTorch.

Para cambiar la versión actual del software de PyTorch, consulta Cambia la versión de PyTorch.

JAX

Debes instalar JAX de forma manual en tu VM de TPU, ya que no hay una versión de software de TPU específica de JAX. Para todas las versiones de TPU, usa tpu-ubuntu2204-base. La versión correcta de libtpu.so se instala automáticamente cuando instalas JAX.

Nodo TPU

TensorFlow

Usa la versión de software de TPU que coincida con la versión de TensorFlow con la que se escribió tu modelo. Por ejemplo, si usas TF 2.12.0, utiliza la versión de software de TPU 2.12.0. Estas son las versiones específicas del software de TPU de TensorFlow:

  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.2
  • 2.7.3

Para obtener más información sobre las versiones de parches de TensorFlow, consulta Versiones de parches de TensorFlow compatibles.

Cuando creas un nodo TPU, la versión más reciente de TensorFlow está preinstalada en el nodo TPU.

PyTorch

Usa la versión de software de TPU que coincida con la versión de PyTorch con la que se escribió tu modelo. Por ejemplo, si usas PyTorch 1.9, usa la versión de software pytorch-1.9.

Estas son las versiones específicas del software de TPU de PyTorch:

  • pytorch‐2.0
  • pytorch-1.13
  • pytorch‐1.12
  • pytorch‐1.11
  • pytorch-1.10
  • pytorch-1.9
  • pytorch-1.8
  • pytorch-1.7

  • pytorch-1.6

  • pytorch-nightly

Cuando creas un nodo TPU, la versión más reciente de PyTorch está preinstalada en el nodo TPU.

JAX

Debes instalar JAX de forma manual en tu VM de TPU para que no haya una versión de software de TPU específica de JAX preinstalada. Puedes usar cualquiera de las versiones de software enumeradas para TensorFlow.

¿Qué sigue?