Scaling MoE reinforcement studying on Amazon EKS with EFA and DeepEP with 40% extra throughput

Whenever you post-train a Combination-of-Specialists (MoE) mannequin with Reinforcement Studying from Human Suggestions (RLHF) or Group Relative Coverage Optimization (GRPO) at scale, three simultaneous challenges emerge. The primary requires coordinating heterogeneous compute for rollout era and coverage coaching. Second, sustaining high-throughput communication throughout a whole lot of accelerators. And third, dynamically orchestrating each subsystem to maintain them in steadiness. On AWS, you possibly can tackle these challenges utilizing Amazon Elastic Kubernetes Service (Amazon EKS), Elastic Cloth Adapter (EFA), and DeepEP.

Combination-of-Specialists (MoE) has develop into an ordinary structure for scaling massive language fashions (LLMs) to a whole lot of billions and even trillions of parameters, whereas sustaining environment friendly inference via sparsity. Nevertheless, sparsity doesn’t take away infrastructure complexity in coaching. As a part of the usual coaching pipeline, these fashions should bear pre-training, mid-training, supervised fine-tuning (SFT), and reinforcement studying (RL). Amongst these levels, large-scale RL coaching locations uncommon calls for on infrastructure as a result of it combines elastic inference work with tightly coupled mannequin coaching that requires high-bandwidth communication. Reward fashions, verifiers, and checkpoint updates add additional reminiscence, networking, and orchestration stress. Any such multi-workload optimization displays a standard infrastructure problem while you run mannequin coaching, inference, and analysis on shared assets with out creating bottlenecks or leaving capability idle.

In contrast with dense fashions, post-training MoE fashions introduce a brand new infrastructure problem: as newer MoE architectures develop into more and more sparse to cut back inference prices, coaching turns into constrained extra by communication than by compute. A key supply of this communication overhead is Skilled Parallelism (EP). EP introduces dynamic all-to-all token routing throughout units, along with the dense, structured communication patterns of Tensor Parallelism (TP), Knowledge Parallelism (DP), and Pipeline Parallelism (PP). In tightly coupled asynchronous RL workloads, the heterogeneous compute and communication calls for of MoE coaching have to be balanced with inference-based era. Sluggish coaching steps stall inference employees, whereas inadequate inference throughput leaves coaching accelerators idle. This problem is frequent throughout large-scale reinforcement studying workloads, together with RLHF pipelines based mostly on Proximal Coverage Optimization (PPO) and newer approaches reminiscent of GRPO.

PPO usually makes use of a critic mannequin to estimate worth throughout coverage optimization, whereas GRPO avoids the necessity for a separate critic mannequin by utilizing group-based relative rewards. Though their algorithmic and mannequin necessities differ, each impose comparable infrastructure calls for: large-scale rollout era, tightly coupled coverage coaching, and high-bandwidth inter-node communication.

On this publish, we describe an structure optimized to speed up MoE coaching that mixes Amazon Elastic Kubernetes Service (Amazon EKS) and EFA to orchestrate and speed up large-scale RL coaching and the way DeepEP optimizes expert-parallel communication over EFA.

Challenges of large-scale RL coaching

Massive-scale RL coaching presents three interrelated challenges:

1. Balancing the competing useful resource calls for of rollout era and coverage coaching.
2. Managing accelerator compute, reminiscence, and community bandwidth concurrently.
3. Dealing with the shift from high-bandwidth intra-node communication to lower-bandwidth inter-node hyperlinks as jobs scale past a single occasion.

The rollout-training loop

Massive-scale asynchronous RL jobs have two distinct, simultaneous workloads to optimize: rollout era and coverage coaching. Throughout rollout era, the system performs large-scale distributed inference targeted on maximizing combination throughput relatively than minimizing time to first token (TTFT) or inter-token latency. In distinction, coverage coaching requires tightly coupled employees that progress in lockstep, very like pre-training or SFT. Any latency spike or straggling employee can stall all the job or set off NVIDIA Collective Communications Library (NCCL) timeouts. RL programs should steadiness these two workloads as a result of any mismatch of their charges can go away {hardware} idle or introduce coaching instability.

Compute, reminiscence, and bandwidth stress

