睡眠模式¶
vLLM 的睡眠模式允许你在不停止服务器或卸载 Docker 容器的情况下,暂时释放模型所占用的大部分 GPU 显存,包括模型权重和 KV 缓存。这对于 RLHF、训练或需要在推理工作负载之间释放 GPU 资源的节省成本场景特别有用。
主要优势
- 释放 GPU 显存:将模型权重卸载到 CPU 内存并丢弃 KV 缓存,可释放高达 90% 以上的 GPU 显存供其他任务使用。
- 快速恢复:无需完整重新加载模型,即可快速唤醒引擎并恢复推理。
- API 端点:通过 HTTP 端点或 Python API 控制睡眠/唤醒状态。
- 支持分布式工作负载:支持张量并行、流水线并行等。
- 精细控制:可选择仅唤醒模型权重或 KV 缓存,以避免在权重更新期间出现内存溢出 (OOM)。
注意
该功能目前在 CUDA 和 ROCm 平台上均受支持。
注意
更多信息,请参阅此 博文。
睡眠级别¶
1 级睡眠将卸载模型权重并丢弃 KV 缓存。KV 缓存的内容会被清除。1 级睡眠适用于挂起并在之后唤醒引擎以再次运行相同模型。模型权重被备份在 CPU 内存中,请确保有足够的 CPU 内存来存储模型权重。2 级睡眠将同时丢弃模型权重和 KV 缓存(同时保留模型在 CPU 中的缓冲区,如 RoPE 缩放张量)。模型权重和 KV 缓存的内容均会被清除。2 级睡眠适用于挂起并在之后唤醒引擎以运行不同的模型或更新模型,此时不再需要之前的模型权重(例如 RLHF 权重更新)。
用法¶
离线推理¶
通过向 LLM 类传递 enable_sleep_mode=True 来启用睡眠模式。
Python API¶
# Sleep level 1
# Put the engine to sleep (level=1: offload weights to CPU RAM, discard KV cache)
llm.sleep(level=1)
# Wake up the engine (restore weights)
llm.wake_up()
# Sleep level 2
# Put the engine to sleep (level=2: discard both weights and KV cache)
llm.sleep(level=2)
# Reallocate weights memory only
llm.wake_up(tags=["weights"])
# Load weights in-place
llm.collective_rpc("reload_weights")
# Reallocate KV cache
llm.wake_up(tags=["kv_cache"])
RLHF 权重更新¶
在 RLHF 训练期间,vLLM 允许你使用 wake_up() 中的 tags 参数有选择地仅唤醒模型权重或 KV 缓存。这种精细控制在更新模型权重时特别有用:通过仅唤醒权重(例如,llm.wake_up(tags=["weights"])),你可以在权重更新完成之前避免为 KV 缓存分配内存。这种方法有助于防止 GPU 内存溢出 (OOM) 错误(尤其是在使用大模型时),通过最小化权重同步和更新操作期间的峰值内存使用量来实现。
使用 tags=["weights"] 或 tags=["kv_cache"] 来控制恢复哪些资源,这对于 RLHF 和权重更新很有用。注意,在所有组件唤醒之前,is_sleeping 将报告为 true。
# Put engine to deep sleep (level=2)
llm.sleep(level=2)
# ... Get the new weights
# Wake up only weights to avoid OOM
llm.wake_up(tags=["weights"])
# ... Update the weights
# wake up KV cache after weights are updated
llm.wake_up(tags=["kv_cache"])
在线服务¶
若要在 vLLM 服务器中启用睡眠模式,你需要使用标志 VLLM_SERVER_DEV_MODE=1 初始化它,并向 vLLM 服务器传递 --enable-sleep-mode 参数。
开发模式下的服务器¶
使用标志 VLLM_SERVER_DEV_MODE=1 时,你将启用开发端点,这些端点不应暴露给用户。
以下是如何以 1 级睡眠和唤醒模型的示例。
以下是如何以 2 级睡眠和唤醒模型的示例。
curl -X POST 'https://:8000/sleep?level=2'
# Reallocate weights memory only
curl -X POST 'https://:8000/wake_up?tags=weights'
# Load weights in-place
curl -X POST 'https://:8000/collective_rpc' -H 'Content-Type: application/json' -d '{"method":"reload_weights"}'
# Reallocate KV cache
curl -X POST 'https://:8000/wake_up?tags=kv_cache'
HTTP 端点¶
POST /sleep?level=1— 使模型进入睡眠状态(level=1)。POST /wake_up— 唤醒模型。支持可选的tags查询参数以进行部分唤醒(例如,?tags=weights)。POST /collective_rpc— 执行集合远程过程调用 (RPC)。GET /is_sleeping— 检查模型是否处于睡眠状态。
注意
这些端点仅在传递 VLLM_SERVER_DEV_MODE=1 时可用。
限制¶
在 ROCm 上,虚拟内存分配是通过分块内存分配完成的。你可以通过 VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE(单位为 MB)控制块大小。默认值设置为 256MB。块大小越大,性能越快。然而,设置得过大将导致 OOM。因此,如果在睡眠模式下遇到 OOM,请尝试减小块大小。建议将块大小定义为 2 的幂。