Scaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughput
aws ·
This article details an optimized architecture for scaling Mixture-of-Experts (MoE) reinforcement learning on Amazon EKS, addressing challenges in heterogeneous compute coordination and high-throughput communication. The solution combines Amazon Elastic Kubernetes Service (Amazon EKS), Elastic Fabric Adapter (EFA), and DeepEP to balance rollout generation with policy training. A key innovation is the use of DeepEP over EFA, which replaces standard all-to-all collectives with specialized kernels to reduce communication overhead in expert parallelism. Testing on 48 P5en instances demonstrated a 40% increase in aggregate RL rollout throughput compared to baseline configurations. The architecture also leverages Amazon EC2 Spot Instances for rollout generation to lower costs without disrupting tightly coupled policy training workloads. Separate node groups are used for GPU-accelerated tasks, CPU-based environments, and memory-optimized experience buffers to ensure independent scaling. This approach allows for efficient co-design across the stack, balancing compute, memory, and bandwidth constraints. The result is a flexible, cost-effective path for large-scale distributed RL workloads involving hundreds or thousands of accelerators.
aws