DeepSeek-V4 Flash 强化学习训练登陆 AMD Instinct MI355X GPU,由 Miles 框架支持
News Bringing DeepSeek-V4 Flash RL Training to AMD Instinct MI355X GPUs with Miles DeepSeek-V4 RL is now supported in Miles on AMD Instinct™ MI355X GPUs with ROCm™! RL requires SGLang rollout and Megatron training to implement the same policy closely enough that token probabilities ... AMD & Miles Team
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 Flash 的 RL 训练完整搬到 AMD MI355X 上,不是画饼,而是给了实测数据和可跑的配置。想用非 NVIDIA 卡做 RL 训练的,可以照着踩坑。
DeepSeek-V4 RL 现已在 AMD Instinct™ MI355X GPU 上通过 ROCm™ 在 Miles 中获得支持!RL 需要 SGLang rollout 与 Megatron 训练,使二者实现的策略足够接近,从而让 token 概率保持一致,即便 Miles 会反复将更新后的权重传回正在运行的 rollout 引擎。
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 worker。
当前的 FP8 路径使用 FP8 Hugging Face checkpoint 进行 rollout,并使用 BF16 Megatron torch_dist checkpoint 进行 actor 训练。因此,这两个引擎以不同的执行格式表示同一个策略,使得转换、重新打分和在线更新行为成为正确性边界的一部分。

图 2. 提示词流向 SGLang rollout;轨迹和对数概率经由 Miles 流向 Megatron actor;更新后的权重在下次 rollout 之前返回 SGLang。
挑战 1:弥合训练与 rollout 之间的对数概率差距
RL 训练依赖于 SGLang 和 Megatron 对同一批生成的 token 赋予相似的概率。我们构建了一套 token 完全一致的对比工作流:SGLang 生成一次序列,两个引擎对相同的 token 打分,token 级别的差异能在代价高昂的多节点验证之前揭示模型层面的不匹配。
这一对比发现了 DeepSeek-V4 特有路径中的两处差异:早期哈希路由 MoE 和 mHC 残差混合。我们将 Megatron 的哈希路由行为与 SGLang 对齐,并修正了 Megatron 侧的 mHC 后混合,使两个引擎保持相同的模型语义。这些针对性的改动让 rollout 与训练在数值上更加一致。
挑战 2:在线更新期间保持量化语义
在 RL 中,rollout 服务器会反复接收更新后的策略权重,而无需重启。对于量化模型而言,仅仅成功传输是不够的:打包权重、scale tensor 以及依赖量化的运行时状态都必须保持其原本的含义。
对于 FP4 和 E8M0 tensor,AMD 让更新路径具备数据类型感知能力,避免了更新后 tensor 被错误解读而产生无效生成。对于 FP8,Miles 已经定义了更新后的生命周期;AMD 在 ROCm 栈中补充了缺失的 SGLang 接口,使 Miles 能够在 rollout 恢复之前执行所需的量化处理。
关键教训很简单:在线更新必须恢复模型的量化状态,而不仅仅是复制其字节。
挑战 3:在 ROCm 上实现稳定的多节点并行策略
在多个 AMD Instinct MI355X 节点上扩展 DeepSeek-V4 Flash RL 时,暴露出两个相互耦合的启动问题:选择一个能容纳 2840 亿参数 MoE、4K 上下文窗口的模型并行策略,以及保持多节点集合通信的稳定。这两者是相互关联的。更重的张量并行会降低每 GPU 显存占用,但会增加集合通信流量,而一些早期的多节点配置在 RCCL 集合通信中卡死——例如某个张量并行的 all-reduce 或专家 all-to-all 未能完成,被通信看门狗捕获,从而中止了运行。
我们最终收敛到一个既显存可行又稳定的布局:张量并行 1、流水线并行 4、专家并行 4,跨四个八 GPU 节点,并配合激活重计算、优化器状态卸载到主机内存,以及有界的每 GPU token 预算。将并行重心从张量并行 all-reduce 转向流水线并行和专家并行,再加上调优后的 RCCL 传输设置,使运行能够端到端推进超过 100 个优化器步而没有集合通信卡死。建立这个稳定运行点是后续更长验证的前提。
在 AMD Instinct MI355X GPU 上的四节点验证
我们在四个八卡 AMD Instinct MI355X 节点上验证了 FP8 路径:两个用于 SGLang rollout,两个用于 Megatron actor 训练。Miles 在一个长上下文数学工作负载(DAPO-Math-17K,4K 上下文)上协调了 GRPO 风格的训练、奖励收集以及反复的在线权重更新,采用张量并行 1 / 流水线并行 4 / 专家并行 4 的模型并行配置。每十步,我们在 AIME-2024 上以每题八个样本进行离线评估。rollout 模型使用 FP8,而 actor 以 BF16 训练。
正确性与在线奖励
一个关键的的正确性检查是:rollout 与训练是否对同一批生成的 token 赋予相近的概率。在记录的各个步骤中,对数概率的平均绝对差约为 ~0.09。如图 3 所示,该差异在超过 100 步以及反复的权重更新过程中始终保持在有界范围内,没有持续上升的漂移,也没有在更新后出现急剧增大。这是一个令人鼓舞的启动阶段结果,而非最终阈值。

