O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

Introdução aos programas no Cloud

O Pathways é um sistema projetado para permitir a criação de sistemas de machine learning (ML) em grande escala, com várias tarefas e ativados esparsamente. Ele permite usar milhares ou dezenas de milhares de aceleradores, com a capacidade de alocar dinamicamente quantidades variadas de computação para diferentes tarefas com base nos requisitos de processamento.

O Pathways simplifica os cálculos de ML em grande escala, permitindo que um único cliente JAX orquestre cargas de trabalho em grandes frações de TPU, que podem abranger milhares de chips de TPU.

O Google usa o Pathways internamente para treinar modelos grandes, como o Gemini. O Pathways no Cloud traz os mesmos benefícios para os clientes do Google Cloud .

Antes de começar

Selecione a guia para como planeja usar as amostras nesta página:
gcloud

Instale a Google Cloud CLI e faça login na CLI gcloud com sua identidade federada. Depois de fazer login, inicialize a Google Cloud CLI executando o seguinte comando:
```
gcloud init
```
REST

Para usar as amostras da API REST desta página em um ambiente de desenvolvimento local, use as credenciais fornecidas para a CLI gcloud.
Saiba mais em Autenticar para usar REST na documentação de autenticação do Google Cloud .

Antes de começar

É possível executar o Pathways no Google Kubernetes Engine (GKE) ou em um cluster do Cluster Director com o Slurm. Selecione a guia da sua plataforma.

GKE

Siga os seguintes pré-requisitos:

Neste documento, apresentamos uma visão geral de como usar TPUs gerenciadas pelo Pathways no GKE para cargas de trabalho em lote, em tempo real e interativas. Ele pressupõe que você já sabe usar TPUs no GKE, incluindo TPUs de fração única e várias frações no GKE, e tem experiência geral com TPUs de várias frações.

Cluster Director

Siga os seguintes pré-requisitos:

Neste documento, mostramos como usar o Pathways com TPUs em um cluster do Cluster Director para cargas de trabalho em lote e interativas. Ele pressupõe que você já conhece o Cluster Director, as máquinas de TPU no Cluster Director e a execução de jobs com o Slurm, além de ter experiência geral com TPUs de várias frações.

Para criar um cluster, consulte Criar um cluster do Cluster Director com Pathways.

Controlador único e multicontrolador

Há duas maneiras principais de gerenciar e organizar cálculos em vários dispositivos:

Recurso	Controlador único (Pathways)	Multicontrolador (padrão do JAX)
Controle	Ponto de controle único: um único programa cliente atua como o controlador central.	Controle distribuído: vários processos participam, cada um com a própria instância do interpretador do Python.
Ver	Visualização unificada: o cliente vê todos os dispositivos como um único sistema unificado.	Visualização localizada: cada processo do Python vê apenas os dispositivos conectados a ele.
Programação	Programação simplificada: você interage com um único cliente, fazendo com que o sistema apareça como uma única máquina grande com muitos aceleradores locais.	Programa único, vários dados (SPMD): usa principalmente o paradigma SPMD, exigindo que todos os dispositivos executem o mesmo programa.
Flexibilidade	Oferece suporte a padrões de computação complexos além do SPMD, incluindo paralelismo de pipeline assimétrico e esparsidade computacional.	Pode ser menos flexível no gerenciamento de recursos, especialmente em diferentes frações de TPU.

Componentes do Pathways

As seções a seguir descrevem os principais componentes da arquitetura do Pathways.

Gerente de recursos do Pathways

Esse componente é o plano de controle central do sistema Pathways. Ele gerencia todos os recursos de acelerador e coordena a alocação de aceleradores para seus jobs. Ele monitora a integridade dos workers e lida com a programação, pausa e retomada de jobs. Ele serve como um ponto de contato único para erros e status do sistema. Esse componente requer apenas recursos de CPU.

Cliente do Pathways

Esse componente é uma implementação do Interim Framework Runtime (IFRT) que serve como ponto de entrada para o sistema Pathways. Ele recebe operações de alto nível (HLOs, na sigla em inglês) do seu programa. O cliente do Pathways se coordena com o gerenciador de recursos do Pathways para determinar onde colocar os programas compilados para execução com base no seu código. Ele apresenta uma visualização unificada do sistema para um determinado cliente JAX. Esse componente requer apenas recursos de CPU.

Trabalhador do Pathways

Os processos de worker do Pathways são executados nas máquinas aceleradoras (VMs de TPU). Eles recebem executáveis compilados do seu programa do servidor proxy do IFRT e realizam os cálculos nas TPUs. Os workers do Pathways enviam dados de volta ao programa pelo servidor proxy do IFRT. Este componente requer recursos de acelerador.

