休眠模式指南

概述

休眠模式是一种用于从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图时，清除ACL图注意力工作空间并使捕获的ACL图缓存失效；

• 重置模型运行器图管理器，以便在唤醒后可以重新捕获ACL图；

• 等待待处理的流水线并行发送工作，同步NPU，并销毁HCCL进程组。

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

备注

额外清理以更长的唤醒延迟为代价，换取更低的休眠期NPU内存使用。特别是，如果启用了ACL图，wake_up()必须在模型状态恢复后再次调用capture_model()。当更低的唤醒延迟比释放HCCL和ACL图工作空间内存更重要时，请保持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图。这避免了在外部重新加载的权重和KV缓存状态就绪之前重新捕获图。

专家权重布局恢复

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

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

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

• w13_weight（门控/上投影）：当其第二维度与hidden_size匹配时，转置回运行时布局；

• w2_weight（下投影）：当其第三维度与hidden_size匹配时，转置回运行时布局。

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

准备模型权重

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

模型	ModelScope 链接
Qwen2.5-0.5B-Instruct	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
  

• 在线服务：

  备注

  考虑到可能存在恶意访问的风险，请确保您处于开发模式下，并明确指定开发环境变量 VLLM_SERVER_DEV_MODE 以暴露这些端点（睡眠/唤醒）。

  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
      }'