RL workloads have heterogeneous compute and communication calls for. Consequently, any RL system should steadiness three useful resource constraints: accelerator compute, reminiscence, and community bandwidth. Coverage coaching is compute-intensive and should maintain tempo with rollout era. On the similar time, distributed inference should handle KV-cache capability and token era. Balancing reminiscence bandwidth and compute is crucial for max throughput as a result of MoE layers add sparse, dynamic all-to-all communication as tokens are routed throughout units. Reward fashions present suggestions throughout coaching, and knowledge motion provides additional stress. All of those subsystems have to be balanced collectively to assist forestall anybody from changing into a bottleneck.

Intra-node versus inter-node communication

As RL coaching jobs scale past a single occasion, mannequin partitions and parallelism teams span a number of nodes, shifting communication from the high-bandwidth intra-node NVLink cloth to lower-bandwidth inter-node hyperlinks. MoE fashions intensify this shift: not like the structured patterns of Tensor Parallelism and Pipeline Parallelism, Skilled Parallelism dynamically routes tokens throughout units via sparse, fine-grained all-to-all communication site visitors that turns into more and more inter-node because the skilled parallelism diploma grows.

Determine 2: Communication domains in multi-node MoE coaching, the place NVLink carries high-bandwidth intra-node site visitors whereas EFA handles inter-node token routing for Skilled Parallelism, Tensor Parallelism, and Knowledge Parallelism

AWS accelerated computing situations reminiscent of P5 and P6 use two major communication domains: an intra-instance NVLink cloth, usually linked via NVSwitch, and inter-instance networking via EFA. EFA supplies high-bandwidth communication site visitors between situations. On supported configurations, EFA works with NVIDIA GPUDirect RDMA and OS bypass to switch knowledge immediately between GPU reminiscence buffers throughout situations, decreasing CPU and operating-system involvement within the communication path. Optimizing bandwidth utilization in RL workloads requires balancing these two communication domains by figuring out which operations can run effectively over EFA and which should stay inside the NVLink cloth.

Structure overview

To scale RL workloads on AWS, we mix Amazon EKS, EFA, and Amazon Easy Storage Service (Amazon S3) in order that orchestration, high-performance communication, and sturdy storage can scale independently. With Amazon EKS, you possibly can handle the lifecycle and placement of heterogeneous employees. With EFA, you get the inter-node knowledge path for communication-intensive GPU workloads. With Amazon S3, you possibly can retailer datasets, mannequin checkpoints, and accomplished coaching artifacts together with the mannequin weights. The next sections describe find out how to map the distinct layers of the RL system, protecting orchestration, high-performance networking, and sturdy storage, and the way every layer scales independently.

EKS cluster topology

The Amazon EKS cluster comprises separate node teams optimized for every stage of the RL workflow. GPU-accelerated situations run rollout era, reward-model inference, and coverage coaching, whereas CPU situations execute environments and preprocessing duties. Reminiscence-optimized situations host expertise buffers and checkpoint caches, permitting producers and customers to change knowledge with out inserting sturdy storage immediately on the crucial path.

Determine 3: EKS cluster topology with GPU node teams for coverage coaching and rollout era, CPU node teams for atmosphere employees and preprocessing, and memory-optimized situations for expertise buffers and checkpoint caches

RL job topology

Throughout rollout, the mannequin generates samples via interactions with CPU-based atmosphere pods, and the ensuing expertise flows right into a memory-optimized buffer. From there, the policy-training step consumes batches, updates mannequin weights, and publishes new checkpoints that feed again into the following spherical of rollout era. Checkpoints and accomplished coaching artifacts are additionally endured to Amazon S3 for sturdy storage, restoration, and downstream use. Coverage coaching, weight updates, and new checkpoint era can all run on the EKS cluster.

Determine 4: RL job knowledge circulation on EKS, the place rollout employees generate expertise via CPU atmosphere interactions and write to a shared reminiscence buffer, and policy-training employees eat batches, publish up to date checkpoints, and persist artifacts to Amazon S3

Community and execution layers

EKS supplies the management airplane for scheduling, scaling, failure restoration, and coordination throughout totally different employee teams. Inside GPU situations, NVLink and NVSwitch carry high-bandwidth intra-node communication. EFA helps latency-sensitive inter-node communication for distributed coverage coaching and different tightly coupled GPU operations. The expertise buffer and Amazon S3 type the info layer, separating high-frequency samples and checkpoint change from long-term artifact storage.

