让视觉语言模型的强化学习训练更快,不改算法也能提速一倍以上。
Rollplex: Cross-Phase GPU Spatial Sharing for Vision Language Model Post-Training

- 把推理和训练的前置计算并行做,不破坏同步更新机制。
- 在32张H800卡上实测提速1.23到2.24倍,内存占用降低超65%。
- 适合做视觉语言模型后训练的工程师,尤其关注推理效率者。
视觉语言模型(VLMs)使具身智能体能基于视觉观察和语言指令进行推理与决策。强化学习(RL)后训练通过任务反馈提升这些能力,但现有在线策略RL运行时将采样、参考评分和策略训练严格串行执行。尽管对纯文本任务有效,这种阶段粒度的执行方式在处理密集视频输入和提示前缀的VLM中效率低下。由于前缀计算与生成响应无关,可与采样解码并行进行,从而更充分地利用GPU算力,且不破坏同步在线策略语义。我们提出Rollplex,一种将参考计算与训练阶段解耦,并将前缀计算移入采样解码窗口的运行时系统。实现该调度需解决多项挑战:直接共置Qwen2.5-VL-32B模型需约165 GiB显存/卡,而采样与训练偏好不同张量并行(TP)度和权重布局。Rollplex通过两种机制应对:阶段感知内存管理根据生产者-消费者生命周期控制高带宽内存驻留;并行感知权重共享使用同一物理存储承载布局兼容的张量,仅重构不兼容部分,避免完整复制第二份演员模型。在32张H800 GPU上,Rollplex相较串行共置实现1.23×–1.30×加速,相较分离部署实现1.57×–2.24×加速,同时保持同步更新。
原文摘要 · Abstract (English)
Vision-language models (VLMs) enable embodied agents to reason and act from visual observations and language instructions. Reinforcement learning (RL) post-training enhances these capabilities using task feedback, but current on-policy RL runtimes execute rollout, reference scoring, and actor training in strict serial phases. While effective for text-only RL, this phase-granular execution is wasteful for VLMs, where processing dense video inputs and prompt prefixes occupies a large fraction of each phase. Because prefix processing is independent of the generated response, it can be run alongside rollout decoding, which leaves GPU compute capacity underutilized, without breaking synchronous on-policy semantics. We present Rollplex, a runtime that decomposes the reference and training phase and moves the prefix computation into the rollout decode window. Realizing this schedule requires more than concurrent kernel launches: naive colocation of Qwen2.5-VL-32\,B requires roughly 165\,GiB per GPU, while rollout and training prefer different tensor-parallel (TP) degrees and weight layouts. Rollplex addresses these constraints with two mechanisms. Phase-aware memory management controls HBM residency according to producer--consumer lifetimes. Parallelism-aware weight sharing uses the same physical storage for layout-compatible tensors across distinct TP degrees and reconstructs only incompatible tensors, avoiding a complete second actor copy. On 32 H800 GPUs, Rollplex achieves $1.23\times$--$1.30\times$ speedup over serial colocation and $1.57\times$--$2.24\times$ over disaggregation under the same GPU budget, while preserving the synchronous RL update.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。