游乐游手机版
首页/AI教程/文章详情

KVCache休眠恢复机制与显存复用

时间:2026-07-06 16:22
训推共卡时,sleep resume并非服务启停,而是显存状态机,用于在生成与训练间切换权重、KVcache等对象。不同后端(vLLM、SGLang、TRT-LLM)的释放粒度各异;非naive权重同步仅释放KVcache以腾出空间。free_cache_engine需权衡显存释放与状态切换成本。

上一篇将 vLLM、SGLang、TRT-LLM 等推理引擎封装到了 rollout backend 的抽象框架中。今天我们来深入探讨:当 rollout engine 与训练 engine 共享同一块 GPU 时,代码中为何频繁出现 sleep()resume()release_kv_cache()resume_kv_cache() 这类看似“运维操作”的调用?

这里先给出核心结论:sleep 和 resume 并非为了让推理服务实现优雅的上下线,它们实际上是训推共卡场景下的显存状态机——简单来说,就是决定 rollout 的权重、KV cache、prefix cache,以及训练参数和优化器状态,哪些留在 GPU 上,哪些在权重同步或训练更新之前需要临时腾出空间。

先看看共卡时显存冲突的具体情况。看下图时需要注意,训练 engine 和 rollout engine 争夺的并非一个笼统的“GPU 内存”,而是多个生命周期截然不同的对象。

\

训推共卡时的显存争用

这张图对应 rollout 配置中几个关键参数:gpu_memory_utilization 控制推理引擎分配给 KV cache 的显存比例;free_cache_engine 决定生成完成后是否释放 engine cache;max_num_batched_tokensmax_num_seqsprompt_length/response_length 则限定了生成阶段的 token 容量上限(可参考 verl/trainer/config/rollout/rollout.yaml:37-79verl/workers/config/rollout.py:188-212)。

1. 生成结束后为何要 sleep

RayPPOTrainer.fit() 的主循环中,执行 generate_sequences() 之后,紧接着会调用 self.checkpoint_manager.sleep_replicas()verl/trainer/ppo/ray_trainer.py:1373-1384)。如果 validation 阶段需要使用 colocated reward model,也会先让 rollout model 进入 sleep,计算完 reward 后,再通过 update_weights() 把 rollout 唤醒(verl/trainer/ppo/ray_trainer.py:553-566)。

下面这张图将这个动作放回一轮 step 中重新审视。它要补充的关键点是:sleep 不是生成阶段的收尾操作,而是从推理状态切换回训练状态的入口。

\

rollout 显存状态生命周期

CheckpointEngineManager.sleep_replicas() 的注释说得非常直接:sleep all rollout replicas,释放 weight 和 KV cache 在 device 上的内存(verl/checkpoint_engine/base.py:409-412)。在 naive 同进程路径中,ActorRolloutRefWorker.update_weights() 会在 free_cache_engine 开启时先 resume(tags=["weights"]),更新 rollout weights,再 resume(tags=["kv_cache"]),将推理状态恢复至下一轮 rollout 可用状态(verl/workers/engine_workers.py:663-740)。

由此可见,sleep/resume 的设计意图根本不是让 server “暂停一下”,而是让同一批 GPU 在生成、reward/logprob、actor update、weight sync 等阶段之间交替承载不同的内存对象。

2. 同一份合同在不同 backend 中的含义并不相同

第 14 篇提到 backend abstraction 是一个薄抽象,sleep/resume 就是最典型的例证。下面这张图将 vLLM、SGLang、TRT-LLM 的语义放在一起对比。看图时需特别留意 tags 和 level 参数:它们决定了释放的是 KV cache、weights,还是两者都释放。

\

不同 backend 的 sleep/resume 语义

vLLM adapter 的 resume(tags) 会转化为 server 的 wake_up(tags=...)release() 会转化为 sleep(level=...);更新权重后还会清除 prefix cache 并设置 global steps(verl/workers/rollout/vllm_rollout/vllm_rollout.py:150-192)。vLLM async server 在 colocated 模式下会执行 wake_up(tags=["kv_cache", "weights"]) 并重置 prefix cache;hybrid sleep 会根据 LoRA 或 NPU 情况选择不同的 sleep level(verl/workers/rollout/vllm_rollout/vllm_async_server.py:583-606verl/workers/rollout/vllm_rollout/vllm_async_server.py:891-906)。

