Suporte do Dataflow para TPUs

As TPUs do Google Cloud Platform são aceleradores de IA personalizados criados pelo Google e otimizados para treinamento e uso de grandes modelos de IA. Elas foram projetadas para escalonar de maneira econômica para uma ampla variedade de cargas de trabalho de IA e oferecem versatilidade para acelerar cargas de trabalho de inferência em frameworks de IA, incluindo PyTorch, JAX e TensorFlow. Para mais detalhes sobre TPUs, consulte Introdução à TPU do Google Cloud Platform.

Pré-requisitos para usar TPUs no Dataflow

  • Seus projetos Google Cloud precisam ser aprovados para usar essa oferta de GA.

Limitações

Esta oferta está sujeita às seguintes limitações:

  • Somente aceleradores de TPU de host único são aceitos: a oferta de TPU do Dataflow aceita apenas configurações de TPU de host único em que cada worker do Dataflow gerencia um ou vários dispositivos de TPU que não estão interconectados com TPUs gerenciadas por outros workers.
  • Só há suporte para pools de workers de TPU homogêneos: recursos como ajuste direito do Dataflow e Dataflow Prime não são compatíveis com cargas de trabalho de TPU.

Preços

Os jobs do Dataflow que usam TPUs são cobrados pelas horas de chip de TPU do worker consumidas, e não pela CPU e memória do worker. Para mais informações, consulte a página de preços do Dataflow.

Disponibilidade

Os seguintes aceleradores de TPU e regiões de processamento estão disponíveis.

Aceleradores de TPU compatíveis

As combinações de aceleradores de TPU compatíveis são identificadas pela tupla (tipo de TPU, topologia de TPU).

  • Tipo de TPU se refere ao modelo do dispositivo de TPU.
  • A topologia de TPU se refere ao número e à disposição física dos chips de TPU em uma fração.

Para configurar o tipo e a topologia de TPUs para workers do Dataflow, use a opção de pipeline worker_accelerator formatada como type:TPU_TYPE;topology:TPU_TOPOLOGY.

As seguintes configurações de TPU são compatíveis com o Dataflow:

Tipo de TPU Topologia worker_machine_type obrigatório
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t
tpu-v6e-slice 1x1 ct6e-standard-1t
tpu-v6e-slice 2x2 ct6e-standard-4t
tpu-v6e-slice 2x4 ct6e-standard-8t
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t

Regiões

Para informações sobre regiões e zonas disponíveis para TPUs, consulte Regiões e zonas de TPU na documentação do Cloud TPU.

A seguir