大规模 RLHF 与 GRPO 训练中，GPU 不只是在执行矩阵计算。对于混合专家模型（Mixture-of-Experts，MoE），每个 token 还要经过路由、跨 GPU 发送到对应专家，再汇总结果。当 rollout 扩展到多节点后，这类 all-to-all 通信很容易成为真正的吞吐瓶颈。
一种有效的扩展方式，是把 Amazon EKS 的编排能力、Elastic Fabric Adapter（EFA）的低延迟网络、DeepEP 的专家并行通信优化，以及 Amazon S3 的持久化能力组合起来。来源案例显示，这套架构将大规模强化学习的 aggregate rollout throughput 提高了 40%。这个数字应被理解为特定模型、实例、批量和网络配置下的实测结果，而不是迁移到任意集群后都能自动获得的固定收益。
为什么 MoE rollout 会被网络卡住
普通稠密模型的数据并行，通信主要集中在梯度同步。MoE 则多出了一条高频通信路径：路由器为 token 选择专家后，token 表示需要在专家并行组内重新分发。
一次典型前向过程可以简化为：
- rollout worker 接收提示词并执行生成；
- MoE 路由器为每个 token 选择一个或多个专家；
- token 通过 all-to-all 发往持有对应专家的 GPU；
- 专家完成计算，结果再次跨 GPU 汇总；
- 生成的序列、奖励和训练样本进入后续 RLHF 或 GRPO 阶段。
当专家分布在多个节点时，问题就不再只是显存或算力不足。网络延迟、有效带宽、消息大小、专家负载偏斜和通信计算重叠程度，都会直接影响每秒生成的 token 数。
EFA 用于改善节点之间的高性能网络通信；DeepEP 面向 MoE 专家并行的数据交换，重点优化 dispatch 与 combine 一类通信操作。两者解决的是同一条关键路径的不同层次：EFA 提供底层传输能力，DeepEP 优化上层专家通信模式。
把控制面、计算面和存储面拆开
在 EKS 上部署这类系统时，可以把整体架构分成三个平面。
计算与通信平面
GPU 节点组成 rollout worker 集群。Pod 申请 GPU 和 EFA 设备，并尽量固定在支持目标网络能力的节点组上。DeepEP 集成在推理或训练运行时中，使 MoE token 在专家并行组内高效交换。
这里最重要的不是简单增加副本，而是保证以下条件同时成立：
- GPU 实例类型和驱动栈满足训练框架要求；
- EFA 设备插件已安装，Pod 能看到 EFA 资源；
- 节点安全组允许参与作业的实例互相通信；
- Pod 的 GPU、EFA 和 CPU 拓扑没有形成明显的跨 NUMA 绕行；
- 专家并行组与节点布局相匹配；
- rollout worker 使用一致的模型权重版本。
编排与弹性平面
EKS 负责调度 rollout、奖励模型、训练控制器和数据处理任务。不同角色最好使用独立的 node group、标签和污点，避免奖励计算或数据预处理抢占高价值的 EFA GPU 节点。
弹性扩缩容也需要谨慎。单个 MoE 作业往往要求完整的 gang 才能启动；只增加一个 worker，未必能组成新的专家并行组。生产环境可以结合队列或 gang scheduling，让调度器按完整拓扑分配资源，而不是让部分 Pod 长时间占着 GPU 等待同伴。
持久化平面
Amazon S3 适合保存模型检查点、rollout 结果、训练数据和评估产物。它不应被放进 token dispatch 的同步热路径中：GPU 间的高频通信交给 EFA，而需要跨任务保存和恢复的数据异步写入 S3。
建议为对象路径加入运行 ID、策略版本和全局步数，例如：
s3://my-rl-bucket/runs/grpo-2025-03-08/
├── checkpoints/step-000100/
├── rollouts/policy-v17/part-00001.jsonl
├── rewards/policy-v17/part-00001.parquet
└── metrics/
这种布局能避免 rollout worker 读取到尚未完整发布的新权重，也方便失败后按明确版本恢复。
一个可改造的 EKS Rollout Worker 清单
下面是一个示意性 Kubernetes 清单。它假设集群已经具备 NVIDIA GPU 支持，EFA device plugin 暴露了 vpc.amazonaws.com/efa 资源，并且镜像内已经安装兼容版本的训练框架、libfabric、EFA 用户态组件和 DeepEP。
运行前需要替换镜像、S3 bucket、节点标签以及分布式作业参数。不同 GPU 实例可提供的 EFA 设备数量并不相同，应先通过 kubectl describe node 检查容量。
apiVersion: v1
kind: Service
metadata:
name: moe-rollout-headless
spec:
clusterIP: None
selector:
app: moe-rollout
ports:
- name: rendezvous
port: 29500
---
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: moe-rollout
spec:
serviceName: moe-rollout-headless
replicas: 2
selector:
matchLabels:
app: moe-rollout
template:
metadata:
labels:
app: moe-rollout
spec:
serviceAccountName: rl-s3-access
terminationGracePeriodSeconds: 120
nodeSelector:
workload: efa-gpu
topologySpreadConstraints:
- maxSkew: 1
topologyKey: kubernetes.io/hostname
whenUnsatisfiable: DoNotSchedule
labelSelector:
matchLabels:
app: moe-rollout
containers:
- name: worker
image: 123456789012.dkr.ecr.us-east-1.amazonaws.com/moe-rl:latest
securityContext:
capabilities:
add: ["IPC_LOCK"]
env:
- name: MASTER_ADDR
value: moe-rollout-0.moe-rollout-headless
- name: MASTER_PORT
value: "29500"
- name: WORLD_SIZE
value: "16"
- name: GPUS_PER_NODE
value: "8"
- name: FI_PROVIDER
value: efa
- name: FI_EFA_USE_DEVICE_RDMA
value: "1"
- name: CHECKPOINT_URI
value: s3://my-rl-bucket/runs/grpo-demo/checkpoints/current/
command: ["/bin/bash", "-lc"]
args:
- |
export NODE_RANK="${HOSTNAME##*-}"
exec python -m my_rl.rollout \
--master-addr "$MASTER_ADDR" \
--master-port "$MASTER_PORT" \
--world-size "$WORLD_SIZE" \
--node-rank "$NODE_RANK" \
--gpus-per-node "$GPUS_PER_NODE" \
--checkpoint "$CHECKPOINT_URI" \
--enable-deepep
resources:
requests:
cpu: "32"
memory: 220Gi
nvidia.com/gpu: "8"
vpc.amazonaws.com/efa: "1"
limits:
cpu: "32"
memory: 220Gi
nvidia.com/gpu: "8"
vpc.amazonaws.com/efa: "1"
volumeMounts:
- name: dshm
mountPath: /dev/shm
volumes:
- name: dshm
emptyDir:
medium: Memory
sizeLimit: 64Gi
应用前可以先确认目标节点确实公布了 GPU 与 EFA 资源：
kubectl get nodes -l workload=efa-gpu
kubectl get nodes -l workload=efa-gpu \
-o custom-columns='NAME:.metadata.name,GPU:.status.allocatable.nvidia\.com/gpu,EFA:.status.allocatable.vpc\.amazonaws\.com/efa'
kubectl apply -f moe-rollout.yaml
kubectl rollout status statefulset/moe-rollout --timeout=15m
kubectl get pods -l app=moe-rollout -o wide
my_rl.rollout 和 --enable-deepep 是示例接口，需要映射到实际使用的 RL 框架。若框架通过 torchrun 启动，也可以让每个 StatefulSet Pod 运行一个节点级 torchrun，并将 ordinal 转换为 node_rank。
不要只看 GPU 利用率：按通信阶段测量
40% 的 aggregate rollout throughput 提升是否能够复现，取决于基线和测量方式。仅比较任务完成时间往往不足以解释优化来自哪里。更可靠的基准测试应固定以下变量：
- 模型、精度、量化策略和专家数量；
- 提示词长度与生成长度分布；
- 每个 token 选择的专家数量；
- GPU 数、每节点 GPU 数与专家并行规模；
- batch size、并发请求数和采样参数；
- rollout 输出落盘与上传策略。
至少记录这些指标：
aggregate_output_tokens_per_second
per_worker_output_tokens_per_second
p50_and_p99_rollout_latency
moe_dispatch_time_ms
moe_combine_time_ms
expert_load_imbalance
network_bytes_per_second
checkpoint_or_weight_load_time
可以先做三组对照：默认通信实现、启用 EFA 的默认实现、启用 EFA 与 DeepEP 的实现。每组预热后运行相同输入，并重复多次。这样才能区分收益来自高速网络、DeepEP 通信优化，还是单纯来自不同的批处理设置。
还要观察专家负载偏斜。如果少数专家持续接收大量 token，再快的网络也无法消除计算热点。此时需要检查路由器的负载均衡损失、容量因子以及 token 丢弃策略，而不是继续堆叠带宽。
上线前的工程检查单
这套架构适合通信占比较高、需要跨节点扩展的 MoE rollout。对于单节点可容纳的小模型，EFA 与 DeepEP 带来的复杂度可能超过收益。正式采用前可以逐项检查：
- 用 profiler 证明 dispatch/combine 或跨节点通信确实处于关键路径；
- 确认所选 GPU 实例、可用区、AMI、驱动与 EFA 软件栈兼容；
- 验证 Pod 能申请并使用
vpc.amazonaws.com/efa； - 通过 IAM Roles for Service Accounts 或 EKS Pod Identity 授予最小化 S3 权限，避免在镜像中保存长期密钥；
- 对模型权重采用版本化发布，防止 worker 混用不同策略版本；
- 为节点中断、Pod 重启和部分 rollout 丢失设计幂等恢复；
- 同时计算每百万 token 的成本，而不是只追求峰值吞吐；
- 用真实提示词长度和专家路由分布完成压测。
EFA 和 DeepEP 的价值并不是让所有 MoE 训练自动加速，而是把最昂贵的跨节点专家通信变成可优化、可观测的工程路径。先定位 rollout 的通信瓶颈，再用 EKS 管理资源边界、用 EFA 改善传输、用 DeepEP 优化 token 交换，并让 S3 承担异步持久化，才更有机会把案例中的 40% 提升转化为自身工作负载上的稳定收益。