Treinamento distribuído com Ray Train DLCs e EFA no EKS

O problema que o Ray Train DLC resolve

Quem já tentou escalar um job de PyTorch para múltiplos nós sabe que o maior obstáculo raramente é o código do modelo. O verdadeiro trabalho está em montar — e manter alinhada — a pilha de dependências dentro do container: o runtime CUDA, o PyTorch compilado contra ele, o NCCL (Biblioteca de Comunicações Coletivas da NVIDIA), as bibliotecas de espaço de usuário do Adaptador de Fabric Elástico (EFA), o plugin AWS OFI NCCL que conecta o NCCL ao EFA, e ainda uma versão do Ray compatível com o operador do cluster. Qualquer desalinhamento entre essas peças não gera um erro claro: o NCCL simplesmente cai para sockets TCP e o job termina em fração da velocidade esperada.

O Ray Train Deep Learning Container (DLC) é a resposta da AWS para esse problema. A empresa monta essa pilha inteira, testa tudo como uma unidade, publica a imagem no Amazon Elastic Container Registry (Amazon ECR) Público e mantém os patches em dia. Nenhuma imagem precisa ser construída pela equipe que vai usá-la.

Como o cluster é organizado

O KubeRay é um operador Kubernetes de código aberto que adiciona o recurso RayCluster: você descreve um grupo de cabeça (head) e um grupo de trabalhadores (workers) em um único manifesto, e o KubeRay constrói e mantém o cluster a partir disso. O Ray opera um job em várias máquinas como um único cluster com dois papéis: o nó cabeça coordena o estado global e serve como ponto de submissão de jobs; os nós workers fornecem as GPUs e executam os processos de treinamento.

No exemplo publicado pela AWS, um único cluster no Amazon Elastic Kubernetes Service (Amazon EKS) é composto por: um nó de sistema m7i.xlarge que hospeda o operador KubeRay e o nó cabeça do Ray (sem GPUs), e dois nós g6.12xlarge com EFA habilitado, ambos na mesma Zona de Disponibilidade. Cada nó g6.12xlarge possui quatro GPUs NVIDIA L4, 48 vCPUs, 192 GiB de memória, 40 Gbps de rede e uma interface EFA. O Ray Train distribui oito workers — um por GPU — com os ranks 0 a 3 no primeiro nó e 4 a 7 no segundo.

Todos os nós ficam em sub-redes privadas de uma única Nuvem Privada Virtual (VPC) e acessam a internet apenas via NAT Gateway — é por ali que cada nó puxa a imagem DLC do Amazon ECR Público e os workers baixam os pesos do modelo do Hugging Face. Como o tráfego EFA não cruza Zonas de Disponibilidade, os scripts fixam ambas as instâncias na mesma zona.

O que vem pronto no DLC, sem precisar construir nada

A tag utilizada no exemplo — public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1 — entrega toda a pilha a seguir sem que a equipe precise montar ou manter nada:

• Runtime distribuído: Ray 2.58.0 com Train, Tune e Data
• Framework: PyTorch 2.13.0, torchvision 0.28.0
• Runtime de GPU: CUDA 13.0.2, cuDNN, NCCL 2.29.7
• Rede: instalador EFA, libfabric, plugin AWS OFI NCCL e GDRCopy — todos compilados contra o NCCL incluído
• Bibliotecas de treinamento: DeepSpeed, flash-attn, Transformer Engine, PyTorch Lightning, bibliotecas Hugging Face
• Diagnóstico: all_reduce_perf já disponível no path para testes de NCCL

Além disso, algumas propriedades eliminam trabalho recorrente por cluster:

• Sem configuração de registro: a imagem é pública, então os nós a puxam sem credenciais.
• Sem correspondência de driver: no startup, o entrypoint verifica o driver NVIDIA do host e ativa automaticamente as bibliotecas de compatibilidade forward do CUDA se o driver for mais antigo que o CUDA da imagem.
• Uma imagem para todos os papéis: a mesma tag serve tanto para o nó cabeça quanto para os workers — o que define o papel é o comando ray start passado na inicialização.
• Família de imagens compartilhada: o Ray Train compartilha o repositório ECR com o Ray Serve DLC e os Ray LLM DLCs, sob o prefixo de tag train-ml.

Configure, não reconstrua

Tudo o que muda de um job para outro é definido fora da imagem. O exemplo usa os seguintes mecanismos de configuração:

