Executar um pipeline com TPUs

Nesta página, explicamos como executar um pipeline do Apache Beam no Dataflow com TPUs. Os jobs que usam TPUs geram cobranças, conforme especificado na página de preços do Dataflow.

Para mais informações sobre o uso de TPUs com o Dataflow, consulte Suporte do Dataflow para TPUs.

Opcional: fazer uma reserva específica para usar aceleradores

Embora seja possível usar TPUs sob demanda, recomendamos usar TPUs do Dataflow com reservas especificamente direcionadas Google Cloud. Isso ajuda a garantir que você tenha acesso aos aceleradores disponíveis e tempos de inicialização rápidos do worker. Os pipelines que consomem uma reserva de TPU não exigem cota adicional de TPU.

Se você não fizer uma reserva e optar por usar TPUs sob demanda, provisione a cota de TPU antes de executar o pipeline.

Opcional: provisionar cota de TPU

É possível usar TPUs em uma capacidade sob demanda ou com uma reserva. Se você quiser usar TPUs sob demanda, provisione a cota de TPU antes. Se você usar uma reserva especificamente direcionada, pule esta seção.

Para usar TPUs sob demanda sem uma reserva, verifique o limite e o uso atual da sua cota da API Compute Engine para TPUs da seguinte maneira:

Console

1. Acesse a página Cotas no console Google Cloud :

   Acessar "Cotas"

2. Na caixa Filtro , faça o seguinte:

   1. Use a tabela a seguir para selecionar e copiar a propriedade da cota com base na versão da TPU e no tipo de máquina. Por exemplo, se você planeja criar nós da TPU v5e sob demanda com o tipo de máquina que começa com ct5lp-, digite Name: TPU v5 Lite PodSlice chips.

      Versão da TPU, o tipo de máquina começa com	Propriedade e nome da cota para instâncias sob demanda
      TPU v5e, ct5lp-	Name: TPU v5 Lite PodSlice chips
      TPU v5p, ct5p-	Name: TPU v5p chips
      TPU v6e, ct6e-	Dimensions (e.g. location): tpu_family:CT6E

   2. Selecione a propriedade Dimensões (por exemplo, local) e insira region: seguido do nome da região em que você planeja iniciar seu pipeline. Por exemplo, insira region:us-west4 se planeja usar a zona us-west4-a. A cota de TPU é regional, portanto, todas as zonas na mesma região consomem a mesma cota de TPU.

Configurar uma imagem de contêiner personalizada

Para interagir com TPUs em pipelines do Dataflow, é necessário fornecer software que possa operar em dispositivos XLA no ambiente de execução do pipeline. Isso exige a instalação de bibliotecas de TPU com base nas necessidades do pipeline e a configuração de variáveis de ambiente com base no dispositivo de TPU usado.

Para personalizar a imagem do contêiner, instale o Apache Beam em uma imagem de base pronta para uso que tenha as bibliotecas de TPU necessárias. Como alternativa, instale o software de TPU nas imagens publicadas com as versões do SDK do Apache Beam.

Para fornecer uma imagem de contêiner personalizada, use a opção sdk_container_image do pipeline. Para mais informações, consulte Usar contêineres personalizados no Dataflow.

Ao usar um acelerador de TPU, é necessário definir as seguintes variáveis de ambiente na imagem do contêiner.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

Dependendo do acelerador usado, as variáveis na tabela a seguir também precisam ser definidas.

tipo	topologia	worker_machine_type obrigatório do Dataflow	variáveis de ambiente adicionais
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	TPU_ACCELERATOR_TYPE=v5litepod-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	TPU_ACCELERATOR_TYPE=v5litepod-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	TPU_ACCELERATOR_TYPE=v5litepod-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice	1x1	ct6e-standard-1t	TPU_ACCELERATOR_TYPE=v6e-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice	2x2	ct6e-standard-4t	TPU_ACCELERATOR_TYPE=v6e-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice	2x4	ct6e-standard-8t	TPU_ACCELERATOR_TYPE=v6e-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	TPU_ACCELERATOR_TYPE=v5p-8 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

