Información sobre las TPU

Las unidades de procesamiento tensorial (TPU) son los circuitos integrados específicos de la aplicación (ASIC) personalizados de Google que se usan para acelerar las cargas de trabajo de aprendizaje automático (AA) y de inteligencia artificial (IA). Ya sea que entrenes modelos de base complejos durante semanas o ejecutes inferencias a gran escala, las TPU ofrecen recursos de procesamiento especializados y escalables que están optimizados para frameworks de IA avanzados.

En este documento, se describe el rol de las TPU en Google Cloud, las especificaciones técnicas de cada versión de TPU disponible, las características de rendimiento, las consideraciones sobre los precios y cómo se asignan a las series de máquinas de Compute Engine.

¿Qué es una TPU?

Una TPU es un ASIC diseñado a medida para satisfacer las demandas de las operaciones de matrices masivas que son fundamentales para los algoritmos de AA.

Entre los componentes arquitectónicos clave de un sistema de TPU, se incluyen los siguientes:

  • TensorCore: Es la unidad de procesamiento de un chip TPU. Cada TensorCore consta de una o más unidades de multiplicación de matrices (MXU) que usan la arquitectura de array sistólico para realizar la multiplicación de matrices con alta eficiencia, junto con unidades escalares y vectoriales para el flujo de control y el cálculo general.

  • SparseCore: Son procesadores de flujo de datos especializados diseñados para acelerar las operaciones dispersas. Los SparseCores funcionan junto con los TensorCores para procesar grandes tablas de incorporación de manera eficiente, lo que los hace ideales para acelerar los modelos de recomendación y las incorporaciones a gran escala.

  • Memoria de alto ancho de banda (HBM): Es una gran memoria física conectada al chip de la TPU que ofrece un ancho de banda de memoria masivo. La HBM permite el entrenamiento y la entrega de modelos más grandes con tamaños de lote más grandes.

  • Instancia de TPU: Es una instancia de procesamiento de Linux que se ejecuta en un host de TPU y tiene acceso directo al hardware de TPU subyacente, lo que ofrece acceso a bibliotecas de alto rendimiento, conectividad SSH y registros de depuración del tiempo de ejecución.

  • Porciones y pods de TPU: Un pod de TPU es un clúster contiguo grande de chips TPU vinculados físicamente. Una porción de TPU es una partición lógica de un Pod (que comprende uno o más hosts conectados con interconexiones de alta velocidad) asignada para ejecutar una sola carga de trabajo.

  • Compilador de XLA: El compilador de álgebra lineal acelerada (XLA) compila grafos de aprendizaje automático en instrucciones de máquina optimizadas que se ejecutan en los TensorCores de la TPU.

Versiones de TPU y series de máquinas

Compute Engine admite varias generaciones y versiones de TPU. En la siguiente tabla, se proporcionan las especificaciones clave de cada versión de TPU compatible:

Versión de TPU Series de máquinas Procesamiento máximo por chip (TFLOPs) Memoria y ancho de banda de la TPU Núcleos por chip Ancho de banda de interconexión (ICI) por chip Ancho de banda de red (DCN) por chip Chips por Pod
TPU7x (Ironwood) tpu7x-standard

BF16: 2,307

FP8: 4,614

192 GiB (7,380 GBps)

2 TensorCores

4 SparseCores

1,200 GBps 100 Gbps 9,216
TPU v6e (Trillium) ct6e-standard

BF16: 918

FP8: 918

32 GiB (1,638 GBps)

1 TensorCore

2 SparseCores

800 GBps 100 Gbps 256
TPU v5p ct5p-hightpu

BF16: 459

FP8: 459

95 GiB (2,765 GBps)

2 TensorCores

4 SparseCores

1,200 GBps 50 Gbps 8,960

Características de rendimiento de la TPU

Las TPU están optimizadas para ejecutar operaciones de álgebra de matrices densas con la máxima eficiencia. Para elegir cuándo es mejor usar una TPU para tus cargas de trabajo de AA, ten en cuenta los siguientes lineamientos.

Casos de uso ideales y requisitos de carga de trabajo

Elegir el modelo de TPU adecuado depende de las características computacionales, los requisitos de memoria y las necesidades de escalabilidad de tu carga de trabajo. Selecciona una de las siguientes pestañas para ver las recomendaciones.

Entrenamiento de IA/AA

  • Preentrenamiento de modelos grandes: Entrenamiento de modelos masivos desde cero Estas cargas de trabajo están vinculadas a la capacidad de procesamiento y la comunicación.

    • Capacidad de TPU requerida: FLOPS máximos altos, compatibilidad con baja precisión y ancho de banda de interconexión alto.
    • Versiones de TPU recomendadas:
      • TPU7x. Esta versión proporciona un alto rendimiento de FP8.
      • TPU v5p Esta versión admite configuraciones de clúster a gran escala.
  • Ajuste y destilación de modelos: Adaptación de modelos existentes con métodos eficientes en cuanto a parámetros o entrenamiento de variantes de modelos más pequeños

    • Capacidad de TPU requerida: capacidad de procesamiento moderada y capacidad de memoria suficiente.
    • Versiones de TPU recomendadas:
      • TPU v6e Esta versión es rentable para las tareas de entrenamiento estándar.
      • TPU7x. Esta versión ofrece un alto rendimiento para el ajuste a gran escala.