SGLang adapter 将 memory occupation 明确拆分成 tags。默认 sleep_level=2 表示释放 weights 和 KV cache;LoRA adapter 模式下 sleep_level=1,只释放 KV cache,保留 base weights(verl/workers/rollout/sglang_rollout/sglang_rollout.py:184-188verl/workers/rollout/sglang_rollout/sglang_rollout.py:258-285)。SGLang async server 也会根据 hybrid/colocated/standalone 模式,选择释放 ["kv_cache", "weights"] 或只释放 ["kv_cache"]verl/workers/rollout/sglang_rollout/async_sglang_server.py:430-494)。

TRT-LLM 路径同样实现了 release_kv_cache()resume_kv_cache();hybrid/colocated sleep 会释放完整 tags,standalone 则跳过 wake/sleep(verl/workers/rollout/trtllm_rollout/trtllm_async_server.py:403-440)。其 adapter 在恢复 KV cache 前还会执行分布式 barrier,避免非 leader rank 的 actor offload 尚未完成就恢复 KV,导致 OOM(verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:395-421)。

3. 非 naive 权重同步仅挪走 KV cache

naive 路径适用于同进程直接同步;但当 backend 走 NCCL/NIXL/HCCL 等 checkpoint engine 时,权重同步策略要精细得多。下面这张图说明了非 naive 路径的执行顺序:先中止并保存未完成请求,再释放 KV cache,保留 weight buffer,完成权重传输,最后恢复 KV cache 和 generation。

checkpoint engine 中的 KV cache 同步路径

CheckpointEngineManager.update_weights() 对这个路径描述得很清楚:非 naive 情况下先 abort_replicas(),再 release_kv_cache_replicas();注释特别说明它不同于 sleep_replicas(),只释放 KV cache,保留 model weights,目的是让传输直接写入已有的 weight buffers。之后构建 process group、trainer/rollout 一起 update weights、finalize,最后 resume_kv_cache_replicas()resume_generation_replicas()verl/checkpoint_engine/base.py:447-492)。

这意味着在 weight sync 过程中,KV cache 就是一块“临时挪开的工作区”。如果直接释放 weights,可能需要重建更大的推理状态;如果完全不释放 KV cache,又可能不给权重同步留足显存。这里的工程取舍并非单纯快慢问题,而是显存峰值、同步成本和请求连续性之间的平衡。

4. 读懂 sleep/resume 代码时需避免两个误区

第一,不要把 free_cache_engine=True 简单理解成“更省显存”。它确实能释放推理侧的 cache,但每次生成和训练之间都要付出状态切换的成本。对于短 response、高频 step 的训练场景,这个成本可能非常明显;反过来,对于长 response 或共卡压力大的场景,不释放又可能直接导致 OOM。

第二,不要以为所有 backend 都能以同等粒度释放 KV cache。vLLM、SGLang、TRT-LLM 虽然都暴露了类似的生命周期接口,但 tags、sleep level、prefix cache reset、standalone 模式和 LoRA adapter 路径的实现各不相同。源码中甚至能看到 vLLM async server 的 release_kv_cache() 仍带有 TODO,说明“只释放 KV、保留 weights”这一能力,并非在所有 backend 中都能完全等价实现(verl/workers/rollout/vllm_rollout/vllm_async_server.py:612-623)。

小结:显存复用是 rollout 抽象的一部分

第 13 篇说明 rollout 是训练中的推理服务;第 14 篇说明不同 backend 被抽象成 rollout contract;本文则补充了这层 contract 的显存含义:

generate 需要 weights + KV cache  
→ sleep 释放 rollout 侧显存  
→ actor/reward/logprob 使用训练侧显存  
→ update_weights 恢复 weights  
→ resume_kv_cache 恢复下一轮生成能力

读到这里,rollout 已经不再只是“生成 response 的服务”,而是一个能在训练和推理状态之间灵活切换的 GPU 内存管理边界。下一篇将继续进入 agent loop:当 rollout 从单轮补全变为多轮工具交互时,系统为何必须坚持 token-in/token-out,而不是仅保存最终文本。

