Como escalonar cargas de trabalho de ML usando o Ray

Este documento explica como executar cargas de trabalho de machine learning (ML) com o Ray e o JAX em TPUs. Há dois modos diferentes para usar TPUs com o Ray: modo centrado no dispositivo (PyTorch/XLA) e modo centrado no host (JAX).

Este documento presume que você já tem um ambiente de TPU configurado. Para saber mais, confira estes recursos:

Cloud TPU: Configurar o ambiente do Cloud TPU e Gerenciar recursos de TPU
Google Kubernetes Engine (GKE): Implantar cargas de trabalho de TPU no GKE Autopilot ou Implantar cargas de trabalho de TPU no GKE Standard

Modo centrado no dispositivo (PyTorch/XLA)

O modo centrado no dispositivo tem grande parte do estilo programático do PyTorch clássico. Nesse modo, você adiciona um tipo de dispositivo XLA, que funciona como qualquer outro dispositivo PyTorch. Cada processo individual interage com um dispositivo XLA.

Esse modo é ideal para pessoas que já sabem utilizar o PyTorch com GPUs e querem usar abstrações de programação semelhantes.

As seções a seguir descrevem como executar uma carga de trabalho do PyTorch/XLA em um ou mais dispositivos sem usar o Ray e como executar a mesma carga de trabalho em vários hosts usando o Ray.

Criar uma TPU

Crie variáveis de ambiente para parâmetros de criação de TPU.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-8
export RUNTIME_VERSION=v2-alpha-tpuv5

Descrições de variáveis de ambiente

Variável	Descrição
`PROJECT_ID`	O ID do projeto do Google Cloud . Use um projeto atual ou crie um novo.
`TPU_NAME`	O nome da TPU.
`ZONE`	A zona em que a VM de TPU será criada. Para mais informações sobre as zonas disponíveis, consulte Zonas e regiões de TPU.
`ACCELERATOR_TYPE`	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores aceitos por cada versão de TPU, consulte Versões de TPU.
`RUNTIME_VERSION`	A versão do software do Cloud TPU.

Use este comando para criar uma VM de TPU v5p com oito núcleos:

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE  \
   --version=$RUNTIME_VERSION

Conecte-se à VM de TPU usando este comando:

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE

Ao usar o GKE, consulte o guia KubeRay no GKE para saber mais sobre a configuração.

Requisitos de instalação

Execute os comandos abaixo na VM de TPU para instalar as dependências necessárias:

Salve as informações a seguir em um arquivo. Por exemplo, requirements.txt.

--find-links https://storage.googleapis.com/libtpu-releases/index.html
--find-links https://storage.googleapis.com/libtpu-wheels/index.html
torch~=2.6.0
torch_xla[tpu]~=2.6.0
ray[default]==2.40.0

Para instalar as dependências necessárias, execute o seguinte:
```
pip install -r requirements.txt
```

Ao executar cargas de trabalho no GKE, a recomendação é criar um Dockerfile que instale as dependências necessárias. Confira um exemplo em Executar uma carga de trabalho em nós de uma fração de TPU na documentação do GKE.

Executar uma carga de trabalho do PyTorch/XLA em um único dispositivo

O exemplo a seguir demonstra como criar um tensor XLA em um único dispositivo, que é um chip de TPU. O PyTorch lida com ele da mesma forma que lida com outros tipos de dispositivos.

Salve o snippet de código a seguir em um arquivo. Por exemplo, workload.py.
```
import torch
import torch_xla
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)
```
A instrução de importação import torch_xla inicia o PyTorch/XLA e a função xm.xla_device() retorna o dispositivo XLA atual, um chip de TPU.
Defina a variável de ambiente PJRT_DEVICE como TPU.
```
export PJRT_DEVICE=TPU
```
Execute o script.
```
python workload.py
```
A saída será parecida com a mostrada abaixo. Verifique se a saída indica que o dispositivo XLA foi encontrado.
```
xla:0
tensor([[ 0.6220, -1.4707],
        [-1.2112,  0.7024]], device='xla:0')
```

Executar o PyTorch/XLA em vários dispositivos

Atualize o snippet de código da seção anterior para execução em vários dispositivos.

import torch
import torch_xla
import torch_xla.core.xla_model as xm

def _mp_fn(index):
    t = torch.randn(2, 2, device=xm.xla_device())
    print(t.device)
    print(t)

if __name__ == '__main__':
    torch_xla.launch(_mp_fn, args=())

Execute o script.

python workload.py

Se você executar o snippet de código em uma TPU v5p-8, a saída será parecida com esta:

xla:0
xla:0
xla:0
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
xla:0
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')
tensor([[ 1.2309,  0.9896],
        [ 0.5820, -1.2950]], device='xla:0')

torch_xla.launch() usa dois argumentos: uma função e uma lista de parâmetros. Ele cria um processo para cada dispositivo XLA disponível e chama a função especificada nos argumentos. Neste exemplo, há quatro dispositivos de TPU disponíveis. Portanto, torch_xla.launch() cria quatro processos e chama _mp_fn() em cada dispositivo. Cada processo só tem acesso a um dispositivo. Portanto, cada dispositivo tem o índice 0, e xla:0 é exibido para todos os processos.

Executar o PyTorch/XLA em vários hosts com o Ray

As seções a seguir mostram como executar o mesmo snippet de código em uma fração de TPU maior com vários hosts. Para mais informações sobre a arquitetura de TPU de vários hosts, consulte Arquitetura do sistema.

Neste exemplo, você configura o Ray manualmente. Se você já sabe configurar o Ray, pule para a última seção, Executar uma carga de trabalho do Ray. Para saber como configurar o Ray para um ambiente de produção, consulte estes recursos:

Criar uma VM de TPU com vários hosts

Crie variáveis de ambiente para parâmetros de criação de TPU.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-16
export RUNTIME_VERSION=v2-alpha-tpuv5

Descrições de variáveis de ambiente

Variável	Descrição
`PROJECT_ID`	O ID do projeto do Google Cloud . Use um projeto atual ou crie um novo.
`TPU_NAME`	O nome da TPU.
`ZONE`	A zona em que a VM de TPU será criada. Para mais informações sobre as zonas disponíveis, consulte Zonas e regiões de TPU.
`ACCELERATOR_TYPE`	O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores aceitos por cada versão de TPU, consulte Versões de TPU.
`RUNTIME_VERSION`	A versão do software do Cloud TPU.

Crie uma TPU v5p de vários hosts com dois hosts (uma v5p-16 com quatro chips de TPU em cada host) usando o seguinte comando:

gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION

Configurar o Ray

Uma TPU v5p-16 tem dois hosts de TPU, cada um com quatro chips de TPU. Neste exemplo, você vai iniciar o nó principal do Ray em um host e adicionar o segundo host como um nó de trabalho ao cluster do Ray.

Conecte-se por SSH ao primeiro host.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0

Instale as dependências com o mesmo arquivo de requisitos da seção Requisitos de instalação.
```
pip install -r requirements.txt
```
Inicie o processo do Ray.
```
ray start --head --port=6379
```
Observação: se você receber um erro como Command 'ray' not found, talvez seja necessário adicionar o diretório que contém o executável ray à variável de ambiente PATH do sistema. Também é possível especificar o caminho completo ao executar comandos ray (por exemplo, ~/.local/bin/ray).
Observação: o comando ray start pede que você ative ou desative a coleta de estatísticas de uso. Para mais informações, consulte Coleta de estatísticas de uso na documentação do Ray.

A saída será assim:
```
Enable usage stats collection? This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup. Confirm [Y/n]: y
Usage stats collection is enabled. To disable this, add
```