Efficiency and price optimizations

This part covers two key optimizations: utilizing DeepEP to cut back expert-parallel communication overhead over EFA, and utilizing Amazon Elastic Compute Cloud (Amazon EC2) Spot Situations to assist decrease the price of rollout era.

DeepEP over EFA

DeepEP, together with different topology-aware expert-parallel communication methods, is a standard optimization for MoE workloads that goals to cut back communication bottlenecks. Commonplace all-to-all collectives are most effective for dense, common communication, however MoE workloads generate sparse, fine-grained, and imbalanced site visitors as tokens are dynamically routed throughout specialists. As Skilled Parallelism spans a number of nodes, synchronization and per-message overhead enhance, making inter-node communication a dominant bottleneck. DeepEP addresses this by changing generic all-to-all collectives with specialised dispatch and mixing kernels. These kernels use NVLink for intra-node communication and an RDMA-capable backend for inter-node communication.

Amazon has contributed a number of options emigrate DeepEP’s communication primitives to libfabric. This makes the transport layer moveable throughout libfabric-supported community materials and optimizes MoE coaching over EFA. With these adjustments, DeepEP v2 features native EFA help. Moreover, NCCL 2.31 incorporates the most recent EFA optimizations for dense collective communication. Within the following part, we describe how DeepEP over EFA improves rollout-generation throughput by decreasing the communication overhead of skilled dispatch and mixing operations.

How DeepEP communicates over EFA

DeepEP replaces normal NCCL all-to-all collectives with two specialised GPU kernels: a dispatch kernel that routes tokens from native GPUs to distant specialists, and a mix kernel that gathers processed tokens again. For intra-node transfers, these kernels use NVLink via NVSwitch. For inter-node transfers, DeepEP makes use of libfabric to ship knowledge over EFA. On supported occasion varieties reminiscent of P5 and P6, EFA works with NVIDIA GPUDirect RDMA to switch knowledge immediately between GPU reminiscence buffers throughout situations, bypassing the CPU and working system. The upstream contributions from Amazon migrate DeepEP’s communication primitives from a CUDA-specific RDMA backend to libfabric. This makes the transport moveable throughout EFA-supported configurations and reduces per-message overhead for the sparse, fine-grained site visitors patterns that Skilled Parallelism generates.

Throughout 48 P5en situations (16 devoted to coaching, 32 to inference) operating a super-sparse MoE model, enabling DeepEP over EFA elevated combination RL rollout throughput by 40 %. Determine 5 exhibits the throughput comparability with and with out DeepEP.

Determine 5: RL rollout throughput with and with out DeepEP over EFA throughout 48 P5en situations

Spot Situations for rollout era

Rollout era is nicely suited to Amazon EC2 Spot Situations as a result of it consists of distributed inference duties that may be partitioned throughout unbiased employees. In contrast to coverage coaching, the place tightly coupled employees should progress collectively, the interruption of a rollout employee doesn’t require all the RL job to cease. Unfinished rollout duties will be returned to the queue and reassigned whereas the remaining employees proceed producing expertise.

With Amazon EKS, you possibly can scale Spot-based rollout node teams in line with rollout demand and queue depth whereas sustaining steady capability for coverage coaching. Rollout employees ought to course of bounded items of labor and publish accomplished samples regularly. When a Spot interruption discover arrives, employees drain lively requests and return unfinished duties to the queue. This separation can assist cut back rollout-generation prices. Coverage-training employees stay insulated from Spot interruptions, delays, or NCCL timeouts.

Placing all of it collectively

This part walks via provisioning the infrastructure described within the earlier sections, from cluster creation via operating an RL job with DeepEP communication enabled.

Stipulations

Earlier than operating RL on Amazon EKS, confirm the next necessities are met:

• AWS account with acceptable AWS Id and Entry Administration (IAM) permissions.
• Amazon EKS 1.31 or later.
• EFA installer 1.49 with AWS OFI NCCL plugin.
• DeepEP 2.0.0, NCCL 2.31.2, SGLang 0.5.17, PyTorch 2.12.1 (CUDA 13.0).
• Supported GPU situations, reminiscent of p5.48xlarge, p5e.48xlarge, or p6-b200.48xlarge.
• Familiarity with Kubernetes and distributed coaching ideas.

