Portar cargas de trabalho do JAX para o Pathways

Devido à natureza distribuída do JAX com Pathways, algumas operações podem não ser escalonadas bem devido a sobrecargas de comunicação. Embora o Pathways minimize esses overheads com recursos como envio assíncrono, há algumas coisas que você precisa saber ao migrar cargas de trabalho do JAX para o Pathways ou escalonar uma carga de trabalho do JAX com o Pathways para um grande número de aceleradores.

Antes de começar

Você precisa ter:

Índice de processo

O JAX com Pathways trata todos os dispositivos no cluster do Pathways como locais. Isso simplifica o gerenciamento de dispositivos e permite que o JAX use todos os recursos disponíveis. Na prática, isso significa:

  • jax.process_index() é sempre 0 para todos os dispositivos.
  • jax.devices() e jax.local_devices() retornam todos os dispositivos de TPU em todo o job.

Tipo de hardware e colocation

Para um melhor desempenho, coloque todos os componentes do programa e o trabalho do usuário na mesma Google Cloud zona da nuvem. Use uma CPU grande, como o proxy IFRT e o gerenciador de recursos. Recomendamos pelo menos um n2-standard-64 dedicado, que vem com 64 vCPUs e 256 GB de memória.

PathwaysUtils

O Pathways-utils é um repositório do GitHub baseado em Python que oferece utilitários e ferramentas essenciais para simplificar a implantação e a execução de cargas de trabalho do JAX na arquitetura Pathways on Cloud. Esse pacote processa as adaptações necessárias para o ambiente de nuvem, permitindo que os desenvolvedores do JAX se concentrem nos principais fluxos de trabalho de machine learning com configuração mínima específica da plataforma. Especificamente, ele oferece:

  • Um back-end JAX "proxy": esse back-end personalizado permite que seu aplicativo JAX use a infraestrutura do Pathways definindo a variável de ambiente JAX_PLATFORMS=proxy.
  • Utilitários de criação de perfil integrados: recursos que permitem entender o desempenho do seu aplicativo. Ao usar APIs de criação de perfil JAX padrão, como jax.profiler.start_trace e jax.profiler.start_server, é possível criar perfis não apenas do seu código JAX, mas também dos componentes do Pathways subjacentes, fornecendo uma visão holística da execução no ambiente de nuvem.
  • Checkpoint distribuído com Orbax: um manipulador de checkpoint Orbax personalizado que permite usar checkpoints distribuídos e restaurar checkpoints ao usar a biblioteca Orbax no ambiente do programa de aprendizado. Essa integração tem como objetivo funcionar sem exigir mudanças no código de checkpointing do Orbax desde que ele importe pathwaysutils.
  • Primitivos de treinamento elástico: fornecem primitivos de treinamento elástico básicos que podem ser usados para criar fluxos de trabalho de treinamento robustos e escalonáveis usando o Pathways. Essas primitivas permitem que seus jobs de treinamento se adaptem dinamicamente às mudanças nos recursos disponíveis, melhorando a eficiência e a resiliência em ambientes de nuvem.

Como estabelecer pontos de verificação

O Orbax é testado minuciosamente com o Pathways para checkpoint distribuído e restauração com o Cloud Storage. Quando você define a variável de ambiente ENABLE_PATHWAYS_PERSISTENCE=1 e chama import pathwaysutils; pathwaysutils.initialize() no seu train.py, um ArrayHandler personalizado é registrado para processar operações de checkpoint de maneira eficiente pelo proxy IFRT, permitindo que os workers do Pathways em aceleradores salvem e restaurem dados diretamente.

Python colocated

O Python alocado é uma API JAX de código aberto que permite executar código Python especificado pelo usuário diretamente nos hosts de TPU ou GPU, o que é mais simples no JAX multicontrolador. Isso permite que tarefas mais intensivas em computação, como carregamento de dados e criação de pontos de verificação, evitem a transferência de dados entre o cliente e as máquinas de TPU. Para configurar seu cluster do Pathways para executar a API JAX Python colocada, siga as instruções no README do Python colocado. Estas instruções explicam como iniciar um sidecar Python alocado junto aos workers do Pathways.

Carregamento de dados

Durante o treinamento, lotes de um conjunto de dados são carregados repetidamente para alimentar o modelo. Ter um carregador de dados assíncrono e eficiente que fragmenta o lote em vários hosts é importante para evitar que os aceleradores fiquem sem trabalho. Ao executar o treinamento com o Pathways, o carregador de dados é executado em uma VM de CPU (ao contrário de uma VM de TPU, que é usada em configurações de vários controladores) e envia dados para VMs de TPU. Isso causa uma latência maior na leitura de dados, mas é parcialmente mitigado pela leitura antecipada de X lotes no host da CPU e pelo envio assíncrono dos dados lidos para as TPUs. Essa solução é suficiente quando executada em pequena ou média escala.

