Sleep Mode Guide

Overview

Sleep Mode is an API designed to offload model weights and discard KV cache from NPU memory. This functionality is essential for reinforcement learning (RL) post-training workloads, particularly in online algorithms such as PPO, GRPO, or DPO. During training, the policy model typically performs autoregressive generation using inference engines like vLLM, followed by forward and backward passes for optimization.

Since the generation and training phases may employ different model parallelism strategies, it becomes crucial to free KV cache and even offload model parameters stored within vLLM during training. This ensures efficient memory utilization and avoids resource contention on the NPU.

Getting started

With enable_sleep_mode=True, the way we manage memory (malloc, free) in vLLM is under a specific memory pool. During model loading and KV cache initialization, we tag the memory as a map: {"weight": data, "kv_cache": data}.

The engine (v0/v1) supports two sleep levels to manage memory during idle periods:

• Level 1 Sleep

  • Action: Offloads model weights and discards the KV cache.
  • Memory: Model weights are moved to CPU memory; KV cache is forgotten.
  • Use Case: Suitable when reusing the same model later.
  • Note: Ensure sufficient CPU memory is available to hold the model weights.

• Level 2 Sleep

  • Action: Discards both model weights and KV cache.
  • Memory: The content of both the model weights and KV cache is forgotten.
  • Use Case: Ideal when switching to a different model or updating the current one.

Since this feature uses the low-level API AscendCL, in order to use sleep mode, you should follow the installation guide and build from source. If you are using < v0.12.0rc1, remember to set export COMPILE_CUSTOM_KERNELS=1.

Optional extra cleanup

By default, sleep mode only releases memory managed by the sleep-mode allocator. For RL workloads that need to return more NPU memory to the trainer, vLLM Ascend also provides an optional extra cleanup path:

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

For online serving, pass the same option through --additional-config:

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

When enable_sleep_mode_extra_cleanup is enabled, sleep() additionally:

• clears ACL graph attention workspaces and invalidates captured ACL graph caches when ACL graph is enabled;
• resets the model runner graph manager so ACL graphs can be captured again after wakeup;
• waits for pending pipeline-parallel send work, synchronizes the NPU, and destroys HCCL process groups.

During wake_up(), vLLM Ascend restores the HCCL process groups, refreshes MoE dispatcher HCCL metadata, restores sleep-mode allocator memory, and recaptures ACL graphs when needed.

Note

Extra cleanup trades lower sleep-time NPU memory usage for longer wakeup latency. In particular, if ACL graph is enabled, wake_up() must call capture_model() again after the model state has been restored. Keep enable_sleep_mode_extra_cleanup disabled when lower wakeup latency is more important than releasing HCCL and ACL graph workspace memory.

For level 2 sleep, wakeup can be split into two phases:

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

With extra cleanup enabled, ACL graphs are recaptured only when tags is None or contains "kv_cache". This avoids recapturing graphs before externally reloaded weights and KV-cache state are ready.

Prepare Model Weights

Use the Qwen2.5-0.5B-Instruct model weights. With VLLM_USE_MODELSCOPE=True, the model will be downloaded automatically from 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)

Usage

The following is a simple example of how to use sleep mode.

• Offline inference:

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