EKS cluster setup

This structure will be deployed on Amazon EKS by separating coverage coaching, rollout era, and supporting providers throughout independently managed node teams. This preserves the isolation between tightly coupled coaching workloads and extra elastic rollout employees whereas permitting every element to make use of the capability mannequin optimized for its execution traits.

Managed node teams streamline provisioning, updates, and occasion lifecycle administration. Coverage-training employees can run on steady GPU capability, whereas rollout-generation capability can scale independently and incorporate Spot Situations the place interruption tolerance permits. CPU-based providers, together with orchestration and supporting parts, will be positioned in separate node teams to keep away from competing with GPU workloads for capability.

The next eksctl ClusterConfig defines a general-purpose CPU node group and a GPU accelerator node group for the cluster:

apiVersion: eksctl.io/v1alpha5
type: ClusterConfig

metadata:
  title: my-eks-cluster
  area: us-west-2
  model: "1.33"

managedNodeGroups:

  # ------------------------------------------------------------
  # Common-purpose CPU node group
  # ------------------------------------------------------------
  - title: general-purpose-ng
    minSize: 3
    desiredCapacity: 3
    maxSize: 6
    capacityType: ON_DEMAND
    privateNetworking: true
    launchTemplate:
      id: lt-xxxxxxxxxxxxxxxxx
      model: "1"
    labels:
      workload-type: general-purpose
    tags:
      Title: eks-general-purpose
      NodeGroup: general-purpose-ng
      Workload: general-purpose
    updateConfig:
      maxUnavailable: 1

  # ------------------------------------------------------------
  # GPU / accelerator node group
  # ------------------------------------------------------------
  - title: accelerator-ng
    minSize: 3
    desiredCapacity: 3
    maxSize: 6
    capacityType: ON_DEMAND
    privateNetworking: true
    launchTemplate:
      id: lt-yyyyyyyyyyyyyyyyy
      model: "1"
    labels:
      workload-type: accelerator
      accelerator: nvidia-b300
    taints:
      - key: nvidia.com/gpu
        worth: "true"
        impact: NoSchedule
    tags:
      Title: eks-accelerator
      NodeGroup: accelerator-ng
      Workload: accelerator
    updateConfig:
      maxUnavailable: 1

The listed Ampere, Hopper, and Blackwell occasion varieties (p4d.24xlarge, p4de.24xlarge, p5.48xlarge, p5e.48xlarge, p6-b200.48xlarge, and p6-b300.48xlarge) help large-scale RL coaching and can be utilized for both rollout or coaching workloads.

Determine 6: Consultant EKS deployment topology exhibiting GPU node teams (P5/P6) for
coaching and rollout, Spot-backed capability for inference employees, and CPU node teams for
orchestration providers.

Establishing EFA drivers and plugins

To optimize asynchronous RL on Amazon EKS, speaking nodes have to be in the identical Availability Zone (AZ) and the EFA Kubernetes system plugin have to be arrange. With EFA, you get quick and environment friendly mannequin updates, high-performance mannequin coaching, and low-latency communication. To use the EKS EFA system plugin, you should utilize kubectl to use it on to the cluster.

apiVersion: apps/v1
type: DaemonSet
metadata:
  title: aws-efa-k8s-device-plugin
  namespace: kube-system
  labels:
    app: aws-efa-k8s-device-plugin
spec:
  selector:
    matchLabels:
      app: aws-efa-k8s-device-plugin
  updateStrategy:
    kind: RollingUpdate
  template:
    metadata:
      labels:
        app: aws-efa-k8s-device-plugin
    spec:
      # Solely set up on the accelerator node group
      nodeSelector:
        eks.amazonaws.com/nodegroup: accelerator-ng
      priorityClassName: system-node-critical
      hostNetwork: true
      automountServiceAccountToken: false
      tolerations:
        - operator: Exists
      containers:
        - title: aws-efa-k8s-device-plugin
          # Present AWS EFA K8s system plugin
          picture: 602401143452.dkr.ecr.us-west-2.amazonaws.com/eks/aws-efa-k8s-device-plugin:v0.5.20
          securityContext:
            privileged: true
            allowPrivilegeEscalation: true
            runAsUser: 0
            runAsNonRoot: false
          assets:
            requests:
              cpu: 10m
              reminiscence: 20Mi
          volumeMounts:
            - title: device-plugin
              mountPath: /var/lib/kubelet/device-plugins
            - title: infiniband
              mountPath: /dev/infiniband
      volumes:
        - title: device-plugin
          hostPath:
            path: /var/lib/kubelet/device-plugins
        - title: infiniband
          hostPath:
            path: /dev/infiniband

