跳到内容

休眠模式

vLLM 的 休眠模式 允许您在不停止服务器或卸载 Docker 容器的情况下,临时释放模型使用的大部分 GPU 内存,包括模型权重和 KV 缓存。这对于 RLHF、训练或成本节省等场景特别有用,在这些场景中,GPU 资源需要在推理工作负载之间释放。

主要优点

  • 释放 GPU 内存:将模型权重卸载到 CPU RAM 并丢弃 KV 缓存,释放高达 90% 以上的 GPU 内存用于其他任务。
  • 快速恢复:无需完全重新加载模型,即可快速唤醒引擎并恢复推理。
  • API 端点:通过 HTTP 端点或 Python API 控制休眠/唤醒状态。
  • 支持分布式工作负载:适用于张量并行、流水线并行等。
  • 细粒度控制:可以选择仅唤醒模型权重或 KV 缓存,以避免在权重更新期间出现 OOM(内存不足)错误。

注意

此功能现在支持 CUDA 和 ROCm 平台。

注意

更多信息,请参阅这篇博客文章

休眠级别

1 级休眠将卸载模型权重并丢弃 KV 缓存。KV 缓存的内容将被遗忘。1 级休眠适用于休眠然后唤醒引擎以再次运行同一模型。模型权重会备份在 CPU 内存中。请确保有足够的 CPU 内存来存储模型权重。2 级休眠将丢弃模型权重和 KV 缓存(但模型的缓冲区会保留在 CPU 中,例如 rope scaling tensors)。模型权重和 KV 缓存的内容都会被遗忘。2 级休眠适用于休眠然后唤醒引擎以运行不同的模型或更新模型,此时不需要之前的模型权重,例如 RLHF 权重更新。

用法

离线推理

通过将 enable_sleep_mode=True 传递给 LLM 类来启用休眠模式。

from vllm import LLM
llm = LLM("Qwen/Qwen3-0.6B", 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 初始化它,并将 --enable-sleep-mode 传递给 vLLM 服务器。

服务器开发模式

使用标志 VLLM_SERVER_DEV_MODE=1 时,您将启用开发端点,这些端点不应暴露给用户。

VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-0.6B \
  --enable-sleep-mode \
  --port 8000

以下是级别 1 中休眠和唤醒模型的示例。

curl -X POST 'https://:8000/sleep?level=1'
curl -X POST 'https://:8000/wake_up'

这是级别 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 上,ROCm 上的虚拟内存分配是通过分块内存分配完成的。您可以通过 VLLM_ROCM_SLEEP_MEM_CHUNK_SIZE(以 MB 为单位)来控制分块大小。默认值为 256MB。分块越大,性能越快。但是,设置得太大将导致 OOM。因此,如果您在使用休眠模式时遇到 OOM。请尝试减小分块大小。建议将分块大小定义为 2 的幂。