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:

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

  1. 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.

  2. 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
  3. 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:

  1. 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
    
  2. 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.

  1. 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.

  2. Defina a variável de ambiente PJRT_DEVICE como TPU.

    export PJRT_DEVICE=TPU
    
  3. 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

  1. 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=())
    
  2. 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

  1. 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.

  2. 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.

  1. Conecte-se por SSH ao primeiro host.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=0
  2. Instale as dependências com o mesmo arquivo de requisitos da seção Requisitos de instalação.

    pip install -r requirements.txt
    
  3. Inicie o processo do Ray.

    ray start --head --port=6379
    

    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 `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.
    
    Local node IP: 10.130.0.76
    
    --------------------
    Ray runtime started.
    --------------------
    
    Next steps
    To add another node to this Ray cluster, run
        ray start --address='10.130.0.76:6379'
    
    To connect to this Ray cluster:
        import ray
        ray.init()
    
    To terminate the Ray runtime, run
        ray stop
    
    To view the status of the cluster, use
        ray status
    

    Esse host de TPU agora é o nó principal do Ray. Anote as linhas que mostram como adicionar outro nó ao cluster do Ray. Elas serão assim:

    To add another node to this Ray cluster, run
        ray start --address='10.130.0.76:6379'
    

    Você vai usar esse comando em uma etapa posterior.

  4. Verifique o status do cluster do Ray:

    ray status
    

    A saída será assim:

    ======== Autoscaler status: 2025-01-14 22:03:39.385610 ========
    Node status
    ---------------------------------------------------------------
    Active:
    1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
    Pending:
    (no pending nodes)
    Recent failures:
    (no failures)
    
    Resources
    ---------------------------------------------------------------
    Usage:
    0.0/208.0 CPU
    0.0/4.0 TPU
    0.0/1.0 TPU-v5p-16-head
    0B/268.44GiB memory
    0B/119.04GiB object_store_memory
    0.0/1.0 your-tpu-name
    
    Demands:
    (no resource demands)
    

    Como você só adicionou o nó principal até agora, o cluster contém apenas quatro TPUs (0.0/4.0 TPU).

    Agora que o nó principal está em execução, é possível adicionar o segundo host ao cluster.

  5. Conecte-se por SSH ao segundo host.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --worker=1
  6. Instale as dependências com o mesmo arquivo de requisitos da seção Requisitos de instalação.

    pip install -r requirements.txt
    
  7. Inicie o processo do Ray. Para adicionar esse nó ao cluster do Ray atual, use o comando da saída do comando ray start. Substitua o endereço IP e a porta neste comando:

    ray start --address='10.130.0.76:6379'

    A saída será assim:

    Local node IP: 10.130.0.80
    [2025-01-14 22:30:07,397 W 75572 75572] global_state_accessor.cc:463: Retrying to get node with node ID 35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
    
    --------------------
    Ray runtime started.
    --------------------
    
    To terminate the Ray runtime, run
    ray stop
    
  8. Verifique novamente o status do Ray:

    ray status
    

    A saída será assim:

    ======== Autoscaler status: 2025-01-14 22:45:21.485617 ========
    Node status
    ---------------------------------------------------------------
    Active:
    1 node_bc0c62819ddc0507462352b76cc06b462f0e7f4898a77e5133c16f79
    1 node_35f9ac0675c91429805cdc1b97c3713422d97eee783ccb0c0304f5c1
    Pending:
    (no pending nodes)
    Recent failures:
    (no failures)
    
    Resources
    ---------------------------------------------------------------
    Usage:
    0.0/416.0 CPU
    0.0/8.0 TPU
    0.0/1.0 TPU-v5p-16-head
    0B/546.83GiB memory
    0B/238.35GiB object_store_memory
    0.0/2.0 your-tpu-name
    
    Demands:
    (no resource demands)
    

    O segundo host de TPU agora é um nó no cluster. A lista de recursos disponíveis agora mostra oito TPUs (0.0/8.0 TPU).

Executar uma carga de trabalho do Ray

  1. Atualize o snippet de código para execução no cluster do Ray:

    import os
    import torch
    import torch_xla
    import torch_xla.core.xla_model as xm
    import ray
    
    import torch.distributed as dist
    import torch_xla.runtime as xr
    from torch_xla._internal import pjrt
    
    # Defines the local PJRT world size, the number of processes per host.
    LOCAL_WORLD_SIZE = 4
    # Defines the number of hosts in the Ray cluster.
    NUM_OF_HOSTS = 4
    GLOBAL_WORLD_SIZE = LOCAL_WORLD_SIZE * NUM_OF_HOSTS
    
    def init_env():
        local_rank = int(os.environ['TPU_VISIBLE_CHIPS'])
    
        pjrt.initialize_multiprocess(local_rank, LOCAL_WORLD_SIZE)
        xr._init_world_size_ordinal()
    
    # This decorator signals to Ray that the `print_tensor()` function should be run on a single TPU chip.
    @ray.remote(resources={"TPU": 1})
    def print_tensor():
        # Initializes the runtime environment on each Ray worker. Equivalent to
        # the `torch_xla.launch call` in the Run PyTorch/XLA on multiple devices section.
        init_env()
    
        t = torch.randn(2, 2, device=xm.xla_device())
        print(t.device)
        print(t)
    
    ray.init()
    
    # Uses Ray to dispatch the function call across available nodes in the cluster.
    tasks = [print_tensor.remote() for _ in range(GLOBAL_WORLD_SIZE)]
    ray.get(tasks)
    
    ray.shutdown()
    
  2. Execute o script no nó principal do Ray. Substitua ray-workload.py pelo caminho do script.

    python ray-workload.py

    A saída será assim:

    WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
    xla:0
    xla:0
    xla:0
    xla:0
    xla:0
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    xla:0
    xla:0
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    xla:0
    tensor([[ 0.6220, -1.4707],
            [-1.2112,  0.7024]], device='xla:0')
    

    A saída indica que a função foi chamada em cada dispositivo XLA (oito dispositivos neste exemplo) na fração de TPU de vários hosts.

Modo centrado no host (JAX)

As seções a seguir descrevem o modo centrado no host com uso do JAX. O JAX usa um paradigma de programação funcional e aceita a semântica de nível superior de programa único e vários dados (SPMD). Em vez de cada processo interagir com um único dispositivo XLA, o código JAX foi projetado para operar em vários dispositivos em um único host simultaneamente.

O JAX é destinado à computação de alto desempenho e pode usar TPUs de maneira eficiente para treinamento e inferência em grande escala. Esse modo é ideal se você já conhece os conceitos de programação funcional, porque é possível aproveitar todo o potencial do JAX.

Estas instruções presumem que você já tem um ambiente do Ray e de TPU configurado, incluindo um ambiente de software com o JAX e outros pacotes relacionados. Para criar um cluster de TPU do Ray, siga as instruções em Iniciar um cluster do GKE do Google Cloud com TPUs para KubeRay. Para saber como usar TPUs com KubeRay, consulte Usar TPUs com KubeRay.

Executar uma carga de trabalho do JAX em uma TPU de host único

O script de exemplo a seguir demonstra como executar uma função JAX em um cluster do Ray com uma TPU de host único, como uma v6e-4. Se você tiver uma TPU de vários hosts, o script vai parar de responder devido ao modelo de execução de vários controladores do JAX. Para saber como executar o Ray em uma TPU de vários hosts, consulte Executar uma carga de trabalho do JAX em uma TPU de vários hosts.

  1. 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-a
    export ACCELERATOR_TYPE=v6e-4
    export RUNTIME_VERSION=v2-alpha-tpuv6e

    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.

  2. Use este comando para criar uma VM de TPU v6e com quatro núcleos:

    gcloud compute tpus tpu-vm create $TPU_NAME \
       --zone=$ZONE \
       --accelerator-type=$ACCELERATOR_TYPE  \
       --version=$RUNTIME_VERSION
  3. Conecte-se à VM de TPU usando este comando:

    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
  4. Instale o JAX e o Ray na TPU.

    pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    
  5. Salve o código a seguir em um arquivo. Por exemplo, ray-jax-single-host.py.

    import ray
    import jax
    
    @ray.remote(resources={"TPU": 4})
    def my_function() -> int:
        return jax.device_count()
    
    h = my_function.remote()
    print(ray.get(h)) # => 4
    

    Se você já sabe executar o Ray com GPUs, há algumas diferenças importantes ao usar TPUs:

    • Em vez de definir num_gpus, especifique TPU como um recurso personalizado e defina o número de chips de TPU.
    • Especifique a TPU usando o número de chips por nó de trabalho do Ray. Por exemplo, se você usar uma v6e-4, executar uma função remota com TPU definido como 4 vai consumir todo o host da TPU.
    • Isso é diferente de como as GPUs normalmente são executadas, com um processo por host. Não recomendamos definir TPU como um número diferente de 4.
      • Exceção: se você tiver uma v6e-8 ou v5litepod-8 de host único, defina esse valor como 8.
  6. Execute o script.

    python ray-jax-single-host.py

Executar uma carga de trabalho do JAX em uma TPU de vários hosts

O script de exemplo a seguir demonstra como executar uma função JAX em um cluster do Ray com uma TPU de vários hosts. O script de exemplo usa uma v6e-16.

  1. 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-a
    export ACCELERATOR_TYPE=v6e-16
    export RUNTIME_VERSION=v2-alpha-tpuv6e

    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.

  2. Use este comando para criar uma VM de TPU v6e com 16 núcleos:

    gcloud compute tpus tpu-vm create $TPU_NAME \
       --zone=$ZONE \
       --accelerator-type=$ACCELERATOR_TYPE  \
       --version=$RUNTIME_VERSION
  3. Instale o JAX e o Ray em todos os workers de TPU.

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --zone=$ZONE \
       --worker=all \
       --command="pip install ray jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
  4. Salve o código a seguir em um arquivo. Por exemplo, ray-jax-multi-host.py.

    import ray
    import jax
    
    @ray.remote(resources={"TPU": 4})
    def my_function() -> int:
        return jax.device_count()
    
    ray.init()
    num_tpus = ray.available_resources()["TPU"]
    num_hosts = int(num_tpus) # 4
    h = [my_function.remote() for _ in range(num_hosts)]
    print(ray.get(h)) # [16, 16, 16, 16]
    

    Se você já sabe executar o Ray com GPUs, há algumas diferenças importantes ao usar TPUs:

    • Semelhanças com as cargas de trabalho do PyTorch em GPUs:
    • Ao contrário das cargas de trabalho do PyTorch em GPUs, o JAX tem uma visão global dos dispositivos disponíveis no cluster.
  5. Copie o script em todos os workers de TPU.

    gcloud compute tpus tpu-vm scp ray-jax-multi-host.py $TPU_NAME: --zone=$ZONE --worker=all
  6. Execute o script.

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --zone=$ZONE \
       --worker=all \
       --command="python ray-jax-multi-host.py"

Executar uma carga de trabalho do JAX de várias frações

Com várias frações, é possível executar cargas de trabalho que abrangem várias frações de TPU em um ou vários Pods de TPU usando a rede do data center.

Você pode usar o pacote ray-tpu para simplificar as interações do Ray com as frações de TPU.

Instale ray-tpu usando pip.

pip install ray-tpu

Para saber como usar o pacote ray-tpu, consulte Introdução no repositório do GitHub. Para um exemplo de como usar várias frações, consulte Execução em várias frações.

Orquestrar cargas de trabalho usando o Ray e o MaxText

Para saber como usar o Ray com o MaxText, consulte Executar um job de treinamento com o MaxText.

Recursos de TPU e do Ray

O Ray trata as TPUs de maneira diferente das GPUs para lidar com a diferença de uso. Neste exemplo, há nove nós do Ray no total:

  • O nó principal do Ray está sendo executado em uma VM n1-standard-16.
  • Os nós de trabalho do Ray estão sendo executados em duas TPUs v6e-16. Cada TPU é composta por quatro workers.
$ ray status
======== Autoscaler status: 2024-10-17 09:30:00.854415 ========
Node status
---------------------------------------------------------------
Active:
 1 node_e54a65b81456cee40fcab16ce7b96f85406637eeb314517d9572dab2
 1 node_9a8931136f8d2ab905b07d23375768f41f27cc42f348e9f228dcb1a2
 1 node_c865cf8c0f7d03d4d6cae12781c68a840e113c6c9b8e26daeac23d63
 1 node_435b1f8f1fbcd6a4649c09690915b692a5bac468598e9049a2fac9f1
 1 node_3ed19176e9ecc2ac240c818eeb3bd4888fbc0812afebabd2d32f0a91
 1 node_6a88fe1b74f252a332b08da229781c3c62d8bf00a5ec2b90c0d9b867
 1 node_5ead13d0d60befd3a7081ef8b03ca0920834e5c25c376822b6307393
 1 node_b93cb79c06943c1beb155d421bbd895e161ba13bccf32128a9be901a
 1 node_9072795b8604ead901c5268ffcc8cc8602c662116ac0a0272a7c4e04
Pending:
 (no pending nodes)
Recent failures:
 (no failures)

Resources
---------------------------------------------------------------
Usage:
 0.0/727.0 CPU
 0.0/32.0 TPU
 0.0/2.0 TPU-v6e-16-head
 0B/5.13TiB memory
 0B/1.47TiB object_store_memory
 0.0/4.0 tpu-group-0
 0.0/4.0 tpu-group-1

Demands:
 (no resource demands)

Descrições de campos de uso de recursos:

  • CPU: o número total de CPUs disponíveis no cluster.
  • TPU: o número de chips de TPU no cluster.
  • TPU-v6e-16-head: um identificador especial para o recurso que corresponde ao worker 0 de uma fração de TPU. Relevante para acessar frações de TPU individuais.
  • memory: memória de heap do worker usada pelo aplicativo.
  • object_store_memory: memória usada quando o aplicativo cria objetos no repositório de objetos usando ray.put e quando retorna valores de funções remotas.
  • tpu-group-0 e tpu-group-1: identificadores exclusivos das frações de TPU individuais. Relevante para executar jobs em frações. Esses campos são definidos como 4 porque há quatro hosts por fração de TPU em uma v6e-16.