For extra directions and efficiency testing, see the EKS EFA setup directions.

Launching DeepEP

The reference implementation combines the CUDA, PyTorch, communication, and inference parts required to help the optimizations described within the earlier sections. The next desk lists the benchmark atmosphere:

• CUDA: 13.0.
• PyTorch: 2.12+cu130.
• NCCL: 2.31.
• EFA: 1.49.
• DeepEP: 2.0.
• SGLang: 0.5.17.
• Miles: 0.1.0.

Holding these parts aligned is crucial. GPU kernels, collective communication libraries, and the underlying EFA transport every contribute to end-to-end efficiency. Utilizing the 763104351884.dkr.ecr.<area>.amazonaws.com/sglang:0.5.17-gpu-py312-cu130-ubuntu24.04-ec2 picture (accessible within the SGLang Deep Learning Containers catalog) is a sensible place to begin. The next Dockerfile can be utilized to arrange the coaching stack.

# syntax=docker/dockerfile:1.7

# Reference base from the AWS Deep Studying Containers SGLang catalog.
# Override AWS_REGION at construct time if the picture is pulled from one other area.
ARG AWS_REGION=us-west-2
ARG SGLANG_DLC_TAG=0.5.17-gpu-py312-cu130-ubuntu24.04-ec2
FROM 763104351884.dkr.ecr.${AWS_REGION}.amazonaws.com/sglang:${SGLANG_DLC_TAG}

ARG BUILD_JOBS=16
ARG NVCC_THREADS=1

# Reference implementation model matrix.
ARG PYTORCH_VERSION=2.12.1
ARG EFA_INSTALLER_VERSION=1.49.0
ARG NCCL_VERSION=2.31.2
ARG NVSHMEM_VERSION=3.7.2
ARG SGLANG_VERSION=0.5.17
ARG DEEPEP_VERSION=2.0.0
ARG DEEPEP_COMMIT=b306af06afd412c88e51e71802951606e40b7358
ARG MILES_VERSION=0.1.0
ARG MILES_REF=v0.1.0

# Present optimization pins retained from the unique Dockerfile.
ARG DEEPGEMM_COMMIT=731e7c7a97d269e4b9f482ea18d0e709a948f293

ENV PIP_RETRIES=20 
    PIP_DEFAULT_TIMEOUT=60 
    LIBFABRIC_HOME=/choose/amazon/efa

RUN apt-get replace && 
    apt-get set up -y --no-install-recommends 
    autoconf automake build-essential ca-certificates cmake curl git 
    libtool ninja-build pkg-config patch

# PyTorch 2.12.1 constructed for CUDA 13.0. That is put in explicitly relatively than
# counting on no matter PyTorch minor model occurred to ship within the base picture.
RUN pip set up --no-cache-dir --force-reinstall 
    --index-url https://obtain.pytorch.org/whl/cu130 
    "torch==${PYTORCH_VERSION}"

# Set up the EFA 1.49 userspace stack and its NGC-compatible OFI NCCL plugin.
# The EFA kernel module is provided by the EKS/EC2 host, not by this container.
RUN apt-get take away -y libnccl-ofi && 
    rm -f /and many others/ld.so.conf.d/aws-ofi-nccl.conf && 
    curl -fsSL "https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz" -o /tmp/aws-efa-installer.tar.gz && 
    mkdir -p /tmp/aws-efa-installer && 
    tar -xzf /tmp/aws-efa-installer.tar.gz -C /tmp/aws-efa-installer --strip-components=1 && 
    cd /tmp/aws-efa-installer && 
    ./efa_installer.sh -y --skip-kmod --no-verify --build-ngc --skip-mpi

