Introdução aos programas no Cloud

O Pathways é um sistema projetado para permitir a criação de sistemas de machine learning em grande escala, com várias tarefas e ativados esparsamente. Ele permite o uso de 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 machine learning em grande escala, permitindo que um único cliente JAX organize cargas de trabalho em várias frações de TPU grandes, que podem abranger milhares de chips de TPU.

O Pathways é usado internamente no Google 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

Você precisa ter:

Neste documento, mostramos como usar TPUs gerenciadas do Pathways no Google Kubernetes Engine (GKE) para cargas de trabalho em lote, em tempo real e interativas. Ele pressupõe que você já esteja familiarizado com o uso de TPUs com o GKE, incluindo TPUs de fração única e várias frações no Google Kubernetes Engine, além de experiência geral com TPUs de várias frações.

Controlador único e multicontrolador

Há duas maneiras principais de gerenciar e orquestrar computações em vários dispositivos:

Recurso

Controlador único (Pathways)

Multicontrolador (padrão do JAX)

Controle

Ponto único de controle: 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: os usuários interagem com um único cliente, fazendo com que o sistema pareça uma única máquina grande com muitos aceleradores locais.

SPMD: usa principalmente o paradigma SPMD, exigindo que todos os dispositivos executem o mesmo programa.

Flexibilidade

Oferece suporte a padrões de computação mais 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

A seção a seguir descreve os principais componentes da arquitetura do Pathways.

Gerente de recursos do Pathways

Esse é o plano de controle central do sistema Pathways. Ele gerencia todos os recursos de acelerador e é responsável por coordenar a alocação de aceleradores para jobs de usuário. 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

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

Trabalhador do Pathways

Esses são os processos 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 seu programa pelo servidor proxy do IFRT. Este componente requer recursos de acelerador.

Cliente proxy do IFRT

Essa é uma implementação de OSS da API Interim Framework Runtime (IFRT), que separa o código do usuário do ambiente de execução 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 o programa e os componentes do Pathways. Ele envia solicitações ao servidor proxy do IFRT e recebe resultados dele. É uma implementação de OSS da API IFRT. Esse componente requer apenas recursos de CPU.

Servidor proxy do IFRT

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

Servidor sidecar

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

Mostra a relação dos componentes do Pathways.
Componentes do Pathways

Componentes do Pathways no GKE

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

As imagens de contêiner do Pathways estão disponíveis nos locais a seguir.

Tipo de contêiner

Local

Servidor proxy do IFRT

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>

Gerente/trabalhador de recursos do Pathways

us-docker.pkg.dev/cloud-tpu-v2-images/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/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

Descrições dos argumentos:

  • --server_port: o gerenciador de recursos do Pathways usa essa porta para se comunicar com outros componentes do Pathways.
  • --node_type: o tipo de nó. Ele precisa ser definido como "resource_manager" para o gerenciador de recursos do Pathways e 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 de tpu{TPU type}:{TPU 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 IFRT:

 - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Descrições dos 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 ser igual ao valor --server_port usado para o contêiner do gerenciador de recursos do Pathways.
  • --server_port: o servidor proxy do IFRT usa essa porta 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/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"

Descrições dos 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 ser igual ao valor --server_port usado para o contêiner do gerenciador de recursos do Pathways.
  • --server_port: os workers usam essa porta 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.

A seguir