Portar cargas de trabalho do JAX para o Pathways

Devido à natureza distribuída do JAX com o Pathways, algumas operações podem não ser escalonadas devido a sobrecargas de comunicação. Embora o Pathways minimize essas sobrecargas com recursos como envio assíncrono, há algumas coisas que você precisa saber ao portar 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 do 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 Pathways e o job do usuário na mesma zona do Cloud do Google Cloud . Use uma CPU grande, como o proxy e o gerenciador de recursos do IFRT. Recomendamos pelo menos um n2-standard-64 dedicado, que vem com 64 vCPUs e 256 GB de memória.

PathwaysUtils

Pathways-utils é um repositório do GitHub baseado em Python que fornece utilitários e ferramentas essenciais para simplificar a implantação e a execução de cargas de trabalho do JAX na arquitetura do Pathways on Cloud. Esse pacote lida com as adaptações necessárias para o ambiente de nuvem, permitindo que os desenvolvedores do JAX se concentrem nos fluxos de trabalho principais de machine learning com o mínimo de configuração 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 de criação de perfil que permitem entender o desempenho do aplicativo. Ao usar APIs de criação de perfil padrão do JAX, como jax.profiler.start_trace e jax.profiler.start_server, é possível criar perfis não apenas do código JAX, mas também dos componentes subjacentes do Pathways, fornecendo uma visão holística da execução no ambiente de nuvem.
  • Checkpoint distribuído com o Orbax: um gerenciador de checkpoint personalizado do Orbax que permite usar checkpoints distribuídos e restaurar seus checkpoints ao usar a biblioteca do Orbax no ambiente do Pathways. Essa integração tem como objetivo funcionar sem exigir mudanças no código de checkpoint do Orbax, desde que ele importe pathwaysutils.
  • Primitivas de treinamento elástico: fornece primitivas de treinamento elástico fundamentais que podem ser usadas para criar fluxos de trabalho de treinamento robustos e escalonáveis usando o Pathways. Essas primitivas permitem que os 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 é totalmente testado 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 com eficiência as operações de checkpoint pelo proxy IFRT, permitindo que os workers do Pathways em aceleradores salvem e restaurem dados diretamente.

Python colocalizado

O Python colocalizado é 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 com vários controladores. Com essa abordagem, é possível evitar a transferência de dados entre o cliente e as máquinas de TPU para tarefas com uso intensivo de computação, como carregamento de dados e checkpointing. Para configurar o cluster do Pathways para executar a API JAX do Python colocada, crie e envie uma imagem de contêiner secundário usando o Dockerfile de amostra do contêiner secundário.

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 gera 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 para operações de pequena a média escala.

Quando você executa um carregador de dados multiprocesso (como o Grain) no contêiner de carga de trabalho principal no GKE, o contêiner compartilha o nó principal com os sidecars pathways-proxy e pathways-rm. Devido a essa arquitetura compartilhada, é necessário configurar o seguinte:

  • Defina uma contagem explícita de workers: dentro de um contêiner, os.cpu_count() informa o total de vCPUs da VM host, em vez do limite de CPU do contêiner. Transmita uma contagem explícita de workers (não mais do que o limite de CPU do contêiner) em vez de derivá-la de os.cpu_count().
  • Monte um volume /dev/shm com suporte de memória: os carregadores de dados de vários processos usam a memória compartilhada POSIX (/dev/shm) para comunicação entre processos. Verifique se um volume emptyDir (medium: Memory) com suporte de memória está montado em /dev/shm para que os lotes pré-buscados não excedam o limite de memória compartilhada padrão do contêiner.

Para otimizar o desempenho em escala, recomendamos colocar o pipeline de dados de entrada no mesmo local usando o Python 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.

Confira 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 com vários controladores e no Pathways de maneira distribuída usando a API Python JAX colocalizada.

Controle de versões do Jax

As versões do Pathways são totalmente integradas às versões do JAX para garantir compatibilidade e estabilidade. Para evitar possíveis problemas, verifique se os artefatos do Pathways e a versão do JAX estão alinhados. Cada versão do Pathways especifica claramente as versões compatíveis do JAX por meio de uma tag do formato jax-<version>.

Cache de compilação

O cache de compilação persistente do Pathways é um recurso que permite que os servidores do Pathways armazenem executáveis XLA compilados em um local persistente, como o Cloud Storage, para evitar compilações redundantes. Esse recurso é ativado por padrão. O local do cache é transmitido como flag --gcs_scratch_location para o gerenciador de recursos e os contêineres de worker do Pathways. Para manter os custos de armazenamento associados ao 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 deixá-la vazia, o Pathways vai usar automaticamente um cache de compilação na memória. Esse cache não mantém os 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 o Pathways configure políticas de ciclo de vida de objetos para aplicar o TTL para remoção de cache.
  • storage.objects.list: para listar os objetos de cache no bucket.
  • storage.objects.create: para gravar novos executáveis compilados no cache.
  • storage.objects.get: para ler executáveis armazenados em cache no bucket.

Criação de perfil

É possível usar o criador de perfil do JAX para gerar traces de um programa JAX. Há duas maneiras comuns de fazer isso com o Pathways:

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

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

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

Esses arquivos de trace 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 perfil, 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 trace
  • PREFIX: um caminho no bucket do Cloud Storage para armazenar os arquivos de trace

Captura de perfil programática

Capturar um perfil de dentro do 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 as informações do 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 perfil estiver em execução, você poderá 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>

É possível encontrar informações de tempo para métodos de cliente de proxy IFRT, como Compile e Execute, no trace do programa. Esses eventos, que detalham as interações com o servidor gRPC do proxy do 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 rastro, você pode ter insights sobre o desempenho dessas operações.

Flags do XLA

Ao usar o Pathways, é necessário definir as sinalizações de XLA no contêiner pathways-proxy.

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

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

Dump de HLO

Para inspecionar as entradas do High Level Optimizer (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