ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

TRL 在线训练如何开启 vLLM sleep 模式(vllm_enable_sleep_mode)释放 GPU 显存?

TRL 在线训练如何开启 vLLM sleep 模式(vllm_enable_sleep_mode)释放 GPU 显存? TRL 在线训练如何开启 vLLM sleep 模式vllm_enable_sleep_mode释放 GPU 显存【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl用 TRL 跑 GRPO、RLOO 这类在线训练时模型要在训练过程中不断生成 completions通常会把 vLLM 作为生成后端。在默认的 colocate 模式下vLLM 引擎和训练模型共享同一块 GPU 显存文档建议优先调整vllm_gpu_memory_utilization来避开内存冲突但如果调整之后仍然出现 out-of-memory 错误TRL 官方给出的下一步就是开启 vLLM sleep 模式——把 vLLM 的参数和缓存卸载到 CPU 内存让优化步骤期间的 GPU 显存降下来。本文基于 Reducing Memory Usage 的 vLLM sleep mode 章节和 GRPO Trainer 文档说明如何开启vllm_enable_sleep_mode、它做了什么、代价是什么。前提你已经用 vLLM 做生成sleep 模式不是独立功能它只在 vLLM 作为生成后端时生效。如果你的 GRPO/RLOO 训练还没有接入 vLLM先完成这一步见 GRPO 文档 Speed up training with vLLM-powered generation 一节pip install trl[vllm]from trl import GRPOConfig training_args GRPOConfig( ..., use_vllmTrue, # vllm_modecolocate by default )这里...代表你训练脚本中已有的其他GRPOConfig参数保持不变即可。colocate 是默认模式无需显式传vllm_mode。文档给出的触发条件是先按 GRPO 文档 的建议调整vllm_gpu_memory_utilization如果推荐的值在你的环境不工作建议加一个小缓冲如 0.05 或 0.1如果仍然报 OOM 错误再开启 sleep 模式。也就是说它是显存紧张场景下的最后手段而不是默认要打开的开关——该参数默认值为False。开启 vllm_enable_sleep_mode在训练配置里加一个参数即可文档为 GRPO 和 RLOO 给出示例from trl import GRPOConfig training_args GRPOConfig(..., vllm_enable_sleep_modeTrue)RLOO 同理from trl import RLOOConfig training_args RLOOConfig(..., vllm_enable_sleep_modeTrue)参数定义见 GRPOConfig官方帮助文本是Enable vLLM sleep mode to offload weights/cache during the optimizer step. Keeps GPU memory usage low, but waking the engine adds host–device transfer latency.开启后实际发生了什么按文档描述开启后 TRL 会在优化步骤期间把 vLLM 的参数和缓存卸载到 CPU RAM在需要权重同步和生成时再加载回 GPU 显存。源码里可以看到这个过程的调用位置vllm_generation.pyvLLM 引擎创建完成后立即self.llm.sleep(level2)即训练开始时 vLLM 不占用显存同步训练权重到 vLLM 前先empty_cache()再wake_up(tags[weights])真正生成 completions 前wake_up(tags[kv_cache])生成结束后再次sleep(level2)回到卸载状态。所以一个完整训练循环里vLLM 只在权重同步和生成的窗口期占用 GPU 显存其余时间显存留给训练侧。代价与限制文档明确写了两点代价开之前要心里有数速度唤醒 vLLM 引擎会引入 host–device 传输延迟可能轻微影响训练速度。显存紧张和训练速度之间是用这一个参数做权衡。适用位置文档示例只覆盖 GRPO 与 RLOO。源码中该参数还存在于 OnlineDPOonline_dpo_config.py、SSD、SDPO、SDFT、Distillation 等配置的 vLLM 相关参数里行为描述一致在优化步骤卸载权重/缓存。另外提醒一点边界sleep 模式解决的是训练与 vLLM 共享显存的紧张如果你的 vLLM 跑在独立 GPU 的 server 模式下通过CUDA_VISIBLE_DEVICES把训练和生成分到不同 GPU见 Speeding Up Training显存压力在 vLLM 自己那几块卡上是否还需要 sleep 模式取决于你那边的显存占用文档对此没有单独展开。验证是否生效文档没有给出固定的检查命令或数值阈值判断依据就是文档描述的行为本身训练能正常跑过此前触发 OOM 的优化步骤不再因 vLLM 占用而爆显存在优化步骤期间 GPU 显存占用保持较低vLLM 参数和缓存已卸载到 CPU RAM文档原话the vllm parameters and cache will be offloaded during the optimization step若发现整体训练时间变长来源就是每次唤醒/休眠的 host–device 传输属于文档声明的预期行为不是故障。如果开启后训练速度明显变慢而显存其实还够用可以关闭该参数改回默认False回到只用vllm_gpu_memory_utilization控制显存的配置。相关文档docs/source/reducing_memory_usage.md — TRL 全部内存优化手段sleep 模式只是其中一种可与其他技术LoRA/PEFT、truncation 等组合docs/source/grpo_trainer.md — vLLM colocate/server 两种模式与vllm_gpu_memory_utilization调参建议docs/source/speeding_up_training.md — vLLM 安装与 server 模式下的 GPU 分配示例。【免费下载链接】trlTrain transformer language models with reinforcement learning.项目地址: https://gitcode.com/GitHub_Trending/tr/trl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表