Scaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughput

來源原標題

Scaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughput

作者

Ashvin Nihalani

原始標籤

Amazon Elastic Kubernetes Service · Technical How-to

原始連結

https://aws.amazon.com/blogs/machine-learning/scaling-moe-reinforcement-learning-on-amazon-eks-with-efa-and-deepep-with-40-more-throughput/

摘要原文

When you post-train a Mixture-of-Experts (MoE) model with Reinforcement Learning from Human Feedback (RLHF) or Group Relative Policy Optimization (GRPO) at scale, three simultaneous challenges emerge. The first requires coordinating heterogeneous compute for rollout generation and policy training. Second, sustaining high-throughput communication across hundreds of accelerators. And third, dynamically orchestrating every subsystem to keep them in balance. On AWS, you can address these challenges using Amazon Elastic Kubernetes Service (Amazon EKS), Elastic Fabric Adapter (EFA), and DeepEP. Mixture-of-Experts (MoE) has become a standard architecture for scaling large language models (LLMs) to hundreds of billions or even trillions of parameters, while maintaining efficient inference through sparsity. However, sparsity doesn’t remove infrastructure complexity in training. As part of the standard training pipeline, these models must undergo pre-training, mid-training, supervised fine-tuning (SFT), and reinforcement learning (RL). Among these stages, large-scale RL training places unusual demands on infrastructure because it combines elastic inference work with tightly coupled model training that requires high-bandwidth communication. Reward models, verifiers, and checkpoint updates add further memory, networking, and orchestration pressure. This type of multi-workload optimization reflects a common infrastructure challenge when you run model training, inference, and evaluation on shared resources without creating bottlenecks or leaving capacity idle. Compared with dense models, post-training MoE models introduce a new infrastructure challenge: as newer MoE architectures become increasingly sparse to reduce inference costs, training becomes constrained more by communication than by compute. A key source of this communication overhead is Expert Parallelism (EP). EP introduces dynamic all-to-all token routing across devices, in addition to the dense, structured communication patterns of Tensor Parallelism (TP), Data Parallelism (DP), and Pipeline Parallelism (PP). In tightly coupled asynchronous RL workloads, the heterogeneous compute and communication demands of MoE training must be balanced with inference-based generation. Slow training steps stall inference workers, while insufficient inference throughput leaves training accelerators idle. This challenge is common across large-scale reinforcement learning workloads, including RLHF pipelines based on Proximal Policy Optimization (PPO) and newer approaches such as GRPO. PPO typically uses a critic model to estimate value during policy optimization, while GRPO avoids the need for a separate critic model by using group-based relative rewards. Although their algorithmic and model requirements differ, both impose similar infrastructure demands: large-scale rollout generation, tightly coupled policy training, and high-bandwidth inter-node communication. In this post, we describe an architecture optimized to accelerate MoE training that combines Amazon Elastic Kubernetes Service (Amazon EKS) and EFA to orchestrate and accelerate large-scale RL training and how DeepEP optimizes expert-parallel communication over EFA. Large-scale RL training presents three interrelated challenges: Large-scale asynchronous RL jobs have two distinct, simultaneous workloads to optimize: rollout generation and policy training. During rollout generation, the system performs large-scale distributed inference focused on maximizing aggregate throughput rather than minimizing time to first token (TTFT) or inter-token latency. In contrast, policy training requires tightly coupled workers that progress in lockstep, much like pre-training or SFT. Any latency spike or straggling worker can stall the entire job or trigger NVIDIA Collective Communications Library (NCCL) timeouts. RL systems must balance these two workloads because any mismatch in their rates can leave hardware idle or introduce training instability. Figure 1: The asynchronous RL loop, where rollout workers generate experience through distributed inference while policy-training workers consume batches and update model weights RL workloads have heterogeneous compute and communication demands. As a result, any RL system must balance three resource constraints: accelerator compute, memory, and network bandwidth. Policy training is compute-intensive and must keep pace with rollout generation. At the same time, distributed inference must manage KV-cache capacity and token generation. Balancing memory bandwidth and compute is critical for maximum throughput because MoE layers add sparse, dynamic all-to-all communication as tokens are routed across devices. Reward models provide feedback during training, and data movement adds further pressure. All of these subsystems must be balanced jointly to help prevent any one from becoming a bottleneck. As RL training jobs scale beyond a single instance, model partitions and parallelism groups span multiple nodes, shifting communication from the high-bandwidth intra-node NVLink fabric to lower-bandwidth inter-node links. MoE models intensify this shift: unlike the structured patterns of Tensor Parallelism and Pipeline Parallelism, Expert Parallelism dynamically routes tokens across devices through sparse, fine-grained all-to-all communication traffic that becomes increasingly inter-node as the expert parallelism degree grows. Figure 2: Communication domains in multi-node MoE training, where NVLink carries high-bandwidth intra-node traffic while EFA handles inter-node token routing for Expert Parallelism, Tensor Parallelism, and Data Parallelism AWS accelerated computing instances such as P5 and P6 use two primary communication domains: an intra-instance NVLink fabric, typically connected through NVSwitch, and inter-instance networking through EFA. EFA provides high-bandwidth communication traffic between instances. On supported configurations , EFA works with NVIDIA GPUDirect RDMA and OS bypass to transfer data directly between GPU memory buffers across instances, reducing CPU and operating-system involvement in the communication path. Optimizing bandwidth utilization in RL workloads requires balancing these two communication domains by determining which operations can run efficiently over EFA and which must remain within the NVLink fabric. To scale RL workloads on AWS, we combine Amazon EKS, EFA, and Amazon Simple Storage Service (Amazon S3) so that orchestration, high-performance communication, and durable storage can scale independently. With Amazon EKS, you can manage the lifecycle and placement of heterogeneous workers. With EFA, you get the inter-node data path for communication-intensive GPU workloads. With Amazon S3, you can store datasets, model checkpoints, and completed training artifacts including the model weights. The following sections describe how to map the distinct layers of the RL system, covering orchestration, high-performance networking, and durable storage, and how each layer scales independently. The Amazon EKS cluster contains separate node groups optimized for each stage of the RL workflow. GPU-accelerated instances run rollout generation, reward-model inference, and policy training, while CPU instances execute environments and preprocessing tasks. Memory-optimized instances host experience buffers and checkpoint caches, allowing producers and consumers to exchange data without placing durable storage directly on the critical path. Figure 3: EKS cluster topology with GPU node groups for policy training and rollout generation, CPU node groups for environment workers and preprocessing, and memory-optimized instances for experience buffers and checkpoint caches During rollout, the model generates samples through interactions with CPU-based environment pods, and the resulting experience flows into a memory-optimized buffer.