Para otimizar a performance em grande escala, recomendamos que você coloque o pipeline de dados de entrada no mesmo local usando Python colocados no mesmo local para executar o pipeline de dados diretamente nos aceleradores. Isso elimina o gargalo da CPU e aproveita as interconexões rápidas da TPU para transferência de dados.

Você pode encontrar uma implementação de referência da migração de um pipeline de entrada baseado em TFDS na implementação RemoteIterator em multihost_dataloading.py. Essa implementação funciona no JAX e no Pathways com vários controladores de maneira distribuída usando a API Python JAX colocada.

Controle de versões do Jax

As versões do Pathways são fortemente acopladas às versões do JAX para garantir compatibilidade e estabilidade. Para evitar problemas, verifique se os artefatos do programa de aprendizado e a versão do JAX estão alinhados. Cada lançamento do Pathways especifica claramente as versões compatíveis do JAX usando uma tag no formato jax-<version>.

Cache de compilação

O cache de compilação persistente do programa de aprendizado é um recurso que permite que os servidores do programa de aprendizado armazenem executáveis XLA compilados em um local persistente, como o Cloud Storage, para evitar compilação redundante. Esse recurso é ativado por padrão. O local do cache é transmitido como uma flag --gcs_scratch_location para o gerenciador de recursos e os contêineres de trabalho do Pathways. Para manter os custos de armazenamento associados no mínimo, o cache anexa uma política de ciclo de vida ao local do Cloud Storage. Há um limite de 50 políticas por bucket do Cloud Storage. Portanto, recomendamos usar um local comum do Cloud Storage em todas as cargas de trabalho.

Se você omitir a flag --gcs_scratch_location ou deixar ela vazia, o programa vai usar automaticamente um cache de compilação na memória. Esse cache não persiste executáveis XLA compilados em implantações do Pathways.

Esse cache é semelhante ao cache de compilação do JAX, que é desativado por pathwaysutils.initialize() para cargas de trabalho do Pathways.

As seguintes permissões do Cloud Storage são necessárias para o cache de compilação:

  • storage.buckets.get: para recuperar metadados do bucket.
  • storage.buckets.update: essencial para que os programas de aprendizado configurem políticas de ciclo de vida de objetos para aplicar o TTL à remoção de cache.
  • storage.objects.list: para listar objetos de cache existentes no bucket.
  • storage.objects.create: para gravar novos executáveis compilados no cache.
  • storage.objects.get: para ler executáveis armazenados em cache do bucket.

Criação de perfil

Você pode usar o criador de perfil do JAX para gerar rastreamentos de um programa JAX. Há duas maneiras comuns de fazer isso com os programas:

  • Programático
    • Capturar perfis de maneira programática no seu código JAX
  • Manual
    • Capturar perfis on demand depois de iniciar o servidor do criador de perfil no código JAX

Em ambos os casos, os perfis são gravados em um bucket do Cloud Storage. Vários arquivos de rastreamento serão criados no bucket do Cloud Storage, possivelmente em pastas de carimbo de data/hora diferentes, por exemplo:

  • Processo principal do Python que invocou o rastreamento (normalmente a VM do notebook): <jax-client-vm-name>.xplane.pb
  • Proxy IFRT de programas: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Gerenciador de recursos de programas: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Trabalhadores de programas de aprendizado: server.*<tpu-node-name>.xplane.pb

Esses arquivos de rastreamento podem ser analisados com o TensorBoard executando o seguinte comando. Para mais informações sobre o TensorBoard e todas as ferramentas de criação de perfis dele, consulte Otimizar o desempenho do TensorFlow usando o Profiler.

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

Substitua:

  • BUCKET : um bucket do Cloud Storage para armazenar os arquivos de rastreamento
  • PREFIX: um caminho no bucket do Cloud Storage para armazenar os arquivos de rastreamento

Captura programática de perfil

Capture um perfil de dentro do seu código. Os perfis são salvos em gs://<bucket>/<prefix> em um diretório de carimbo de data/hora

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Captura manual de perfil

Para capturar manualmente informações de perfil, inicie o servidor do criador de perfil no código Python:

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

Enquanto o servidor do criador de perfis está em execução, é possível capturar um perfil e exportar os dados para o local de destino do Cloud Storage:

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

Você pode encontrar informações de tempo para métodos de cliente proxy IFRT, como Compile e Execute, no rastreamento do seu programa. Esses eventos, que detalham as interações com o servidor gRPC do proxy IFRT durante a compilação e a execução, aparecem na linha de execução chamada GrpcClientSessionUserFuturesWorkQueue. Ao examinar essa linha de execução no seu rastreamento, você pode receber insights sobre a performance dessas operações.

Flags do XLA

Ao usar os programas, é necessário definir as flags de XLA no contêiner pathways-proxy.

Ao usar o Cluster Toolkit (comando gcluster job submit), defina as flags do XLA usando --pathways-proxy-args:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

HLO Dump

Para inspecionar as entradas do otimizador de alto nível (HLO) fornecidas ao compilador XLA, configure o Pathways para despejar o HLO em um local especificado do Cloud Storage usando --pathways-proxy-args:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

A seguir