DeepSeek-V4 Flash RL训练成功迁移至AMD MI355X,非NVIDIA路线可复现

在过去两年里,大模型强化学习(RL)训练几乎被视为NVIDIA GPU的“专属领地”——CUDA生态的成熟度和工具链的深度绑定,让AMD、Intel等替代方案长期停留在推理优化或小规模微调层面。但LMSYS Chatbot Arena团队最新披露的工程报告,正在打破这种认知惯性:他们将2840亿参数的稀疏MoE模型DeepSeek-V4 Flash(每token实际激活130亿参数)的RL训练完整搬到了AMD Instinct MI355X GPU上,并且给出了端到端的可复现配置。

核心技术栈并不“省心”:此次训练依赖Miles框架(负责异步循环与权重同步)作为调度中枢,rollout生成使用SGLang,策略更新则调用Megatron。这并非现成集成套件,而是跨框架的“拼图”——三个工具分别来自不同社区,且SGLang与Megatron在模型分片、量化状态管理上存在天然对齐难点。团队重点攻克了三大工程挑战:SGLang与Megatron之间的模型参数一致性维护、量化状态(如INT8/FP8)在训练过程线性更新时的数值稳定性、以及多节点并行下梯度通信的延迟抖控制。

实测数据证明了有效性与可控性:在四个八GPU节点(共32张MI355X)上,训练完成了超过100个优化器步骤。关键指标方面,训练与rollout之间的对数概率差始终保持在可接受范围内,表明模型对齐策略准确;在线奖励曲线持续上升,离线AIME-2024基准分数也同步上涨。这意味着非NVIDIA硬件承载大规模RL训练不再只是“能跑起来”,而是可以产生有效收益。

行业背景值得深挖:目前大规模RL训练主要依赖NVIDIA的Megatron-LM + NeMo框架,而AMD的ROCm软件栈虽然在推理端有了不少落地案例(如Llama系列),但在训练端尤其RL这类需要频繁rollout与策略更新的场景中,始终缺乏标杆级验证。LMSYS这次选择的DeepSeek-V4 Flash恰恰是一个“高难度样本”——MoE结构带来的动态稀疏性与RL的非稳态特性叠加,对内存带宽和通信拓扑提出了极高要求。MI355X的HBM3内存与Infinity Fabric互连在此次测试中展现了足够支撑力,没有出现明显的显存溢出或通信死锁。

实用建议:对于考虑非NVIDIA路线做RL训练的团队,该配置的核心参考价值并非硬件本身,而是Miles + SGLang + Megatron的组合踩坑经验。直接复现时需特别注意三点:一是确保两个推理引擎使用相同的量化校准集以维持参数对齐;二是建议将SGLang的rollout async timeout设为保守值以避免单节点掉队;三是多节点环境下建议用NCCL的AMI插件替代默认后端以压缩通信量。从趋势看,AMD MI355X若能持续获得类似的高难度工作负载验证,其在大模型RL训练市场将不再扮演“备胎”角色。