diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 31a2670ba..86a0e9201 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -179,7 +179,134 @@ class Flux2Transformer2DModel(nn.Module): - Support single GPU only for now -### Implementation Notes +## Distributed Layerwise Offloading + +### How It Works + +Distributed layerwise offloading extends single-GPU layerwise offloading to +multi-device deployments. Each DP rank stores only **1/dp_size** of the model +weights in host memory; full layer weights are reconstructed at runtime via +**AllGather** on a dedicated communication stream, overlapped with computation +via a fixed double-buffer scheme. + +**Key features:** +- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights, + reconstructed per-layer via `all_gather_into_tensor` +- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any + time, regardless of model size +- **DP multi-concurrency**: N concurrent requests processed in parallel + (AllGather only gathers weight shards, which are request-independent) +- **mmap weight loading**: weights loaded as mmap views pointing to shared OS + page cache, eliminating O(dp_size × model_size) RSS during model creation +- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU + (CANN/HCCL) via vLLM-Omni's platform abstraction + +**Execution Flow (per device):** + +``` +Compute Stream: [Layer N] [Layer N+1] [Layer N+2] +H2D Stream: [H2D Shard N+1] [H2D Shard N+2] +AllGather: [AG N+1] [AG N+2] +Slot Usage: Slot0: Layer N Slot1: Layer N+1 +``` + +### Usage + +**CLI:** +```bash +# 4× GPU/NPU with AllGather (recommended) +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 + +# Without AllGather (each rank loads full weights, no sharding) +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 \ + --dlo-no-use-allgather + +# With SP instead of DP (long sequences) +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --usp 4 +``` + +**Python API:** +```python +from vllm_omni import Omni + +m = Omni( + model="/path/to/model", + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, # default +) +``` + +### CLI Flags + +| Flag | Description | Default | +|------|-------------|---------| +| `--enable-distributed-layerwise-offload` | Enable DLO | `false` | +| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` | +| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` | +| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` | + +### How model weights are loaded (mmap path) + +When DLO + AllGather is active, the offloader: + +1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`) + from the normally-created transformer +2. **Converts to meta device** via `to_empty(device="meta")`, releasing random + initialization weights +3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`, + which return views into the OS page cache (shared across all ranks, 0 RSS) +4. **Calls `post_load_weights()`** to apply model-specific dtype conversions + (e.g. Cosmos3's `time_embedder` → FP32) +5. **Restores non-persistent buffers** from saved copies + +This approach requires **zero model-specific code changes** — no pipeline or +transformer modifications are needed. + +### OffloadPlan (declarative topology metadata) + +Models can optionally declare an `OffloadPlan` class variable to provide +topology metadata (block attribute names, submodules to offload) without +any offload-specific logic: + +```python +from vllm_omni.diffusion.offloader import OffloadPlan + +class MyPipeline(nn.Module): + _dit_modules = ["transformer"] + _offload_plan = OffloadPlan( + block_attrs={"transformer": ("blocks",)}, + offload_submodules={"context_encoder": "layers"}, + ) +``` + +When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs` +and heuristic attribute search. This is backward-compatible — existing models +work without any changes. + +### DP Multi-concurrency + +When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches +up to `dp_size` requests per denoise step. Each DP rank processes a different +request while AllGather synchronizes only weight shards (request-independent). + +**Requirements for concurrent requests:** +- `num_inference_steps` must be specified explicitly (None is not allowed) +- All concurrent requests must have the same `num_inference_steps` value +- These constraints exist because AllGather is a collective that requires + every rank to participate at each denoise step + +### Limitations +- Online quantization (FP8) is incompatible with mmap loading — falls back to + regular `load_weights()` automatically +- Tensor Parallel is not supported (DLO uses DP-based sharding) +- HSDP + AllGather is rejected (would double-shard weights) +- `num_inference_steps=None` is not allowed in DP multi-concurrency mode **Module Discovery** @@ -208,7 +335,10 @@ OffloadBackend (base class) ├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap) │ (delegates to a pipeline's enable_omni_model_cpu_offload │ for split models like Cosmos3) -└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook +├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook +│ (single-GPU, full weights on host) +└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook + (multi-GPU, 1/dp_size sharded weights + AllGather) ``` Factory function `get_offload_backend()` selects the appropriate backend based on @@ -221,20 +351,22 @@ reasoner/generator components inside the model forward pass. ## Supported Models -| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) | -|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------| -| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` | -| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` | -| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` | -| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` | -| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` | -| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` | -| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` | -| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | -| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | -| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` | -| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` | +| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) | +|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------| +| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | +| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` | +| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` | +| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` | +| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` | +| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` | +| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` | +| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | +| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` | +| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` | +| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` | +| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` | **Notes:** - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks +- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples. diff --git a/recipes/cosmos3/Cosmos3-DistOffload.md b/recipes/cosmos3/Cosmos3-DistOffload.md new file mode 100644 index 000000000..27d1c77bd --- /dev/null +++ b/recipes/cosmos3/Cosmos3-DistOffload.md @@ -0,0 +1,191 @@ +# Cosmos3 — Distributed Layerwise Offload + +> Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment + +## When to use this recipe + +Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency. + +**Key benefits:** +- Each rank stores only 1/dp_size of the model weights on host memory +- Only 2 transformer blocks reside on each device at any time (fixed double-buffer) +- DP multi-concurrency: N concurrent requests with near-linear throughput scaling +- mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS + +## Prerequisites + +- 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.) +- vLLM-Omni with distributed layerwise offload support +- Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID) + +## Serve commands + +### Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput) + +```bash +vllm serve /path/to/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 +``` + +### Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM) + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 2 +``` + +### Cosmos3-Super with DP=4 + AllGather + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 +``` + +### Without AllGather (full weights per rank, no sharding) + +Use `--dlo-no-use-allgather` when you want each rank to load full weights independently — +no AllGather synchronization, no concurrent request requirement, but N× host memory. + +```bash +vllm serve /path/to/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --data-parallel-size 4 \ + --dlo-no-use-allgather +``` + +### With SP (sequence parallel) instead of DP + +For long-sequence workloads where SP parallelism is preferred: + +```bash +vllm serve /path/to/Cosmos3-Super \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --no-guardrails --init-timeout 3600 \ + --vae-use-tiling --vae-patch-parallel-size 1 \ + --enable-distributed-layerwise-offload \ + --usp 4 +``` + +## Request examples + +### Text-to-Image (T2I) + +```bash +curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \ + -H "Content-Type: multipart/form-data" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, cinematic" \ + -F "size=1024x1024" \ + -F "num_inference_steps=50" \ + -F "guidance_scale=6.0" \ + -F "seed=42" +``` + +### Text-to-Video (T2V) + +```bash +curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \ + -H "Accept: video/mp4" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, cinematic shot" \ + -F "negative_prompt=blurry" \ + -F "size=832x480" \ + -F "num_frames=29" \ + -F "fps=24" \ + -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" \ + -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" +``` + +> **Important:** When using DP multi-concurrency (AllGather mode), `num_inference_steps` +> must be specified explicitly and must be the same for all concurrent requests. +> `num_inference_steps=None` is not allowed because it may resolve differently +> per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather. + +### Multiple concurrent requests (DP multi-concurrency) + +With `--data-parallel-size 4`, send 4 concurrent requests to maximize throughput: + +```bash +for i in 1 2 3 4; do + curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \ + -H "Accept: video/mp4" \ + -F "model=/path/to/Cosmos3-Nano" \ + -F "prompt=A robot arm cleaning a plate, shot ${i}" \ + -F "negative_prompt=blurry" \ + -F "size=832x480" \ + -F "num_frames=29" \ + -F "fps=24" \ + -F "num_inference_steps=35" \ + -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" \ + -F "flow_shift=10.0" \ + -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \ + -F "seed=17" & +done +wait +``` + +## CLI flags + +| Flag | Description | Default | +|------|-------------|---------| +| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` | +| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` | +| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` | +| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` | +| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` | + +## Mode comparison + +| Mode | CLI flags | CPU/rank | HBM/card | Throughput | Use case | +|------|----------|---------|---------|-----------|----------| +| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences | +| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory | +| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) | +| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU | +| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM | + +## Memory expectations + +| Model | Config | Host RAM (cgroup) | HBM/card | +|-------|--------|-------------------|----------| +| Nano (33 GB) | DP2 + AllGather | ~38 GB | ~10 GB | +| Nano (33 GB) | DP4 + AllGather | ~47 GB | ~10 GB | +| Super (124 GB) | DP2 + AllGather | ~157 GB | ~15 GB | +| Super (124 GB) | DP4 + AllGather | ~172 GB | ~10 GB | + +> Host RAM is dominated by page cache (1× model_size, shared across ranks). +> HBM only holds 2 transformer blocks + framework + activations, independent of model size. + +## Limitations + +- `--dlo-use-allgather` (default) requires all concurrent requests to have the same + `num_inference_steps` (AllGather is a collective that needs all ranks synchronized) +- Online quantization (FP8) is incompatible with mmap loading — falls back to regular + `load_weights()` automatically +- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension) +- HSDP + AllGather is rejected (would double-shard weights) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py new file mode 100644 index 000000000..c02fdaabc --- /dev/null +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -0,0 +1,1027 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for DistributedLayerwiseOffloadHook and backend utilities.""" + +import gc +import json +import os +import socket +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from safetensors.torch import save_file +from torch import nn +from torch.distributed.tensor import DeviceMesh, DTensor, Replicate + +import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.block_discovery import ( + get_blocks_attr_names, + get_blocks_from_dit, + set_blocks_attr_names, +) +from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( + DistributedLayerwiseOffloadBackend, + DistributedLayerwiseOffloadHook, +) +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class DummyStream: + def wait_stream(self, _stream) -> None: + return None + + def wait_event(self, _event) -> None: + return None + + +class DummyEvent: + def record(self, _stream) -> None: + return None + + +@contextmanager +def dummy_stream(_stream): + yield None + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + +def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None: + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + + +def _cleanup_distributed() -> None: + if dist.is_initialized(): + dist.destroy_process_group() + + for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]: + os.environ.pop(key, None) + + gc.collect() + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + + +@pytest.fixture(scope="module") +def dist_group(): + master_port = _find_free_port() + _set_dist_env(rank=0, world_size=1, master_port=master_port) + + dist.init_process_group("gloo", rank=0, world_size=1) + try: + yield + finally: + _cleanup_distributed() + + +@pytest.fixture +def patched_offload_runtime(monkeypatch): + monkeypatch.setattr(dist_backend_module.current_omni_platform, "Stream", DummyStream) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream()) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream) + + +class TinyBlock(nn.Module): + def __init__(self, values: torch.Tensor): + super().__init__() + mesh = DeviceMesh("cpu", [0]) + dtensor = DTensor.from_local(values, mesh, [Replicate()]) + self.weight = nn.Parameter(dtensor) + + +def _make_values(start: float) -> torch.Tensor: + return torch.arange(start, start + 4, dtype=torch.float32) + + +class TestDistributedLayerwiseOffloadHook: + def test_shard_and_pin_single_rank(self, dist_group, patched_offload_runtime): + """With dp_size=1, the shard should equal the full weights.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + + hook.initialize_hook(current_block) + + # With dp_size=1, the CPU shard should contain all 4 elements + assert next_block.weight.dtype in hook.cpu_shards + shard = hook.cpu_shards[next_block.weight.dtype] + assert shard.numel() == 4 + assert torch.equal(shard, _make_values(10.0)) + + def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime): + """DTensor wrapper should be preserved through prefetch/offload cycle.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + + hook.initialize_hook(current_block) + + # After init, next_block weights should be placeholders + assert isinstance(next_block.weight, DTensor) + assert next_block.weight.to_local().is_meta + + # Prefetch into slot 0 + hook.prefetch_layer(slot=0, non_blocking=False) + + # After prefetch, next_block weights should be materialized + assert isinstance(next_block.weight, DTensor) + assert torch.equal(next_block.weight.to_local(), _make_values(10.0)) + + # Offload current block + hook.offload_layer() + assert isinstance(current_block.weight, DTensor) + assert current_block.weight.to_local().is_meta + assert not hook.is_materialized + + def test_double_buffer_slot_swapping(self, dist_group, patched_offload_runtime): + """Verify slot swapping works correctly after each layer.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + + hook.initialize_hook(current_block) + + assert hook.current_slot == 0 + + # Prefetch into slot 1 (next slot) + hook.prefetch_layer(slot=1, non_blocking=False) + + # Simulate post_forward: offload + swap + hook.offload_layer() + hook.current_slot = 1 - hook.current_slot + + assert hook.current_slot == 1 + + # Next iteration: prefetch into slot 0 + hook.prefetch_layer(slot=0, non_blocking=False) + hook.offload_layer() + hook.current_slot = 1 - hook.current_slot + + assert hook.current_slot == 0 + + def test_device_buffers_allocated(self, dist_group, patched_offload_runtime): + """Verify exactly two device buffers are allocated.""" + current_block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(10.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + + hook.initialize_hook(current_block) + + # Two slots should be allocated + assert hook.gpu_buffers[0] is not None + assert hook.gpu_buffers[1] is not None + + # Each slot should have one dtype entry (float32) + assert len(hook.gpu_buffers[0]) == 1 + assert len(hook.gpu_buffers[1]) == 1 + + # Buffer size should match total numel + dtype = next_block.weight.dtype + assert hook.gpu_buffers[0][dtype].numel() == 4 + assert hook.gpu_buffers[1][dtype].numel() == 4 + + def test_sharding_multiple_ranks(self, dist_group, patched_offload_runtime): + """Verify weight sharding splits correctly across ranks.""" + block = nn.Module() + block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32)) + + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(100, 108, dtype=torch.float32)) + + # Simulate dp_size=4, rank=1 + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=4, + rank=1, + pin_memory=False, + ) + + hook.initialize_hook(block) + + shard = hook.cpu_shards[torch.float32] + # 8 elements / 4 ranks = 2 elements per shard + assert shard.numel() == 2 + # Rank 1 should have elements [102, 103] + assert torch.equal(shard, torch.tensor([102.0, 103.0])) + + def test_sharding_with_remainder(self, dist_group, patched_offload_runtime): + """Verify sharding handles non-even division with padding.""" + block = nn.Module() + block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + + # 3 elements, dp_size=2: both ranks get ceil(3/2)=2 elements (padded) + hook0 = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + # Need fresh next_block for each hook since init replaces tensors + next_block0 = nn.Module() + next_block0.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + hook0.next_block = next_block0 + hook0.initialize_hook(block) + + shard0 = hook0.cpu_shards[torch.float32] + assert shard0.numel() == 2 + assert torch.equal(shard0, torch.tensor([100.0, 101.0])) + + next_block1 = nn.Module() + next_block1.weight = nn.Parameter(torch.arange(100, 103, dtype=torch.float32)) + block1 = nn.Module() + block1.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + hook1 = DistributedLayerwiseOffloadHook( + next_block=next_block1, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=1, + pin_memory=False, + ) + hook1.initialize_hook(block1) + + shard1 = hook1.cpu_shards[torch.float32] + # Equal-sized shards: rank 1 gets [102, 0] (zero-padded) + assert shard1.numel() == 2 + assert torch.equal(shard1, torch.tensor([102.0, 0.0])) + + +class _DummyBlock(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.randn(10, 10)) + + +class _SingleBlockModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self, num_blocks: int = 3): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) + + +class _MultiBlockModel(nn.Module): + _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] + + def __init__(self, num_transformer: int = 2, num_single: int = 2): + super().__init__() + self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)]) + self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)]) + + +class _EmptyBlocksModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([]) + + +class _NoAttrsModel(nn.Module): + def __init__(self, num_blocks: int = 2): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)]) + + +class _MmapPostLoadModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.time_embedder = nn.Linear(2, 2, bias=False) + self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False) for _ in range(2)]) + self.post_load_calls = 0 + + def post_load_weights(self) -> None: + self.time_embedder.to(torch.float32) + self.post_load_calls += 1 + + +class _MmapPostLoadPipeline(nn.Module): + def __init__(self): + super().__init__() + with torch.device("meta"): + self.transformer = _MmapPostLoadModel() + + @staticmethod + def _remap_ckpt_key(key: str) -> str: + return key + + +class TestMmapWeightLoading: + def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime): + pipeline = _MmapPostLoadPipeline() + weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()} + weight_file = tmp_path / "model.safetensors" + save_file(weights, str(weight_file)) + weight_map = {name: weight_file.name for name in weights} + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + model_path=str(tmp_path), + ), + torch.device("cpu"), + ) + modules = SimpleNamespace( + dits=[pipeline.transformer], + dit_names=["transformer"], + ) + + backend._load_weights_via_mmap(pipeline, modules) + + assert pipeline.transformer.post_load_calls == 1 + assert pipeline.transformer.time_embedder.weight.dtype == torch.float32 + assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16 + + +class TestGetBlocksFromDit: + def test_get_blocks_from_dit_single_block_attr(self): + model = _SingleBlockModel(num_blocks=3) + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == ["blocks"] + assert len(blocks) == 3 + assert all(isinstance(b, _DummyBlock) for b in blocks) + + def test_get_blocks_from_dit_multi_block_attrs(self): + model = _MultiBlockModel(num_transformer=2, num_single=3) + attr_names, blocks = get_blocks_from_dit(model) + assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"} + assert len(blocks) == 5 + + def test_get_blocks_from_dit_empty_blocks(self): + model = _EmptyBlocksModel() + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == [] + assert blocks == [] + + def test_get_blocks_from_dit_no_attrs_defined(self): + model = _NoAttrsModel(num_blocks=3) + attr_names, blocks = get_blocks_from_dit(model) + assert attr_names == [] + assert blocks == [] + + +class TestGetBlocksAttrNames: + def test_get_blocks_attr_names_new_format(self): + model = _MultiBlockModel() + attrs = get_blocks_attr_names(model) + assert attrs == ["transformer_blocks", "single_transformer_blocks"] + + def test_get_blocks_attr_names_no_attrs(self): + model = _NoAttrsModel() + attrs = get_blocks_attr_names(model) + assert attrs == [] + + def test_set_blocks_attr_names(self): + model = _NoAttrsModel() + set_blocks_attr_names(model, ["new_blocks"]) + assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs") + assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"] + + +class TestCrossGroupSharedBuffer: + """Regression test: when multiple DiT groups share the same 2 GPU + buffers, the first block of each group must always sync-prefetch on + entry, because another group may have overwritten the shared slot. + """ + + def test_group_first_hook_forces_sync_prefetch(self, dist_group, patched_offload_runtime): + """Verify _is_group_first=True forces sync-prefetch in pre_forward + even when is_materialized returns True.""" + block_a = TinyBlock(_make_values(1.0)) + block_b = TinyBlock(_make_values(10.0)) + + hook_a = DistributedLayerwiseOffloadHook( + next_block=block_b, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + ) + hook_a.initialize_hook(block_a) + + # Prefetch to make block_a "materialized" (non-empty params) + hook_a.prefetch_layer(hook_a.current_slot, non_blocking=False) + assert hook_a.is_materialized + + # Track sync-prefetch calls + sync_called = [False] + orig = hook_a.prefetch_layer + + def tracking(slot, non_blocking=True): + if not non_blocking: + sync_called[0] = True + orig(slot, non_blocking=non_blocking) + + hook_a.prefetch_layer = tracking + + # Case 1: _is_group_first=False, materialized → no sync-prefetch + hook_a._is_group_first = False + sync_called[0] = False + hook_a.pre_forward(block_a) + assert not sync_called[0], "Materialized block should not sync-prefetch without _is_group_first" + + # Case 2: _is_group_first=True, materialized, slot contaminated + # (no _shared_slot_group set → defaults to contaminated) → MUST sync-prefetch + hook_a._is_group_first = True + hook_a._prev_hook = hook_a + sync_called[0] = False + hook_a.pre_forward(block_a) + assert sync_called[0], "_is_group_first must force sync-prefetch when slot is contaminated" + + # Case 3: _is_group_first=True, slot NOT contaminated (same group owns slot) + # → skip sync-prefetch, just wait for async event + hook_a._shared_slot_group = [hook_a._group_id, -1] + sync_called[0] = False + hook_a.pre_forward(block_a) + assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)" + + +class TestOffloadPlan: + """Test declarative OffloadPlan metadata.""" + + def test_get_offload_plan_returns_none_when_not_declared(self): + """Pipelines without _offload_plan should return None.""" + model = _SingleBlockModel(num_blocks=3) + assert get_offload_plan(model) is None + + def test_get_offload_plan_returns_declared_plan(self): + """Pipelines with _offload_plan should return it.""" + plan = OffloadPlan( + block_attrs={"transformer": ("gen_layers",)}, + offload_submodules={"context_encoder": "layers"}, + ) + + class PipelineWithPlan(nn.Module): + _offload_plan = plan + + model = PipelineWithPlan() + result = get_offload_plan(model) + assert result is plan + assert result.block_attrs == {"transformer": ("gen_layers",)} + assert result.offload_submodules == {"context_encoder": "layers"} + + def test_offload_plan_defaults_to_empty(self): + """OffloadPlan with no arguments should have empty dicts.""" + plan = OffloadPlan() + assert plan.block_attrs == {} + assert plan.offload_submodules == {} + + def test_offload_plan_is_frozen(self): + """OffloadPlan should be immutable (frozen=True).""" + plan = OffloadPlan() + with pytest.raises(Exception): + plan.block_attrs = {"x": ("y",)} # type: ignore + + +class TestMmapValidation: + """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" + + def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch): + """DLO+AllGather should reject when TP world size > 1.""" + import json + from types import SimpleNamespace + + from safetensors.torch import save_file + + from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy + + # Mock TP world size = 2 + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", + lambda: 2, + ) + + class SimpleModel(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)]) + + class SimplePipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = SimpleModel() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = SimplePipeline() + weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta} + save_file(weights, str(tmp_path / "model.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) + ) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + model_path=str(tmp_path), + ), + torch.device("cpu"), + ) + modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + + with pytest.raises(ValueError, match="Tensor Parallel"): + backend._load_weights_via_mmap(pipeline, modules) + + def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch): + """DLO+AllGather should work when TP world size = 1, even if params + have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1).""" + import json + from types import SimpleNamespace + + from safetensors.torch import save_file + + from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy + + # Mock TP world size = 1 (default) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", + lambda: 1, + ) + + def custom_weight_loader(param, weight): + param.data.copy_(weight) + + class ModelWithCustomLoader(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4, bias=False) + # Simulate QKVParallelLinear: custom weight_loader even at TP=1 + self.linear.weight.weight_loader = custom_weight_loader + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False)]) + + class PipelineWithLoader(nn.Module): + def __init__(self): + super().__init__() + self.transformer = ModelWithCustomLoader() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = PipelineWithLoader() + weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta} + save_file(weights, str(tmp_path / "model.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) + ) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + model_path=str(tmp_path), + ), + torch.device("cpu"), + ) + modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + + # Should NOT raise — TP=1 is allowed even with custom weight_loader + backend._load_weights_via_mmap(pipeline, modules) + + def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime): + """validate_loaded_weights should be called after mmap loading.""" + import json + from types import SimpleNamespace + + from safetensors.torch import save_file + + from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy + + class ModelWithValidate(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + validate_called = False + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(4, 4, bias=False) for _ in range(2)]) + + def post_load_weights(self): + pass + + def validate_loaded_weights(self, loaded): + self.validate_called = True + assert len(loaded) > 0, "Should have loaded weights" + + class PipelineWithValidate(nn.Module): + def __init__(self): + super().__init__() + self.transformer = ModelWithValidate() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = PipelineWithValidate() + # Save on meta to trigger mmap path + with torch.device("meta"): + pipeline.transformer = ModelWithValidate() + + weights = { + name: torch.ones(param.shape, dtype=param.dtype if not param.is_meta else torch.float32) + for name, param in pipeline.named_parameters() + } + save_file(weights, str(tmp_path / "model.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) + ) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path) + ), + torch.device("cpu"), + ) + modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + + backend._load_weights_via_mmap(pipeline, modules) + assert pipeline.transformer.validate_called, "validate_loaded_weights should be called" + + +class TestConfigValidation: + """Tests for configuration validation in OffloadConfig / execute_request.""" + + def test_hsdp_with_allgather_rejected(self): + """HSDP + DLO + AllGather should raise ValueError (double sharding).""" + from vllm_omni.diffusion.offloader.base import OffloadConfig + + class FakePC: + data_parallel_size = 1 + use_hsdp = True + hsdp_shard_size = 2 + hsdp_replicate_size = 1 + sequence_parallel_size = 1 + tensor_parallel_size = 1 + cfg_parallel_size = 1 + pipeline_parallel_size = 1 + ulysses_degree = 1 + ring_degree = 1 + allgather_degree = 1 + enable_expert_parallel = False + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = True + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + with pytest.raises(ValueError, match="incompatible with HSDP"): + OffloadConfig.from_od_config(FakeODConfig()) + + def test_hsdp_without_allgather_allowed(self): + """HSDP + DLO + no-AllGather should be allowed (full weights per rank).""" + from vllm_omni.diffusion.offloader.base import OffloadConfig + + class FakePC: + data_parallel_size = 1 + use_hsdp = True + hsdp_shard_size = 2 + hsdp_replicate_size = 1 + sequence_parallel_size = 1 + tensor_parallel_size = 1 + cfg_parallel_size = 1 + pipeline_parallel_size = 1 + ulysses_degree = 1 + ring_degree = 1 + allgather_degree = 1 + enable_expert_parallel = False + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = False # no AllGather → should be allowed + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + config = OffloadConfig.from_od_config(FakeODConfig()) + assert config.dp_size == 1 # forced to 1 when no AllGather + + def test_num_inference_steps_none_rejected(self): + """DP multi-concurrency should reject None num_inference_steps.""" + from types import SimpleNamespace + + # Mock requests with None steps + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id=f"req-{i}", + sampling_params=SimpleNamespace(num_inference_steps=None), + ) + ) + for i in range(2) + ] + + # We can't easily instantiate the full executor, but we can test + # the validation logic by checking that the code path raises. + # The validation is in execute_request, which needs self._ensure_open(). + # Instead, test the validation logic directly: + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + assert has_none, "Test setup: should have None steps" + assert len(step_counts) == 0, "Test setup: no explicit steps" + + # The validation condition: (len(step_counts) > 1) or has_none → should reject + should_reject = (len(step_counts) > 1) or has_none + assert should_reject, "None steps should trigger rejection" + + def test_num_inference_steps_same_explicit_allowed(self): + """DP multi-concurrency should allow same explicit num_inference_steps.""" + from types import SimpleNamespace + + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id=f"req-{i}", + sampling_params=SimpleNamespace(num_inference_steps=35), + ) + ) + for i in range(4) + ] + + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + + should_reject = (len(step_counts) > 1) or has_none + assert not should_reject, "Same explicit steps should be allowed" + assert step_counts == {35} + + def test_num_inference_steps_different_explicit_rejected(self): + """DP multi-concurrency should reject different explicit steps.""" + from types import SimpleNamespace + + reqs = [ + SimpleNamespace( + req=SimpleNamespace( + request_id="req-0", + sampling_params=SimpleNamespace(num_inference_steps=35), + ) + ), + SimpleNamespace( + req=SimpleNamespace( + request_id="req-1", + sampling_params=SimpleNamespace(num_inference_steps=30), + ) + ), + ] + + step_counts = { + r.req.sampling_params.num_inference_steps + for r in reqs + if r.req.sampling_params.num_inference_steps is not None + } + has_none = any(r.req.sampling_params.num_inference_steps is None for r in reqs) + + should_reject = (len(step_counts) > 1) or has_none + assert should_reject, "Different steps should trigger rejection" + assert len(step_counts) == 2 + + +class TestDynamicSlotTracking: + """Tests for dynamic slot tracking with odd and even block counts. + + With i%2 initialization alone, odd block counts cause the circular + tail→head prefetch to collide with the head's read slot (both map + to slot 0). Dynamic slot tracking via _prefetched_slot corrects + this by reading the actual slot from the previous hook's prefetch. + """ + + @staticmethod + def _trace_slots(num_blocks: int, num_iterations: int = 1): + """Trace slot assignments through forward passes (pure logic, no hooks). + + Returns a list of (read_slot, write_slot) per block per iteration. + """ + # Initial i%2 assignment (same as enable()) + current_slots = [i % 2 for i in range(num_blocks)] + prefetched_slots: list[int | None] = [None] * num_blocks + + # Sync-prefetch: tail hook prefetches into slot 0 + prefetched_slots[-1] = 0 + + results = [] + for _ in range(num_iterations): + for i in range(num_blocks): + prev_idx = (i - 1) % num_blocks + + # Dynamic slot tracking: read from prev's prefetched slot + if prefetched_slots[prev_idx] is not None: + current_slots[i] = prefetched_slots[prev_idx] + + read_slot = current_slots[i] + + # Prefetch into other slot + next_slot = 1 - read_slot + current_slots[i] = next_slot + prefetched_slots[i] = next_slot + + results.append((read_slot, next_slot)) + + return results + + def test_odd_block_count_no_collision(self): + """3 blocks (odd): no read/write slot collision.""" + results = self._trace_slots(3) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 3} (step {i}): read={read} == write={write}" + + def test_even_block_count_no_collision(self): + """4 blocks (even): no read/write slot collision.""" + results = self._trace_slots(4) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 4} (step {i}): read={read} == write={write}" + + def test_large_odd_block_count(self): + """7 blocks (odd): no collision across all blocks.""" + results = self._trace_slots(7) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 7} (step {i}): read={read} == write={write}" + + def test_5_blocks_two_iterations(self): + """5 blocks (odd), 2 iterations: no collision in either pass.""" + results = self._trace_slots(5, num_iterations=2) + n = 5 + for iteration in range(2): + for i in range(n): + idx = iteration * n + i + read, write = results[idx] + assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}" + + def test_3_blocks_three_iterations(self): + """3 blocks (odd), 3 iterations: no collision across all passes.""" + results = self._trace_slots(3, num_iterations=3) + n = 3 + for iteration in range(3): + for i in range(n): + idx = iteration * n + i + read, write = results[idx] + assert read != write, f"Iter {iteration} block {i}: read={read} == write={write}" + + def test_72_blocks_cosmos3(self): + """72 blocks (Cosmos3, even): no collision.""" + results = self._trace_slots(72) + for i, (read, write) in enumerate(results): + assert read != write, f"Block {i % 72} (step {i}): read={read} == write={write}" + + def test_slots_alternate_within_iteration(self): + """Within one iteration, consecutive blocks should use alternating slots.""" + results = self._trace_slots(6) + for i in range(6): + read, write = results[i] + # Each block reads from one slot and writes to the other + assert {read, write} == {0, 1}, f"Block {i}: slots {read},{write} don't cover both" + + def test_prefetched_slot_set_after_prefetch(self, dist_group, patched_offload_runtime): + """prefetch_layer should set _prefetched_slot to the slot argument.""" + block = TinyBlock(_make_values(1.0)) + next_block = TinyBlock(_make_values(2.0)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook.initialize_hook(block) + + assert hook._prefetched_slot is None # initially None + + hook.prefetch_layer(slot=0, non_blocking=False) + assert hook._prefetched_slot == 0 + + hook.prefetch_layer(slot=1, non_blocking=False) + assert hook._prefetched_slot == 1 + + def test_dynamic_slot_overrides_initial_assignment(self, dist_group, patched_offload_runtime): + """When _prev_hook._prefetched_slot is set, pre_forward should use + that value instead of the initial i%2 assignment.""" + block_a = TinyBlock(_make_values(1.0)) + block_b = TinyBlock(_make_values(2.0)) + + hook_a = DistributedLayerwiseOffloadHook( + next_block=block_b, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook_a.initialize_hook(block_a) + + hook_b = DistributedLayerwiseOffloadHook( + next_block=block_a, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + copy_stream=DummyStream(), + comm_stream=DummyStream(), + pin_memory=False, + ) + hook_b.initialize_hook(block_b) + + # Wire: hook_b._prev_hook = hook_a + hook_b._prev_hook = hook_a + + # hook_a prefetched into slot 1 + hook_a.prefetch_layer(slot=1, non_blocking=False) + assert hook_a._prefetched_slot == 1 + + # hook_b's initial slot is 0 (from i%2) + hook_b.current_slot = 0 + + # pre_forward should override to slot 1 + hook_b.pre_forward(block_b) + assert hook_b.current_slot == 1, "pre_forward should read _prev_hook._prefetched_slot=1, not keep initial 0" diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 192675908..e26211e9c 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -92,6 +92,21 @@ class MockScheduler: # assume all new req finished return [req.request_id for req in sched_output.scheduled_new_reqs] + def get_request_state(self, request_id): + return None + + def pop_request_state(self, request_id): + return None + + def finish_requests(self, request_id, status): + pass + + def close(self): + pass + + def initialize(self, od_config): + pass + class _BatchCapablePipeline: supports_request_batch = True @@ -644,7 +659,7 @@ async def test_async_add_req_and_stream_response(): engine.step_execution = False engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH - def _finalize(rid, out, err=None): + def _finalize(rid, out, err=None, **kwargs): # Stream consumers stop on ``finished``; keep result_data for assertions. return SimpleNamespace(result_data=out.result.result_data, finished=True) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 10fe7e86c..932630a23 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -102,7 +102,10 @@ def _start_worker(req_q, res_q, count=2): def _run(): for _ in range(count): - req = req_q.get(timeout=10) + try: + req = req_q.get(timeout=10) + except queue.Empty: + break method = req.get("method", "") args = req.get("args", ()) if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput): @@ -146,7 +149,7 @@ def _inject_interleave(executor): orig_enqueue(item) if threading.current_thread().name == "thread_a": a_enqueued.set() # tell B: "A has enqueued" - b_complete.wait(5) # block A until B finishes + b_complete.wait(1) # block A until B finishes executor._broadcast_mq.enqueue = _controlled return a_enqueued, b_complete @@ -173,8 +176,8 @@ class TestConcurrentRequestExecution: results["B"] = engine.add_req_and_wait_for_response(_mock_request("B")) b_complete.set() # release A - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -274,8 +277,8 @@ class TestConcurrentCollectiveRpc: ) b_complete.set() - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -311,8 +314,8 @@ class TestConcurrentRequestExecutionAndCollectiveRpc: ) b_complete.set() - ta = threading.Thread(target=_a, name="thread_a") - tb = threading.Thread(target=_b, name="thread_b") + ta = threading.Thread(target=_a, name="thread_a", daemon=True) + tb = threading.Thread(target=_b, name="thread_b", daemon=True) ta.start() tb.start() ta.join(10) @@ -683,7 +686,10 @@ class TestMultiprocExecutorStepStreamingOutput: def _worker(): for runner_output in runner_outputs: - req_q.get(timeout=10) + try: + req_q.get(timeout=10) + except queue.Empty: + break res_q.put(runner_output) thread = threading.Thread(target=_worker, daemon=True) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 422e8d784..c8fa51448 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -474,6 +474,8 @@ class _DiffusionConfigProjection: output_type: str = "pil" enable_cpu_offload: bool = False enable_layerwise_offload: bool = False + enable_distributed_layerwise_offload: bool = False + dlo_use_allgather: bool = True pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index b97ffc311..31617ccdb 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -672,6 +672,12 @@ class OmniDiffusionConfig: enable_cpu_offload: bool = False # Layer-wise offloading (block-level offloading) parameters enable_layerwise_offload: bool = False + # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1) + enable_distributed_layerwise_offload: bool = False + # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires + # concurrent requests in DP mode). If False: each rank loads full weights + # via H2D only (N× CPU memory, but no AllGather synchronization needed). + dlo_use_allgather: bool = True pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 44ca9da89..e89b95571 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -156,6 +156,10 @@ class DiffusionEngine: #: remains a pure explicit user override (never mutated by engines). default_diffusion_model_runner_cls: str | None = None + # Class-level default so tests using object.__new__ (without __init__) + # don't hit AttributeError when _busy_loop accesses self.dp_concurrent. + dp_concurrent: bool = False + def __init__( self, od_config: OmniDiffusionConfig, @@ -224,6 +228,29 @@ class DiffusionEngine: self.scheduler.initialize(od_config) def _init_runtime_state(self) -> None: + # DP multi-concurrency: allow batching dp_size requests so each + # worker processes a different request in parallel. Only enabled + # for distributed layerwise offload (which shards weights and + # needs all ranks active simultaneously). Ordinary DP with a + # non-batch pipeline should not schedule multiple requests. + dp_size = 1 + if getattr(self.od_config, "parallel_config", None) is not None: + dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1) + dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + if dp_size > 1 and dist_offload: + self.scheduler.max_num_running_reqs = dp_size + self.dp_concurrent = True + # Set batch admission wait so concurrent requests accumulate + # before scheduling. The scheduler reads this from od_config. + if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0: + self.od_config.request_batch_max_wait_ms = 500.0 + logger.info( + "dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms", + dp_size, + self.od_config.request_batch_max_wait_ms, + ) + else: + self.dp_concurrent = False self.main_loop: asyncio.AbstractEventLoop | None = None self.stop_event: threading.Event | None = None self.worker_thread: threading.Thread | None = None @@ -402,7 +429,7 @@ class DiffusionEngine: # Only RPC / abort work pending; loop back to drain it. continue - if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH: + if self.supports_request_batch or self.dp_concurrent: self._wait_for_request_batch_admission_locked() sched_output = self.scheduler.schedule() @@ -442,7 +469,7 @@ class DiffusionEngine: Caller must hold ``self._cv``. """ - if self.step_execution or not self.supports_request_batch: + if self.step_execution or (not self.supports_request_batch and not self.dp_concurrent): return max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0 @@ -460,9 +487,12 @@ class DiffusionEngine: deadline = start + max_wait_s last_waiting = -1 stable_since = start - # Require a short idle period with no queue growth so bursty HTTP - # ingress can land before the first schedule() of a wave. - stable_window_s = min(0.05, max_wait_s / 5.0) + # For dp_concurrent, use a longer stable window so all dp_size + # HTTP requests have time to land before scheduling. + if self.dp_concurrent: + stable_window_s = min(0.3, max_wait_s / 2.0) + else: + stable_window_s = min(0.05, max_wait_s / 5.0) while not self.stop_event.is_set(): waiting = self.scheduler.num_waiting_requests() @@ -783,6 +813,27 @@ class DiffusionEngine: raise RuntimeError(f"Could not {action} profiler: {e}") from e def run_startup_warmup(self) -> None: + dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True) + # Skip dummy run when AllGather is used with more than 1 rank, + # because the dummy run sends only 1 request but AllGather requires + # all ranks to participate simultaneously. This covers both DP > 1 + # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig). + pc = getattr(self.od_config, "parallel_config", None) + dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1 + sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1 + effective_shard_size = max(dp_size, sp_size) + skip_dummy = ( + getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and dlo_use_allgather + and effective_shard_size > 1 + ) + if skip_dummy: + logger.info( + "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)", + dp_size, + sp_size, + ) + return try: self._dummy_run() except Exception as e: diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 84a460c53..fa0213fca 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -1,6 +1,7 @@ from __future__ import annotations import concurrent.futures +import json import multiprocessing as mp import multiprocessing.connection import queue @@ -64,12 +65,17 @@ class _ExecutorShutdownCleaner: class MultiprocDiffusionExecutor(DiffusionExecutor): uses_multiproc: bool = True + # Class-level defaults so tests using object.__new__ (without _init_executor) + # don't hit AttributeError when collective_rpc accesses these. + _rpc_wave_id: int = 0 + def _init_executor(self) -> None: self._processes: list[mp.Process] = [] self._closed = False self._is_failed = False self._failure_callbacks: list[Callable[[], None]] = [] self._result_mq: MessageQueue | None = None + self._rpc_wave_id: int = 0 num_workers = cast(int, self.od_config.num_gpus) self.wake_events = [mp.Event() for _ in range(num_workers)] @@ -215,6 +221,30 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response) return response + _MAX_STALE_DISCARDS = 16 + + def _validate_wave_id(self, response: Any, expected_wave_id: int, deadline: float | None, method: str) -> Any: + """Discard stale RPC responses from a previous wave.""" + discards = 0 + while discards < self._MAX_STALE_DISCARDS: + if not isinstance(response, dict): + return response + resp_wave_id = response.get("wave_id") + if resp_wave_id is None or resp_wave_id == expected_wave_id: + return response + logger.warning( + "Discarding stale RPC response (wave_id=%s, expected=%s, method=%s)", + resp_wave_id, + expected_wave_id, + method, + ) + discards += 1 + response = self._dequeue_one_with_failure_polling(deadline, method) + raise TimeoutError( + f"Discarded {self._MAX_STALE_DISCARDS} stale RPC responses " + f"without finding wave_id={expected_wave_id} for method={method}." + ) + def _launch_workers( self, broadcast_handle: Handle, @@ -355,9 +385,84 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput self._ensure_open() + new_reqs = scheduler_output.scheduled_new_reqs runner_outputs: list[RunnerOutput] = [] - for new_req in scheduler_output.scheduled_new_reqs: + # DP multi-concurrency: when DLO+AllGather is active and multiple + # requests are scheduled, send ALL requests in one broadcast RPC. + # Each rank picks req[rank % len(reqs)] and computes independently. + # All ranks reply (unique_reply_rank=None) so we collect dp_size + # responses and match by dp_rank. + if ( + len(new_reqs) > 1 + and getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and getattr(self.od_config, "dlo_use_allgather", True) + ): + # Validate: all concurrent requests must share the same + # num_inference_steps and identical extra_args, because + # AllGather is a collective that requires every rank to + # participate at each step. + step_counts = { + nr.req.sampling_params.num_inference_steps + for nr in new_reqs + if nr.req.sampling_params.num_inference_steps is not None + } + has_none = any(nr.req.sampling_params.num_inference_steps is None for nr in new_reqs) + if (len(step_counts) > 1) or has_none: + raise ValueError( + "DP multi-concurrency requires all concurrent requests to have " + "the same explicit num_inference_steps (None is not allowed), got " + f"{[nr.req.sampling_params.num_inference_steps for nr in new_reqs]}." + ) + extra_args_signatures: set = set() + for nr in new_reqs: + ea = getattr(nr.req, "extra_args", None) + if ea and isinstance(ea, dict): + extra_args_signatures.add(json.dumps(ea, sort_keys=True)) + else: + extra_args_signatures.add(None) + if len(extra_args_signatures) > 1: + raise ValueError( + "DP multi-concurrency requires all concurrent requests to " + "share identical extra_args. Different extra_args can change " + "the forward schedule and cause AllGather deadlock." + ) + + reqs_list = [nr.req for nr in new_reqs] + try: + results = self.collective_rpc( + "execute_model", + args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job), + unique_reply_rank=None, + exec_all_ranks=True, + ) + results = results if isinstance(results, list) else [results] + for i, new_req in enumerate(new_reqs): + res = results[i] if i < len(results) else results[0] + if isinstance(res, DiffusionOutput): + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=res, + ) + ) + else: + raise RuntimeError(f"Unexpected response type [{i}]: {type(res)!r}") + except Exception as exc: + for new_req in new_reqs: + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=DiffusionOutput(error=str(exc)), + ) + ) + return BatchRunnerOutput.from_list(runner_outputs) + + for new_req in new_reqs: req = new_req.req try: result = self.collective_rpc( @@ -511,24 +616,61 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # ── Path 2: step-execution or other non-execute_model RPCs ── # _dequeue_one_with_failure_polling reads from _sync_result_buffer # when pump is active (request-mode), or result_mq directly otherwise. + self._rpc_wave_id += 1 + wave_id = self._rpc_wave_id + rpc_request["wave_id"] = wave_id + try: # Broadcast RPC request to all workers via unified message queue self._broadcast_mq.enqueue(rpc_request) # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown - num_responses = 1 + # Determine number of responses to collect: + # - unique_reply_rank=None + exec_all_ranks=True: all DP ranks reply + # (N responses, one per DP worker). + # - Otherwise: 1 response (only rank 0 or specified rank) + if unique_reply_rank is None and exec_all_ranks: + dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1) + num_responses = max(1, dp_size) + else: + num_responses = 1 - responses = [] - for _ in range(num_responses): - response = self._dequeue_one_with_failure_polling(deadline, method) + responses: list = [] + if unique_reply_rank is None and exec_all_ranks and num_responses > 1: + # DP multi-concurrency: collect num_responses replies, sort by dp_rank. + tagged: list[tuple[int, Any]] = [] + collected_errors: list[str] = [] + for _ in range(num_responses): + response = self._dequeue_one_with_failure_polling(deadline, method) + response = self._validate_wave_id(response, wave_id, deadline, method) + try: + unpack_diffusion_output_shm(response) + except Exception as e: + logger.warning("SHM unpack failed (data may already be inline): %s", e) + if isinstance(response, dict) and response.get("status") == "error": + collected_errors.append(str(response.get("error", "unknown"))) + else: + response = MultiprocDiffusionExecutor._handle_rpc_response(response) + if isinstance(response, dict) and "dp_rank" in response: + tagged.append((response["dp_rank"], response["output"])) + else: + tagged.append((len(tagged), response)) + if collected_errors: + raise RuntimeError(f"Worker error: {collected_errors[0]}") + tagged.sort(key=lambda x: x[0]) + responses = [r for _, r in tagged] + else: + for _ in range(num_responses): + response = self._dequeue_one_with_failure_polling(deadline, method) + response = self._validate_wave_id(response, wave_id, deadline, method) - try: - unpack_diffusion_output_shm(response) - except Exception as e: - logger.warning("SHM unpack failed (data may already be inline): %s", e) + try: + unpack_diffusion_output_shm(response) + except Exception as e: + logger.warning("SHM unpack failed (data may already be inline): %s", e) - response = MultiprocDiffusionExecutor._handle_rpc_response(response) + response = MultiprocDiffusionExecutor._handle_rpc_response(response) - responses.append(response) + responses.append(response) return responses[0] if unique_reply_rank is not None else responses except Exception as e: diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index ca6bc858d..5b7465af0 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -172,8 +172,9 @@ def pack_diffusion_output_shm( Supports a bare ``DiffusionOutput``, a wrapper object carrying one in ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying - the diffusion output in ``["result"]``, or a batch wrapper carrying - ``RunnerOutput`` objects in ``.runner_outputs``. + the diffusion output in ``["result"]``, a batch wrapper carrying + ``RunnerOutput`` objects in ``.runner_outputs``, or a DP-tagged dict + ``{"dp_rank": int, "output": DiffusionOutput}`` used by DP multi-concurrency. If *d2h_stream* is provided, D2H copies use that stream (non-blocking on the default stream). The caller must synchronize *d2h_stream* afterward. @@ -181,6 +182,13 @@ def pack_diffusion_output_shm( if isinstance(output, DiffusionOutput): return _pack_diffusion_fields(output, d2h_stream=d2h_stream) + # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} + if isinstance(output, dict) and "dp_rank" in output and "output" in output: + inner = output["output"] + if isinstance(inner, DiffusionOutput): + output["output"] = _pack_diffusion_fields(inner) + return output + if _is_rpc_result_envelope(output): result = output.get("result") if isinstance(result, DiffusionOutput): @@ -211,6 +219,13 @@ def unpack_diffusion_output_shm(output: object) -> object: if isinstance(output, DiffusionOutput): return _unpack_diffusion_fields(output) + # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} + if isinstance(output, dict) and "dp_rank" in output and "output" in output: + inner = output["output"] + if isinstance(inner, DiffusionOutput): + output["output"] = _unpack_diffusion_fields(inner) + return output + if _is_rpc_result_envelope(output): result = output.get("result") if isinstance(result, DiffusionOutput): diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d7770c266..df7aa2523 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -362,17 +362,42 @@ class DiffusersPipelineLoader: ) else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - logger.debug("Loading weights on %s ...", load_device) - if load_format == "diffusers": - # DiffusersAdapterPipeline.load_weights() calls - # DiffusionPipeline.from_pretrained() internally; it does - # not use our native customized pipeline classes. - cast(DiffusersAdapterPipeline, model).load_weights() + + # Skip load_weights only for DLO+AllGather when the model + # supports mmap loading and no online quantization is active. + # This condition MUST match the gate in + # DistributedLayerwiseOffloadBackend.enable() so that the + # loader skips ⟺ enable() uses mmap. + from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading + + _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + _use_ag = getattr(self.od_config, "dlo_use_allgather", True) + _has_online_quant = self._has_online_quant(model) + _supports_mmap = supports_mmap_loading(model) + + _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant + + if _skip_load: + logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())") else: - self.load_weights(model) - # HSDP processes quantized weights before wrapping parameters as - # DTensors. The non-HSDP path can process them here as usual. - self._process_weights_after_loading(model, target_device) + if _dist_offload and _use_ag and _has_online_quant: + raise ValueError( + "Online quantization is incompatible with DLO+AllGather: " + "the sharding + AllGather mechanism flattens weights by " + "dtype, which breaks quantized weight/scale layouts. " + "Please use --dlo-no-use-allgather or disable online " + "quantization." + ) + if _dist_offload and _use_ag and not _supports_mmap: + logger.info( + "DLO+AllGather active but model does not support mmap: loading weights via regular loader." + ) + logger.debug("Loading weights on %s ...", load_device) + if load_format == "diffusers": + cast(DiffusersAdapterPipeline, model).load_weights() + else: + self.load_weights(model) + self._process_weights_after_loading(model, target_device) if offload_after_quant: model.to("cpu") diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 7987da887..807fad7ed 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -8,24 +8,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, OffloadStrategy +from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names +from .distributed_layerwise_backend import ( + DistributedLayerwiseOffloadBackend, + DistributedLayerwiseOffloadHook, + apply_distributed_block_hook, + remove_distributed_block_hook, +) from .layerwise_backend import LayerWiseOffloadBackend +from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, remove_sequential_offload, ) +from .tensor_utils import ( + dtype_size, + is_dtensor, + is_materialized_tensor, + make_offload_placeholder, + set_tensor_storage, +) logger = init_logger(__name__) __all__ = [ "OffloadBackend", "OffloadConfig", + "OffloadPlan", "OffloadStrategy", "LayerWiseOffloadBackend", + "DistributedLayerwiseOffloadBackend", + "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", "apply_sequential_offload", "remove_sequential_offload", + "apply_distributed_block_hook", + "remove_distributed_block_hook", "get_offload_backend", + "get_offload_plan", + "supports_mmap_loading", + "get_blocks_attr_names", + "get_blocks_from_dit", + "set_blocks_attr_names", + "dtype_size", + "is_dtensor", + "is_materialized_tensor", + "make_offload_placeholder", + "set_tensor_storage", ] @@ -75,6 +105,8 @@ def get_offload_backend( return ModelLevelOffloadBackend(config, device) elif config.strategy == OffloadStrategy.LAYER_WISE: return LayerWiseOffloadBackend(config, device) + elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE: + return DistributedLayerwiseOffloadBackend(config, device) else: logger.error("Unknown offload strategy: %s", config.strategy) return None diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index e10438765..20b2f5e14 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -18,6 +18,7 @@ class OffloadStrategy(Enum): NONE = "none" MODEL_LEVEL = "model_level" # Sequential offloading between DiT and encoders LAYER_WISE = "layer_wise" # Block-level + DISTRIBUTED_LAYER_WISE = "distributed_layer_wise" # Block-level with DP sharding + H2D/AllGather overlap @dataclass @@ -25,13 +26,21 @@ class OffloadConfig: strategy: OffloadStrategy pin_cpu_memory: bool = True use_hsdp: bool = False + dp_size: int = 1 # derived from parallel_config, not user-configurable + dlo_use_allgather: bool = True # if True: shard + AllGather; if False: full weights H2D only + model_path: str | None = None # checkpoint path for mmap weight loading @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": """Extract and validate offload settings from OmniDiffusionConfig. - For now, enforces mutual exclusion between model-level and layer-wise offloading. - Layer-wise takes priority if both are enabled. + Enforces mutual exclusion among the three offload strategies. + Distributed layer-wise takes the highest priority, then layer-wise, + then model-level. + + The ``dp_size`` is automatically derived from ``parallel_config`` — + it is NOT a user-configurable parameter. The distributed layerwise + offload works with whatever DP/SP parallelism is already set up. Args: od_config: OmniDiffusionConfig with offload settings @@ -41,13 +50,39 @@ class OffloadConfig: """ enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False) enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) + enable_distributed_layerwise_offload = getattr(od_config, "enable_distributed_layerwise_offload", False) pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) parallel_config = getattr(od_config, "parallel_config", None) use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False - # Determine strategy (mutual exclusion, layer-wise takes priority) - if enable_layerwise_offload: + # Derive dp_size from parallel_config — not user-configurable. + # The offload adapts to whatever DP/SP is already configured. + dp_size = 1 + if parallel_config is not None: + dp_size = getattr(parallel_config, "data_parallel_size", 1) + # HSDP's fully_shard_degree also contributes to effective DP + hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1 + hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1 + if use_hsdp and hsdp_shard_size > 0: + dp_size = hsdp_shard_size * hsdp_replicate_size + + # When there is no DP but SP > 1, shard weights across SP ranks. + # AllGather reconstructs full weights per layer; each rank then + # computes on its SP portion of the sequence. This gives N× + # compute parallelism with 1/N H2D transfer, reusing the exact + # same AllGather code path — only the process group changes. + if dp_size <= 1: + sp_size = getattr(parallel_config, "sequence_parallel_size", 1) + if sp_size and sp_size > 1: + dp_size = sp_size + + # Determine strategy (mutual exclusion, distributed layer-wise takes priority) + if enable_distributed_layerwise_offload: + strategy = OffloadStrategy.DISTRIBUTED_LAYER_WISE + if enable_layerwise_offload or enable_cpu_offload: + logger.info("Distributed layer-wise offloading takes priority, disabling other offloading strategies.") + elif enable_layerwise_offload: strategy = OffloadStrategy.LAYER_WISE if enable_cpu_offload: logger.info( @@ -59,10 +94,37 @@ class OffloadConfig: else: strategy = OffloadStrategy.NONE + # When dlo_use_allgather=False, each rank loads full weights (no sharding, + # no AllGather). This avoids AllGather synchronization requirements + # (concurrent requests, dummy run skip) at the cost of N× CPU memory. + dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) + + # If dlo_use_allgather=False, force dp_size=1 (each rank independent) + if enable_distributed_layerwise_offload and not dlo_use_allgather: + dp_size = 1 + logger.info( + "Distributed layerwise offload: dlo_use_allgather=False, " + "each rank loads full weights (no shard, no AllGather)" + ) + + # HSDP already shards parameters into DTensors. Running distributed + # layerwise offload on top would shard each to_local() again, producing + # incorrect reconstruction after AllGather. Reject this combination. + if enable_distributed_layerwise_offload and use_hsdp and dlo_use_allgather: + raise ValueError( + "Distributed layerwise offload with AllGather is incompatible with " + "HSDP: HSDP parameters are already sharded DTensors, and the offloader " + "would double-shard them. Use --dlo-no-use-allgather (full weights per " + "rank) or disable HSDP." + ) + return cls( strategy=strategy, pin_cpu_memory=pin_cpu_memory, use_hsdp=use_hsdp, + dp_size=dp_size, + dlo_use_allgather=dlo_use_allgather, + model_path=getattr(od_config, "model", None), ) diff --git a/vllm_omni/diffusion/offloader/block_discovery.py b/vllm_omni/diffusion/offloader/block_discovery.py new file mode 100644 index 000000000..643485ed6 --- /dev/null +++ b/vllm_omni/diffusion/offloader/block_discovery.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Block discovery for layerwise offload. + +Shared between LayerWiseOffloadBackend and DistributedLayerwiseOffloadBackend. +""" + +from __future__ import annotations + +from torch import nn +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def get_blocks_attr_names(model: nn.Module) -> list[str]: + """Get block attribute names from model class.""" + attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", []) + + if not attrs: + old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None) + if old_attr is not None: + logger.warning( + "'_layerwise_offload_blocks_attr' is deprecated, " + "please use '_layerwise_offload_blocks_attrs' instead. " + "Example: _layerwise_offload_blocks_attrs = ['blocks']" + ) + attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr) + + return attrs + + +def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None: + if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"): + setattr(model.__class__, "_layerwise_offload_blocks_attrs", names) + + +def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]: + """Retrieve blocks and attribute names from provided DiT model.""" + blocks_attr_names = get_blocks_attr_names(model) + if not blocks_attr_names: + logger.warning( + f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, " + "skipping distributed layerwise offloading" + ) + return [], [] + + blocks: list[nn.Module] = [] + for name in blocks_attr_names: + attr = getattr(model, name, None) + if attr is None: + raise AttributeError( + f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs " + f"does not exist on model {model.__class__.__name__}" + ) + try: + attr_iter = iter(attr) + except TypeError: + if isinstance(attr, nn.Module): + logger.warning( + "Attribute '%s' on %s is not iterable; treating it as one block.", + name, + model.__class__.__name__, + ) + blocks.append(attr) + continue + + logger.warning( + "Attribute '%s' on %s is not iterable (got %s); skipping it.", + name, + model.__class__.__name__, + type(attr).__name__, + ) + else: + blocks.extend(attr_iter) + + if not blocks: + logger.warning( + "No blocks found in %s for %s, skipping distributed layerwise offloading", + blocks_attr_names, + model.__class__.__name__, + ) + return [], [] + + return blocks_attr_names, blocks diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py new file mode 100644 index 000000000..797209781 --- /dev/null +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -0,0 +1,1491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Distributed Layerwise Offload backend with H2D + AllGather overlap. + +This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: + +* Shards model weights across DP ranks and stores only the local shard + (1/DP_size of full model) in host pinned memory per rank. +* Uses a fixed double-buffer scheme that keeps only two layers' worth of + weights on each device at any time. +* Asynchronously pipelines both H2D transfers and AllGather communications + on dedicated streams, fully overlapping them with computation. +* Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU + (CANN) platforms via vLLM-Omni's platform abstraction layer. +""" + +from __future__ import annotations + +import os +from itertools import chain +from typing import Any + +import torch +import torch.distributed +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.platforms import current_omni_platform + +from .base import OffloadBackend, OffloadConfig +from .block_discovery import get_blocks_from_dit +from .module_collector import ModuleDiscovery +from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .tensor_utils import ( + dtype_size as _dtype_size, +) +from .tensor_utils import ( + is_materialized_tensor, + make_offload_placeholder, + set_tensor_storage, +) + +logger = init_logger(__name__) + +# Threshold (in MB) for deciding whether a non-block DiT submodule should +# use layerwise offload (streaming hooks) or be moved to GPU as a resident +# module. Submodules larger than this are offloaded to save HBM; smaller +# ones stay resident for lower latency. +_ON_DEMAND_THRESHOLD_MB = 1024 + + +class DistributedLayerwiseOffloadHook(ModelHook): + """Hook for distributed layerwise offloading with fixed double-buffer. + + Each rank stores only a shard of each block's weights on host memory. + Two device slots alternate: one holds current weights, one holds next weights. + H2D and AllGather run asynchronously on dedicated streams, overlapped + with computation. + + Supports both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms. + """ + + _HOOK_NAME = "distributed_layerwise_offload" + + def __init__( + self, + next_block: nn.Module, + device: torch.device, + dp_group: torch.distributed.ProcessGroup | None, + dp_size: int, + rank: int, + copy_stream: Any | None = None, + comm_stream: Any | None = None, + pin_memory: bool = True, + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + ): + assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" + + self.next_block = next_block + self.device = device + self.dp_group = dp_group + self.dp_size = dp_size + self.rank = rank + self.pin_memory = pin_memory + + self.copy_stream = copy_stream or current_omni_platform.Stream() + self.comm_stream = comm_stream or current_omni_platform.Stream() + + # Double buffers: either shared (from backend) or self-allocated (lazy) + if shared_buffers is not None: + self.gpu_buffers: list[dict[torch.dtype, torch.Tensor] | None] = shared_buffers + self._owns_buffers = False + else: + self.gpu_buffers = [None, None] + self._owns_buffers = True + self.ready_events: list[Any | None] = [None, None] + + # Sharded host weights for the next block, keyed by dtype + self.cpu_shards: dict[torch.dtype, torch.Tensor] = {} + self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + # Current slot index (0 or 1). Updated dynamically by the previous + # hook's prefetch_layer call via _prefetched_slot. This ensures + # correct slot tracking for ALL block counts (including odd N). + self.current_slot = 0 + self._prefetched_slot: int | None = None + + # Backward link to previous hook for fallback (cache-dit skip) + self._prev_hook: DistributedLayerwiseOffloadHook | None = None + + # Marks the first hook in a shared-buffer group. When multiple DiT + # groups share the same 2 GPU buffers, another group may have + # overwritten this group's slot between forwards. The first block + # must sync-prefetch on entry to ensure it loads the correct + # weights, even if is_materialized sees a non-empty tensor left by + # the other group. + self._is_group_first: bool = False + + # Group ID for per-slot contamination tracking. Shared across all + # hooks in the same group. When a hook prefetches into a slot, it + # stamps the slot with its group ID. On group entry, the first + # hook checks whether the slot was last written by its own group + # (tail hook's async prefetch) — if so, skip the sync-prefetch. + self._group_id: int = -1 + self._shared_slot_group: list[int] | None = None # [-1, -1], shared + + # Parameters/buffers of the current and next blocks + self.block_parameters: dict[str, nn.Parameter] = {} + self.block_buffers: dict[str, torch.Tensor] = {} + self.next_block_parameters: dict[str, nn.Parameter] = {} + self.next_block_buffers: dict[str, torch.Tensor] = {} + + # Per-block synchronization primitive: set after H2D copy completes. + self._prefetch_done: Any | None = None + + # Shared shard (AllGather input) buffers — assigned by backend. + self.gpu_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + + # Pending async AllGather work (prevent GC before completion). + self._pending_work: Any | None = None + self._cached_repoint: list | None = None + + # ------------------------------------------------------------------ # + # DTensor helpers (shared with LayerwiseOffloadHook) # + # ------------------------------------------------------------------ # + + def initialize_hook(self, module: nn.Module) -> nn.Module: + module = super().initialize_hook(module) + + self.block_parameters = dict(module.named_parameters()) + self.block_buffers = dict(module.named_buffers()) + + self.next_block_parameters = dict(self.next_block.named_parameters()) + self.next_block_buffers = dict(self.next_block.named_buffers()) + + # Shard next block's weights and store local shard in pinned CPU memory + self.cpu_shards, self.metadata = self._shard_and_pin( + self.next_block_parameters, + self.next_block_buffers, + self.dp_size, + self.rank, + self.pin_memory, + ) + + # Allocate device buffers only if not using shared buffers from backend + if self._owns_buffers: + self._allocate_device_buffers() + + # Cache parameter re-pointing metadata to avoid per-layer dict lookups. + self._cached_repoint = [] + for slot in range(2): + repoint = [] + for dtype, metas in self.metadata.items(): + for m in metas: + target = ( + self.next_block_parameters[m["name"]] + if m["name"] in self.next_block_parameters + else self.next_block_buffers[m["name"]] + ) + repoint.append((target, dtype, m["offset"], m["numel"], m["shape"])) + self._cached_repoint.append(repoint) + + # Pre-compute AG output sizes (avoid sum() per layer). + self._ag_output_sizes: dict[torch.dtype, int] = {} + for dtype, metas in self.metadata.items(): + total_numel = sum(m["numel"] for m in metas) + shard_numel = self.cpu_shards[dtype].numel() + self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel + + return module + + @staticmethod + def _shard_and_pin( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + dp_size: int, + rank: int, + pin_memory: bool, + ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: + """Flatten params+buffers by dtype, split into DP shards, store local shard. + + Each rank stores only ``1/dp_size`` of the total weights. The full + tensor is reconstructed at runtime via AllGather. + """ + dtype_grouped: dict[torch.dtype, dict[str, torch.Tensor]] = {} + dtype_metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + + for name, param_or_buf in chain(params.items(), bufs.items()): + dtype = param_or_buf.dtype + if dtype not in dtype_grouped: + dtype_grouped[dtype] = {} + dtype_grouped[dtype][name] = param_or_buf + + cpu_shards: dict[torch.dtype, torch.Tensor] = {} + + for dtype, name2weights in dtype_grouped.items(): + # Resolve local tensors (handle DTensor via to_local) + weights_with_local = [] + for name, t in name2weights.items(): + local_t = t.to_local() if hasattr(t, "to_local") else t + weights_with_local.append((name, t, local_t)) + + total_numel = sum(local.numel() for _, _, local in weights_with_local) + + # Equal-sized shards (ceil division) for all_gather_into_tensor + shard_size = (total_numel + dp_size - 1) // dp_size # ceil + shard_start = rank * shard_size + shard_end = min(shard_start + shard_size, total_numel) + + # Allocate ONLY the shard (1/dp_size), zero-padded to ceil. + # Avoids materialising the full block on CPU. + shard = torch.zeros(shard_size, dtype=dtype, device="cpu") + + current_offset = 0 + for name, original_tensor, local_tensor in weights_with_local: + numel = local_tensor.numel() + if dtype not in dtype_metadata: + dtype_metadata[dtype] = [] + # Offsets remain relative to the FULL flattened buffer + # (needed for correct AllGather reconstruction). + dtype_metadata[dtype].append( + { + "name": name, + "offset": current_offset, + "numel": numel, + "shape": local_tensor.shape, + } + ) + + # Copy ONLY the portion within [shard_start, shard_end) + overlap_start = max(current_offset, shard_start) + overlap_end = min(current_offset + numel, shard_end) + if overlap_start < overlap_end: + src_start = overlap_start - current_offset + src_end = overlap_end - current_offset + dst_start = overlap_start - shard_start + dst_end = overlap_end - shard_start + shard[dst_start:dst_end].copy_(local_tensor.flatten()[src_start:src_end]) + + # Replace original tensor with placeholder (frees CPU storage) + set_tensor_storage( + original_tensor, + make_offload_placeholder(original_tensor), + ) + current_offset += numel + + if pin_memory: + shard = shard.pin_memory() + + cpu_shards[dtype] = shard + + return cpu_shards, dtype_metadata + + def _allocate_device_buffers(self) -> None: + """Pre-allocate exactly two device buffers (one per slot).""" + for slot in range(2): + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + for dtype, metas in self.metadata.items(): + total_numel = sum(m["numel"] for m in metas) + # AllGather output = dp_size * shard_size (padded) + padded = total_numel + if self.dp_size > 1: + shard_sz = (total_numel + self.dp_size - 1) // self.dp_size + padded = shard_sz * self.dp_size + gpu_weights[dtype] = torch.empty(padded, dtype=dtype, device=self.device) + self.gpu_buffers[slot] = gpu_weights + + @property + def is_materialized(self) -> bool: + """Check whether this block's parameters hold real data on device.""" + for param in self.block_parameters.values(): + return is_materialized_tensor(param) + return True + + # ------------------------------------------------------------------ # + # Prefetch: H2D + AllGather (overlapped on dedicated streams) # + # ------------------------------------------------------------------ # + + @torch.compiler.disable + def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None: + """Prepare next block's weights into the shared device buffer for *slot*. + + Uses the pre-allocated ``self.gpu_buffers[slot]`` instead of + allocating fresh tensors every layer. This enforces the fixed + double-buffer memory bound (exactly 2 blocks on device). + """ + self.copy_stream.wait_stream(current_omni_platform.current_stream()) + + evt = current_omni_platform.Event() + gpu_weights = self.gpu_buffers[slot] + assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated" + + if self.dp_size <= 1 or self.dp_group is None: + with current_omni_platform.stream(self.copy_stream): + for dtype, cpu_shard in self.cpu_shards.items(): + gw = gpu_weights[dtype] + gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=non_blocking) + evt.record(self.copy_stream) + else: + gpu_shards: dict[torch.dtype, torch.Tensor] = {} + shard_bufs = self.gpu_shard_buffers[slot] + assert shard_bufs is not None, f"gpu_shard_buffers[{slot}] not allocated" + with current_omni_platform.stream(self.copy_stream): + for dtype, cpu_shard in self.cpu_shards.items(): + gpu_shard = shard_bufs[dtype][: cpu_shard.numel()] + gpu_shard.copy_(cpu_shard, non_blocking=non_blocking) + gpu_shards[dtype] = gpu_shard + + self.comm_stream.wait_stream(self.copy_stream) + with current_omni_platform.stream(self.comm_stream): + for dtype, local_shard in gpu_shards.items(): + gw = gpu_weights[dtype] + torch.distributed.all_gather_into_tensor( + gw, + local_shard, + group=self.dp_group, + ) + evt.record(self.comm_stream) + + self.ready_events[slot] = evt + self._prefetch_done = evt + self._prefetched_slot = slot + + # Stamp the slot with this hook's group ID so that group-first + # hooks can detect whether another group has overwritten the slot. + if self._shared_slot_group is not None: + self._shared_slot_group[slot] = self._group_id + + # Re-point using cached metadata (avoids per-layer dict lookups). + for target, dtype, offset, numel, shape in self._cached_repoint[slot]: + set_tensor_storage( + target, + gpu_weights[dtype][offset : offset + numel].view(shape), + ) + + def get_weights(self, slot: int) -> dict[torch.dtype, torch.Tensor] | None: + """Wait for AllGather completion and return full weights for the slot. + + The ready event for this slot was set by the *previous* hook's + prefetch_layer (which prefetched THIS block's weights into the + shared buffer). This hook's own ready_events[slot] is None because + it never prefetched into this slot itself. Fall back to the + previous hook's event to ensure the compute stream waits for the + AllGather (or H2D) to complete before reading the buffer. + """ + evt = self.ready_events[slot] + if evt is None and self._prev_hook is not None: + evt = self._prev_hook.ready_events[slot] + if evt is not None: + current_omni_platform.current_stream().wait_event(evt) + return self.gpu_buffers[slot] + + # ------------------------------------------------------------------ # + # Offload: free device memory for current block # + # ------------------------------------------------------------------ # + + @torch.compiler.disable + def offload_layer(self) -> None: + """Free GPU memory for current block by replacing tensors with placeholders.""" + evt = self._prefetch_done + if evt is not None: + current_omni_platform.current_stream().wait_event(evt) + self._prefetch_done = None + + for _, param in self.block_parameters.items(): + set_tensor_storage(param, make_offload_placeholder(param)) + for _, buf in self.block_buffers.items(): + set_tensor_storage(buf, make_offload_placeholder(buf)) + + # ------------------------------------------------------------------ # + # ModelHook interface # + # ------------------------------------------------------------------ # + + def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]: + # Dynamic slot tracking: read current_slot from the previous hook's + # _prefetched_slot (the slot where this block's data was actually + # loaded). This corrects the initial i%2 assignment for odd block + # counts, where the circular tail→head prefetch would otherwise + # collide with the head's read slot. + if self._prev_hook is not None and self._prev_hook._prefetched_slot is not None: + self.current_slot = self._prev_hook._prefetched_slot + + # Group-first hook: check whether the shared buffer slot was + # overwritten by another group since our last forward. If the + # slot still contains our own data (from the tail hook's async + # prefetch), skip the sync-prefetch and just wait for the event. + if self._is_group_first and self._prev_hook is not None: + slot_contaminated = True + if self._shared_slot_group is not None: + slot_contaminated = self._shared_slot_group[self.current_slot] != self._group_id + if slot_contaminated: + # Another group (or no group) wrote to our slot — re-fetch + self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) + # Always wait for data to be ready (handles both sync and async paths) + self._prev_hook.get_weights(self.current_slot) + elif not self.is_materialized and self._prev_hook is not None: + # Previous hook was skipped (e.g. by cache-dit), sync-prefetch + self._prev_hook.prefetch_layer(self.current_slot, non_blocking=False) + self._prev_hook.get_weights(self.current_slot) + + # Prefetch next layer into the other slot (overlapped with compute). + # No explicit get_weights() here — offload_layer() in the previous + # hook's post_forward already enqueued wait_event for the current + # layer's H2D/AllGather, which is sufficient to ensure compute_stream + # waits for weights before reading them. The redundant wait_event + # caused cascading NPU stream synchronization stalls (~10ms/layer). + next_slot = 1 - self.current_slot + self.prefetch_layer(next_slot, non_blocking=True) + + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + self.offload_layer() + return output + + +# ---------------------------------------------------------------------- # +# Module-level helpers # +# ---------------------------------------------------------------------- # + + +def apply_distributed_block_hook( + module: nn.Module, + next_block: nn.Module, + device: torch.device, + dp_group: torch.distributed.ProcessGroup | None, + dp_size: int, + rank: int, + copy_stream: Any | None = None, + comm_stream: Any | None = None, + pin_memory: bool = True, + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, +) -> DistributedLayerwiseOffloadHook: + """Register a DistributedLayerwiseOffloadHook on *module*.""" + registry = HookRegistry.get_or_create(module) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=device, + dp_group=dp_group, + dp_size=dp_size, + rank=rank, + copy_stream=copy_stream, + comm_stream=comm_stream, + pin_memory=pin_memory, + shared_buffers=shared_buffers, + ) + registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) + return hook + + +def remove_distributed_block_hook(module: nn.Module) -> None: + """Remove the distributed layerwise offload hook from *module*.""" + registry: HookRegistry | None = getattr(module, "_hook_registry", None) + if registry is not None: + registry.remove_hook(DistributedLayerwiseOffloadHook._HOOK_NAME) + logger.debug("Removed distributed offload hook from %s", module.__class__.__name__) + + +# ---------------------------------------------------------------------- # +# Backend # +# ---------------------------------------------------------------------- # + + +class DistributedLayerwiseOffloadBackend(OffloadBackend): + """Distributed layer-wise (block-level) offloading backend. + + Supports both GPU (CUDA) and NPU (CANN) platforms. + Device type is determined by the device passed to enable(). + + Each rank stores only a shard of each block's weights on host memory. + Two device slots alternate: one holds current weights, one holds next + weights. H2D and AllGather run asynchronously on dedicated streams, + overlapped with computation. + """ + + def __init__(self, config: OffloadConfig, device: torch.device): + super().__init__(config, device) + + self.copy_stream = current_omni_platform.Stream() + self.comm_stream = current_omni_platform.Stream() + self.dp_group: torch.distributed.ProcessGroup | None = None + self.dp_size = config.dp_size + self.rank = 0 + self._blocks: list[list[nn.Module]] = [] + self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] + + def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: + """Load DiT weights from safetensors via mmap views (no RSS). + + When the transformer is created on meta device, this method + replaces meta params with mmap views of the checkpoint files. + The views point to OS page cache (shared across ranks), so + no private copy is created. _shard_and_pin then copies only + the 1/dp_size shard from the mmap view to a private buffer. + + Non-DiT modules (VAE, encoders) are NOT affected — they were + created on CPU with real weights via from_pretrained. + """ + import glob + import json + + from safetensors import safe_open + + model_path = self.config.model_path + if not model_path: + logger.warning("No model_path for mmap weight loading, skipping") + return + + # Resolve HF repo ID to local snapshot path. + # If model_path is a local directory, use it directly; otherwise + # download the safetensors index + weight files via HuggingFace Hub. + if not os.path.isdir(model_path): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + logger.info("model_path %s is not local, downloading from HF", model_path) + model_path = download_weights_from_hf( + model_name_or_path=model_path, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + ) + + # Build {checkpoint_key: file_path} from safetensors index + weight_map: dict[str, str] = {} + for idx_file in glob.glob(os.path.join(model_path, "**", "*.safetensors.index.json"), recursive=True): + idx_dir = os.path.dirname(idx_file) + with open(idx_file) as f: + idx = json.load(f) + for key, filename in idx.get("weight_map", {}).items(): + weight_map[key] = os.path.join(idx_dir, filename) + + if not weight_map: + single_files = glob.glob(os.path.join(model_path, "**", "model.safetensors"), recursive=True) + for sf in single_files: + from safetensors import safe_open as _safe_open + + f = _safe_open(sf, framework="pt", device="cpu") + for key in f.keys(): + weight_map[key] = sf + logger.info("Found single-file safetensors: %s (%d keys)", sf, len(f.keys())) + + if not weight_map: + raise RuntimeError( + "Distributed layerwise offload could not find safetensors " + f"weights at {model_path}. Expected *.safetensors.index.json " + "or model.safetensors. Ensure the checkpoint path is correct." + ) + + logger.info("Loading DiT weights via mmap (meta → page cache, no RSS): %d keys", len(weight_map)) + + # --- Convert DiT modules to meta device --- + # The transformer was created normally (with random weights and + # correct non-persistent buffers from __init__). We convert it + # to meta device to release the random weights (they're useless), + # then load real weights via mmap. + # + # Non-persistent buffers (e.g. RoPE inv_freq, timestep freqs) are + # NOT in the checkpoint — they are computed from formulas in + # __init__. Save them before meta conversion and restore after + # mmap loading, so we don't need model-specific buffer rebuild code. + # + # Fail closed for unsupported configurations BEFORE to_empty. + # The mmap path bypasses AutoWeightsLoader, which means TP-aware + # weight_loader callbacks (fused QKV, row-parallel, etc.) are not + # invoked. When TP > 1, these callbacks shard weights across TP + # ranks — skipping them produces incorrect weights. + # Check the actual TP world size from parallel_state (initialized + # before enable()). Note: params may have custom weight_loader + # attributes even at TP=1 (e.g. QKVParallelLinear), so we cannot + # reject based on attribute presence alone. + try: + from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_world_size, + ) + + tp_world = get_tensor_model_parallel_world_size() + except Exception: + tp_world = 1 + if tp_world > 1: + raise ValueError( + "Distributed layerwise offload with mmap loading does not " + "support Tensor Parallel (TP > 1). TP-aware weight_loader " + "callbacks are bypassed by the mmap path. Use DP or SP " + "instead of TP, or disable distributed layerwise offload." + ) + + # + # Important: when DiT modules are nested (e.g. transformer contains + # transformer.language_model), to_empty("meta") on the parent + # converts the child's buffers too. So we must save ALL buffers + # from ALL DiT modules BEFORE any to_empty call. + saved_buffers: dict[int, dict[str, torch.Tensor]] = {} + for dit_module in modules.dits: + bufs: dict[str, torch.Tensor] = {} + for name, buf in dit_module.named_buffers(): + # Check if this buffer is non-persistent on its OWNING module + owner = dit_module + parts = name.split(".") + for part in parts[:-1]: + owner = getattr(owner, part) + buf_name = parts[-1] + is_non_persistent = buf_name in owner._non_persistent_buffers_set + # Save if non-persistent OR already meta (shouldn't happen + # before to_empty, but be safe) + if is_non_persistent: + bufs[name] = buf.detach().clone() + saved_buffers[id(dit_module)] = bufs + + # Now convert all DiT modules to meta (after saving all buffers). + # Skip modules that are already meta (happens when a parent DiT + # module contains a child DiT module — to_empty on the parent + # already converted the child). + for dit_module in modules.dits: + if any(p.is_meta for p in dit_module.parameters()): + logger.info( + "%s already on meta device (skipping to_empty, %d buffers saved)", + dit_module.__class__.__name__, + len(saved_buffers.get(id(dit_module), {})), + ) + continue + dit_module.to_empty(device="meta") + logger.info( + "Converted %s to meta device (%d non-persistent buffers saved)", + dit_module.__class__.__name__, + len(saved_buffers.get(id(dit_module), {})), + ) + + # Build reverse mapping {model_param_name: (ckpt_key, file_path)} using + # the pipeline's _remap_ckpt_key (if available). This handles the + # model-specific checkpoint key remapping (e.g., Cosmos3's + # layers.0.mlp_moe_gen.gate_proj → transformer.gen_layers.0.mlp.gate_proj). + remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) + model_to_ckpt: dict[str, tuple[str, str]] = {} + if remap_fn is not None: + for ckpt_key, file_path in weight_map.items(): + model_name = remap_fn(ckpt_key) + if model_name is not None: + model_to_ckpt[model_name] = (ckpt_key, file_path) + logger.info( + "Built reverse remap: %d model params from %d checkpoint keys", len(model_to_ckpt), len(weight_map) + ) + else: + # No remap function — use checkpoint keys directly + for ckpt_key, file_path in weight_map.items(): + model_to_ckpt[ckpt_key] = (ckpt_key, file_path) + + # Collect all meta params in DiT modules + dit_ids = set() + for dit_module in modules.dits: + dit_ids.update(id(m) for m in dit_module.modules()) + + # Cache open file handles + file_cache: dict[str, Any] = {} + loaded = 0 + skipped = 0 + loaded_names: set[str] = set() + + # Discover blocks first so we can distinguish block params (deferred + # to _shard_and_pin) from non-block params (loaded here as mmap views). + block_module_ids: set[int] = set() + for dit_module in modules.dits: + blocks_attr_names, blocks = get_blocks_from_dit(dit_module) + for block in blocks: + block_module_ids.update(id(m) for m in block.modules()) + + for name, param in pipeline.named_parameters(): + if not (hasattr(param, "is_meta") and param.is_meta): + continue + + # Check if this param is inside a block (will be handled by + # _shard_and_pin via mmap views in section 2 below). + parent = pipeline + parts = name.split(".") + is_block_param = False + for part in parts[:-1]: + try: + if part.isdigit(): + parent = parent[int(part)] + else: + parent = getattr(parent, part) + except (AttributeError, IndexError, TypeError): + break + if id(parent) in block_module_ids: + is_block_param = True + break + + if is_block_param: + # Block param — will be loaded as mmap view in section 2 + continue + + # Non-block param (merger, patch_embed, norm, proj_out, etc.) + # Look up in reverse remap mapping (handles Cosmos3 key remapping). + entry = model_to_ckpt.get(name) + if entry is None: + skipped += 1 + continue + + ckpt_key, file_path = entry + if file_path not in file_cache: + file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") + f = file_cache[file_path] + tensor = f.get_tensor(ckpt_key) + + # Replace meta param with mmap view (no copy!) + parent = pipeline + for part in parts[:-1]: + if part.isdigit(): + parent = parent[int(part)] + else: + parent = getattr(parent, part) + parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + loaded += 1 + loaded_names.add(name) + + logger.info("Mmap weight loading: %d non-block params loaded, %d skipped", loaded, skipped) + + # Now load DiT block params from safetensors using mmap views. + # Each block param is assigned an mmap view (no RSS). + # _shard_and_pin will later copy the shard portion to a private buffer. + block_loaded = 0 + for dit_idx, dit_module in enumerate(modules.dits): + dit_name = modules.dit_names[dit_idx] + blocks_attr_names, blocks = get_blocks_from_dit(dit_module) + if not blocks: + continue + + # Determine the blocks_attr for this DiT module (e.g. "layers" or "gen_layers") + blocks_attr = None + for attr_name in blocks_attr_names: + if hasattr(dit_module, attr_name): + candidate = getattr(dit_module, attr_name) + if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: + blocks_attr = attr_name + break + + for block_idx, block in enumerate(blocks): + # Build the full model param prefix for this block + # e.g. "transformer.language_model.layers.0" or "transformer.gen_layers.0" + block_full_prefix = f"{dit_name}.{blocks_attr}.{block_idx}" + + for bname, bparam in block.named_parameters(): + if not (hasattr(bparam, "is_meta") and bparam.is_meta): + continue + + # Use reverse remap to find checkpoint key + full_param_name = f"{block_full_prefix}.{bname}" + entry = model_to_ckpt.get(full_param_name) + if entry is None: + skipped += 1 + continue + + ckpt_key, file_path = entry + if file_path not in file_cache: + file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") + f = file_cache[file_path] + tensor = f.get_tensor(ckpt_key) + + # Replace meta param with mmap view + parent = block + bparts = bname.split(".") + for part in bparts[:-1]: + if part.isdigit(): + parent = parent[int(part)] + else: + parent = getattr(parent, part) + parent._parameters[bparts[-1]] = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + block_loaded += 1 + loaded_names.add(full_param_name) + + logger.info("Mmap block loading: %d params loaded, %d skipped", block_loaded, skipped) + + # Strict validation: check that all meta params were loaded. + # This replaces the validation that AutoWeightsLoader.load_weights() + # performs in the regular load path. + remaining_meta = [ + name for name, param in pipeline.named_parameters() if hasattr(param, "is_meta") and param.is_meta + ] + if remaining_meta: + # Filter out params inside submodules that will be loaded later + # by _load_module_weights_from_mmap (non-block submodules) + dit_param_names = set() + for dit_module in modules.dits: + for pname, _ in dit_module.named_parameters(): + dit_param_names.add(pname) + truly_missing = [n for n in remaining_meta if n in dit_param_names] + if truly_missing: + logger.warning( + "Mmap loading: %d params still on meta device after " + "loading (first 5: %s). These may be loaded later by " + "_load_module_weights_from_mmap or are expected to be " + "meta (e.g. DTensor placeholders).", + len(truly_missing), + truly_missing[:5], + ) + + # Keep file handles open — _shard_and_pin will read from the mmap views. + # They will be released after _shard_and_pin completes (when params are + # replaced with offload placeholders). + self._mmap_file_cache = file_cache + self._mmap_model_to_ckpt = model_to_ckpt + + # The regular loader runs model-specific post-load transforms after + # assigning checkpoint tensors. The mmap path bypasses that loader, so + # preserve the same lifecycle for transforms such as Cosmos3's fp32 + # timestep embedder. + for dit_module in modules.dits: + post_load_weights = getattr(dit_module, "post_load_weights", None) + if callable(post_load_weights): + post_load_weights() + + # Call validate_loaded_weights if the model defines it. + # This preserves the sound-weight and action-weight validation + # that AutoWeightsLoader.load_weights() triggers in the regular + # load path (e.g. Cosmos3 checks for missing audio/action weights). + validate = getattr(dit_module, "validate_loaded_weights", None) + if callable(validate): + validate(loaded_names) + + # Restore non-persistent buffers that were saved before meta + # conversion. These buffers (e.g. RoPE inv_freq, timestep + # freqs) are computed from formulas in __init__ and are not + # present in the checkpoint. By saving and restoring them + # generically, we avoid model-specific buffer rebuild code. + bufs = saved_buffers.get(id(dit_module), {}) + for name, buf in bufs.items(): + parent = dit_module + parts = name.split(".") + for part in parts[:-1]: + if part.isdigit(): + parent = parent[int(part)] + else: + parent = getattr(parent, part) + # Place on the same device as the module's parameters + # Place on the offloader's target device (not the module's + # current device, which may still be meta at this point). + parent._buffers[parts[-1]] = buf.to(self.device) + + def _load_module_weights_from_mmap(self, module: nn.Module, dit_name: str, child_name: str) -> None: + """Load a non-block DiT submodule's weights from safetensors via mmap.""" + from safetensors import safe_open + + if not hasattr(self, "_mmap_file_cache") or not self._mmap_file_cache: + return + + file_cache = self._mmap_file_cache + model_to_ckpt = getattr(self, "_mmap_model_to_ckpt", {}) + loaded = 0 + + for pname, param in module.named_parameters(): + if not (hasattr(param, "is_meta") and param.is_meta): + continue + + # Build full model param name and look up in reverse remap + full_name = f"{dit_name}.{child_name}.{pname}" + entry = model_to_ckpt.get(full_name) + if entry is None: + # Try without dit_name prefix + full_name = f"{child_name}.{pname}" + entry = model_to_ckpt.get(full_name) + if entry is None: + continue + + ckpt_key, file_path = entry + if file_path not in file_cache: + file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") + f = file_cache[file_path] + tensor = f.get_tensor(ckpt_key) + + # Replace meta param with mmap view + parent = module + parts = pname.split(".") + for part in parts[:-1]: + if part.isdigit(): + parent = parent[int(part)] + else: + parent = getattr(parent, part) + parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + loaded += 1 + + if loaded > 0: + logger.info("Loaded %d params for submodule %s.%s via mmap", loaded, dit_name, child_name) + + def _init_dp_group(self) -> None: + """Reuse the process group initialized by parallel_state. + + When DP > 1, uses the DP group (handles strided groups with SP/CFG/PP). + When DP = 1 but SP > 1 (and dp_size was set to sp_size in OffloadConfig), + uses the SP group so weights are sharded across SP ranks. + """ + if self.dp_size <= 1: + logger.info("Distributed layerwise offload: dp_size=1, running without AllGather") + self.dp_group = None + return + + if not torch.distributed.is_initialized(): + raise RuntimeError( + "torch.distributed is not initialized. " + "Distributed layerwise offload with dp_size > 1 requires " + "an initialized process group." + ) + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_data_parallel_world_size, + get_dp_group, + ) + + # Determine which parallel group to use for sharding. + # When data_parallel_size > 1, use the DP group. + # When data_parallel_size == 1 but dp_size > 1 (set from sp_size in + # OffloadConfig), use the SP group for weight sharding. + dp_world = get_data_parallel_world_size() + if dp_world > 1: + coord = get_dp_group() + else: + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + coord = get_sp_group() + logger.info( + "Distributed layerwise offload: DP=1, using SP group (world_size=%d) for weight sharding", + coord.world_size, + ) + + self.dp_group = coord.device_group + self.rank = coord.rank_in_group + self.dp_size = coord.world_size + + logger.info( + "Distributed layerwise offload: dp_size=%d, rank_in_group=%d, global_rank=%d, group_ranks=%s", + self.dp_size, + self.rank, + coord.rank, + coord.ranks, + ) + + def _register_on_demand_hook(self, module: nn.Module, label: str) -> None: + """Move *module* to GPU permanently. + + Previously, on-demand hooks moved the module to GPU before forward + and back to CPU after forward (with synchronize + empty_cache). + However, _post_forward's global synchronize() and empty_cache() + disrupt the async prefetch pipeline (copy_stream / comm_stream) + used by distributed layerwise offload, causing garbled output. + + Moving the module to GPU permanently avoids this issue at the + cost of slightly higher HBM usage (VAE/encoders are small). + """ + module.to(self.device) + logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + + def _try_layerwise_offload_submodule(self, module: nn.Module, name: str, plan: OffloadPlan | None = None) -> bool: + """Try to apply layerwise offload to a large submodule's blocks. + + Resolution order: + 1. OffloadPlan.offload_submodules (declarative, if plan is provided) + 2. Heuristic search for common block-list attributes + + Returns True if layerwise offload was applied, False otherwise. + """ + from operator import attrgetter + + blocks = None + blocks_attr = None + + # 1. Check OffloadPlan first (declarative — no guessing) + if plan is not None and name in plan.offload_submodules: + attr_name = plan.offload_submodules[name] + try: + candidate = attrgetter(attr_name)(module) + if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: + blocks = candidate + blocks_attr = attr_name + except AttributeError: + logger.warning( + "OffloadPlan declared block attr '%s' for submodule '%s' " + "but attribute not found — falling back to heuristic", + attr_name, + name, + ) + + # 2. Fallback: heuristic search + if blocks is None: + for attr_name in ("layers", "blocks", "h", "model.layers"): + try: + candidate = attrgetter(attr_name)(module) + except AttributeError: + continue + if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: + blocks = candidate + blocks_attr = attr_name + break + + if blocks is None: + return False + + num_blocks = len(blocks) + logger.info( + "Distributed layerwise offload for submodule '%s.%s' (%d blocks, %.0f MB total, dp_size=%d)", + name, + blocks_attr, + num_blocks, + sum(p.nelement() * p.element_size() for p in module.parameters()) / 1048576, + self.dp_size, + ) + + # Move non-block parts of the submodule to GPU (small: embeddings, norms) + for child_name, child in module.named_children(): + if child_name != blocks_attr: + child.to(self.device) + + # Apply distributed hooks (1/4 sharding + AllGather, same as DiT) + # Pass shared_buffers=[None,None] to defer per-hook allocation + # (prevents OOM on large models where N hooks × 2 buffers >> HBM) + last_block, first_block = blocks[-1], blocks[0] + last_hook = apply_distributed_block_hook( + last_block, + first_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + sub_hooks = [last_hook] + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + hook = apply_distributed_block_hook( + block, + next_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + sub_hooks.append(hook) + + # Wire backward references + slot alternation + for i in range(len(sub_hooks)): + sub_hooks[i]._prev_hook = sub_hooks[i - 1] + # Assign slots in list order: sub_hooks = [last_hook, block0, ..., blockN-2] + # This ensures last_hook.current_slot != block0_hook.current_slot, + # so the circular prefetch (last_hook -> block0) writes to a + # different slot than block0 reads from. Correct for ALL N. + for i, hook in enumerate(sub_hooks): + hook.current_slot = i % 2 + # Mark block0_hook (index 1) as group-first — it must sync-prefetch + # on entry because another group may have overwritten the shared slot. + if len(sub_hooks) > 1: + sub_hooks[1]._is_group_first = True + + # Defer buffer allocation and prefetch to enable() unified allocation + self._all_hook_groups.append(sub_hooks) + self._blocks.append(blocks) + return True + + def enable(self, pipeline: nn.Module) -> None: + if self.enabled: + logger.warning("DistributedLayerwiseOffloadBackend already enabled") + return + + self._on_demand_shard_infos: list[dict] = [] + self._on_demand_handles: list[Any] = [] + + # Initialize DP group (if not already done by early init) + if self.dp_group is None and self.dp_size > 1: + self._init_dp_group() + + modules = ModuleDiscovery.discover(pipeline) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading") + return + + # Retrieve optional declarative OffloadPlan from the pipeline. + # When present, replaces heuristic block discovery. + plan = get_offload_plan(pipeline) + + # Load weights via mmap for DLO+AllGather. + # Gate condition MUST match diffusers_loader.py: + # supports_mmap_loading(pipeline) and not _has_online_quant + # When the gate is False, the loader has already loaded weights via + # regular load_weights() + _process_weights_after_loading(). + if self.config.dlo_use_allgather and self.dp_size > 1: + _has_online_quant = any( + getattr(getattr(module, "quant_method", None), "uses_meta_device", False) + for module in pipeline.modules() + ) + if _has_online_quant: + raise ValueError( + "Online quantization is incompatible with DLO+AllGather: " + "the sharding + AllGather mechanism flattens weights by " + "dtype, which breaks quantized weight/scale layouts. " + "Please use --dlo-no-use-allgather or disable online " + "quantization." + ) + if supports_mmap_loading(pipeline): + self._load_weights_via_mmap(pipeline, modules) + else: + logger.info("Weights loaded via regular loader — skipping mmap (model does not support mmap)") + + # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. + # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) + # during the DiT forward pass. They are only needed briefly for + # text-encoding (before DiT) and VAE-decode (after DiT). + for enc in modules.encoders: + self._register_on_demand_hook(enc, "encoder") + for vae in modules.vaes: + self._register_on_demand_hook(vae, "vae") + + # Move resident modules to GPU (small modules needed every forward) + for name, module in zip(modules.resident_names, modules.resident_modules): + try: + module.to(self.device) + except Exception as exc: + logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + + logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + + # Collect all DiT module objects to detect submodules that are + # already handled as a separate DiT module (avoids duplicate hooks). + all_dit_modules = set(id(m) for m in modules.dits) + + # Apply hooks for each DiT module + for i, dit_module in enumerate(modules.dits): + dit_name = modules.dit_names[i] + logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})") + + blocks_attr_names, blocks = get_blocks_from_dit(dit_module) + + if not blocks: + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + num_blocks = len(blocks) + if num_blocks <= 1: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + dit_name, + dit_module.__class__.__name__, + ) + dit_module.to(self.device) + continue + + # Move non-block modules to GPU (they stay resident). + # Large modules (> 1 GB) use on-demand hooks instead — they are + # only needed briefly (e.g. language_model for text encoding) + # and keeping them resident wastes HBM during the DiT loop. + _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB # alias for readability + for name, m in dit_module.named_children(): + if name not in blocks_attr_names: + _has_meta = any((p.is_meta if hasattr(p, "is_meta") else False) for p in m.parameters()) + if _has_meta: + # Meta params: load from safetensors before moving to GPU + self._load_module_weights_from_mmap(m, dit_name, name) + _mb = ( + sum( + (p.nelement() * p.element_size() if not (hasattr(p, "is_meta") and p.is_meta) else 0) + for p in m.parameters() + ) + / 1048576 + ) + if _mb > _ON_DEMAND_THRESHOLD: + # Skip if this submodule is already a separate DiT + # module (e.g. transformer.language_model when both + # "transformer.language_model" and "transformer" are + # in _dit_modules) — hooks already applied. + if id(m) in all_dit_modules: + logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) + elif self._try_layerwise_offload_submodule(m, name, plan): + pass # layerwise hooks applied + else: + self._register_on_demand_hook(m, name) + else: + try: + m.to(self.device) + except (NotImplementedError, RuntimeError): + # Meta params: load from safetensors via mmap, then move to GPU + self._load_module_weights_from_mmap(m, dit_name, name) + # Materialize any remaining meta buffers (e.g. RoPE + # inv_freq, timestep freqs) on device and + # recompute them from their original formulas. + # These are non-persistent buffers not present in + # the checkpoint — they must be reconstructed, not + # left as torch.empty. + # NOTE: Do NOT call reset_parameters() here — it + # would re-initialize ALL parameters (including + # mmap-loaded weights) with random values. + _has_meta_buf = any(hasattr(b, "is_meta") and b.is_meta for b in m.buffers(recurse=True)) + if _has_meta_buf: + # Save mmap-loaded param values before + # reset_parameters (which re-initializes + # ALL params with random values). + saved_params = { + n: p.data.clone() + for n, p in m.named_parameters() + if not (hasattr(p, "is_meta") and p.is_meta) + } + m.to_empty(device=self.device) + for submod in m.modules(): + if hasattr(submod, "reset_parameters"): + submod.reset_parameters() + # Restore mmap-loaded weights + for n, p in m.named_parameters(): + if n in saved_params: + p.data.copy_(saved_params[n]) + try: + m.to(self.device) + except Exception: + logger.warning("Module %s still has meta params after mmap load", name) + else: + logger.debug(f"Skipped blocks module {name}") + + # Move top-level params/buffers to GPU + for param in dit_module._parameters.values(): + if param is not None and not (hasattr(param, "is_meta") and param.is_meta): + param.data = param.data.to(self.device, non_blocking=True) + for buffer in dit_module._buffers.values(): + if buffer is not None: + buffer.data = buffer.data.to(self.device, non_blocking=True) + + # Register hooks in a circular sliding window: + # last block prefetches first block, block i prefetches block (i+1) + # All hooks share 2 global device buffers (RFC: "exactly two layers on device") + last_block, first_block = blocks[-1], blocks[0] + + # Pass 1: create hooks with deferred buffer allocation + # (shared_buffers=[None,None] prevents per-hook OOM on large models) + last_hook = apply_distributed_block_hook( + last_block, + first_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + + block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] + for i, block in enumerate(blocks[:-1]): + next_block = blocks[(i + 1) % num_blocks] + hook = apply_distributed_block_hook( + block, + next_block, + self.device, + self.dp_group, + self.dp_size, + self.rank, + self.copy_stream, + self.comm_stream, + self.config.pin_cpu_memory, + shared_buffers=[None, None], + ) + block_hooks.append(hook) + + # Wire backward references for cache-dit fallback + for i in range(len(block_hooks)): + block_hooks[i]._prev_hook = block_hooks[i - 1] + + # Assign slots in list order: block_hooks = [last_hook, block0, ..., blockN-2] + # This ensures last_hook.current_slot != block0_hook.current_slot, + # so the circular prefetch (last_hook -> block0) writes to a + # different slot than block0 reads from. Correct for ALL N. + for i, hook in enumerate(block_hooks): + hook.current_slot = i % 2 + + # Mark block0_hook (index 1) as group-first + if len(block_hooks) > 1: + block_hooks[1]._is_group_first = True + + # Defer buffer allocation — collected for unified allocation below + self._all_hook_groups.append(block_hooks) + self._blocks.append(blocks) + + if not self._all_hook_groups: + return + + # Unified allocation: 2 shared output buffers + 2 shared shard buffers + # sized to the max block across ALL module groups (gen_layers + + # language_model). Groups execute sequentially, so 2 buffers suffice. + all_hooks: list[DistributedLayerwiseOffloadHook] = [] + for group in self._all_hook_groups: + all_hooks.extend(group) + + unified_buffers = self._allocate_shared_buffers(all_hooks) + unified_shard_buffers = None + if self.dp_size > 1: + unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) + + # Shared slot-group tracker: _shared_slot_group[slot] = group_id + # that last wrote to that slot. Group-first hooks use this to + # skip sync-prefetch when the slot still contains their own data. + shared_slot_group = [-1, -1] + + for group_idx, group in enumerate(self._all_hook_groups): + for hook in group: + hook.gpu_buffers = unified_buffers + hook._owns_buffers = False + if unified_shard_buffers is not None: + hook.gpu_shard_buffers = unified_shard_buffers + hook._group_id = group_idx + hook._shared_slot_group = shared_slot_group + + # Prefetch first block of the FIRST module group only. + # Subsequent groups share the same 2 device buffers; prefetching + # them now would overwrite the first group's data in the shared + # buffer (both groups default to slot 0). Instead, subsequent + # groups' first blocks remain as meta placeholders, and their + # pre_forward will sync-prefetch on-demand via the is_materialized + # check — by which point the first group's forward has completed + # and its buffer slots are free. + if self._all_hook_groups: + group = self._all_hook_groups[0] + first_slot = group[0].current_slot + group[-1].prefetch_layer(slot=first_slot, non_blocking=False) + group[-1].get_weights(first_slot) + + total_blocks = sum(len(b) for b in self._blocks) + logger.info( + f"Distributed layer-wise offloading enabled on {total_blocks} blocks " + f"across {len(self._all_hook_groups)} group(s), dp_size={self.dp_size}, " + f"unified shared_buffers=2" + ) + + self.enabled = True + + # Release mmap file handles — _shard_and_pin has copied all shards, + # params now point to offload placeholders (not mmap views). + self._release_mmap_handles() + + # Assign GPU buffers to sharded on-demand modules (VAE/encoders). + # Each module gets a dedicated input (shard-sized) and output + # (full-sized) buffer. VAE/encoders run before/after DiT, never + # concurrently, so peak HBM = max(DiT, VAE) not sum. + for si in self._on_demand_shard_infos: + out_bufs: dict[torch.dtype, torch.Tensor] = {} + in_bufs: dict[torch.dtype, torch.Tensor] = {} + for dtype, shard in si["cpu_shards"].items(): + # AllGather output = dp_size * shard_size + full_size = shard.numel() * self.dp_size if self.dp_size > 1 else shard.numel() + out_bufs[dtype] = torch.empty(full_size, dtype=dtype, device=self.device) + in_bufs[dtype] = torch.empty(shard.shape, dtype=dtype, device=self.device) + si["gpu_output"] = out_bufs + si["gpu_input"] = in_bufs + _mb = sum(t.nelement() * t.element_size() for t in out_bufs.values()) / 1048576 + logger.info("Allocated %.0f MB GPU buffer for sharded on-demand module", _mb) + + self._cleanup_after_loading() + + def _release_mmap_handles(self) -> None: + """Release safetensors mmap file handles.""" + if hasattr(self, "_mmap_file_cache"): + self._mmap_file_cache.clear() + del self._mmap_file_cache + if hasattr(self, "_mmap_model_to_ckpt"): + del self._mmap_model_to_ckpt + logger.info("Released safetensors mmap file handles") + + def _cleanup_after_loading(self) -> None: + """Synchronize and release freed device/CPU memory after sharding.""" + current_omni_platform.synchronize() + current_omni_platform.empty_cache() + import ctypes as _ctypes + import gc as _gc + + _gc.collect() + try: + _ctypes.CDLL("libc.so.6").malloc_trim(0) + except Exception: + pass + + def disable(self) -> None: + if not self.enabled: + return + + for blocks in self._blocks: + for block in blocks: + remove_distributed_block_hook(block) + + for h in getattr(self, "_on_demand_handles", []): + h.remove() + self._on_demand_handles = [] + self._on_demand_shard_infos = [] + + self._blocks.clear() + self._all_hook_groups.clear() + self.enabled = False + logger.info("Distributed layer-wise offloading disabled") + + @staticmethod + def _allocate_shared_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate exactly 2 shared device buffers sized to the largest block. + + All hooks share these 2 buffers. At any time, slot 0 holds the current + layer's weights and slot 1 holds the next layer's weights (or vice + versa). This ensures only 2 layers' worth of weights reside on device, + regardless of the total number of blocks. + """ + max_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + dp = hook.dp_size + for dtype, metas in hook.metadata.items(): + total = sum(m["numel"] for m in metas) + # AllGather output = dp * ceil(total/dp) (padded for equal shards) + if dp > 1: + total = ((total + dp - 1) // dp) * dp + if dtype not in max_sizes or total > max_sizes[dtype]: + max_sizes[dtype] = total + + device = hooks[0].device + shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + gpu_weights: dict[torch.dtype, torch.Tensor] = {} + for dtype, total_numel in max_sizes.items(): + gpu_weights[dtype] = torch.empty(total_numel, dtype=dtype, device=device) + shared_buffers[slot] = gpu_weights + + logger.info( + "Allocated 2 shared device buffers (max block size: %s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()}, + ) + return shared_buffers + + @staticmethod + def _allocate_shared_shard_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate 2 shared shard (AllGather input) buffers sized to the + largest per-rank shard. + + All hooks share these 2 input buffers — the same device address is + reused for every layer's AllGather so HCCL can reuse its internal + communication buffers (preventing rank-dependent HBM imbalance). + Cost: 2 × max_shard_size (~184 MB) instead of 2 × N_hooks × shard. + """ + max_shard_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + for dtype, shard in hook.cpu_shards.items(): + numel = shard.numel() + if dtype not in max_shard_sizes or numel > max_shard_sizes[dtype]: + max_shard_sizes[dtype] = numel + + device = hooks[0].device + shared_shard_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + shard_bufs: dict[torch.dtype, torch.Tensor] = {} + for dtype, numel in max_shard_sizes.items(): + shard_bufs[dtype] = torch.empty(numel, dtype=dtype, device=device) + shared_shard_buffers[slot] = shard_bufs + + logger.info( + "Allocated 2 shared shard buffers (max shard size: %s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_shard_sizes.items()}, + ) + return shared_shard_buffers diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index 17d4ca2fd..d6fcb2569 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -80,11 +80,13 @@ class ModuleDiscovery: "text_encoder_2", "text_encoder_3", "image_encoder", + "vision_encoder", "mllm", ] _FALLBACK_VAE_ATTRS = [ "vae", "audio_vae", + "sound_tokenizer", ] @staticmethod diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py new file mode 100644 index 000000000..575194ea0 --- /dev/null +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Declarative OffloadPlan for distributed layerwise offload. + +Models declare this as a class attribute ``_offload_plan`` on the +pipeline class. When present, the offloader uses it instead of +heuristic block discovery. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from torch import nn + + +@dataclass(frozen=True) +class OffloadPlan: + """Optional declarative metadata for distributed layerwise offload. + + Models declare this as a class attribute ``_offload_plan`` on the + pipeline class. When present, the offloader uses it instead of + heuristic block discovery, making it easier to support new models + without offload-specific logic in constructor or forward code. + + If not declared, the offloader falls back to: + 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. + 2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes. + + Attributes: + block_attrs: Maps DiT path → tuple of block-list attribute names. + e.g. ``{"transformer": ("gen_layers",), + "transformer.language_model": ("layers",)}`` + offload_submodules: Maps child name → block-list attribute name, + for large non-DiT submodules within a DiT that should be + independently offloaded with their own hooks. + e.g. ``{"context_encoder": "layers"}`` + """ + + block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) + offload_submodules: dict[str, str] = field(default_factory=dict) + + +def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: + """Retrieve the OffloadPlan declared by the pipeline, if any.""" + return getattr(pipeline, "_offload_plan", None) + + +def supports_mmap_loading(pipeline: nn.Module) -> bool: + """Whether the pipeline supports mmap weight loading. + + A pipeline supports mmap if any module defines ``_remap_ckpt_key``, + which maps checkpoint keys to model parameter names. Models without + this method must use the regular ``load_weights()`` path. + + This check is shared between ``diffusers_loader.py`` (to decide whether + to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()`` + (to decide whether to call ``_load_weights_via_mmap``), ensuring both + paths use the same gate condition. + """ + return any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules()) diff --git a/vllm_omni/diffusion/offloader/tensor_utils.py b/vllm_omni/diffusion/offloader/tensor_utils.py new file mode 100644 index 000000000..ef2d96011 --- /dev/null +++ b/vllm_omni/diffusion/offloader/tensor_utils.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared tensor utilities for distributed layerwise offload. + +These helpers are used by both DistributedLayerwiseOffloadHook and +DistributedLayerwiseOffloadBackend, and can be reused by other +offload backends. +""" + +from __future__ import annotations + +import torch +from torch.distributed.tensor import DTensor + + +def dtype_size(dtype: torch.dtype) -> int: + """Return element size in bytes for a torch.dtype.""" + return torch.empty(1, dtype=dtype).element_size() + + +def is_dtensor(t: torch.Tensor) -> bool: + """Check if tensor is a DTensor.""" + return isinstance(t, DTensor) + + +def set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None: + """Replace target's underlying storage with value (zero-copy).""" + if is_dtensor(target): + target._local_tensor = value + else: + target.data = value + + +def make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor: + """Create a zero-element placeholder to free GPU memory.""" + if is_dtensor(tensor): + local_shape = tuple(tensor.to_local().shape) + return torch.empty(local_shape, device="meta", dtype=tensor.dtype) + return torch.empty((0,), device=tensor.device, dtype=tensor.dtype) + + +def is_materialized_tensor(t: torch.Tensor) -> bool: + """Check if tensor holds real data (not meta or empty placeholder).""" + if is_dtensor(t): + local_t = t.to_local() + return not local_t.is_meta + return not t.is_meta and t.data.numel() > 0 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index be68f071f..3255e59db 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -212,7 +212,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) load_device = ( - "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device) + "cpu" + if self.od_config.enable_cpu_offload + or self.od_config.enable_layerwise_offload + or getattr(self.od_config, "enable_distributed_layerwise_offload", False) + else str(self.device) ) def get_memory_context() -> AbstractContextManager[Any]: @@ -478,7 +482,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # better perf. HSDP2's fully_shard pre-forward hooks need tensor version # counters, which inference tensors do not track. use_hsdp = od_config.parallel_config.use_hsdp - grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() + use_distributed_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) + grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode() with grad_context: for req in reqs: self._prepare_request_for_forward( diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 58c8fd1eb..9a2958fcf 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -417,12 +417,34 @@ class DiffusionWorker: def execute_model( self, - req: OmniDiffusionRequest, + req: OmniDiffusionRequest | list[OmniDiffusionRequest], od_config: OmniDiffusionConfig, kv_prefetch_job: KVPrefetchJob | None = None, ) -> DiffusionOutput: - """Execute a forward pass by delegating to the model runner.""" + """Execute a forward pass by delegating to the model runner. + + If *req* is a list (DP multi-concurrency), each rank picks one + request based on its distributed rank. AllGather in the layerwise + offload only gathers weight shards (request-independent), so all + ranks stay synchronised at each AllGather call while computing + different activations. + + Each rank returns its OWN DiffusionOutput (no gather). The executor + collects N responses via result_mq (all ranks share one queue). + """ assert self.model_runner is not None, "Model runner not initialized" + + # DP multi-concurrency: pick one request per DP rank. + # Use rank_in_group (DP rank) not global rank, so that SP/TP/CFG + # ranks within the same DP replica select the same request. + is_batch = isinstance(req, list) + if is_batch: + from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank + + dp_rank = get_data_parallel_rank() + idx = dp_rank % len(req) + req = req[idx] + if self.lora_manager is not None: try: self.lora_manager.set_active_adapter(req.sampling_params.lora_request, req.sampling_params.lora_scale) @@ -436,6 +458,14 @@ class DiffusionWorker: output = self.model_runner.execute_model(req, kv_prefetch_job=kv_prefetch_job) if profiler: profiler.step() + + # Each primary rank returns its own output tagged with its DP rank + # so the executor can match results to requests by dp_rank. + # Non-primary ranks (SP/TP > 0) do not reply. + if is_batch: + from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank + + return {"dp_rank": get_data_parallel_rank(), "output": output} return output def execute_model_batch( @@ -781,17 +811,12 @@ class WorkerProc: self.result_mq = None self.result_mq_handle = None - # Setup result sender (only for rank 0) - if gpu_id == 0: - self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) - self.result_mq_handle = self.result_mq.export_handle() - WorkerProc._shared_result_handle = self.result_mq_handle - logger.info(f"Worker {gpu_id} created result MessageQueue") - else: - handle = getattr(WorkerProc, "_shared_result_handle", None) - if handle: - self.result_mq = MessageQueue.create_from_handle(handle, gpu_id) - logger.info(f"Worker {gpu_id} attached to shared result MessageQueue") + # Each worker creates its own result MessageQueue (as writer). + # The executor creates one reader per worker to collect responses. + # This supports DP multi-concurrency where all ranks reply independently. + self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0]) + self.result_mq_handle = self.result_mq.export_handle() + logger.info(f"Worker {gpu_id} created result MessageQueue") assert od_config.master_port is not None @@ -920,12 +945,33 @@ class WorkerProc: output_rank = rpc_request.get("output_rank") exec_all_ranks = rpc_request.get("exec_all_ranks", False) collect_rank_status = rpc_request.get("collect_rank_status", False) + wave_id = rpc_request.get("wave_id") if collect_rank_status and not exec_all_ranks: raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather") should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id - should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None + # For DP multi-concurrency (output_rank=None), only the primary rank + # within each DP replica should reply. This prevents SP/TP/CFG/PP + # ranks from enqueuing extra replies that the executor doesn't drain. + if output_rank is None and exec_all_ranks: + from vllm.distributed.parallel_state import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_pipeline_parallel_rank, + get_sequence_parallel_rank, + ) + + is_primary_in_replica = ( + get_sequence_parallel_rank() == 0 + and get_classifier_free_guidance_rank() == 0 + and get_tensor_model_parallel_rank() == 0 + and get_pipeline_parallel_rank() == 0 + ) + should_reply = is_primary_in_replica and self.result_mq is not None + else: + should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None if not should_execute: return None, False @@ -968,6 +1014,7 @@ class WorkerProc: "method": method, "result": result, "rank_statuses": rank_statuses, + "wave_id": wave_id, }, True, ) @@ -975,6 +1022,8 @@ class WorkerProc: if rpc_exception is not None: raise rpc_exception + if isinstance(result, dict) and wave_id is not None: + result["wave_id"] = wave_id return result, should_reply def _worker_busy_loop(self) -> None: @@ -1016,10 +1065,16 @@ class WorkerProc: self._return_result(result, rpc_id=rpc_id) except Exception as e: logger.error(f"Error processing RPC: {e}", exc_info=True) + # Apply the same reply gate as the success path so + # non-output ranks don't enqueue stale error replies + # that compete with the expected responder's message. + output_rank = msg.get("output_rank") + exec_all_ranks = msg.get("exec_all_ranks", False) + wave_id = msg.get("wave_id") if self.result_mq is not None: if rpc_id is not None: - # Async RPC: must complete the executor''s pending - # future so collective_rpc() doesn''t hang. + # Async RPC: must complete the executor's pending + # future so collective_rpc() doesn't hang. self.result_mq.enqueue( AsyncDiffusionOutput( kind=AsyncOutputKind.RPC_RESULT, @@ -1027,8 +1082,36 @@ class WorkerProc: error=str(e), ) ) - else: - self._return_result({"status": "error", "error": str(e)}) + elif output_rank is None and exec_all_ranks: + # DP multi-concurrency: primary ranks reply, tagged + from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + ) + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_rank, + get_data_parallel_rank, + get_pipeline_parallel_rank, + get_sequence_parallel_rank, + ) + + is_primary = ( + get_sequence_parallel_rank() == 0 + and get_classifier_free_guidance_rank() == 0 + and get_tensor_model_parallel_rank() == 0 + and get_pipeline_parallel_rank() == 0 + ) + if is_primary: + try: + dp_rank = get_data_parallel_rank() + except Exception: + dp_rank = self.gpu_id + self._return_result( + {"status": "error", "error": str(e), "dp_rank": dp_rank, "wave_id": wave_id} + ) + elif output_rank is None or output_rank == self.gpu_id: + # Normal RPC: only the expected rank replies + self._return_result({"status": "error", "error": str(e), "wave_id": wave_id}) elif isinstance(msg, dict) and msg.get("type") == "shutdown": logger.info("Worker %s: Received shutdown message", self.gpu_id) @@ -1103,7 +1186,7 @@ class WorkerProc: pipe_writer.send( { "status": "ready", - "result_handle": worker_proc.result_mq_handle if rank == 0 else None, + "result_handle": worker_proc.result_mq_handle, } ) worker_proc._worker_busy_loop() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ed3f5464d..096649b1f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -480,6 +480,8 @@ class OrchestratorArgs: num_weight_load_threads: int = 4 enable_cpu_offload: bool = False enable_layerwise_offload: bool = False + enable_distributed_layerwise_offload: bool = False + dlo_use_allgather: bool = True boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a01d905b1..1fd747df5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1008,6 +1008,8 @@ class AsyncOmniEngine: "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), + "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), + "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 83e975528..2d014891e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -641,6 +641,30 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + omni_config_group.add_argument( + "--enable-distributed-layerwise-offload", + action="store_true", + help="Enable distributed layerwise offloading with H2D + AllGather overlap. " + "Shards weights across DP ranks, stores only 1/DP_size on each host, " + "and overlaps H2D transfers and AllGather with computation. " + "DP size is automatically derived from the parallel configuration.", + ) + omni_config_group.add_argument( + "--dlo-use-allgather", + dest="dlo_use_allgather", + action="store_true", + default=True, + help="Use shard + AllGather for weight reconstruction (default: True). " + "When disabled (--dlo-no-use-allgather), each rank loads full weights " + "via H2D only — no sharding, no AllGather, no concurrent request " + "requirement, but N× CPU memory.", + ) + omni_config_group.add_argument( + "--dlo-no-use-allgather", + dest="dlo_use_allgather", + action="store_false", + help="Disable AllGather: each rank loads full weights independently.", + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio",