vime 对 AMD Instinct GPU 的 ROCm 支持

vime 对 AMD Instinct GPU 的 ROCm 支持

vLLM 已将其 ROCm 支持集成到 vime（其强化学习 (RL) 框架）中，使得大规模 RL 后训练能够原生运行在 AMD Instinct MI300X 和 MI355X GPU 上。此次集成提供了一个经过验证的端到端流水线，包括一个预构建的容器，从而消除了从源码构建的需求。

vime 架构与设计

vime 基于 slime 框架，采用三阶段、解耦的训练-推理设计。它将 RL 过程分为三个不同的组件：

  • Training (Megatron): 管理主训练循环、参数更新以及向 rollout 端同步权重。
  • Rollout (vLLM + Router): 处理推理采样，根据奖励或验证器信号生成训练样本。
  • Data Buffer: 作为训练和 rollout 之间的桥梁，管理自定义的 rollout 逻辑和 prompt 注入。

AMD Instinct GPU 在 RL 中的硬件优势

RL 后训练是高度内存密集型的，因为它需要同时存储训练端权重 (Megatron 格式) 和推理端 KV caches (vLLM rollouts)。AMD Instinct GPU 针对这一特性进行了优化：

  • 高 HBM 容量: MI300X 提供 192 GB 的 HBM3，而 MI355X 提供 288 GB。这种大容量减少了对激进张量并行 (tensor parallelism) 的需求，简化了拓扑结构并提高了集群利用率。
  • 卓越的内存带宽: MI300X 提供超过 5 TB/s 的聚合带宽，MI355X 提供 8 TB/s。由于 RL rollouts 受内存带宽限制，更高的吞吐量降低了自回归 token 生成阶段的步长延迟。
  • 原生生态系统集成: ROCm 的开源特性使 vime 能够继承现有的 vLLM rollout 栈，而无需单独的代码路径，充分利用了 PyTorch 和 vLLM 对 ROCm 的原生支持。

在 ROCm 上的技术实现

将 vime 引入 AMD 硬件涉及多个关键软件组件的集成：

Megatron-LM Backend

vime 使用了 Megatron-LM 的 ROCm 兼容分支。为了确保兼容性，实现了一个补丁来保护非 CUDA 构建版本中的 CUDA fused-kernel 初始化。该流水线支持在单个 GPU 上进行 HuggingFace 到 torch_dist 的 checkpoint 转换，并利用原生 PyTorch 路径进行梯度累积。

Colocated Weight Synchronization

在共存 (colocated) 模式下，Megatron 和 vLLM 共享同一个 GPU 池。vime 使用进程间通信 (IPC) 在每个优化器步长后将更新后的权重从 Megatron 同步到 vLLM。这在 ROCm 上通过 torch.cuda.get_device_properties(i).uuid 接口实现，该接口为路由提供了稳定且进程一致的设备 UUID。

资源管理与 Ray 集成

GPU 分配通过 HIP_VISIBLE_DEVICES 进行管理。vime 通过设置此变量以及 CUDA_VISIBLE_DEVICES 来确保 Megatron 训练 actor 与 vLLM 子进程之间的一致性。Ray 的 AMD GPU 管理器被配置为遵循这些掩码，并且系统需要提高文件描述符限制 (--ulimit nofile=1048576:1048576) 以大规模生成 actor workers。

MI355X 上的性能基准测试

在 MI355X GPU 上使用 Qwen3-8B 模型进行的测试显示了以下 100 个训练步长的结果：

  • 吞吐量: 持续保持在约 4,100 tokens_per_gpu_per_second。随着策略学习到更可预测的输出，吞吐量呈上升趋势，从而允许 vLLM 进行更高效的批处理。
  • Logprob 稳定性: train_rollout_logprob_abs_diff 稳定在 0.012 左右。训练与 rollout 对数概率之间如此低的差异表明权重同步成功，且与 NVIDIA 硬件上的结果相当。
  • 奖励收敛: 使用 dapo-math-17k 数据集，raw_reward 从接近 0 上升到约 0.5–0.6，这标志着策略成功学会了偏好验证器所奖励的推理模式。

支持的功能与路线图

当前支持

  • GRPO 训练
  • 共存模式及异步（不相交 GPU 池）训练与 rollout
  • Megatron-LM 训练后端和 vLLM rollout 后端
  • 支持 Qwen3 Dense 和 MoE 模型

未来路线图

  • 完整的 vLLM Router 和 PD (Prefill-Decode) 解耦支持
  • FP8 流水线优化
  • 针对 AMD MoE 工作负载的 R3 (Rollout Routing Replay)
  • 异步训练的性能优化，以解决内存泄漏和 logprob 分歧问题
  • 用于多智能体设置和多轮工具调用的 Agentic RL

Sources