图 3. 前 100 个训练步骤中训练与 rollout 的绝对对数概率差。
除了对数概率一致性保持在有界范围内,在线奖励也随训练而提升。在一次延长运行中,在线原始奖励呈现出明显的上升趋势,而非保持平稳:其均值从运行的前三分之一到后三分之一有所上升(图 4)。这表明 actor 在持续的 GRPO 训练和反复的在线权重更新下正在改进,而不仅仅是维持奖励。

图 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 步评估一次,每个问题八个样本)。单次准确率(pass@1)和覆盖率(pass@8)均呈上升趋势。
我们的经验总结
离线基准评估才是诚实的信号。训练工作负载上的在线通过率会受到动态采样和过度采样的偏差影响;每隔十步进行一次留出的 AIME-2024 评估,才能给出可信的模型质量衡量。我们建议每次 RL 运行都搭配周期性的离线评估。
强化学习同时提升了准确率和覆盖率。在 100 步中,离线 AIME pass@1 从 0.39 升至 0.49,pass@8 从 0.53 升至 0.67。pass@1 与 pass@k 的同时上升表明策略获得了能力提升,而不仅仅是在已有解附近变得更加锐利。
跨引擎一致性在长时间运行中保持稳定。训练-推演对数概率差在 100+ 步和反复的权重更新中始终保持有界,证实了启动对齐在初始验证窗口之外依然良好成立。
响应截断是绝对评测分数的主要天花板。约 55-60% 的 AIME 响应触及了 4,096 token 的生成上限;提高评测响应预算,是提升绝对准确率最具杠杆效应的手段。
看趋势,而非单步。每步的在线奖励和每次评测的通过率都充满噪声(在线奖励逐步波动范围为 0.36-0.77);移动平均和周期性离线评测才是可靠的进展信号。
前路展望
- 启用 FP8 actor 训练。将 Megatron actor 从 BF16 扩展到 FP8,并评估其对训练-推演对齐和训练质量的影响。
- 性能剖析与差距分析。识别端到端 RL 流水线中最大的性能差距,并优先处理影响最大的瓶颈。
- 性能优化。 提升 rollout 吞吐量、训练效率,以及 rollout 与 actor 执行之间的重叠。
- 扩展。 在更大规模集群上评估吞吐量、效率和正确性,然后调优分布式执行以保持扩展效率。
启动命令
这些实验使用一个外部 Ray 集群,并通过单条启动命令在 ROCm 容器内运行。
Docker 镜像:rlsys/miles:rocm7.2-mi35x-dsv4
关键设置如下所示;完整的 RCCL 传输和 flight-recorder 环境变量在启动脚本中设置。
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 rollout 与 Megatron 训练之间对齐模型行为,并在在线权重更新期间保留量化状态,在 ROCm 上启用了端到端的 DeepSeek-V4 Flash RL 训练工作流。
在一个四节点 AMD Instinct MI355X 验证中,Miles 在超过 100 个优化器步骤中协调了 FP8 rollout、BF16 actor 训练、奖励收集以及反复的策略更新。训练与 rollout 的对数概率差异始终保持在有界范围内,在线奖励有所提升,离线 AIME-2024 基准分数从 pass@1 0.39 升至 0.49(pass@8 0.53 升至 0.67)。接下来,我们将启用 FP8 actor 训练,推进性能优化,并在更大规模上评估该工作流。
致谢
本工作建立在 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、Yueming Yuan、Mao Cheng、Zhichen Zeng、Shi Don、Yanbin Jiang、Ying Sheng,以及 miles 团队
附录
ROCm 运行时与内核路径映射
所报告的运行使用了以下路径来加载那些最直接影响跨引擎一致性和在线更新的模型组件。
| 模型组件 | 选定的运行时路径 | 为何重要 |
|---|---|---|
| mHC 残差混合 | Rollout: AITER mHC pre/post。 Training: TileKernels pre;显式 PyTorch/HIP post-mix。 | 使残差流映射在各引擎之间保持可见且可比较。 |
| 混合注意力 | Rollout: ROCm 融合 MLA decode;Triton 滑动窗口准备;融合 compressor 与 paged-compressor 路径;TileLang indexer。 Training: Miles DeepSeek-V4 attention,BF16。 | 通过不同的执行栈覆盖本地注意力与压缩注意力。 |
| MoE 与路由 | Rollout: Triton FP8 MoE;融合 hash top-k。 Training: Megatron MoE 与 router 路径。 | 要求两侧具备相同的确定性 hash-routing 语义。 |
| 在线权重更新 | Rollout: 分布式权重更新加 SGLang post_process_weights。 Training: 从 BF16 actor 进行 Miles 广播更新。 | 在生成恢复之前,重建依赖量化的运行时状态。 |
表 1. 所报告的 ROCm 配置中使用的运行时路径。
启动器使所选后端显式化,而 Docker 作用域内的补丁则移除了依赖的 Megatron 和 Transformer Engine 路径中残留的仅限 CUDA 的假设。这使所测试的配置保持可复现且可审查。
来源:LMSYS:Blog(Chatbot Arena 团队) · lmsys.org