• Papel do pod (cabeça ou worker): via parâmetros rayStartParams do KubeRay. O cabeça roda com num-gpus: "0".
• GPUs e EFA por pod: via resource requests do Kubernetes (nvidia.com/gpu: 4, vpc.amazonaws.com/efa: 1).
• Transporte de rede: variáveis de ambiente no pod — FI_PROVIDER, NCCL_NET_PLUGIN, NCCL_SOCKET_IFNAME, NCCL_DEBUG.
• Pacotes Python extras: via runtime_env do Ray no momento da submissão, sem alterar a imagem.
• Versão da imagem: tag versionada ou fixada por digest em DLC_IMAGE.
• Forma do cluster: variáveis em scripts/env.sh — tipo de instância, contagem de nós, GPUs por nó, região e zona.

Um detalhe importante: solicitar vpc.amazonaws.com/efa: 1 no manifesto é o que faz o device plugin do EFA montar a interface dentro do container. Sem isso, o NCCL não enxerga o EFA e cai para sockets TCP.

O código de treinamento: uma função e um trainer

O arquivo code/train.py é o job inteiro e não contém nada específico ao container, ao cluster ou à rede. O Ray Train executa a função train_func uma vez por GPU. Ela consulta o Ray para obter seu rank, carrega o modelo e o passa para prepare_model para o FSDP (Paralelismo de Dados Totalmente Fragmentado), que divide pesos, gradientes e estado do otimizador entre todas as oito GPUs:

def train_func(config):
    ctx = ray.train.get_context()
    world_size, rank = ctx.get_world_size(), ctx.get_world_rank()
    device = ray.train.torch.get_device()
    model = AutoModelForCausalLM.from_pretrained(config["model_id"], torch_dtype=torch.bfloat16)
    # Any submodule above 10M parameters becomes its own FSDP unit. Wrapping by
    # size, not by a model-specific layer class, keeps the script model-agnostic.
    auto_wrap_policy = partial(size_based_auto_wrap_policy, min_num_params=int(1e7))
    model = ray.train.torch.prepare_model(
        model,
        parallel_strategy="fsdp",
        parallel_strategy_kwargs={"auto_wrap_policy": auto_wrap_policy},
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=config["learning_rate"])
    for step in range(config["steps"]):
        input_ids = torch.randint(0, model.config.vocab_size, (config["batch_size"], config["seq_len"]), device=device)
        optimizer.zero_grad()
        loss = model(input_ids=input_ids, labels=input_ids).loss
        loss.backward()
        # parameters gathered again, gradients reduce-scattered
        optimizer.step()
        # local: each rank updates only its own slice
        ray.train.report({"step": step, "loss": loss.item()})

O driver se conecta ao cluster em execução, conta as GPUs disponíveis e inicia o trainer. Como a contagem de workers vem do cluster, o mesmo script funciona sem alteração em um ou vários nós:

ray.init(address="auto")  # join the RayCluster
num_workers = int(ray.cluster_resources().get("GPU", 0))  # 8 here
trainer = TorchTrainer(
    train_func,
    train_loop_config={...},
    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=True),
    run_config=RunConfig(name="ray-train-fsdp-llm"),
)
result = trainer.fit()

O Qwen2.5-1.5B em bfloat16 precisa de cerca de 12 GB para pesos, gradientes e estado do otimizador. Fragmentado entre oito GPUs, cada L4 segura aproximadamente 1,5 GB. O custo é comunicação: o FSDP reúne os pesos de cada camada antes de usá-la e dispersa os gradientes depois — e é exatamente esse tráfego entre instâncias que a pilha EFA do DLC transporta.

O manifesto: onde a imagem encontra o cluster

O arquivo manifest/raycluster.yaml é curto justamente porque a imagem precisa de pouca configuração. O grupo de workers concentra as decisões principais:

workerGroupSpecs:
  - groupName: gpu-workers
    replicas: 2
    rayStartParams:
      num-gpus: "${GPUS_PER_NODE}"  # keeps Ray's GPU count equal to the pod's
    template:
      spec:
        nodeSelector: { role: gpu-worker }
        containers:
          - name: ray-worker
            image: "${DLC_IMAGE}"  # the DLC as published; same tag as the head
            env:
              - { name: NCCL_DEBUG,        value: "INFO"  }  # print the chosen transport
              - { name: NCCL_SOCKET_IFNAME, value: "eth0" }  # pod network on EKS
              - { name: FI_PROVIDER,        value: "efa"  }  # libfabric: EFA, nothing else
              - { name: NCCL_NET_PLUGIN,    value: "ofi"  }  # load the AWS OFI plugin
            resources:
              limits:   { nvidia.com/gpu: "${GPUS_PER_NODE}", vpc.amazonaws.com/efa: "1", cpu: "40", memory: 160Gi }
              requests: { nvidia.com/gpu: "${GPUS_PER_NODE}", vpc.amazonaws.com/efa: "1", cpu: "40", memory: 160Gi }