# DeepEP V2 makes use of NCCL Gin. Preserve NCCL on the requested 2.31 line and set up it
# earlier than compiling DeepEP so build-time headers and runtime libraries agree.
RUN pip set up --no-cache-dir --force-reinstall --no-deps 
    "nvidia-nccl-cu13==${NCCL_VERSION}" 
    "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}"

ENV DEEPEP_REPO=https://github.com/deepseek-ai/DeepEP.git 
    DEEPEP_SRC_DIR=/choose/amazon/DeepEP 
    TORCH_CUDA_ARCH_LIST="9.0 10.0 10.3 12.0" 
    CUDASTDCXX_INCLUDE=/usr/native/lib/python3.12/dist-packages/flashinfer/knowledge/cccl/libcudacxx/embody 
    DEEPGEMM_REPO=https://github.com/sgl-project/DeepGEMM.git 
    DEEPGEMM_SRC_DIR=/choose/amazon/DeepGEMM 
    TVM_FFI_CUDA_ARCH_LIST="9.0 10.0 10.3 12.0" 
    DG_FORCE_BUILD=1

ENV CPATH="${CUDASTDCXX_INCLUDE}" 
    CPLUS_INCLUDE_PATH="${CUDASTDCXX_INCLUDE}"

# Construct DeepEP from the pinned public-release commit.
RUN export MAX_JOBS="${BUILD_JOBS}" CMAKE_BUILD_PARALLEL_LEVEL="${BUILD_JOBS}" && 
    mkdir -p /choose/amazon && 
    rm -rf "${DEEPEP_SRC_DIR}" && 
    git clone --filter=blob:none --no-checkout "${DEEPEP_REPO}" "${DEEPEP_SRC_DIR}" && 
    cd "${DEEPEP_SRC_DIR}" && 
    git fetch --depth 1 --no-tags origin "${DEEPEP_COMMIT}" && 
    git checkout FETCH_HEAD && 
    pip set up --no-cache-dir wheel && 
    python setup.py bdist_wheel && 
    pip set up --no-cache-dir --force-reinstall --no-deps dist/*.whl

# Preserve the prevailing DeepGEMM optimization layer.
RUN export MAX_JOBS="${BUILD_JOBS}" CMAKE_BUILD_PARALLEL_LEVEL="${BUILD_JOBS}" && 
    mkdir -p /choose/amazon && 
    rm -rf "${DEEPGEMM_SRC_DIR}" && 
    git clone --filter=blob:none --no-checkout "${DEEPGEMM_REPO}" "${DEEPGEMM_SRC_DIR}" && 
    cd "${DEEPGEMM_SRC_DIR}" && 
    git fetch --depth 1 --no-tags origin "${DEEPGEMM_COMMIT}" && 
    git checkout FETCH_HEAD && 
    pip set up --no-cache-dir --upgrade "apache-tvm-ffi==0.1.11" && 
    ./build_sgl_deep_gemm.sh && 
    pip set up --force-reinstall --no-deps --no-cache-dir dist/sgl_deep_gemm*.whl

# Miles 0.1.0 is put in with out dependency decision so it can't change
# the benchmark-critical PyTorch/NCCL/SGLang variations above.
ENV MILES_REPO=https://github.com/radixark/miles.git 
    MILES_SRC_DIR=/choose/amazon/miles

RUN rm -rf "${MILES_SRC_DIR}" && 
    git clone --branch "${MILES_REF}" --depth 1 "${MILES_REPO}" "${MILES_SRC_DIR}" && 
    pip set up --no-cache-dir --no-deps -e "${MILES_SRC_DIR}"

# Re-assert NCCL after each supply/package deal set up, then validate the entire
# reference matrix declared within the documentation paragraph.
RUN pip set up --no-cache-dir --force-reinstall --no-deps 
    "nvidia-nccl-cu13==${NCCL_VERSION}" 
    "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}"

RUN grep -Fx "# EFA installer model: ${EFA_INSTALLER_VERSION}" /choose/amazon/efa_installed_packages && 
    check -f /choose/amazon/efa/lib/libfabric.so.1 && 
    check -f /choose/amazon/ofi-nccl/lib/libnccl-net-ofi.so

Launch a job

With TorchX, you possibly can submit RL workloads to Amazon EKS whereas retaining utility configuration separate from infrastructure configuration. TorchX interprets job necessities into Kubernetes assets, together with compute requests, storage mounts, and runtime configuration. This separation makes it simpler to differ mannequin and coaching configurations with out coupling them to cluster provisioning.

The next command makes use of TorchX to submit an RL coaching job to your Amazon EKS cluster. The reference train_rl.py script, accessible within the Miles repository, orchestrates the rollout-training loop described earlier, launching rollout employees and policy-training employees as separate Kubernetes pods. The deepseek_v3_moe configuration specifies the MoE mannequin structure, together with the variety of specialists and parallelism technique. TorchX retains utility configuration separate from the underlying Kubernetes infrastructure, so you possibly can regulate mannequin parameters, employee counts, and parallelism settings with out modifying cluster-level assets.

A RL job will be launched with the next command:

torchx run -s kubernetes -cfg queue=default 
    dist.ddp -j 4x8 
    --script train_rl.py 
    --model_conf deepseek_v3_moe 
    --num_rollout_workers 32 
    --checkpoint_dir s3://my-bucket/checkpoints

Clear up

To keep away from ongoing expenses, delete the assets you created throughout this walkthrough. First, delete any operating coaching jobs and their related pods. Subsequent, take away the EKS managed node teams for coverage coaching, rollout era, and atmosphere employees. Then, delete the EKS cluster itself. Lastly, confirm that any related Amazon EC2 Spot capability, Elastic Community Interfaces for EFA, and Amazon S3 buckets used for checkpoints are additionally eliminated or emptied as wanted.

# Delete coaching jobs
kubectl delete jobs --all -n rl-training

# Delete node teams
eksctl delete nodegroup --cluster my-eks-cluster --name general-purpose-ng
eksctl delete nodegroup --cluster my-eks-cluster --name accelerator-ng

# Delete the cluster
eksctl delete cluster --name my-eks-cluster

For extra details about managing Amazon EKS assets, see Deleting an Amazon EKS cluster within the Amazon EKS documentation.

Conclusion and key takeaways

Scaling RLHF and GRPO for MoE fashions calls for infrastructure that balances elastic rollout era with tightly coupled coverage coaching. In our inner workloads, the structure improved combination rollout throughput by 40 % (see Determine 5), whereas additionally decreasing end-to-end coverage iteration time and scaling to roughly a thousand accelerators. This comparability makes use of two configurations on the identical 48 P5en situations (16 devoted to coaching, 32 to inference) operating the identical super-sparse MoE mannequin. The baseline used a Slime-based stack (CUDA 12.9, PyTorch 2.9.1, NCCL 2.27, SGLang 0.5.9, Slime 0.2.4) with out EFA-accelerated skilled parallelism. The improved configuration used the DeepEP-over-EFA stack described on this publish (CUDA 13.0, PyTorch 2.12, NCCL 2.31, EFA 1.49, DeepEP 2.0, SGLang 0.5.17, Miles 0.1.0). See Determine 5 for the corresponding throughput comparability.

This publish demonstrates find out how to construct and optimize a scalable structure that mixes Amazon EKS, EFA, and Amazon S3 to supply a versatile and cost-effective path for large-scale distributed RL workloads.

Key takeaways

• Scale parts independently: The heterogeneous nature of RL workloads means totally different subsystems can scale independently. Price optimizations reminiscent of Spot Situations can subsequently be utilized selectively.
• Steadiness useful resource constraints: RL programs place simultaneous calls for on compute, reminiscence, and bandwidth. Finish-to-end features require co-design throughout the stack, combining optimizations from transport and communication primitives to mannequin design.

AWS, together with open supply parts, can assist you run large-scale MoE RL workloads with sturdy efficiency and price effectivity. To get began, evaluate the structure patterns described on this publish and adapt them to your personal MoE coaching workloads on Amazon EKS with EFA. For extra info, see the Amazon EKS documentation, the EFA setup information for EKS, and the DeepEP repository on GitHub. In case you have questions or need assist adapting this structure to your workloads, contact your AWS account crew or publish within the AWS HPC forum.

In regards to the authors