Um exemplo de Dockerfile para a imagem de contêiner personalizada pode ter esta aparência:

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

Executar o job com TPUs

As considerações para executar um job do Dataflow com TPUs incluem:

• Como os contêineres de TPU podem ser grandes, aumente o tamanho padrão do disco de inicialização para 50 gigabytes ou um tamanho adequado, conforme exigido pela imagem do contêiner, usando a opção de pipeline --disk_size_gb.
• Limite o paralelismo intraworker.

TPUs e paralelismo de worker

Na configuração padrão, os pipelines do Python do Dataflow iniciam um processo do SDK do Apache Beam por núcleo de VM. Os tipos de máquinas com TPU têm um grande número de núcleos de vCPU, mas apenas um processo pode realizar cálculos em um dispositivo de TPU. Além disso, um dispositivo TPU pode ser reservado por um processo durante toda a vida útil dele. Portanto, é necessário limitar o paralelismo entre workers ao executar um pipeline de TPU do Dataflow. Para limitar o paralelismo de workers, use as seguintes orientações:

• Se o caso de uso envolver a execução de inferências em um modelo, use a API RunInference do Beam. Para mais informações, consulte Inferência de modelo de linguagem grande no Beam.
• Se não for possível usar a API RunInference do Beam, use os objetos compartilhados multiprocesso do Beam para restringir determinadas operações a um único processo.
• Se você não puder usar as recomendações anteriores e preferir iniciar apenas um processo do Python por worker, defina a opção de pipeline --experiments=no_use_multiple_sdk_containers.
• Para workers com mais de 100 vCPUs, reduza o número de linhas de execução usando a opção de pipeline --number_of_worker_harness_threads. Use a tabela a seguir para saber se o tipo de TPU usa mais de 100 vCPUs.

A tabela a seguir lista o total de recursos de computação por worker para cada configuração de TPU.

Tipo de TPU	topologia	tipo de máquina	Chips de TPU	vCPU	RAM (GB)
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	1	24	48
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	4	112	192
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	8	224	384
tpu-v6e-slice	1x1	ct6e-standard-1t	1	44	176
tpu-v6e-slice	2x2	ct6e-standard-4t	4	180	720
tpu-v6e-slice	2x4	ct6e-standard-8t	8	360	1440
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	4	208	448

Executar um pipeline com TPUs

Para executar um job do Dataflow com TPUs, use o seguinte comando:

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

Substitua:

• PIPELINE: o arquivo de código-fonte do pipeline.
• PROJECT: o nome do projeto Google Cloud .
• BUCKET: o bucket do Cloud Storage.
• REGION: uma região do Dataflow, por exemplo, us-central1.
• TPU_TYPE: um tipo de TPU compatível, por exemplo, tpu-v5-lite-podslice. Para uma lista completa de tipos e topologias, consulte Aceleradores de TPU compatíveis.
• TPU_TOPOLOGY: a topologia da TPU, por exemplo, 1x1.
• MACHINE_TYPE: o tipo de máquina correspondente, por exemplo, ct5lp-hightpu-1t.
• DISK_SIZE_GB: o tamanho do disco de inicialização de cada VM de worker. Por exemplo, 100.
• IMAGE: o caminho do Artifact Registry da imagem do Docker.
• NUMBER_OF_THREADS: opcional. O número de linhas de execução de uso de worker.

Verificar o job do Dataflow

Para confirmar que o job usa VMs de worker com TPUs, siga estas etapas:

1. No Google Cloud console, acesse a página Dataflow > Jobs.

   Acessar "Jobs"

2. Selecione um job.

3. Clique na guia Métricas do job.

4. Na seção Escalonamento automático, confirme se há pelo menos uma VM Workers atuais.

5. No painel lateral Informações do job, verifique se o machine_type começa com ct. Por exemplo, ct6e-standard-1t. Isso indica o uso da TPU.

Resolver problemas do job do Dataflow

Se você encontrar problemas ao executar o job do Dataflow com TPUs, consulte Resolver problemas do job da TPU do Dataflow.

A seguir

• Saiba mais sobre o suporte a TPUs no Dataflow.
• Saiba mais sobre a inferência de modelos grandes no Beam.