vLLM × TRL在线强化学习训练中 vLLM 双模式集成与内存优化实战【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm本文基于 vLLM 官方文档 TRL 集成指南 展开讲解如何在 Transformers Reinforcement LearningTRL训练框架中启用 vLLM 作为在线 RL 训练的 rollout 生成引擎。读完后你将能够区分并配置 vLLM 的 server 模式与 colocate 模式、理解use_vllm/vllm_mode两个关键参数的作用、掌握 vLLM sleep mode 对训练期间显存占用的优化原理并结合 vLLM 仓库中的权重传输Weight Transfer与暂停/恢复Pause/Resume机制构建一套完整的“训练 高速生成”RLHF/GRPO 工作流。为什么 TRL 的在线训练方法需要 vLLMTRL 是一套完整的强化学习训练工具库覆盖监督微调SFT、Group Relative Policy OptimizationGRPO、Direct Preference OptimizationDPO、奖励建模Reward Modeling等多种方法并深度集成 Hugging Face transformers 生态。在线 RL 方法Online GRPO、Online DPO 等的核心特征在于每一步训练都需要模型实时生成 completion——即先由当前策略模型采样出回答再对回答打分/计算优势最后更新策略。生成阶段往往是整个训练循环中最耗时的部分而 vLLM 正是为高吞吐推理场景设计的引擎用它来生成这些 completion 可以显著加速 rollout 阶段。这就是 docs/training/trl.md 所述的核心集成关系TRL 负责训练vLLM 负责高速生成两者通过权重同步配合完成在线 RL 循环。TRL 目前支持以下与 vLLM 集成的在线 TrainerTrainer方法类型GRPOGroup Relative Policy Optimization组内相对优势估计的策略优化Online DPO在线 Direct Preference Optimization训练中实时生成偏好对RLOOReinforcement Learning with Leave-One-Out baselineNash-MD基于纳什均衡的偏好对齐方法XPO交叉偏好优化eXtended Preference Optimization启用方式use_vllm标志在 TRL 的 Trainer 配置中将use_vllm标志设为True即可让 Trainer 通过 vLLM 完成 rollout 生成from trl import GRPOConfig training_args GRPOConfig( ..., use_vllmTrue, )启用后GRPO 等在线 Trainer 的生成阶段不再走 transformers 的逐 token 解码而是经由 vLLM 的调度与批处理能力完成批量 rollout训练侧只负责前向/反向更新。两种集成模式server 模式与 colocate 模式TRL 支持两种模式在训练期间集成 vLLM通过 Trainer 配置中的vllm_mode参数控制。选择哪种模式本质上取决于你的 GPU 资源布局是“训练卡与推理卡分离”还是“单卡/同卡混布”。Server 模式默认在server 模式下vLLM 作为独立进程运行在专用 GPU上通过 HTTP 请求与 Trainer 通信。这种配置适合拥有独立推理 GPU 的场景生成负载与训练负载物理隔离互不抢占显存训练性能稳定也更容易横向扩展推理侧资源。from trl import GRPOConfig training_args GRPOConfig( ..., use_vllmTrue, vllm_modeserver, # default value, can be omitted )Colocate 模式在colocate 模式下vLLM 运行在 Trainer 进程内部与训练模型共享 GPU 显存。这种方式无需启动独立服务可以复用训练 GPU、提升整体利用率但代价是可能产生显存竞争——vLLM 的权重与 KV cache 要和训练参数、梯度、优化器状态共用同一块显存。from trl import GRPOConfig training_args GRPOConfig( ..., use_vllmTrue, vllm_modecolocate, )原文档提醒更细粒度的配置项与标志位如量化、KV cache 比例等请查阅你所使用的具体 Trainer 的文档。两种模式的取舍可以从一句话概括server 模式换稳定性与可扩展性colocate 模式换部署简洁性与 GPU 利用率但需要更精细的显存管理——这正是下面 sleep mode 与权重传输机制发挥作用的地方。显存优化vLLM Sleep Mode 在训练循环中的角色TRL 文档指出部分 Trainer 支持vLLM sleep mode训练期间把参数与缓存从 GPU 上卸载offload到 CPU RAM从而降低训练阶段的显存占用。这一能力在 vLLM 侧的完整实现在 Sleep Mode 文档 中有详细说明其核心机制值得理解Level 1 睡眠将模型权重 offload 到 CPU 内存丢弃 KV cache。适合“同一模型睡下去再唤醒继续跑”的场景前提是 CPU 内存足够存放权重。Level 2 睡眠权重与 KV cache 全部丢弃仅保留 rope 缩放张量等缓冲区在 CPU。特别适合 RLHF 权重更新场景——因为下一步马上要用新权重覆盖旧权重旧权重没有保留价值当 colocate Trainer 已经用 CPU 内存做自身 offload、CPU 内存吃紧时Level 2 更是唯一选择唤醒后可通过collective_rpc(reload_weights)原地恢复权重。在 Python API 中sleep mode 的开启与控制方式如下离线场景from vllm import LLM llm LLM(Qwen/Qwen3-0.6B, enable_sleep_modeTrue) # Level 2 睡眠丢弃权重与 KV cache为权重更新腾出显存 llm.sleep(level2) # 只唤醒权重显存避免权重同步期间同时分配 KV cache 导致 OOM llm.wake_up(tags[weights]) # ... 在此完成新权重写入 ... # 权重更新完成后再分配 KV cache llm.wake_up(tags[kv_cache])这种wake_up(tags[weights])的细粒度唤醒正是 RLHF 权重同步的关键技巧先只恢复权重显存完成参数覆写最后才恢复 KV cache把同步过程中的显存峰值压到最低。在线服务场景下设置VLLM_SERVER_DEV_MODE1并加--enable-sleep-mode启动后可通过POST /sleep?level2、POST /wake_up?tagsweights等开发端点完成同样的控制详见 Sleep Mode 文档。权重如何从训练侧同步到推理侧无论 server 还是 colocate 模式在线 RL 的核心闭环都依赖一件事每个训练 step 之后把更新后的权重同步到 vLLM让下一轮 rollout 使用新策略。vLLM 为此提供了可插拔的 Weight Transfer 系统它由训练进程侧的TrainerWeightTransferEngine与推理侧的WeightTransferEngine对称组成每轮同步都执行同一套四阶段协议初始化init_weight_transfer_engine建立训练与推理 worker 间的通信通道训练循环开始前调用一次开始start_weight_update推理引擎准备接收权重更新传输update_weights实际搬运权重可多次调用以支持分块传输收尾finish_weight_update完成更新并做后处理如提交新的weight_version。vLLM 当前提供的后端与适用场景后端传输方式适用场景NCCLNCCL broadcast训练与推理在不同 GPU对应 TRL 的 server 模式IPCCUDA IPC handles训练与推理同卡 colocate对应 TRL 的 colocate 模式sparse_ncclNCCL broadcastcheckpoint 坐标下的稀疏权重补丁sharded_rdtNIXL / Ray Direct Transport拉取式超大模型如 MoE 专家并行每个 worker 只需自己的分片可以看到后端选择与 TRL 的vllm_mode天然对应server 模式推理在别的 GPU 上通常走 NCCLcolocate 模式同进程同卡通常走 IPC。推理侧的配置只需给出后端名其余传输细节由训练侧在初始化握手中下发from vllm import LLM from vllm.config import WeightTransferConfig llm LLM( modelmy-model, weight_transfer_configWeightTransferConfig(backendnccl), # 或 ipc、sparse_nccl、sharded_rdt )在线服务则可用vllm serve my-model \ --weight-transfer-config {backend: nccl}训练侧只需构建一次引擎、每轮同步调用一次send_weights()它会自动驱动推理侧的 start → update → finish 全流程from vllm.distributed.weight_transfer import ( ModuleSource, HTTPVLLMWeightSyncClient, WeightTransferTrainerFactory, ) from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo # 训练循环开始前调用一次 engine WeightTransferTrainerFactory.trainer_init( init_infoNCCLTrainerInitInfo( master_addressmaster_address, master_portmaster_port, world_sizeworld_size, # trainer 所有推理 worker rank0, # 该 trainer rankrank 0 是发送方 packedTrue, ), clientHTTPVLLMWeightSyncClient(http://localhost:8000), sourceModuleSource(model), ) # 每次权重同步调用一次 for step in range(num_steps): train_one_step(model) engine.send_weights()一个值得注意的设计约束FSDP、TP/PP/EP 等分片式 Trainer 要求每个训练 rank 都构建引擎并调用send_weights()。rank 0 是发送方只有它持有通信器并向 wire 上写数据但非发送 rank 也必须迭代WeightSource——因为物化参数本身往往是一次集合通信如 FSDP 的full_tensor()all-gather若某些 rank 跳过就会死锁。边生成边更新Pause/Resume 机制与典型异步 RL 循环在线 RL 还有一个微妙的问题权重更新发生时推理引擎里往往还有正在进行的请求。vLLM 提供pause_generation/resume_generationAPI 来解决这一点详见 Async RL 指南await engine.pause_generation(modekeep, clear_cacheTrue) # ... 同步权重 ... await engine.resume_generation()mode参数决定在途请求的处理方式Mode行为abort立即中止所有在途请求并返回部分结果默认wait等待所有在途请求完成后才暂停keep请求冻结在队列中resume_generation时继续clear_cache参数控制暂停后是否清除 KV cache 与前缀缓存设为True时恢复后生成的 token 完全由新权重计算设为False时保留旧 KV 缓存意味着部分上下文 token 仍反映旧权重stale KV cache。把 TRL 双模式、Weight Transfer 与 Pause/Resume 串起来一个典型的权重同步 RL 循环是由当前策略开始生成 rolloutTrainer 产生新权重后用modekeep暂停生成将新权重从 Trainer 同步到推理引擎NCCL/IPC 后端 sleep mode 精细唤醒恢复生成——在途请求以新权重继续循环往复。仓库中的示例 rlhf_async_new_apis.py 演示了完整模式基于vllm.AsyncLLMEngine、NCCL 权重传输配合在途暂停/恢复与校验。此外 examples/rl/ 目录下还有 IPC、NCCL、FSDPEP 等更多组合的训练侧示例可供参考HTTP 权重传输端点/init_weight_transfer_engine、/start_weight_update、/update_weights、/finish_weight_update、/pause、/resume、/weight_info等需设置VLLM_SERVER_DEV_MODE1后开启HTTPVLLMWeightSyncClient会自动替你调用其中大部分端点。小结启用开关TRL Trainer 配置中设use_vllmTrue即可让 GRPO、Online DPO、RLOO、Nash-MD、XPO 等在线 Trainer 的 rollout 走 vLLM模式选择vllm_modeserver默认独立进程 专用 GPUHTTP 通信稳定易扩展或vllm_modecolocate训练进程内共享显存利用率高但有显存竞争显存治理借助 vLLM sleep mode 的 Level 1/Level 2 睡眠与wake_up(tags[weights])细粒度唤醒把权重同步期间的显存峰值压到最低同步通道Weight Transfer 系统的四阶段协议 NCCL/IPC 等后端与 server/colocate 两种部署形态一一对应在途安全Pause/Resume API 保证权重可以在生成进行中安全更新。更多上下文可参考仓库内 RLHF 总览文档、Weight Transfer 文档 与 Async RL 指南。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考