Dimensionar o aprendizado por reforço MoE no Amazon EKS com EFA e DeepEP com 40% mais throughput
aws ·
Este artigo detalha uma arquitetura otimizada para dimensionar o aprendizado por reforço com Mistura de Especialistas (MoE) no Amazon EKS, abordando desafios na coordenação de computação heterogênea e na comunicação de alta taxa de transferência. A solução combina o Amazon Elastic Kubernetes Service (Amazon EKS), o Elastic Fabric Adapter (EFA) e o DeepEP para equilibrar a geração de rollouts com o treinamento da política. Uma inovação chave é o uso do DeepEP sobre o EFA, que substitui as coletivas all-to-all padrão por kernels especializados para reduzir a sobrecarga de comunicação no paralelismo de especialistas. Testes em 48 instâncias P5en demonstraram um aumento de 40% na taxa de transferência agregada de rollouts de RL em comparação com configurações de referência. A arquitetura também utiliza Instâncias Spot da Amazon EC2 para a geração de rollouts, reduzindo custos sem interromper as cargas de trabalho de treinamento de políticas fortemente acopladas. Grupos de nós separados são usados para tarefas aceleradas por GPU, ambientes baseados em CPU e buffers de experiência otimizados para memória, garantindo escalabilidade independente. Essa abordagem permite um co-design eficiente em toda a pilha, equilibrando restrições de computação, memória e largura de banda. O resultado é um caminho flexível e econômico para cargas de trabalho de RL distribuído em grande escala, envolvendo centenas ou milhares de aceleradores.
aws