Suporte do Dataflow para TPUs

As TPUs da Google Cloud Platform são aceleradores de IA personalizados criados pela Google que são otimizados para a preparação e a utilização de modelos de IA grandes. Estão concebidas para serem rentáveis para uma vasta gama de cargas de trabalho de IA e oferecem versatilidade para acelerar as cargas de trabalho de inferência em frameworks de IA, incluindo PyTorch, JAX e TensorFlow. Para mais detalhes sobre as TPUs, consulte o artigo Introdução à TPU da Google Cloud Platform.

Pré-requisitos para usar as TPUs no Dataflow

  • Os seus Google Cloud projetos têm de ser aprovados para usar esta oferta do GA.

Limitações

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

  • Apenas são suportados aceleradores de TPU de anfitrião único: a oferta de TPU do Dataflow suporta apenas configurações de TPU de anfitrião único em que cada trabalhador do Dataflow gere um ou vários dispositivos de TPU que não estão interligados com TPUs geridas por outros trabalhadores.
  • Apenas são suportados pools de trabalhadores de TPU homogéneos: funcionalidades como o ajuste certo do Dataflow e o Dataflow Prime não suportam cargas de trabalho de TPU.

Preços

Os trabalhos do Dataflow que usam TPUs são faturados em função das horas de chip de TPU do trabalhador consumidas e não em função da CPU e da memória do trabalhador. Para mais informações, consulte a página de preços do Dataflow.

Disponibilidade

Estão disponíveis os seguintes aceleradores de TPU e regiões de processamento.

Aceleradores de TPU suportados

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

  • O tipo de TPU refere-se ao modelo do dispositivo TPU.
  • A topologia da TPU refere-se ao número e à disposição física dos chips da TPU numa fatia.

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

As seguintes configurações de TPU são suportadas com o Dataflow:

Tipo de TPU Topologia Obrigatório worker_machine_type
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 obter informações sobre as regiões e as zonas disponíveis para TPUs, consulte o artigo Regiões e zonas de TPUs na documentação do Cloud TPU.

O que se segue?