Três decisões estão explícitas aqui: qual imagem usar (${DLC_IMAGE}, sem modificação, para cabeça e workers); qual hardware alocar (quatro GPUs e uma interface EFA por pod); e qual transporte usar (quatro variáveis de ambiente, sem depender de auto-detecção). Dois detalhes de estabilidade: requests iguais a limits colocam os workers na classe de Qualidade de Serviço Garantida (QoS) do Kubernetes — um rank throttled ou despejado paralisa os outros sete. E num-gpus deriva da mesma variável que o request de GPU, então a visão do Ray sobre capacidade sempre bate com a do Kubernetes.

Subindo a infraestrutura e rodando o job

Toda a configuração fica em scripts/env.sh: nome do cluster, região, versão do Kubernetes, tipos e contagens de nós, a imagem, versões fixadas do KubeRay e dos device plugins, e os parâmetros de treinamento.

O script deploy_all.sh é o ponto de entrada. Ele começa executando find_gpu_capacity.sh, que verifica por região candidata: se o g6.12xlarge está disponível em cada Zona de Disponibilidade, se a cota de vCPUs On-Demand da família G/VT tem espaço para a contagem de nós solicitada, e o score de placement de Spot como sinal de capacidade. Com a flag --probe, ele obtém uma resposta definitiva criando e cancelando imediatamente uma Reserva de Capacidade real em cada zona candidata.

Em seguida, os quatro passos são executados em ordem: deploy_cluster.sh (plano de controle EKS, VPC, provedor OIDC, add-ons e node group de sistema), deploy_node_group.sh (dois g6.12xlarge com EFA, rotulados role=gpu-worker), install_gpu_plugins.sh (device plugins NVIDIA e EFA) e install_kuberay.sh (operador KubeRay no nó de sistema). O conjunto leva de 20 a 30 minutos em uma região nova.

O job de treinamento é um passo separado, executado por deploy_ray_train_job.sh. Ele renderiza o manifesto com envsubst, aplica, verifica os recursos alocáveis de cada nó GPU, aguarda os pods ficarem prontos e então submete o job de dentro do pod cabeça via kubectl exec:

ray job submit --address http://localhost:8265 --working-dir /tmp/ray-train-code \
    python3 train.py --model_id Qwen/Qwen2.5-1.5B

A chamada de submit transmite os logs de todos os oito ranks e retorna quando o job termina. O script verifica qual transporte o NCCL usou e sai com erro se o job falhou ou se o EFA não estava no caminho — o que o torna utilizável como gate de Integração Contínua (CI).

Adaptando para outros workloads

Como a imagem é usada como publicada e o job aprende sua escala do cluster, a maioria das mudanças toca exatamente um arquivo:

• Trocar o modelo: altere apenas MODEL_ID. Código, manifesto e imagem ficam intactos.
• Ajustar steps, sequence length ou batch size: altere STEPS, SEQ_LEN, BATCH_SIZE. Código, manifesto e imagem ficam intactos.
• Adicionar pacote Python: use --runtime-env-json com lista pip. A imagem não muda.
• Fixar versão do DLC: aponte DLC_IMAGE para uma tag versionada. Todo o resto permanece.
• Trocar família de instâncias GPU: altere GPU_NODE_TYPE e GPUS_PER_NODE. A família precisa ter suporte a EFA. Migrar para p5 ou outra família com GPUDirect não exige mudança na imagem ou no código de treinamento.
• Escalar o número de instâncias: altere GPU_NODE_COUNT mais replicas, minReplicas e maxReplicas no manifesto. O código conta as GPUs automaticamente.

O repositório de exemplo contém toda a solução e usa a mesma base KubeRay no EKS do walkthrough multi-nó do Ray Serve.

Derrubando o ambiente

As instâncias GPU custam em torno de US$ 9,50 por hora junto com o plano de controle e o NAT Gateway. A remoção segue a ordem inversa da criação:

./delete_ray_train_job.sh        # RayCluster
./install_kuberay.sh cleanup     # KubeRay operator
./install_gpu_plugins.sh cleanup # NVIDIA and EFA device plugins
./delete_node_group.sh           # GPU instances
./delete_cluster.sh              # EKS cluster, VPC and system node

Próximos passos sugeridos

A AWS sugere como evoluções naturais a partir desse exemplo:

• Substituir os batches aleatórios por um dataset real e fazer checkpoint no Amazon S3 ou Amazon FSx for Lustre.
• Fixar DLC_IMAGE em uma tag versionada para runs reproduzíveis.
• Ativar FailureConfig para que runs longas sobrevivam à perda de uma instância.
• Migrar para um modelo maior em uma família com GPUDirect RDMA, como p5. A imagem e o código permanecem os mesmos.

Fonte

Distributed training with Ray Train DLCs and EFA on EKS (https://aws.amazon.com/blogs/machine-learning/distributed-training-with-ray-train-dlcs-and-efa-on-eks/)