Treinamento distribuído com Ray Train DLC e EFA no EKS
aws ·
Este artigo detalha um método para realizar treinamento distribuído do PyTorch no Amazon EKS usando Deep Learning Containers (DLCs) do Ray Train e o Elastic Fabric Adapter (EFA). Os autores explicam que o principal gargalo ao escalar o treinamento é geralmente o gerenciamento de dependências complexas de contêineres, em vez do próprio código do modelo. Para resolver isso, a AWS oferece DLCs pré-construídos que agrupam componentes necessários, como bibliotecas CUDA, NCCL e EFA, eliminando a necessidade de construção de imagens personalizadas. A solução utiliza o KubeRay para gerenciar a estrutura do cluster Ray dentro do Kubernetes, definindo pods de cabeçalho e de trabalho por meio de um único manifesto. Uma implementação de exemplo demonstra o ajuste fino do modelo Qwen2.5-1.5B em oito GPUs NVIDIA L4 usando Parallelismo de Dados Totalmente Fragmentado (Fully Sharded Data Parallelism). Ao configurar variáveis de ambiente e solicitações de recursos externamente, os usuários podem aproveitar o EFA para comunicação inter-nós de alta velocidade sem modificar a imagem do contêiner. Essa abordagem garante um treinamento reprodutível e escalável, reduzindo a sobrecarga operacional associada à manutenção da infraestrutura.
aws