DeepSeek-V4 Flash 强化学习训练登陆 AMD Instinct MI355X GPU，由 Miles 框架支持

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 亿参数、每个 token 激活 130 亿参数的 MoE 模型。本工作中使用的配置包含 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 之间的对数概率差距

RL 训练依赖于 SGLang 和 Megatron 对相同生成的 token 分配相似的概率。我们构建了一个 token 级对比工作流：SGLang 生成一次序列，两个引擎对相同的 token 进行评分，token 级别的差异能在代价高昂的多节点验证之前揭示模型层面的不匹配。

这种对比识别出 DeepSeek-V4 特定路径中的两个差异：早期哈希路由 MoE 和 mHC 残差混合。我们使 Megatron 的哈希路由行为与 SGLang 对齐，并修正了 Megatron 侧的 mHC 后混合，使两个引擎保持相同的模型语义。这些有针对性的改动使 rollout 和训练在数值上更加接近。

挑战 2：在线更新期间保持量化语义

在 RL 中，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 协调了在长上下文数学任务（DAPO-Math-17K，4K 上下文窗口）上的 GRPO 风格训练、奖励收集以及重复的在线权重更新，使用了张量并行 1 / 流水线并行 4 / 专家并行 4 的模型并行配置。每十步，我们在 AIME-2024 上运行一次离线评估，每个问题采样八次。推理生成模型使用 FP8，而演员模型使用 BF16 进行训练。

正确性与在线奖励

一项关键的准确性校验是，推理生成（rollout）与训练是否对相同的生成 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、Yueming Yuan、Mao Cheng、Zhichen Zeng、Shi Don、Yanbin Jiang、Ying Sheng 以及 miles 团队

附录

ROCm 运行时与内核路径映射

所报告的运行使用了以下路径来处理对跨引擎一致性和在线更新影响最直接的模型组件。

表 1. 所报告的 ROCm 配置中使用的运行时路径。

启动器会明确指定所选的后端，而 Docker 作用域内的补丁则移除了依赖 Megatron 和 Transformer Engine 路径中残留的仅限 CUDA 的假设。这使得所测试的配置具有可复现性和可审查性。