Prhub

#24854 [RL] Call torch.cuda.empty_cache() for `in-place` pause mode to avoid OOM

原始 PR 作者 ByronHsu 合并时间 2026-05-10 14:36 文件变更 2 提交数 5 评论 7 代码增减 +15 / -1

执行摘要

修复 in-place pause 模式因缺少 empty_cache 导致的 OOM

Post-weight-update processing (e.g. DeepSeek MLA w_kc/w_vc derivation, FP8 scale rebuild) creates transient CUDA allocations that fragment PyTorch's block cache. Without empty_cache(), reserved memory grows each iteration and eventually OOMs. The in_place path never calls flush_cache (to preserve KV cache), so empty_cache() was never triggered — this PR closes that gap.

该 PR 修复了一个明确的 OOM 问题,代码简洁,建议合入。讨论中关于跨平台兼容的取舍值得记录,未来如果有更多后端需求可考虑重构。

讨论亮点

Review 中 hebiao064 建议使用 empty_device_cache 兼容 AMD 等其他加速器,但 ByronHsu 认为增加复杂性且与已有 torch_empty_cache 命名约定冲突,最终保持仅 CUDA 支持。此外,hebiao064 建议字段名改为 torch_empty_cache 以对齐已有命名,被接受。

实现拆解

  1. ContinueGenerationReqInput 数据类中添加 torch_empty_cache 字段,默认 True,允许调用方跳过。
  2. Scheduler.continue_generation 方法中,当 torch_empty_cacheTrue 时,调用 torch.cuda.empty_cache() 并记录前后 reserved 内存量。
  3. 由于 empty_cache 在 engine 仍处于 paused 状态时执行,没有活跃流竞争,线程安全。
  4. 命名与现有 UpdateWeightFromDiskReqInput 中的 torch_empty_cache 保持一致。
  5. 测试:在 in-place pause + weight update 场景中确认日志触发、内存下降且不改变 loss。
文件 模块 状态 重要度
python/sglang/srt/managers/scheduler.py 调度器 modified 6.09
python/sglang/srt/managers/io_struct.py 数据结构 modified 5.23

关键符号

continue_generation ContinueGenerationReqInput

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

使用 empty_device_cache 实现跨设备兼容 设计

hebiao064 建议使用 empty_device_cache 兼容 AMD 等其他加速器 (https://github.com/sgl-project/sglang/blob/12f42f2e7e7520bef573230cadbdf6c10241aa61/python/sglang/srt/utils/common.py#L498)

结论:ByronHsu 认为太混乱且与已有 torch_empty_cache 命名矛盾,决定仅支持 torch,保持简单。 · 已解决

字段命名对齐已有约定 style

hebiao064 建议将 empty_cache 重命名为 torch_empty_cache 以对齐已有 update_* 请求的字段

结论:接受并重命名。 · 已解决

安全性确认 正确性

ByronHsu 询问 @guapisolo @yueming-yuan 是否安全

结论:未直接回应,但 PR 最终合并,可能无安全问题。 · 已解决

风险与影响

仅调用 torch.cuda.empty_cache(),没有考虑 AMD ROCm 或其他后端,但 PR 明确仅针对 CUDA,且已有类似用法。empty_cache 调用开销低,但日志记录可能在高频场景下产生大量 INFO 日志,需注意日志级别配置。另外,重复在 abort/retract 中调用 empty_cache(已在 flush_cache 中调用一次)可能会导致轻微性能开销,但 PR 说明第二次调用通常是 no-op。

修复 RL 训练中 weight update 后 in-place pause 模式的 OOM 问题,使模型更新过程更稳定。日志监控内存趋势。影响范围:使用 in-place pause mode 进行权重更新的 RL 训练场景。其他场景不受影响(新字段默认 True,但 empty_cache 只在 continue_generation 中调用,若不使用 pause 不会触发)。

仅支持 CUDA 日志可能高频输出

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论