Cliente proxy do IFRT

O cliente proxy do IFRT é uma implementação de código aberto da API Interim Framework Runtime (IFRT) que separa o código do usuário do ambiente de execução subjacente e melhora a portabilidade e a transparência do código. O JAX usa essa implementação como alternativa ao ambiente de execução padrão de vários controladores. O cliente proxy do IFRT atua como uma ponte de comunicação entre seu programa e os componentes do Pathways. Ele envia solicitações ao servidor proxy do IFRT e recebe resultados dele. Esse componente requer apenas recursos de CPU.

Servidor proxy do IFRT

O servidor proxy do IFRT é um servidor gRPC que recebe solicitações do cliente proxy do IFRT e as encaminha para o cliente do Pathways, que processa a distribuição do trabalho. Esse componente requer apenas recursos de CPU.

Servidor sidecar

O servidor sidecar é um servidor gRPC que fica no mesmo local que o worker do Pathways na VM do acelerador. Ele executa o código Python especificado pelo usuário diretamente na VM do acelerador para reduzir a latência de transferência de dados entre o controlador e os aceleradores. O servidor sidecar interage com o worker do Pathways por um protocolo personalizado com controle de versões no transporte gRPC.

Diagrama de arquitetura que mostra como o cliente JAX, o proxy IFRT,
o gerenciador de recursos do Pathways e os workers do Pathways interagem. — Componentes do Pathways

Componentes do Pathways no GKE

Nesta seção, mapeamos os componentes do Pathways para componentes do GKE, como contêineres e pods.

As imagens de contêiner do Pathways estão disponíveis nos seguintes locais:

Tipo de contêiner	Local
Servidor proxy do IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION`
Gerenciador/worker de recursos do Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION`

Gerente de recursos do Pathways

Depois de criar um cluster do GKE, use o seguinte containerSpec para implantar o gerenciador de recursos do Pathways:

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

O contêiner usa os seguintes argumentos:

--server_port: a porta que o gerenciador de recursos do Pathways usa para se comunicar com outros componentes do Pathways.
--node_type: o tipo de nó. Defina esse argumento como resource_manager para o gerenciador de recursos do Pathways. Ele não é necessário para os outros contêineres.
--instance_count: o número de frações de TPU.
--instance_type: o tipo de TPU e a topologia da fração, no formato tpuTPU_TYPE:TOPOLOGY, por exemplo, tpuv5e:4x4.
--gcs_scratch_location: um bucket do Cloud Storage usado para arquivos temporários.

Servidor proxy do IFRT

Use o seguinte containerSpec para implantar um servidor proxy do IFRT:

  - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

O contêiner usa os seguintes argumentos:

--resource_manager_address: o nome do host e a porta que o servidor proxy usa para se comunicar com o gerenciador de recursos do Pathways. A porta precisa corresponder ao valor --server_port usado para o contêiner do gerenciador de recursos do Pathways.
--server_port: a porta que o servidor proxy do IFRT usa para se comunicar com o cliente proxy do IFRT.
--gcs_scratch_location: um bucket do Cloud Storage usado para arquivos temporários.

Trabalhador do Pathways

Use o seguinte containerSpec para implantar workers do Pathways:

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

O contêiner usa os seguintes argumentos:

--resource_manager_address: o nome do host e a porta que os workers da TPU usam para se comunicar com o gerenciador de recursos do Pathways. A porta precisa corresponder ao valor --server_port usado para o contêiner do gerenciador de recursos do Pathways.
--server_port: a porta que os workers usam para se comunicar com o servidor proxy e o gerenciador de recursos do Pathways.
--gcs_scratch_location: um bucket do Cloud Storage usado para arquivos temporários.

O gerenciador de recursos do Pathways, o servidor proxy do IFRT e os workers do Pathways podem ter portas diferentes, mas, neste exemplo, o gerenciador de recursos do Pathways e o worker do Pathways compartilham a mesma porta.

Componentes do Pathways no Cluster Director

Esta seção mapeia os componentes do Pathways para jobs do Slurm em um cluster do Cluster Director.

Os scripts do Pathways executam os componentes do Pathways com o Enroot. O script setup_env.sh usa as seguintes imagens de contêiner, que correspondem à versão do JAX (0.11.1) do ambiente do Python que ele cria:

