睡眠模式指南¶

概述¶

睡眠模式是一个旨在从 NPU 内存中卸载模型权重并丢弃 KV 缓存的 API。此功能对于强化学习（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={"enable_sleep_mode_extra_cleanup": True},
)

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

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

当启用 enable_sleep_mode_extra_cleanup 时，sleep() 还会额外执行以下操作：

当启用 ACL graph 时，清除 ACL graph attention 工作空间并使已捕获的 ACL graph 缓存失效；
重置模型运行器 graph 管理器，以便在唤醒后可以重新捕获 ACL graphs；
等待待处理的流水线并行发送工作完成，同步 NPU，并销毁 HCCL 进程组。

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

Note

额外清理以更长的唤醒延迟为代价换取更低的睡眠期 NPU 内存使用。特别是，如果启用了 ACL graph，wake_up() 必须在模型状态恢复后再次调用 capture_model()。当更低的唤醒延迟比释放 HCCL 和 ACL graph 工作空间内存更重要时，请保持 enable_sleep_mode_extra_cleanup 禁用。

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

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

启用额外清理后，仅当 tags 为 None 或包含 "kv_cache" 时才会重新捕获 ACL graphs。这避免了在外部重新加载的权重和 KV-cache 状态准备就绪之前重新捕获 graphs。

专家权重布局恢复¶

对于稠密模型，wake_up() 仅将模型权重恢复到 NPU 内存；张量布局保持不变。

对于**未量化的 MoE 模型**（quant_config is None），融合的专家权重以转置布局存储以提高 NPU matmul 效率。该布局在模型加载时由 process_weights_after_loading() 生成一次：加载权重后，该方法对 transpose(1, 2) 和 w13_weight 的第二个和第三个维度进行转置（w2_weight），将标准检查点布局转换为 torch_npu.npu_grouped_matmul 算子所需的格式。

在睡眠模式分配器恢复原始（未转置）内存后，wake_up() 在恢复 "weights" 标签时对受影响的专家权重重新应用相同的转置：

w13_weight（gate/up projection）：当其第二个维度匹配 hidden_size 时，转置回运行时布局；
w2_weight（down projection）：当其第三个维度匹配 hidden_size 时，转置回运行时布局。

对于稠密模型（没有专家权重）和量化模型（其权重由量化方法处理），此步骤完全跳过。

准备模型权重¶

使用 Qwen2.5-0.5B-Instruct 模型的权重。若设置 VLLM_USE_MODELSCOPE=True，模型将自动从模型Scope 下载。

: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"
os.environ["VLLM_ASCEND_ENABLE_NZ"] = "0"

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)
    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

Online serving: !!! note

    Considering there may be a risk of malicious access, please make sure you are under a dev-mode, and explicitly specify the dev environment `VLLM_SERVER_DEV_MODE` to expose these endpoints (sleep/wake up).

export VLLM_SERVER_DEV_MODE="1"
export VLLM_WORKER_MULTIPROC_METHOD="spawn"
export VLLM_USE_MODELSCOPE="True"
export VLLM_ASCEND_ENABLE_NZ="0"

vllm serve Qwen/Qwen2.5-0.5B-Instruct --enable-sleep-mode

# after serving is up, post to these endpoints

# sleep level 1
curl -X POST http://127.0.0.1:8000/sleep \
    -H "Content-Type: application/json" \
    -d '{"level": "1"}'

curl -X GET http://127.0.0.1:8000/is_sleeping

# sleep level 2
curl -X POST http://127.0.0.1:8000/sleep \
    -H "Content-Type: application/json" \
    -d '{"level": "2"}'

# wake up
curl -X POST http://127.0.0.1:8000/wake_up

# wake up with tag, tags must be in ["weights", "kv_cache"]
curl -X POST "http://127.0.0.1:8000/wake_up?tags=weights"

curl -X GET http://127.0.0.1:8000/is_sleeping

# after sleep and wake up, the serving is still available
curl http://localhost: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
    }'