本文源码索引

  • verl/trainer/config/rollout/rollout.yaml:37-79:KV cache、free_cache_engine、batch token 和序列容量配置。
  • verl/workers/config/rollout.py:188-212:rollout dataclass 中长度、KV 显存比例、free_cache_engine 和容量字段。
  • verl/trainer/ppo/ray_trainer.py:553-566:validation 中 rollout sleep 后为 colocated reward 腾显存。
  • verl/trainer/ppo/ray_trainer.py:1373-1384:训练主循环生成后调用 sleep_replicas()
  • verl/checkpoint_engine/base.py:409-445:sleep/wake、abort、release/resume KV cache 的 manager 方法。
  • verl/checkpoint_engine/base.py:447-492:非 naive checkpoint engine 更新权重时释放 KV、同步权重、恢复 KV 和 generation。
  • verl/workers/engine_workers.py:663-740:naive 权重同步路径中的 resume(weights)resume(kv_cache)
  • verl/workers/rollout/vllm_rollout/vllm_rollout.py:150-192:vLLM adapter 的 resume/release/update weights。
  • verl/workers/rollout/vllm_rollout/vllm_async_server.py:583-623:vLLM server wake/sleep/cache 控制。
  • verl/workers/rollout/vllm_rollout/vllm_async_server.py:891-906:vLLM hybrid sleep level 选择。
  • verl/workers/rollout/sglang_rollout/sglang_rollout.py:184-188:SGLang sleep_level 的语义。
  • verl/workers/rollout/sglang_rollout/sglang_rollout.py:258-285:SGLang adapter 的 resume/release tags。
  • verl/workers/rollout/sglang_rollout/async_sglang_server.py:430-494:SGLang async server 的 wake/sleep/KV cache 控制。
  • verl/workers/rollout/trtllm_rollout/trtllm_async_server.py:403-440:TRT-LLM async server 的 cache 和 sleep/wake 控制。
  • verl/workers/rollout/trtllm_rollout/trtllm_rollout.py:395-421:TRT-LLM adapter 恢复 KV 前的 barrier 和 tags 映射。
来源:https://cloud.tencent.com.cn/developer/article/2703893
上一篇企业AI智能体真正门槛不是模型能力而是工作底座 下一篇比利时股票行情数据API云上部署:BEL20指数弹性拉取与成本控制
本站内容用于信息整理与展示,如有侵权或内容问题请及时联系处理。

相关推荐

补充同频道和同主题内容,方便继续浏览更多相关内容。

同类最新

继续查看同栏目最近更新的文章。

更多
阿里云Qoder CN灵码AI助手免费版及credits计费指南
AI教程 · 2026-07-06

阿里云Qoder CN灵码AI助手免费版及credits计费指南

阿里云QoderCN(原通义灵码)是一款AI智能编码助手,提供IDE插件、独立IDE等形态,覆盖编码及日常办公场景。产品分个人社区版(免费)、个人专业版、企业标准版和企业VPC版,采用Credits计费模式,支持多种AI模型。

基于大模型的城市文旅知识图谱构建与内容分发
AI教程 · 2026-07-06

基于大模型的城市文旅知识图谱构建与内容分发

大模型构建城市知识图谱时优先采信权威信源。贵港西江传媒联合《度假旅游》杂志,通过本地采编、期刊发布、阿里云多平台分发模式,产出产业文旅等结构化内容,提升AI知识库收录权重,为城市品牌长效传播提供可复制路径。

贵州文旅AIGEO内容运营本地媒体落地实践
AI教程 · 2026-07-06

贵州文旅AIGEO内容运营本地媒体落地实践

针对贵州文旅行业在阿里云平台的内容运营痛点,总结合规发文规则,包括弱化营销、避免引流信息与极限词。以《度假旅游杂志》在贵阳设立本地化运营站点为例,为黔域文旅商家产出合规原创内容,提升AI平台收录权重。同时指出商家常踩的审核红线,强调以干货分享获取自然流量。

阿里内部禁用Claude Code OpenCode成替代方案
AI教程 · 2026-07-06

阿里内部禁用Claude Code OpenCode成替代方案

ClaudeCode对国内用户定向封禁,网传阿里内部已全面禁用。OpenCode作为开源替代,支持接入多种AI模型,通过CCSwitch或手动配置可无缝迁移原有MCP与AgentSkill,规避账号风险。

年实测Homebrew安装配置国内源多种安装方式一篇搞定
AI教程 · 2026-07-06

年实测Homebrew安装配置国内源多种安装方式一篇搞定

Homebrew是macOS上流行的包管理工具,提供pkg安装包、脚本安装和Git克隆等多种安装方式。国内用户推荐使用镜像脚本安装并配置中科大或清华大学镜像源以加速下载。常用命令包括brewinstall安装工具、brewinstall--cask安装图形应用、brewupdate更新及brewdoctor诊断环境。