Inferencia de IA/AA

  • Entrega de modelos en línea: Implementa modelos para la interacción en tiempo real en la que la baja latencia es fundamental. Estas cargas de trabajo están limitadas por el ancho de banda de la memoria durante la generación secuencial de tokens.

    • Capacidad de TPU requerida: Gran ancho de banda de HBM y SRAM en chip grande para minimizar la latencia de recuperación de datos.
    • Versiones de TPU recomendadas:
      • TPU v6e Esta versión está optimizada para la publicación rentable a gran escala.
      • TPU7x. Esta versión proporciona un alto ancho de banda de HBM y una gran SRAM para acelerar la generación de tokens.
  • Ejecución de inferencia por lotes: Procesamiento sin conexión de grandes conjuntos de datos en los que el objetivo principal es una alta capacidad de procesamiento Estas cargas de trabajo están vinculadas a la capacidad de procesamiento durante la fase de relleno previo de la instrucción.

    • Capacidad de TPU requerida: Alta densidad de procesamiento para maximizar la capacidad de procesamiento por dólar.
    • Versiones de TPU recomendadas:
      • TPU v6e Esta versión es rentable para el procesamiento por lotes de alto rendimiento.
      • TPU7x. Esta versión ofrece una alta densidad de procesamiento para procesar grandes lotes rápidamente.

Sistemas de recomendación

  • Procesamiento de embeddings grandes: Entrenamiento y publicación de modelos de recomendación que usan tablas de embeddings masivas Estas cargas de trabajo están limitadas por la capacidad y la comunicación.
    • Capacidad de TPU requerida: Hardware SparseCore especializado para procesar operaciones dispersas en paralelo, gran capacidad de HBM y compatibilidad con la descarga en la memoria del host.
    • Versiones de TPU recomendadas:
      • TPU7x. Esta versión incluye cuatro SparseCores por chip y ofrece la mayor capacidad de HBM.
      • TPU v6e Esta versión incluye dos SparseCores por chip para una publicación rentable.
      • TPU v5p Esta versión incluye cuatro SparseCores por chip para el entrenamiento a gran escala.

Aprendizaje por refuerzo

  • Ejecución de bucles de aprendizaje por refuerzo: Administración de cargas de trabajo híbridas que requieren bucles ajustados de generación de muestras y actualización de pesos de políticas.

    • Capacidad de TPU requerida: ICI de baja latencia para transferencias rápidas de pesos y activaciones entre chips, y conexiones de host de alto ancho de banda.
    • Versiones de TPU recomendadas:

      Estas versiones de TPU usan una interconexión de toro 3D de alta velocidad.

Consideraciones sobre la carga de trabajo

Las TPU podrían no ser la opción óptima para las siguientes cargas de trabajo:

  • Programas de álgebra lineal que requieren ramificaciones lógicas frecuentes o contienen muchas operaciones algebraicas a nivel de los elementos
  • Cargas de trabajo que dependen de operadores personalizados en JAX o PyTorch que se ejecutan en CPUs.
  • Cargas de trabajo científicas que requieren aritmética de punto flotante de doble precisión (FP64).

Consideraciones sobre el precio

Los siguientes factores determinan los precios de las TPU en Google Cloud:

  • Versión de TPU

  • Ubicación

  • Modelo de consumo

Para obtener detalles completos sobre los precios, consulta la página Precios de TPU.

Modelos de compra y consumo

Puedes aprovisionar instancias de TPU con cualquiera de las siguientes opciones de consumo, según tus requisitos de disponibilidad y tolerancia al costo:

  • A pedido: Precios estándar de pago por uso (PAYG) para la ejecución inmediata. Ofrece la máxima flexibilidad en el ciclo de vida, pero la disponibilidad está sujeta a las limitaciones de recursos físicos.
  • Spot: Capacidad Google Cloud excedente con grandes descuentos (hasta un 70%). Dado que Google Cloud puede detener o borrar las VMs Spot para recuperar capacidad con un aviso de 30 segundos, estas instancias de procesamiento son ideales para cargas de trabajo tolerantes a errores y con puntos de control.
  • Inicio flexible: Permite solicitar instancias de procesamiento desde un grupo exclusivo por hasta siete días con mayor disponibilidad.
  • Reservas futuras: Son reservas que garantizan la capacidad y ofrecen un precio con descuento en comparación con los precios según demanda. Reserva capacidad de TPU para una fecha y hora futuras especificadas, ya sea por hasta 90 días (modo de calendario) o por 1 año o más. SiGoogle Cloud aprueba tu solicitud, podrás comenzar a usar la capacidad durante el período de la reserva.

Descuentos por compromiso de uso (CUD)

Cuando solicitas una reserva futura de 1 año o más, puedes obtener un CUD en la capacidad de TPU reservada. Aplicar un CUD a tus instancias de procesamiento te brinda una reducción significativa del precio en comparación con las tarifas a pedido a cambio de un período de uso comprometido.

Para obtener más información, consulta CUDs para Compute Engine.

¿Qué sigue?