290 lines
10 KiB
Markdown
290 lines
10 KiB
Markdown
# Sleep Mode & ACK Protocol
|
||
|
||
vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches—**without stopping the server or unloading the Docker container**.
|
||
|
||
This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html) and extended with the **Omni ACK Protocol** to support multi-stage pipelines and heterogeneous hardware backends (NVIDIA, AMD, Intel, Huawei). It is especially useful in **RLHF**, **dynamic model switching**, or **cost-saving scenarios**.
|
||
|
||
---
|
||
|
||
## 1. Feature Documentation
|
||
|
||
### Overview
|
||
|
||
Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization.
|
||
|
||
### Sleep Levels
|
||
|
||
We support two levels of hibernation to balance recovery speed and memory efficiency:
|
||
|
||
| Level | Name | Mechanism | Recovery Speed | Memory Freed |
|
||
| :--- | :--- | :--- | :--- | :--- |
|
||
| **Level 1** | **Weight Offloading** | Offloads weights to Host CPU RAM. | **Fast** (DMA) | Substantial |
|
||
| **Level 2** | **Full De-mapping** | Physically releases memory pages via VRAM scavenging. | **Moderate** | **Maximum** (up to 95%+) |
|
||
|
||
### Supported Platforms
|
||
|
||
Omni Sleep Mode is optimized for high-performance computing backends:
|
||
|
||
* **NVIDIA**: Supported via Virtual Memory Management (VMM).
|
||
* **AMD (ROCm)**: Fully supported with physical page de-mapping.
|
||
* **Intel XPU**: Supported via Level Zero memory management.
|
||
* **Huawei NPU**: Supported via Ascend memory scavenging.
|
||
|
||
### Hardware Requirements
|
||
|
||
* **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep.
|
||
* **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers.
|
||
|
||
### Quiescing a busy diffusion stage: `pause_generation(mode="keep")`
|
||
|
||
`sleep()` on a diffusion stage is a worker-level RPC: it offloads memory but does not stop the
|
||
stage's scheduler, so a stage that still has queued requests must be made quiet first. Aborting
|
||
discards the batch that is running. `pause_generation(mode="keep")` keeps it instead:
|
||
|
||
* The batch that is already running finishes on the weights it started with and its outputs are
|
||
delivered normally.
|
||
* Requests that have not started stay queued and run after `resume_generation()`, on whatever
|
||
weights are loaded by then. A request never continues denoising across a weight change.
|
||
* The call returns only after every replica has acknowledged: the batch that was running has
|
||
finished executing, its asynchronous output copies are done, and the device is synchronized; its
|
||
outputs are then delivered on the normal path. On a non-preemptible backend that acknowledgement
|
||
arrives at the next batch boundary, so expect it to take as long as the remaining batch.
|
||
* Control RPCs (`list_loras`, `sleep`, `wake_up`, `resume_generation`) keep working while paused.
|
||
|
||
Recommended order for a weight update on a diffusion stage:
|
||
|
||
```python
|
||
await engine.pause_generation(mode="keep", clear_cache=False)
|
||
await engine.sleep(level=1) # optional: free memory for the trainer
|
||
await engine.wake_up() # admission stays closed until resume
|
||
# install new weights here
|
||
await engine.resume_generation()
|
||
```
|
||
|
||
Support matrix and limits:
|
||
|
||
| Configuration | `mode="keep"` |
|
||
| :--- | :--- |
|
||
| Request-level execution (`step_execution=False`) with the in-process (`uni`) executor | Supported |
|
||
| Request-level execution with the multi-process (`mp`) executor, including asynchronous output | Supported; the acknowledgement waits for the output copies |
|
||
| Request-level execution with asynchronous KV prefetch | Supported; the acknowledgement joins any prefetch still in flight |
|
||
| Step-level execution (`step_execution=True` / streaming output) | Not supported: the call raises `NotImplementedError` |
|
||
| `mode="abort"` / `mode="wait"` / `wait_for_inflight_requests=True` on a diffusion stage | Unchanged: frontend admission is paused only |
|
||
|
||
Behavior notes:
|
||
|
||
* Callers that already passed `mode="keep"` for a diffusion stage used to get an immediate return
|
||
that only closed frontend admission. They now also stop the backend and wait for the ACK.
|
||
* A timeout or error from `pause_generation` means the pause did **not** complete; the scheduler
|
||
stays closed and the caller must retry the pause or call `resume_generation()` explicitly. Do not
|
||
sleep or mutate weights after a failed pause.
|
||
* `resume_generation(stage_ids=...)` that leaves another paused stage closed (an AR stage, or a
|
||
keep-paused diffusion stage) keeps frontend admission closed as well, so new requests cannot
|
||
queue on a stage that will not schedule them. Admission reopens once every stage paused by
|
||
`pause_generation` has been resumed.
|
||
* `clear_cache` has no diffusion-specific effect.
|
||
* Without a pause call nothing changes: no barrier runs and batching, abort and output ordering are
|
||
as before.
|
||
|
||
---
|
||
|
||
## 2. Usage Examples
|
||
|
||
### Python API Example
|
||
|
||
You can programmatically control the lifecycle of stages using the `AsyncOmni` engine.
|
||
|
||
```python
|
||
|
||
import asyncio
|
||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||
|
||
async def run_sleep_demo():
|
||
# 1. initialization
|
||
engine = AsyncOmni(
|
||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
||
enable_sleep_mode=True
|
||
)
|
||
|
||
# 2. sleep mode level2
|
||
acks = await engine.sleep(stage_ids=[0], level=2)
|
||
print(f"Freed {acks[0].freed_bytes / 1024**3:.2f} GiB on Stage 0")
|
||
|
||
# 3. wake up
|
||
await engine.wake_up(stage_ids=[0])
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(run_sleep_demo())
|
||
|
||
```
|
||
|
||
### server command Example
|
||
|
||
Start the server with sleep mode enabled:
|
||
|
||
The first method
|
||
|
||
```
|
||
|
||
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
|
||
--omni \
|
||
--enable-sleep-mode \
|
||
--trust-remote-code \
|
||
--gpu-memory-utilization 0.7
|
||
|
||
```
|
||
|
||
The second method
|
||
|
||
```
|
||
|
||
python3 -m vllm_omni.entrypoints.openai.api_server \
|
||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||
--omni \
|
||
--enable-sleep-mode \
|
||
--trust-remote-code \
|
||
--gpu-memory-utilization 0.7
|
||
|
||
```
|
||
|
||
### Test Scenarios & Commands
|
||
|
||
#### Scenario 1: LLM Engine Sleep
|
||
|
||
Objective: Verify VRAM reclamation for Stage 0 (Thinker).
|
||
|
||
Trigger sleep (Level 1 or Level 2) via client:
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [0], "level": 2}'
|
||
|
||
```
|
||
|
||
Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage.
|
||
|
||
#### Scenario 2: Diffusion Sleep
|
||
|
||
Objective: Verify VRAM reclamation for Stage 1 (Diffusion).
|
||
|
||
Trigger sleep (Level 1 or Level 2) via client:
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [1], "level": 2}'
|
||
|
||
```
|
||
|
||
#### Scenario 3: Multi-Stage Coordinated Stress Test
|
||
|
||
Objective: Test concurrent sleep and rapid wake-up across multiple stages.
|
||
|
||
Concurrent Sleep (Stage 0 & 1):
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [0, 1], "level": 2}'
|
||
|
||
```
|
||
|
||
Rapid Wake-up:
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/wakeup \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [0, 1]}'
|
||
|
||
```
|
||
|
||
#### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity
|
||
|
||
Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation.
|
||
|
||
Check Initial State: Observe baseline VRAM usage.
|
||
|
||
Trigger Deep Sleep (Level 2):
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/sleep \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [0], "level": 2}'
|
||
|
||
```
|
||
|
||
Wake-up Model:
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/omni/wakeup \
|
||
-H "Content-Type: application/json" \
|
||
-d '{"stage_ids": [0]}'
|
||
|
||
```
|
||
|
||
Verify Functional Integrity (Inference):
|
||
Ensure the model still generates valid output after reloading weights.
|
||
|
||
```
|
||
|
||
curl -X POST http://localhost:8000/v1/images/generations \
|
||
-H "Content-Type: application/json" \
|
||
-d '{
|
||
"prompt": "A huge swimming pool, with many people swimming.",
|
||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||
"response_format": "b64_json",
|
||
"extra_body": {"sampling_params": {"num_inference_steps": 4, "seed": 42}}
|
||
}' > post.json
|
||
|
||
```
|
||
|
||
## 3. API Reference
|
||
|
||
### Methods
|
||
|
||
| Method | Arguments | Return Type | Description |
|
||
| :--- | :--- | :--- | :--- |
|
||
| **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. |
|
||
| **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. |
|
||
| **pause_generation** | `mode: str, stage_ids: List[int]` | `None` | Stops admission; with `mode="keep"` also stops diffusion schedulers and returns after their ACK. |
|
||
| **resume_generation** | `stage_ids: List[int]` | `None` | Reopens paused schedulers, then admission. |
|
||
|
||
### OmniACK Dataclass Fields
|
||
|
||
| Field | Type | Description |
|
||
| :--- | :--- | :--- |
|
||
| **task_id** | `str` | Unique identifier for the operation. |
|
||
| **status** | `str` | `SUCCESS` or `ERROR`. |
|
||
| **stage_id** | `int` | The ID of the stage that responded. |
|
||
| **rank** | `int` | The rank ID within the Tensor Parallel group. |
|
||
| **freed_bytes** | `int` | Actual amount of physical VRAM reclaimed. |
|
||
| **metadata** | `dict` | Additional platform-specific metrics. |
|
||
|
||
Metadata Field Analysis
|
||
The metadata field is a dynamic dictionary containing hardware-specific telemetry and audit data, primarily used for verifying memory reclamation on various backends (e.g., AMD ROCm, NVIDIA CUDA).
|
||
|
||
```
|
||
"metadata": {
|
||
"source": "Platform_AMD_Instinct_MI300X",
|
||
"total_freed_gib": "78.57",
|
||
"rank_residual_gib": "2.07"
|
||
}
|
||
```
|
||
|
||
#### Core Utility
|
||
|
||
**VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM.
|
||
|
||
**Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines.
|
||
|
||
**Backend Traceability (source)**: Identifies the underlying hardware driver or audit source. This is critical for debugging synchronization issues in multi-stage, distributed environments.
|
||
|
||
**Performance Analytics (Roadmap)**: Future updates will include latency_ms (context-switch overhead) and cuda_graph_recalled (graph engine status) to optimize performance in high-frequency sleep/wake scenarios.
|