睡眠模式指南¶

概述¶

睡眠模式是一个 API，用于将模型权重从 NPU 内存中卸载，并丢弃 KV 缓存。此功能对于强化学习（RL）后训练工作负载至关重要，尤其是在 PPO、GRPO 或 DPO 等在线算法中。在训练期间，策略模型通常使用 vLLM 等推理引擎执行自回归生成，然后进行前向和反向传播以进行优化。

由于生成阶段和训练阶段可能采用不同的模型并行策略，在训练期间释放 KV 缓存，甚至卸载存储在 vLLM 中的模型参数变得至关重要。这确保了高效的内存利用率，并避免 NPU 上的资源争用。

快速开始¶

当 enable_sleep_mode=True 时，vLLM 中的内存管理（malloc、free）在特定内存池下进行。在模型加载和 KV 缓存初始化期间，我们将内存标记为一个映射：{"weight": data, "kv_cache": data}。

引擎（v0/v1）支持两种睡眠级别，用于在空闲期间管理内存：

• 一级睡眠

  • 操作：卸载模型权重并丢弃 KV 缓存。
  • 内存：模型权重被移动到 CPU 内存；KV 缓存被丢弃。
  • 用例：适合之后重新使用同一模型。
  • 注意：确保有足够 CPU 内存来保存模型权重。

• 二级睡眠

  • 操作：丢弃模型权重和 KV 缓存。
  • 内存：模型权重和 KV 缓存的内容均被丢弃。
  • 用例：适合切换到不同模型或更新当前模型。

由于此功能使用底层 API AscendCL，若要使用睡眠模式，请遵循安装指南并从源码构建。如果你使用的是 < v0.12.0rc1，请记得设置 export COMPILE_CUSTOM_KERNELS=1。

可选的额外清理¶

默认情况下，睡眠模式仅释放由睡眠模式分配器管理的内存。对于需要将更多 NPU 内存返回给训练器的 RL 工作负载，vLLM Ascend 还提供了一条可选的额外清理路径：

llm = LLM(
    "Qwen/Qwen2.5-0.5B-Instruct",
    enable_sleep_mode=True,
    additional_config={
        "rl_config": {
            "enabled": True,
            "sleep_mode_extra_cleanup": True,
        }
    },
)

对于在线服务，请通过 --additional-config 传递相同选项：

vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enable-sleep-mode \
    --additional-config '{"rl_config": {"enabled": true, "sleep_mode_extra_cleanup": true}}'

当启用 rl_config.sleep_mode_extra_cleanup 时，sleep() 还会：

• 当启用 ACL graph 时，清除 ACL graph 注意力工作区并使已捕获的 ACL graph 缓存失效；
• 重置模型运行器图管理器，以便唤醒后能够再次捕获 ACL graph；
• 等待待处理的流水线并行发送工作，同步 NPU，并销毁 HCCL 进程组。

在 wake_up() 期间，vLLM Ascend 会恢复 HCCL 进程组，刷新 MoE 分发器 HCCL 元数据，恢复睡眠模式分配器内存，并在需要时重新捕获 ACL graph。

对于二级睡眠，唤醒可以拆分为两个阶段：

llm.wake_up(tags=["weights"])
# Reload or update model weights here.
llm.wake_up(tags=["kv_cache"])

启用额外清理后，仅当 tags 为 None 或包含 "kv_cache" 时，才会重新捕获 ACL graph。这可以避免在外部重新加载的权重和 KV 缓存状态就绪之前重新捕获图。

准备模型权重¶

使用 Qwen2.5-0.5B-Instruct 模型权重。当 VLLM_USE_MODELSCOPE=True 时，模型将从 ModelScope 自动下载。

:header-rows: 1

* - Model
  - ModelScope Link
* - Qwen2.5-0.5B-Instruct
  - [Qwen/Qwen2.5-0.5B-Instruct](https://www.modelscope.cn/models/Qwen/Qwen2.5-0.5B-Instruct)

使用¶

以下是如何使用睡眠模式的简单示例。

• 离线推理：

  import os
  
  import torch
  from vllm import LLM, SamplingParams
  from vllm.utils.mem_constants import GiB_bytes
  
  os.environ["VLLM_USE_MODELSCOPE"] = "True"
  os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
  
  if __name__ == "__main__":
      prompt = "How are you?"
  
      free, total = torch.npu.mem_get_info()
      print(f"Free memory before sleep: {free / 1024 ** 3:.2f} GiB")
      # record npu memory use baseline in case other process is running
      used_bytes_baseline = total - free
      llm = LLM(
          "Qwen/Qwen2.5-0.5B-Instruct",
          enable_sleep_mode=True,
          additional_config={"weight_nz_mode": 0},
      )
      sampling_params = SamplingParams(temperature=0, max_tokens=10)
      output = llm.generate(prompt, sampling_params)
  
      llm.sleep(level=1)
  
      free_npu_bytes_after_sleep, total = torch.npu.mem_get_info()
      print(f"Free memory after sleep: {free_npu_bytes_after_sleep / 1024 ** 3:.2f} GiB")
      used_bytes = total - free_npu_bytes_after_sleep - used_bytes_baseline
      # now the memory usage should be less than the model weights
      # (0.5B model, 1GiB weights)
      assert used_bytes < 1 * GiB_bytes
  
      llm.wake_up()
      output2 = llm.generate(prompt, sampling_params)
      # cmp output
      assert output[0].outputs[0].text == output2[0].outputs[0].text
  

• 在线服务： !!! note

由于可能存在恶意访问的风险，请确保处于开发模式，并显式设置开发环境变量 VLLM_SERVER_DEV_MODE，以暴露这些端点（sleep/wake up）。

```bash
export VLLM_SERVER_DEV_MODE="1"
export VLLM_WORKER_MULTIPROC_METHOD="spawn"
export VLLM_USE_MODELSCOPE="True"

vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --enable-sleep-mode \
    --additional-config='{"weight_nz_mode": 0}'

# after serving is up, post to these endpoints.
# /sleep reads level from the query string (JSON body is ignored).

# --- Level 1: offload weights, discard KV cache ---
curl -X POST "http://127.0.0.1:8000/sleep?level=1"
curl -X GET http://127.0.0.1:8000/is_sleeping

# wake all tags (weights + kv_cache)
curl -X POST http://127.0.0.1:8000/wake_up
curl -X GET http://127.0.0.1:8000/is_sleeping

# serving is available again after Level-1 wake_up
curl http://127.0.0.1:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen2.5-0.5B-Instruct",
        "prompt": "The future of AI is",
        "max_tokens": 7,
        "temperature": 0
    }'

# --- Level 2: discard weights and KV cache ---
# tags must be in ["weights", "kv_cache"]. After waking weights,
# reload the checkpoint on every worker before waking kv_cache.
curl -X POST "http://127.0.0.1:8000/sleep?level=2"
curl -X POST "http://127.0.0.1:8000/wake_up?tags=weights"
curl -X POST http://127.0.0.1:8000/collective_rpc \
    -H "Content-Type: application/json" \
    -d '{
        "method": "reload_weights",
        "kwargs": {"weights_path": "Qwen/Qwen2.5-0.5B-Instruct"}
    }'
curl -X POST "http://127.0.0.1:8000/wake_up?tags=kv_cache"
```