Tipo de contêiner	Local
Servidor proxy do IFRT	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1`
Gerenciador/worker de recursos do Pathways	`us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1`

Para iniciar os componentes do Pathways, envie o script start_pathways.sh como o job de serviço do Pathways. Não é necessário escrever os comandos a seguir: os scripts derivam os argumentos das variáveis de ambiente que você define ao enviar start_pathways.sh, como SLICE_SHAPE, MIN_SLICES, MAX_SLICES e GCS_BUCKET.

Gerente de recursos do Pathways

O job de serviço do Pathways executa o gerenciador de recursos do Pathways em um nó de CPU. Ele executa o seguinte comando no contêiner do servidor do Pathways:

/pathways/server \
    --server_port=29001 \
    --node_type=resource_manager \
    --instance_type=tpuTPU_TYPE:TOPOLOGY \
    --instance_count=MAX_SLICES \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

O comando usa os seguintes argumentos:

--server_port: a porta em que o gerenciador de recursos do Pathways detecta. O servidor proxy do IFRT e os workers do Pathways se conectam a essa porta.
--node_type: o tipo de nó, resource_manager. Somente o gerenciador de recursos do Pathways define esse argumento.
--instance_type: o tipo de TPU e a topologia de cada fração, de SLICE_SHAPE. Por exemplo, SLICE_SHAPE=2x4 define tpuv6e:2x4.
--instance_count: o número de frações de TPU, de MAX_SLICES.
--gcs_scratch_location: um caminho no bucket do Cloud Storage, de GCS_BUCKET, usado para arquivos temporários, como o cache de compilação.
--enable_gcs_olm: se o Pathways cria regras de ciclo de vida de objetos do Cloud Storage para limpar entradas antigas do cache de compilação. O script desativa isso porque requer permissões adicionais no bucket.

Servidor proxy do IFRT

O job de serviço do Pathways executa o servidor proxy do IFRT no mesmo nó de CPU que o gerenciador de recursos do Pathways. Ele executa o seguinte comando no contêiner do servidor proxy do IFRT:

/pathways/proxy_server \
    --server_port=29000 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --num_elastic_instances=MAX_SLICES

O comando usa os seguintes argumentos:

--server_port: a porta em que o servidor proxy do IFRT detecta conexões. Seu programa JAX se conecta a essa porta.
--resource_manager_address: o endereço do gerenciador de recursos do Pathways. RM_IP é o endereço IP interno do nó da CPU.
--gcs_scratch_location: o mesmo caminho do Cloud Storage que o gerenciador de recursos do Pathways.
--num_elastic_instances: o número de frações que podem ficar indisponíveis sem interromper a sessão de proxy do IFRT, de MAX_SLICES. Sua carga de trabalho aplica MIN_SLICES.

Quando pelo menos MIN_SLICES frações estão em execução, o job de serviço grava o endereço do servidor proxy IFRT em ~/pathways/cluster.env. Seu programa JAX, o cliente proxy do IFRT, lê esse arquivo para se conectar ao Pathways. O programa JAX é executado como sua carga de trabalho, por exemplo, run_maxtext.sh, ou em um notebook do Jupyter.

Trabalhador do Pathways

O job de serviço do Pathways envia uma matriz de jobs do Slurm (run_worker_slice.sh) com uma tarefa de matriz para cada fração de TPU. Cada tarefa de matriz executa o seguinte comando no contêiner do servidor Pathways em cada VM de TPU na fração:

/pathways/server \
    --server_port=29001 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

O comando usa os seguintes argumentos:

--server_port: a porta que o worker do Pathways detecta. O gerenciador de recursos do Pathways e o servidor proxy do IFRT se conectam a essa porta.
--resource_manager_address: o endereço do gerenciador de recursos do Pathways, o mesmo do servidor proxy do IFRT.
--gcs_scratch_location, --enable_gcs_olm: o mesmo que para o gerenciador de recursos do Pathways.

O script também define a topologia da TPU e as variáveis de ambiente MEGASCALE_* para cada worker dos nós da fração e o índice de tarefas da matriz.

Se uma fração de TPU for interrompida ou falhar, o Slurm vai colocar a tarefa de matriz na fila novamente. Se você tiver sudo sem senha nos nós do cluster, o job de serviço também vai desligar os nós da fração para que o Slurm possa recriá-los. O gerenciador de recursos do Pathways e o servidor proxy do IFRT continuam em execução, então o programa JAX não precisa se conectar a um novo endereço.

Introdução aos programas no Cloud Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Antes de começar

gcloud

REST

Antes de começar

GKE

Cluster Director

Controlador único e multicontrolador

Componentes do Pathways

Gerente de recursos do Pathways

Cliente do Pathways

Trabalhador do Pathways

Cliente proxy do IFRT

Servidor proxy do IFRT

Servidor sidecar

Componentes do Pathways no GKE

Gerente de recursos do Pathways

Servidor proxy do IFRT

Trabalhador do Pathways

Componentes do Pathways no Cluster Director

Gerente de recursos do Pathways

Servidor proxy do IFRT

Trabalhador do Pathways

A seguir

Introdução aos programas no Cloud