DeepSeek-V4 Flash 强化学习训练登陆 AMD Instinct MI355X GPU,由 Miles 框架支持
阅读原文· lmsys.org把 DeepSeek-V4 Flash 的 RL 训练完整搬到 AMD MI355X 上,不是画饼,而是给了实测数据和可跑的配置。想用非 NVIDIA 卡做 RL 训练的,可以照着踩坑。
DeepSeek-V4 Flash 的强化学习训练现已在 AMD Instinct MI355X GPU 上通过 Miles 框架获得支持,基于 ROCm 软件栈运行。该 2840 亿参数 MoE 模型(每 token 激活 130 亿参数)需 SGLang 进行 rollout 生成、Megatron 进行策略更新,Miles 负责异步循环与权重同步。团队解决了 SGLang 与 Megatron 间模型对齐、量化状态在线更新及多节点并行稳定性三大挑战,最终在四个八 GPU 节点上完成端到端验证:超过 100 个优化器步骤中训练-rollout 对数概率差可控,在线奖励持续提升,离线 AIME-2024 基准分数同步上涨。
DeepSeek-V4 RL 现已在搭载 ROCm™ 的 AMD Instinct™ MI355X GPU 上通过 Miles 获得支持!RL 需要借助 SGLang 进行 rollout 并使用 Megatron 进行训练,以便在 Miles 反复将更新后的权重传回实时 rollout 引擎时,仍能足够紧密地实现相同的策略,从而保持模型 token 概率的对齐。
DeepSeek-V4 Flash 通过混合压缩注意力、mHC 残差混合以及 MoE 路由,使得这一过程颇具挑战性。我们的适配工作使 SGLang 和 Megatron 之间的模型行为保持一致,在在线权重更新期间保留了量化状态,并建立了一个端到端的四节点工作流。我们通过一个有界的训练与 rollout 对数概率差异验证了准确性,并且在一次长时间运行中,离线 AIME-2024 基准测试分数持续上升,同时在线奖励也在不断改善。
关键要点
- DeepSeek-V4 Flash RL 现已在 AMD Instinct MI355X GPU 上通过 Miles 运行。我们解决了在 ROCm 上进行端到端执行所需的模型对齐和在线更新问题。
- 四节点验证已完成。成功完成了超过 100 个优化步骤的端到端运行,训练与 rollout 的对数概率差异有界,在线奖励持续改善,离线 AIME-2024 评估分数也在上升。
- 性能优化是下一步工作。未来的工作包括低精度训练、端到端优化,以及在更大集群上进行扩展。
DeepSeek-V4 Flash 一图概览
DeepSeek-V4 Flash 是一个拥有 2840 亿参数的 MoE 模型,每个 token 激活 130 亿参数。本工作中使用的配置包含 43 个解码器层、256 个路由专家(采用 top-6 选择)、四个 mHC 残差流,以及结合了 128 token 滑动窗口与压缩长上下文注意力的混合注意力机制。
C4 层从经过 4:1 压缩的 KV 序列中选择前 512 个条目,而 C128 层则密集关注经过 128:1 压缩的序列。与 mHC 和 MoE 路由一起,这些是 SGLang 和 Megatron 必须一致实现的主要架构特定路径。
图 1. 简化的 DeepSeek-V4 模块,展示了 mHC 残差混合、混合压缩注意力以及 top-6 MoE 路由。
Miles 技术栈
Miles 负责编排异步循环。SGLang 生成候选回复和 rollout 对数概率;Megatron 对相同序列进行评分,计算策略更新,并训练 actor 模型;随后 Miles 将更新后的权重传回正在运行的 SGLang 工作节点。
当前的 FP8 路径使用 FP8 格式的 Hugging Face 检查点进行 rollout,并使用 BF16 格式的 Megatron torch_dist 检查点进行 actor 训练。因此,两个引擎以不同的执行格式表示同一个策略,这使得格式转换、重新评分和在线更新行为成为正确性验证边界的一部分。
图 2. 提示词流向 SGLang 进行 rollout;轨迹和对数概率通过 Miles 流向 Megatron actor 模型;更新后的权重在下一轮 rollout 之前返回 SGLang。
挑战 1:缩小训练与 rollout 之间的对数概率差距
强化学习训练依赖于 SGLang 和 Megatron 对相同生成的 token 赋予相似的概率。我们构建了一个 token 级比对工作流:SGLang 生成一次序列,两个引擎对相同的 token 进行评分,token 级别的差异能在代价高昂的多节点验证之前揭示模型层面的不匹配。
这种比对识别出 DeepSeek-V4 特有两条路径上的差异:早期哈希路由 MoE 和 mHC 残差混合。我们使 Megatron 的哈希路由行为与 SGLang 对齐,并修正了 Megatron 侧的 mHC 后混合步骤,使两个引擎保持相同的模型语义。这些针对性的改动使 rollout 和训练在数值上更加一致。
挑战 2:在线更新期间保持量化语义
在强化学习中,rollout 服务器会反复接收更新后的策略权重而无需重启。对于量化模型,仅成功传输权重是不够的:打包后的权重、缩放张量和依赖于量化的运行时状态都必须保持其预期含义。
对于 FP4 和 E8M0 张量,AMD 使更新路径能够感知数据类型,从而防止更新后张量被错误解读导致生成无效内容。对于 FP8,Miles 已经定义了更新后的生命周期;AMD 在 ROCm 软件栈中补充了缺失的 SGLang 接口,使得 Miles 能在 rollout 恢复前执行所需的量化处理。
关键教训很简单:在线更新必须恢复模型的量化状态,而不仅仅是复制其字节。
挑战三:在 ROCm 上实现稳定的多节点并行策略
将 DeepSeek-V4 Flash RL 扩展到多个 AMD Instinct MI355X 节点时,出现了两个相互关联的启动问题:选择一种适合 2840 亿参数 MoE 且上下文窗口为 4K 的模型并行策略,以及保持多节点集体通信的稳定性。这两个问题相互关联。更重的张量并行会降低每 GPU 内存,但会增加集体通信流量,并且一些早期的多节点配置会在 RCCL 集体通信中停滞——例如,一个张量并行的 all-reduce 或专家 all-to-all 未能完成,被通信看门狗捕获,从而导致运行中断。
我们最终确定了一种既在内存上可行又稳定的布局:在四个八 GPU 节点上采用张量并行 1、流水线并行 4 和专家并行 4,并配合激活重计算、将优化器状态卸载到主机内存以及限制每 GPU token 预算。将并行策略从张量并行 all-reduce 转向流水线并行和专家并行,并结合调整后的 RCCL 传输设置,使得运行能够端到端地持续超过 100 个优化器步骤,而不会出现集体通信停滞。建立这个稳定的运行点是后续更长时间验证的先决条件。
在 AMD Instinct MI355X GPU 上进行四节点验证
我们在四个八 GPU 的 AMD Instinct MI355X 节点上验证了 FP8 路径:两个节点用于 SGLang 部署,两个节点用于 Megatron 演员模型训练。Miles 协调了基于 GRPO 风格的训练,使用长上下文数学工作负载(DAPO-Math-17K,上下文窗口 4K)、奖励收集以及重复的在线权重更新,采用了张量并行 1 / 流水线并行 4 / 专家并行 4 的模型并行配置。每十步,我们使用 AIME-2024 进行一次离线评估,每个问题采样八次。部署模型使用 FP8,而演员模型使用 BF16 进行训练。
正确性与在线奖励
一个关键的准确性检查是,推理阶段和训练阶段是否对相同的生成 token 赋予了相似的概率。在记录的步骤中,平均绝对对数概率差异约为 0.09。如图 3 所示,该差异在超过 100 个步骤和重复的权重更新中保持有界,没有出现持续的上行漂移或更新后的急剧增加。这是一个令人鼓舞的初步结果,而非最终阈值。
图 3. 前 100 个训练步骤中训练与推理的绝对对数概率差异。
除了有界的对数概率一致性外,在线奖励也随着训练过程得到改善。在一次扩展运行中,在线原始奖励呈现出明显的上升趋势,而非保持平稳:其均值从运行的前三分之一段上升到后三分之一段(图 4)。这表明,在持续的 GRPO 训练和重复的在线权重更新下,actor 模型正在改进,而不仅仅是维持奖励水平。
图 4. 100 个训练步骤中的在线原始奖励(每步值、移动平均线和线性拟合)。奖励在运行过程中增加,呈现正线性斜率。
AIME-2024 离线评估
在线通过率是在训练负载上测量的,并受到动态采样的影响,因此我们还每十步运行一次留出的离线基准测试 AIME-2024,每个问题采样八次。这是对模型质量的诚实衡量。在前 100 个步骤中,离线 AIME pass@1 从 0.39 提升至 0.49,pass@8 从 0.53 提升至 0.67,同时因 4,096 token 上限导致的响应截断率从 60% 下降至 55%。单次准确率和多样本覆盖率同步提升,表明在 GRPO 下获得了真正的能力增益,而不仅仅是结果的锐化。在仅有 30 道题的基准测试上,每次评估的值存在噪声,因此趋势而非任何单个数据点才是信号所在。
图 5. 前 100 个 RL 训练步骤中的离线 AIME-2024 pass@1/2/4/8(每 10 步评估一次,每个问题采样 8 次)。单次准确率(pass@1)和覆盖率(pass@8)均呈上升趋势。
我们的收获
离线基准评测是可靠的信号。训练任务上的在线通过率会受到动态采样和过采样的影响而产生偏差;每十步进行一次留出集 AIME-2024 评测,能为模型质量提供可信的衡量标准。我们建议在每次强化学习运行中配合定期的离线评测。
强化学习同时提升了准确率和覆盖率。在超过 100 步的训练中,离线 AIME pass@1 从 0.39 提升至 0.49,pass@8 从 0.53 提升至 0.67。pass@1 和 pass@k 的同时提升表明,策略获得了能力增长,而不仅仅是围绕已有解决方案变得更加精准。
跨引擎的一致性在长时间运行中保持稳定。训练-推理对数概率差异在超过 100 步和多次权重更新中始终保持在有限范围内,这证实了启动阶段的对齐效果在初始验证窗口之外依然有效。
响应截断是限制绝对评测分数的主要因素。大约 55-60% 的 AIME 响应达到了 4096 token 的生成上限;提高评测响应预算,是提升绝对准确率最高杠杆率的措施。
关注趋势,而非单一步骤。单步的在线奖励和单次评测的通过率存在噪声(在线奖励在步与步之间波动范围为 0.36-0.77);移动平均线和定期的离线评测才是可靠的进展信号。
未来方向
- 启用 FP8 演员模型训练。将 Megatron 演员模型从 BF16 扩展到 FP8,并评估其对训练-推理对齐和训练质量的影响。
- 性能分析与差距识别。识别端到端强化学习流程中最大的性能差距,并优先处理影响最大的瓶颈。
- 性能优化。提升推理吞吐量、训练效率,以及推理与演员模型执行之间的重叠程度。
- 规模扩展。在更大规模的集群上评估吞吐量、效率和正确性,然后调整分布式执行以保持扩展效率。
启动命令
这些实验使用一个外部 Ray 集群和一条启动命令,在 ROCm 容器内运行。
Docker 镜像:rlsys/miles:rocm7.2-mi35x-dsv4
基本设置如下所示;完整的 RCCL 传输和飞行记录器环境变量在启动脚本中设置。
ray start --head --node-ip-address=$HEAD_IP --port=6379 --num-gpus=8
ray start --address=$HEAD_IP:6379 --node-ip-address=$WORKER_IP --num-gpus=8
export MASTER_ADDR=xxx
export MILES_SCRIPT_EXTERNAL_RAY=1
export RAY_ADDRESS=xxx
export PYTHONUNBUFFERED=1
export NCCL_SOCKET_IFNAME=xxx
export GLOO_SOCKET_IFNAME=xxx
export TP_SOCKET_IFNAME=xxx
export NCCL_IB_HCA=xxx
export NCCL_IB_GID_INDEX=1
RUN_ID=dsv4-fp8-4node-2roll-tp1-pp4-ep4-$(date +%Y%m%d_%H%M%S)
LOG=/workspace/miles/${RUN_ID}.log
/opt/venv/bin/python3 scripts/amd/run_deepseek_v4.py train \
--run-id "${RUN_ID}" \
--mode normal \
--enable-eval \
--num-nodes 4 \
--actor-num-nodes 2 \
--rollout-num-nodes 2 \
--num-rollout 200 \
--num-steps-per-rollout 1 \
--rollout-batch-size 32 \
--n-samples-per-prompt 8 \
--context-length 16384 \
--rollout-max-response-len 4096 \
--max-tokens-per-gpu 8192 \
--sglang-max-running-requests 48 \
--sglang-max-total-tokens 524288 \
--tensor-model-parallel-size 1 \
--pipeline-model-parallel-size 4 \
--decoder-last-pipeline-num-layers 10 \
--context-parallel-size 1 \
--expert-model-parallel-size 4 \
--expert-tensor-parallel-size 1 \
--extra-args '--wandb-team xxx --use-tis' \
--extra-env-vars 'TORCH_NCCL_DUMP_ON_TIMEOUT=1 TORCH_NCCL_TRACE_BUFFER_SIZE=200000
TORCH_FR_BUFFER_SIZE=200000 TORCH_NCCL_DESYNC_DEBUG=1 TORCH_NCCL_ASYNC_ERROR_HANDLING=1
TORCH_NCCL_DEBUG_INFO_TEMP_FILE=/workspace/miles/nccl_fr_trace_ GPU_MAX_HW_QUEUES=2
NCCL_P2P_NET_CHUNKSIZE=262144' \
2>&1 | tee "${LOG}"
总结
我们通过在 SGLang 推理生成与 Megatron 训练之间对齐模型行为,并在在线权重更新过程中保留量化状态,在 ROCm 上实现了一套端到端的 DeepSeek-V4 Flash RL 训练工作流。
在一个四节点 AMD Instinct MI355X 验证环境中,Miles 协调了 FP8 推理生成、BF16 演员模型训练、奖励收集以及超过 100 个优化器步骤的重复策略更新。整个过程中,训练与推理生成的对数概率差保持有界,在线奖励得到提升,离线 AIME-2024 基准测试分数从 pass@1 0.39 提升至 0.49(pass@8 从 0.53 提升至 0.67)。接下来,我们将启用 FP8 演员模型训练,进行性能优化,并在更大规模上评估该工作流。
致谢
本工作建立在 SGLang 和 Miles 社区对 DeepSeek-V4 的支持之上。我们感谢与 AMD 合作的 Miles 团队,以及 Megatron、AITER、Triton、TileLang、Transformer Engine 和 ROCm 的贡献者,他们的软件构成了端到端的技术栈。
AMD 贡献者:Xinyu Kang, Liz Li, Yuankai Chen, Zhiyao Jiang, Kailesh Gogineni, Yao Fu, Wen Xie, Gowtham Ramesh, Cheng Yao, Xiaobo Chen, Shekhar Pandey, Sree Rohith Pulipaka, Wen Chen, Yuzhen Zhou, Xinyu Jiang, Hai Xiao, Andy Luo, Zhenyu Gu。
Miles 贡献者:Yusheng Su, Jiajun Li, Banghua Zhu 以及 miles 团队
附录
ROCm 运行时与内核路径映射
所报告的运行使用了以下路径来处理对跨引擎一致性和在线更新影响最直接的模型组件。
| 模型组件 | 选定的运行时路径 | 重要性说明 |
|---|---|---|
| mHC 残差混合 | 推理生成:AITER mHC 前/后处理。训练:TileKernels 前处理;显式 PyTorch/HIP 后混合。 | 保持残差流映射在各引擎间可见且可比较。 |
| 混合注意力 | 推理生成:ROCm 融合 MLA 解码;Triton 滑动窗口准备;融合压缩器与分页压缩器路径;TileLang 索引器。训练:Miles DeepSeek-V4 BF16 注意力机制。 | 覆盖了通过不同执行栈实现的局部注意力与压缩注意力。 |
| MoE 与路由 | 推理生成:Triton FP8 MoE;融合哈希 top-k。训练:Megatron MoE 与路由器路径。 | 要求两侧具有相同的确定性哈希路由语义。 |
| 在线权重更新 | 推出:分布式权重更新加上 SGLang 的 post_process_weights。训练:从 BF16 智能体进行 Miles 广播更新。 | 在生成恢复之前,重建依赖于量化(quantization)的运行时状态。 |
表 1. 所报告的 ROCm 配置中使用的运行时路径。
启动器使所选的后端显式化,而 Docker 作用域内的补丁则移除了依赖的 Megatron 和 Transformer Engine 路径中剩余的仅限 CUDA 的假设。这确保了所测试的配置是可复现且可审查的。