基于自适应张量并行加速同步RLHF训练中的长尾生成
人工智能
2026-05-26 v1 分布式、并行与集群计算
摘要
强化学习从人类反馈(RLHF)已成为提高模型质量的关键后训练范式。然而，同步三阶段RLHF管道常因生成阶段而成为瓶颈，其中响应长度偏斜导致解码期间有效 batch size 快速缩小，GPU在进行解码时利用率低下，而少数长响应仍未完成。主流框架采用静态张量并行(TP)配置，无法适应变化的 batch 特征，留下了 substantial performance headroom。我们提出了PAT方法，一种在每个RLHF迭代的生成阶段动态重新配置TP的自适应方法。PAT引入了两项关键技术。首先，一种基于预测器引导的在线重新配置方法决定重新配置点和目标 TP 配置，基于离线轮廓，当预测的延迟收益超过重新配置开销时触发重新配置。其次，一种轻量级的在线重新配置机制仅更新受 TP 更改影响的状态和布局：它通过基于成本模型的选择来适应未完成的解码状态，选择 KV-cache 迁移或重计算，执行原位权重重分片，并重用已缓存的通信组。我们在SGLang之上实现PAT并将其集成到VeRL框架中。在LLaMA3.1-8B和Qwen3-14B上使用DeepScaleR的评估显示，PAT较原始VeRL设置将生成延迟降低最高可达34.6%，端到端RLHF训练迭代延迟降低最高可达27.2%。
引用
@article{arxiv.2605.23945,
title = {Accelerating Long-Tail Generation in Synchronous RLHF Training via Adaptive Tensor Parallelism},
author = {Long Zhao and Qinghe Wang and Jiaan Zhu and Youhui Bai and Zewen Jin and Chaoyi Ruan and Shengnan Wang and Cheng Li},
journal= {arXiv preprint arXiv:2605.23945},
year = {2026}
}
备注
11page, 14 figures