From 1c2a81f6d84aea4fff53bd2f894c2a287c237245 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:02:29 +0800 Subject: [PATCH] feat(minimax-h3): enable 2x RTX 4090/5090 with DLO (#5764) Signed-off-by: lishunyang12 Signed-off-by: lishunyang12 Co-authored-by: lishunyang12 --- .../test_examples/l4_performance_tests.inc.md | 2 +- recipes/MiniMaxAI/MiniMax-H3-5090.md | 61 +++ recipes/MiniMaxAI/MiniMax-H3.md | 118 ++++- tests/diffusion/attention/test_cudnn_attn.py | 64 +++ .../minimax_h3/test_minimax_h3_contract.py | 285 +++++++++++ .../test_distributed_layerwise_backend.py | 231 ++++++++- .../test_multiproc_engine_concurrency.py | 50 ++ .../test_async_omni_diffusion_config.py | 30 ++ vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 3 + .../diffusion/attention/backends/abstract.py | 6 + .../attention/backends/cudnn_attn.py | 22 +- vllm_omni/diffusion/data.py | 8 +- vllm_omni/diffusion/diffusion_engine.py | 2 +- .../diffusion/executor/multiproc_executor.py | 26 +- .../diffusion/models/minimax_h3/encoder.py | 68 +++ .../minimax_h3/minimax_h3_transformer.py | 35 +- .../models/minimax_h3/pipeline_minimax_h3.py | 135 ++++- vllm_omni/diffusion/models/minimax_h3/vae.py | 51 +- vllm_omni/diffusion/offloader/__init__.py | 2 + vllm_omni/diffusion/offloader/base.py | 27 +- .../distributed_layerwise_backend.py | 465 ++++++++++++++---- .../diffusion/offloader/module_collector.py | 4 +- .../diffusion/offloader/module_residency.py | 89 ++++ vllm_omni/diffusion/offloader/offload_plan.py | 15 +- vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 18 +- 28 files changed, 1658 insertions(+), 162 deletions(-) create mode 100644 recipes/MiniMaxAI/MiniMax-H3-5090.md create mode 100644 tests/diffusion/attention/test_cudnn_attn.py create mode 100644 vllm_omni/diffusion/offloader/module_residency.py diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index 5595732de..48d327334 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -200,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par | num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. | | request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. | | max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). | -| baseline | object | No | see below | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS. +| baseline | object | No | see below | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS. diff --git a/recipes/MiniMaxAI/MiniMax-H3-5090.md b/recipes/MiniMaxAI/MiniMax-H3-5090.md new file mode 100644 index 000000000..fb8b0f118 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-5090.md @@ -0,0 +1,61 @@ +# MiniMax-H3 on RTX 5090 + +This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a +second GPU is available, and distributed layerwise offload (DLO). It is a +memory-first serving configuration; lower resident counts reduce HBM use and +increase CPU-to-GPU transfer time. + +## Capacity requirements + +| Resource | One RTX 5090 | Two RTX 5090s | +| --- | ---: | ---: | +| GPU HBM | 32 GiB | 32 GiB per GPU | +| Checkpoint storage | 135 GiB per partition | 135 GiB per partition | +| Available system RAM | 200 GiB minimum | 200 GiB minimum | +| Recommended system RAM | 384 GiB | 384 GiB | + +`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one +server at a time. DLO keeps rank-local weights in pinned host memory; increasing +`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the +current implementation because resident layers retain pinned CPU master copies. + +> **Modular H3:** after #5720 lands, preserve this recipe's one-partition +> behavior with `--task-type fl2va` or `--task-type ref2va`. + +## One RTX 5090: 1344x768, 5 seconds + +Use 12 resident DiT layers. A 50-step B300 allocation test with this exact +single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target +card before increasing the resident count. + +```bash +CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \ + --omni --trust-remote-code --host 0.0.0.0 --port 8000 \ + --num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \ + --usp 1 --ring 1 --vae-patch-parallel-size 1 \ + --vae-parallel-mode tile --vae-use-tiling \ + --enable-distributed-layerwise-offload --dlo-no-use-allgather \ + --dlo-resident-layers 12 --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +## Two RTX 5090s: 1344x768, 5 seconds + +Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at +27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness +proxy, not a consumer-GPU latency claim. + +```bash +CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \ + --omni --trust-remote-code --host 0.0.0.0 --port 8000 \ + --num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \ + --usp 1 --ring 1 --vae-patch-parallel-size 2 \ + --vae-parallel-mode tile --vae-use-tiling \ + --enable-distributed-layerwise-offload --dlo-no-use-allgather \ + --dlo-resident-layers 20 --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +For Ref2VA, stop the FL2VA server and restart the same command with +`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length +can increase activation memory; begin with one request at a time. diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md index fcdc836e4..9fc59628d 100644 --- a/recipes/MiniMaxAI/MiniMax-H3.md +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -33,8 +33,10 @@ export MODEL_ROOT=/path/to/MiniMax-H3 hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}" ``` -Install vLLM-Omni from the checkout containing MiniMax H3 support. On -Blackwell, install the optional FlashAttention-4 dependency: +Install vLLM-Omni from the checkout containing MiniMax H3 support. The +two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need +FlashAttention-4. Install the optional dependency only for the four-GPU +B300/GB200 `FLASH_ATTN` profile: ```bash uv venv @@ -50,6 +52,26 @@ reference-video preparation and MP4 output. One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA and FL2VA requests, or to `Ref2VA` for either Ref2VA request. +### Memory and storage requirements + +Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each +H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of +BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally +therefore needs roughly **270 GiB** of model storage, although the examples +serve only one partition at a time. + +CPU offload and distributed layerwise offload reduce GPU residency; they do +not make the model weights disappear. With `--dlo-no-use-allgather`, each +worker retains its standard-loader rank-local weights in host memory, including +pinned CPU buffers used for H2D streaming. Use at least **200 GiB available +system RAM** before starting the two-GPU recipe; a **384 GiB host is +recommended** to leave room for the OS, CUDA/PyTorch allocations, request +inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the +same time on a host sized for this minimum. + +The consumer-GPU profiles below are HBM budgets only. They still require the +host-RAM budget above. + ### Single GPU: accuracy and memory first The single-GPU configuration uses model-level CPU offload. @@ -77,6 +99,98 @@ Use a GPU with enough memory for the active H3 component and enough system RAM for the offloaded components. CPU offload reduces GPU memory pressure but adds PCIe/NVLink transfer latency. +### Two 24/32 GB GPUs: TP2 distributed layerwise offload + +For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload +(DLO). The standard loader first creates the rank-local TP shard. DLO keeps +that shard in pinned host memory and streams the 30 tail DiT blocks through a +shared two-buffer window without a DP AllGather. The first 20 DiT blocks are +copied to the GPUs once per denoise stage, reused by every sampling step, and +released before VAE decode so the decoder can reuse their HBM. + +```bash +export MODEL="${MODEL_ROOT}/FL2VA" +export PORT=8091 + +CUDA_VISIBLE_DEVICES=0,1 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 2 \ + --tensor-parallel-size 2 \ + --usp 1 \ + --ring 1 \ + --text-encoder-tp-size 2 \ + --vae-patch-parallel-size 2 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --dlo-resident-layers 20 \ + --enforce-eager \ + --diffusion-attention-backend CUDNN_ATTN +``` + +Use the profile that matches the per-GPU memory capacity: + +| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status | +|---|---:|---:|---:|---|---|---| +| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated | +| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point | + +This topology uses all available parallel capacity: TP2 shards both the DiT +and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard +without reconstructing full blocks, and VAE patch parallelism splits tiled +decode across both GPUs. cuDNN attention is selected explicitly for the RTX +consumer path; the server stays eager to avoid an unqualified compile path. + +The resident count changes placement and transfer frequency only; it does not +quantize or change the BF16/FP32 denoise math. Re-measure peak memory before +increasing it on a different request shape. + +### RTX 5090 target-hardware validation + +At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on +2 x RTX 5090 without OOM: + +| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation | +|---:|---:|---:|---:|---| +| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed | + +This is a single end-to-end validation run, not a warmed multi-run latency +benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator +high-water mark. The environment used vLLM 0.26.0, vLLM-Omni +`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The +[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png) +captures the environment, output contract, elapsed time, and sampled peak. + +Before the target run, both profiles were exercised on two B300 ranks as an +allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the +20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer +profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident +and fully streamed placements produced identical decoded video-frame and audio +hashes for the same shape, step count, prompt, and seed. The B300 result does +not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative +starting point until it is measured on that GPU. + +To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate +every MP4's H.264/AAC streams, and retain live server and GPU-memory logs: + +```bash +RUN_ROOT=/path/to/run-root \ +MODEL_ROOT=/path/to/MiniMax-H3 \ +GPU_IDS=0,1 \ +PROFILE=rtx5090 \ +bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh +``` + +The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for +`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default. + ### Four GPUs: measured best-practice throughput The best validated four-GPU configuration on four NVIDIA B300 GPUs is: diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py new file mode 100644 index 000000000..de1adf8c7 --- /dev/null +++ b/tests/diffusion/attention/test_cudnn_attn.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import nullcontext + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): + observed = {} + + def fake_sdpa(query, key, value, **kwargs): + observed.update(query=query, key=key, value=value, kwargs=kwargs) + return query + + monkeypatch.setattr( + "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel", + lambda _backends: nullcontext(), + ) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, 8, 2, 4) + key = torch.randn_like(query) + value = torch.randn_like(query) + + output = impl.forward_cuda( + query, + key, + value, + AttentionMetadata(extra={"valid_kv_length": 5}), + ) + + assert output.shape == query.shape + assert observed["query"].shape == (1, 2, 8, 4) + assert observed["key"].shape == (1, 2, 5, 4) + assert observed["value"].shape == (1, 2, 5, 4) + assert observed["kwargs"]["attn_mask"] is None + + +def test_cudnn_rejects_invalid_valid_kv_length(): + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=4, + softmax_scale=0.5, + ) + query = torch.randn(1, 8, 2, 4) + + with pytest.raises(ValueError, match="valid_kv_length"): + impl.forward_cuda( + query, + query, + query, + AttentionMetadata(extra={"valid_kv_length": 9}), + ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 95c379d7b..460aceaa7 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -30,6 +30,13 @@ def test_pipeline_import_registry_and_component_discovery(): assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] +def test_dlo_offload_plan_includes_token_refiner(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"} + assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"}) + + def test_joint_postprocess_is_multiprocessing_picklable(): from vllm_omni.diffusion.models.minimax_h3 import ( get_minimax_h3_post_process_func, @@ -87,6 +94,87 @@ def test_shifted_sigma_schedule_matches_reference_values(): ) +def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = True + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + # Model-side Ulysses shards rows before Attention gathers them again. + # The Python packed_total must therefore remain global even though q is + # local here. + q = torch.randn(4, 2, 4) + cu_seqlens = torch.tensor([0, 5, 8], dtype=torch.int32) + + output = attention._run_packed_attention( + q, + q, + q, + cu_seqlens=cu_seqlens, + max_seqlen=5, + packed_total=8, + ) + + assert output.shape == q.shape + assert attention.attention.metadata.attn_mask is None + assert attention.attention.metadata.extra["valid_kv_length"] == 5 + + +def test_packed_attention_keeps_padding_mask_for_other_backends(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = False + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + q = torch.randn(8, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + assert torch.equal( + attention.attention.metadata.attn_mask, + torch.tensor([[True, True, True, True, True, False, False, False]]), + ) + + def test_reference_image_resize_contract(): from PIL import Image @@ -293,6 +381,127 @@ def test_layerwise_offload_releases_text_encoder(): pipeline.text_encoder.offload_to_cpu.assert_called_once_with() +def test_distributed_layerwise_offload_releases_text_encoder(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + +def test_distributed_layerwise_offload_stages_vae_component(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + ) + component = Mock() + + with pipeline._component_on_device(component): + component.load_to_device.assert_called_once_with() + component.offload_to_cpu.assert_not_called() + + component.offload_to_cpu.assert_called_once_with() + + +def test_distributed_layerwise_resident_blocks_are_stage_scoped(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + controller = Mock() + pipeline._dlo_residency_controller = controller + + with pipeline._resident_dit_layers_on_device(): + controller.load_resident_layers.assert_called_once_with() + controller.offload_resident_layers.assert_not_called() + + controller.offload_resident_layers.assert_called_once_with() + + +def test_distributed_layerwise_resident_blocks_release_on_failure(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + controller = Mock() + pipeline._dlo_residency_controller = controller + + with pytest.raises(RuntimeError, match="denoise failed"): + with pipeline._resident_dit_layers_on_device(): + raise RuntimeError("denoise failed") + + controller.load_resident_layers.assert_called_once_with() + controller.offload_resident_layers.assert_called_once_with() + + +def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + class Stack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + class TextStack(torch.nn.Module): + def __init__(self, count): + super().__init__() + self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)]) + + hooks = [] + + def fake_apply(block, next_block, device, stream, pin_memory): + hook = SimpleNamespace( + block=block, + next_block=next_block, + device=device, + stream=stream, + pin_memory=pin_memory, + _prev_hook=None, + offload_layer=Mock(), + ) + hooks.append(hook) + return hook + + monkeypatch.setattr( + "vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook", + fake_apply, + ) + monkeypatch.setattr( + "vllm_omni.platforms.current_omni_platform.Stream", + Mock(return_value="copy-stream"), + ) + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + encoder.device_target = torch.device("cpu") + encoder.vision = Stack(2) + encoder.text_model = TextStack(3) + + encoder.enable_omni_layerwise_offload(pin_memory=False) + + assert len(hooks) == 5 + assert all(hook._prev_hook is not None for hook in hooks) + assert all(hook.device == torch.device("cpu") for hook in hooks) + + def test_video_vae_keeps_reference_fp32_weights(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module @@ -324,6 +533,82 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch): assert next(video_vae.parameters()).dtype == torch.float32 +def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("meta"), + load_device=torch.device("cpu"), + ) + + assert next(video_vae.parameters()).device.type == "cpu" + assert video_vae._device_target.type == "meta" + + +def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1) + + stager = Mock() + stager_cls = Mock(return_value=stager) + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("cuda"), + load_device=torch.device("cpu"), + ) + video_vae.load_to_device() + video_vae.offload_to_cpu() + + stager_cls.assert_called_once_with( + video_vae.remote, + torch.device("cuda"), + pin_memory=True, + ) + stager.load.assert_called_once_with() + stager.offload.assert_called_once_with() + + def test_video_vae_encode_uses_configured_parallel_tiling(): from vllm_omni.diffusion.models.minimax_h3.vae import ( MiniMaxH3VideoVAE, diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 5a72ddeb0..fd004d056 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import ( from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( DistributedLayerwiseOffloadBackend, DistributedLayerwiseOffloadHook, + PinnedResidentLayerGroup, ) +from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan from vllm_omni.platforms import current_omni_platform @@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch): 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) + monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None) class TinyBlock(nn.Module): @@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook: assert shard1.numel() == 2 assert torch.equal(shard1, torch.tensor([102.0, 0.0])) + def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) + next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) + next_block.weight.mmap_weight_transform_pending = True + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + + assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0])) + + def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32)) + next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + + assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0])) + + +class TestPinnedResidentLayerGroup: + def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): + blocks = [nn.Linear(2, 2), nn.Linear(2, 2)] + expected = [] + for index, block in enumerate(blocks): + block.weight.data.fill_(index + 1) + block.bias.data.fill_(index + 3) + expected.append((block.weight.detach().clone(), block.bias.detach().clone())) + + group = PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + ) + + assert all(block.weight.numel() == 0 for block in blocks) + + group.load() + for block, (weight, bias) in zip(blocks, expected): + assert torch.equal(block.weight, weight) + assert torch.equal(block.bias, bias) + + # Device-side changes must not overwrite the persistent host master. + blocks[0].weight.data.zero_() + group.offload() + assert all(block.weight.numel() == 0 for block in blocks) + + group.load() + for block, (weight, bias) in zip(blocks, expected): + assert torch.equal(block.weight, weight) + assert torch.equal(block.bias, bias) + + group.offload() + + def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): + module = nn.Linear(2, 2, bias=False) + + def transform(tensor): + return tensor.flip(0) + + module.weight.mmap_weight_transform = transform + + backend = object.__new__(DistributedLayerwiseOffloadBackend) + backend._remember_mmap_param_attrs(module) + module.to_empty(device="meta") + + replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float()) + backend._attach_mmap_param_attrs("weight", replacement, module.weight) + + assert replacement.mmap_weight_transform is transform + assert replacement.mmap_weight_transform_pending is True + + def test_shared_allgather_output_is_narrowed_to_current_block( + self, dist_group, patched_offload_runtime, monkeypatch + ): + block = nn.Module() + block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32)) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32)) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=object(), + dp_size=2, + rank=0, + pin_memory=False, + ) + hook.initialize_hook(block) + # Simulate unified buffers sized for a larger heterogeneous block. + hook.gpu_buffers = [ + {torch.float32: torch.empty(12)}, + {torch.float32: torch.empty(12)}, + ] + hook.gpu_shard_buffers = [ + {torch.float32: torch.empty(6)}, + {torch.float32: torch.empty(6)}, + ] + output_sizes = [] + + def fake_allgather(output, local_shard, *, group): + del group + output_sizes.append(output.numel()) + output.zero_() + output[: local_shard.numel()].copy_(local_shard) + + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather) + + hook.prefetch_layer(slot=0, non_blocking=False) + + assert output_sizes == [4] + + +class TestPinnedModuleStager: + def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime): + module = nn.Linear(2, 2) + expected_weight = module.weight.detach().clone() + expected_bias = module.bias.detach().clone() + stager = PinnedModuleStager( + module, + torch.device("cpu"), + pin_memory=False, + copy_stream=DummyStream(), + ) + + stager.load() + module.weight.data.zero_() + module.bias.data.zero_() + stager.offload() + + assert torch.equal(module.weight, expected_weight) + assert torch.equal(module.bias, expected_bias) + stager.load() + assert torch.equal(module.weight, expected_weight) + assert torch.equal(module.bias, expected_bias) + stager.offload() + class _DummyBlock(nn.Module): def __init__(self): @@ -592,6 +753,7 @@ class TestOffloadPlan: plan = OffloadPlan( block_attrs={"transformer": ("gen_layers",)}, offload_submodules={"context_encoder": "layers"}, + resident_dit_paths=frozenset({"transformer"}), ) class PipelineWithPlan(nn.Module): @@ -602,12 +764,14 @@ class TestOffloadPlan: assert result is plan assert result.block_attrs == {"transformer": ("gen_layers",)} assert result.offload_submodules == {"context_encoder": "layers"} + assert result.resident_dit_paths == frozenset({"transformer"}) 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 == {} + assert plan.resident_dit_paths == frozenset() def test_offload_plan_is_frozen(self): """OffloadPlan should be immutable (frozen=True).""" @@ -615,6 +779,50 @@ class TestOffloadPlan: with pytest.raises(Exception): plan.block_attrs = {"x": ("y",)} # type: ignore + def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime): + class TokenRefiner(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) + self.token_refiner = TokenRefiner() + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan( + offload_submodules={"token_refiner": "blocks"}, + resident_dit_paths=frozenset({"transformer"}), + ) + + def __init__(self): + super().__init__() + self.transformer = Transformer() + + pipeline = Pipeline() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + dlo_resident_layers=2, + ), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert len(backend._resident_blocks) == 2 + assert len(backend._all_hook_groups) == 1 + assert len(backend._all_hook_groups[0]) == 2 + assert backend.enabled + + backend.disable() + class TestMmapValidation: """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" @@ -821,7 +1029,7 @@ class TestConfigValidation: OffloadConfig.from_od_config(FakeODConfig()) def test_hsdp_without_allgather_allowed(self): - """HSDP + DLO + no-AllGather should be allowed (full weights per rank).""" + """HSDP + DLO + no-AllGather uses standard-loader rank-local weights.""" from vllm_omni.diffusion.offloader.base import OffloadConfig class FakePC: @@ -843,12 +1051,33 @@ class TestConfigValidation: enable_layerwise_offload = False enable_distributed_layerwise_offload = True dlo_use_allgather = False # no AllGather → should be allowed + dlo_resident_layers = 20 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 + assert config.dlo_resident_layers == 20 + + def test_resident_layers_with_allgather_rejected(self): + class FakePC: + data_parallel_size = 2 + use_hsdp = False + sequence_parallel_size = 1 + + class FakeODConfig: + enable_cpu_offload = False + enable_layerwise_offload = False + enable_distributed_layerwise_offload = True + dlo_use_allgather = True + dlo_resident_layers = 20 + pin_cpu_memory = True + parallel_config = FakePC() + model = "/fake/path" + + with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"): + OffloadConfig.from_od_config(FakeODConfig()) def test_num_inference_steps_none_rejected(self): """DP multi-concurrency should reject None num_inference_steps.""" diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 932630a23..f1c48ba30 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -247,6 +247,56 @@ class TestRequestModeDispatch: results = {ro.request_id: ro.result.error for ro in out.runner_outputs} assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids} + def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self): + executor, _, _ = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + enable_distributed_layerwise_offload=True, + dlo_use_allgather=True, + ) + executor.execute_request = Mock(return_value="dlo-dp") + executor.collective_rpc = Mock() + scheduler_output = _make_sched_output("A", "B") + + result = executor.execute_batch(scheduler_output) + + assert result == "dlo-dp" + executor.execute_request.assert_called_once_with(scheduler_output) + executor.collective_rpc.assert_not_called() + + def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self): + executor, req_q, res_q = _make_executor(num_gpus=2) + executor.od_config = SimpleNamespace( + step_execution=False, + parallel_config=SimpleNamespace(data_parallel_size=2), + ) + executor._sync_result_buffer = res_q + + def worker(): + request = req_q.get(timeout=2) + assert "rpc_id" not in request + wave_id = request["wave_id"] + for dp_rank in range(2): + res_q.put( + { + "dp_rank": dp_rank, + "output": _tagged_output(str(dp_rank)), + "wave_id": wave_id, + } + ) + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + results = executor.collective_rpc( + "execute_model", + unique_reply_rank=None, + exec_all_ranks=True, + ) + thread.join(timeout=2) + + assert [result.error for result in results] == ["0", "1"] + # ───────────────── concurrent collective RPC ───────────────── diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 8987ae1d8..d31e9f28a 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -262,6 +262,36 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True +def test_serve_cli_forwards_distributed_offload_residency(): + """Ensure the two-GPU DLO placement controls reach the diffusion stage.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--enable-distributed-layerwise-offload", + "--dlo-no-use-allgather", + "--dlo-resident-layers", + "20", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + engine_args = stage_cfg["engine_args"] + + assert args.enable_distributed_layerwise_offload is True + assert args.dlo_use_allgather is False + assert args.dlo_resident_layers == 20 + assert engine_args["enable_distributed_layerwise_offload"] is True + assert engine_args["dlo_use_allgather"] is False + assert engine_args["dlo_resident_layers"] == 20 + + def test_serve_cli_accepts_diffusion_compile_controls(): """Ensure both compile controls reach the diffusion stage.""" parser = TrackingArgumentParser() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 0f523eb48..002352cd2 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -479,6 +479,7 @@ class _DiffusionConfigProjection: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_resident_layers: int = Field(default=0, ge=0) 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/config/stage_config.py b/vllm_omni/config/stage_config.py index 4831e6f92..602ca1e20 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -413,6 +413,9 @@ class StageDeployConfig: enable_cpu_offload: bool | None = None enable_layerwise_offload: bool | None = None + enable_distributed_layerwise_offload: bool | None = None + dlo_use_allgather: bool | None = None + dlo_resident_layers: int | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 2e30877ca..fd5a8e01e 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -15,6 +15,10 @@ class AttentionBackend(ABC): accept_output_buffer: bool = False supports_piecewise_spans: bool = False + # The backend can represent a contiguous valid K/V prefix by slicing the + # tensors instead of materializing a padding mask. Models may use this to + # avoid a slower masked-attention plan when tail padding is not semantic. + supports_prefix_kv_slicing: bool = False @classmethod def supports_attention_mask(cls) -> bool: @@ -83,6 +87,8 @@ class AttentionMetadata: # variable-length query/key sequences for FlashAttention. # "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with # the packed cu_seqlens tensors. + # "valid_kv_length": int — contiguous valid K/V prefix length for a + # backend that advertises supports_prefix_kv_slicing. # Piecewise attention metadata (mixed causal/full masks). # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index 78e428796..ffcc2fb44 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -17,6 +17,7 @@ logger = init_logger(__name__) class CuDNNAttentionBackend(AttentionBackend): accept_output_buffer: bool = True + supports_prefix_kv_slicing: bool = True @classmethod def supports_attention_mask(cls) -> bool: @@ -61,7 +62,26 @@ class CuDNNAttentionImpl(AttentionImpl): ) -> torch.Tensor: attention_mask = None if attn_metadata: - attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k") + valid_kv_length = attn_metadata.extra.get("valid_kv_length") + if attn_metadata.attn_mask is None and isinstance(valid_kv_length, int): + if not 0 < valid_kv_length <= key.shape[1]: + raise ValueError( + "valid_kv_length must be within the K/V sequence, " + f"got {valid_kv_length} for length {key.shape[1]}" + ) + # A contiguous valid prefix is mathematically equivalent to a + # broadcast key-padding mask. Slicing K/V keeps Q/output in the + # checkpoint's aligned layout while letting cuDNN select its + # much faster mask-free FMHA plan. + key = key[:, :valid_kv_length] + value = value[:, :valid_kv_length] + else: + attention_mask = _maybe_reshape_attn_mask( + query, + key, + attn_metadata.attn_mask, + mask_mode="broadcast_k", + ) enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 24260a5b7..e47b642ff 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -687,9 +687,13 @@ class OmniDiffusionConfig: # 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). + # concurrent requests in DP mode). If False: each rank streams the standard + # loader's rank-local tensors (including TP-local shards) via H2D only. + # This avoids AllGather synchronization, while host memory follows the + # loader's existing rank-local layout instead of adding a second DP shard. dlo_use_allgather: bool = True + # Leading main-DiT blocks kept resident by distributed layerwise offload. + dlo_resident_layers: int = 0 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 e89b95571..086206fa0 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -237,7 +237,7 @@ class DiffusionEngine: 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: + if dp_size > 1 and dist_offload and getattr(self.od_config, "dlo_use_allgather", True): self.scheduler.max_num_running_reqs = dp_size self.dp_concurrent = True # Set batch admission wait so concurrent requests accumulate diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index fa0213fca..fa7130e40 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -416,11 +416,8 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): ) 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) + ea = getattr(nr.req.sampling_params, "extra_args", None) + extra_args_signatures.add(json.dumps(ea, sort_keys=True, default=repr) if ea is not None else None) if len(extra_args_signatures) > 1: raise ValueError( "DP multi-concurrency requires all concurrent requests to " @@ -521,6 +518,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): if len(scheduler_output.scheduled_new_reqs) <= 1: return self.execute_request(scheduler_output) + parallel_config = getattr(self.od_config, "parallel_config", None) + dp_size = getattr(parallel_config, "data_parallel_size", 1) + if ( + dp_size > 1 + and getattr(self.od_config, "enable_distributed_layerwise_offload", False) + and getattr(self.od_config, "dlo_use_allgather", True) + ): + # DLO DP uses one independent request per DP replica. It is not a + # fused pipeline request batch, so models such as MiniMax-H3 do not + # need to advertise supports_request_batch=True. + return self.execute_request(scheduler_output) + result = self.collective_rpc( "execute_model_batch", args=(scheduler_output, self.od_config), @@ -599,7 +608,12 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # so the D2H copy runs on a side stream in the worker background # thread while the default stream is free for the next forward. # pump routes the result back to a Future via rpc_id. - if not self.od_config.step_execution and method in ("execute_model", "execute_model_batch"): + multi_rank_reply = unique_reply_rank is None and exec_all_ranks + if ( + not self.od_config.step_execution + and method in ("execute_model", "execute_model_batch") + and not multi_rank_reply + ): rpc_id = self._next_rpc_id() rpc_request["rpc_id"] = rpc_id fut: concurrent.futures.Future = concurrent.futures.Future() diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py index 438ff80a2..f811eeecb 100644 --- a/vllm_omni/diffusion/models/minimax_h3/encoder.py +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -1029,16 +1029,84 @@ class MiniMaxH3Qwen3VLEncoder(nn.Module): def load_to_device(self) -> None: if not self.is_loaded: return + if getattr(self, "_omni_layerwise_enabled", False): + for name, child in self.vision.named_children(): + if name != "blocks": + child.to(self.device_target) + for name, child in self.text_model.named_children(): + if name != "layers": + child.to(self.device_target) + return self.vision.to(self.device_target) self.text_model.to(self.device_target) def offload_to_cpu(self) -> None: if not self.is_loaded: return + if getattr(self, "_omni_layerwise_enabled", False): + for hook in self._omni_layerwise_hooks: + hook.offload_layer() + for name, child in self.vision.named_children(): + if name != "blocks": + child.to("cpu") + for name, child in self.text_model.named_children(): + if name != "layers": + child.to("cpu") + torch.accelerator.empty_cache() + return self.vision.to("cpu") self.text_model.to("cpu") torch.accelerator.empty_cache() + def enable_omni_layerwise_offload(self, *, pin_memory: bool = True) -> None: + """Stream the TP-local Qwen vision/text blocks for low-HBM serving. + + The encoder has its own TP process group, so these blocks must remain + rank-local. Reusing the DiT AllGather group would concatenate + different TP shards and corrupt the encoder weights. + """ + if not self.is_loaded or getattr(self, "_omni_layerwise_enabled", False): + return + + from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook + from vllm_omni.platforms import current_omni_platform + + self._omni_layerwise_hooks = [] + self._omni_layerwise_block_groups = [] + copy_stream = current_omni_platform.Stream() + for blocks in (self.vision.blocks, self.text_model.layers): + if len(blocks) <= 1: + continue + last_hook = apply_block_hook( + blocks[-1], + blocks[0], + self.device_target, + copy_stream, + pin_memory, + ) + hooks = [last_hook] + for index, block in enumerate(blocks[:-1]): + hooks.append( + apply_block_hook( + block, + blocks[index + 1], + self.device_target, + copy_stream, + pin_memory, + ) + ) + for index, hook in enumerate(hooks): + hook._prev_hook = hooks[index - 1] + self._omni_layerwise_hooks.extend(hooks) + self._omni_layerwise_block_groups.append(blocks) + + self._omni_layerwise_enabled = bool(self._omni_layerwise_hooks) + logger.info( + "MiniMax H3 encoder layerwise offload enabled on %d blocks across %d stacks", + sum(len(blocks) for blocks in self._omni_layerwise_block_groups), + len(self._omni_layerwise_block_groups), + ) + def _encode( self, input_ids: torch.Tensor, diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index fd5ca67e9..7bd397a15 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -376,18 +376,30 @@ class MiniMaxH3Attention(nn.Module): *, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int, ) -> torch.Tensor: """Run packed attention as a small eager island. - The scalar packed-layout metadata and the CuTe FlashAttention-4 DSL - are intentionally opaque to Dynamo. Keeping this boundary narrow lets - regional compile fuse projections, norms, RoPE, and the surrounding - DiT block without repeatedly graph-breaking inside the FA4 compiler. + The scalar packed-layout metadata and backend-specific attention + kernels are intentionally opaque to Dynamo. Keeping this boundary + narrow lets regional compile fuse projections, norms, RoPE, and the + surrounding DiT block without repeated graph breaks. """ - used = int(cu_seqlens[1].item()) - packed_total = int(cu_seqlens[-1].item()) + # max_seqlen is already the first (real) packed document length. Do + # not read the CUDA cu_seqlens scalars here: this function runs once + # per layer and .item() would serialize every attention launch. + if not 0 < max_seqlen <= packed_total: + raise ValueError( + f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}" + ) + used = min(max_seqlen, packed_total) attn_mask = None - if used < packed_total: + # Ring attention can dispatch to a different implementation from the + # configured backend, so this no-mask fast path is local-only. + prefix_slice = ( + not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing + ) + if used < packed_total and not prefix_slice: attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, @@ -396,6 +408,7 @@ class MiniMaxH3Attention(nn.Module): "cu_seqlens_k": cu_seqlens, "max_seqlen_q": max_seqlen, "max_seqlen_k": max_seqlen, + "valid_kv_length": used, }, ) return self.attention( @@ -412,6 +425,7 @@ class MiniMaxH3Attention(nn.Module): rope_freqs: torch.Tensor | None, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int | None = None, sp_seq_lens: list[int] | None = None, ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. @@ -448,6 +462,10 @@ class MiniMaxH3Attention(nn.Module): v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + # Before Ulysses, q contains only this rank's row shard. The + # backend receives the global sequence after all-to-all, so carry + # its Python length explicitly instead of inferring it from q. + packed_total=packed_total if packed_total is not None else q.shape[0], ) out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) @@ -655,6 +673,7 @@ class MiniMaxH3DiTBlock(nn.Module): rope_freqs: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, + packed_total: int, sp_seq_lens: list[int] | None = None, ) -> torch.Tensor: """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] @@ -681,6 +700,7 @@ class MiniMaxH3DiTBlock(nn.Module): rope_freqs=rope_freqs, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + packed_total=packed_total, sp_seq_lens=sp_seq_lens, ) x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) @@ -1127,6 +1147,7 @@ class MiniMaxH3DiTModel(nn.Module): rope_freqs=block_rope, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, + packed_total=seq_len, ) hidden = self.sp_gather(hidden) diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 9ab54765a..9923b8d1c 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,6 +8,7 @@ import math import os import tempfile from collections.abc import Iterable, Mapping, Sequence +from contextlib import contextmanager from pathlib import Path from typing import Any, ClassVar @@ -36,6 +37,7 @@ from vllm_omni.diffusion.models.interface import ( SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.offloader import OffloadPlan from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) @@ -454,8 +456,25 @@ class MiniMaxH3Pipeline( """CFG-distilled joint video/audio generation for MiniMax H3.""" _dit_modules: ClassVar[list[str]] = ["transformer"] + # TODO(offload): Reconcile this static plan with the modular H3 pipeline + # (#5720). Its combined mode has ``transformer`` and ``transformers_ref``; + # build the resident/staging plan from the selected task partition so a + # consumer-GPU single-task server never stages both DiTs. _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"] + _offload_plan: ClassVar[OffloadPlan] = OffloadPlan( + offload_submodules={"token_refiner": "blocks"}, + resident_dit_paths=frozenset({"transformer"}), + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "video_vae", "audio_vae"}), + ) + # H3's regular loader performs checkpoint-layout conversions (grouped QKV + # and fused MLP). Keep it on that path until mmap can run the same loader + # callbacks for every affected parameter. + _supports_mmap_loading: ClassVar[bool] = False + # TODO(offload): Re-enable after the generic rank-local mmap path can run + # this model's grouped-QKV reorder and fused-MLP packing before TP sharding. + # Do not bypass the regular loader until that equivalence is tested. _PROFILER_TARGETS: ClassVar[list[str]] = [ "_prepare_reference_videos", "encode_prompt", @@ -542,13 +561,19 @@ class MiniMaxH3Pipeline( load_model=rank < text_encoder_tp_size, encoder_group=self.text_encoder_group, ) + stage_components = bool( + od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) + ) + component_load_device = torch.device("cpu") if stage_components else self.device self.video_vae = MiniMaxH3VideoVAE( os.path.join(model_path, "video_vae"), device=self.device, + load_device=component_load_device, ) self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, + load_device=component_load_device, ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae @@ -935,12 +960,16 @@ class MiniMaxH3Pipeline( input_ids: torch.Tensor, vision_kwargs: dict[str, torch.Tensor], ) -> torch.Tensor: - if self.od_config.enable_cpu_offload: + if self.od_config.enable_cpu_offload and not getattr( + self.od_config, "enable_distributed_layerwise_offload", False + ): # Invoke nn.Module.__call__ so the generic model-level offloader # swaps the resident DiT and encoder. return self.text_encoder(input_ids, **vision_kwargs) - if self.od_config.enable_layerwise_offload: + if self.od_config.enable_layerwise_offload or getattr( + self.od_config, "enable_distributed_layerwise_offload", False + ): # Layerwise DiT offload already provides the low-residency encoder # phase used by the checkpoint reference. self.text_encoder.load_to_device() @@ -955,12 +984,33 @@ class MiniMaxH3Pipeline( self.text_encoder.load_to_device() return self.text_encoder.encode_ids(input_ids, **vision_kwargs) + def _uses_manual_component_offload(self) -> bool: + od_config = getattr(self, "od_config", None) + return bool( + getattr(od_config, "enable_layerwise_offload", False) + or getattr(od_config, "enable_distributed_layerwise_offload", False) + ) + + @contextmanager + def _component_on_device(self, component: nn.Module): + staged = self._uses_manual_component_offload() + if staged: + component.load_to_device() + try: + yield + finally: + if staged: + component.offload_to_cpu() + def _encode_visual_condition( self, image: Image.Image, ) -> torch.Tensor: _, rank, _ = _dit_rank_world() - rows = self.video_vae.encode_image(image) if rank == 0 else None + rows = None + if rank == 0: + with self._component_on_device(self.video_vae): + rows = self.video_vae.encode_image(image) return _broadcast_tensor( rows, dtype=torch.float32, @@ -996,7 +1046,8 @@ class MiniMaxH3Pipeline( rows = None audio_t = 0 if rank == 0: - rows, audio_t = self.audio_vae.encode_waveform(*audio) + with self._component_on_device(self.audio_vae): + rows, audio_t = self.audio_vae.encode_waveform(*audio) audio_t_tensor = torch.tensor( [audio_t], dtype=torch.long, @@ -1042,6 +1093,15 @@ class MiniMaxH3Pipeline( prepared_videos: list[dict[str, Any]] | None, *, count: int, + ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: + with self._component_on_device(self.video_vae): + return self._encode_video_conditions_resident(prepared_videos, count=count) + + def _encode_video_conditions_resident( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + count: int, ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: group, rank, world_size = _dit_rank_world() distributed_encode = self.video_vae.is_distributed_enabled() @@ -1089,6 +1149,15 @@ class MiniMaxH3Pipeline( prepared_videos: list[dict[str, Any]] | None, *, has_audio: list[bool], + ) -> tuple[torch.Tensor | None, list[int]]: + with self._component_on_device(self.audio_vae): + return self._encode_video_audio_conditions_resident(prepared_videos, has_audio=has_audio) + + def _encode_video_audio_conditions_resident( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + has_audio: list[bool], ) -> tuple[torch.Tensor | None, list[int]]: _, rank, _ = _dit_rank_world() count = sum(has_audio) @@ -1156,6 +1225,17 @@ class MiniMaxH3Pipeline( ) return video_rows, audio_rows + @contextmanager + def _resident_dit_layers_on_device(self): + controller = getattr(self, "_dlo_residency_controller", None) + if controller is not None: + controller.load_resident_layers() + try: + yield + finally: + if controller is not None: + controller.offload_resident_layers() + def diffuse( self, *, @@ -1277,21 +1357,22 @@ class MiniMaxH3Pipeline( num_steps=num_steps, shift_scale=audio_shift, ) - with self.progress_bar(total=len(video_sigmas) - 1) as progress: - video_rows, audio_rows = minimax_h3_denoise_loop( - model=self.transformer, - positive=branch, - initial_video_rows=initial_video, - initial_audio_rows=initial_audio, - keyframe_cond_rows=visual_anchor, - audio_ref_rows=audio_anchor, - sigmas_video=video_sigmas, - sigmas_audio=audio_sigmas, - device=self.device, - imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), - audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), - on_step=lambda step, video, audio: progress.update(), - ) + with self._resident_dit_layers_on_device(): + with self.progress_bar(total=len(video_sigmas) - 1) as progress: + video_rows, audio_rows = minimax_h3_denoise_loop( + model=self.transformer, + positive=branch, + initial_video_rows=initial_video, + initial_audio_rows=initial_audio, + keyframe_cond_rows=visual_anchor, + audio_ref_rows=audio_anchor, + sigmas_video=video_sigmas, + sigmas_audio=audio_sigmas, + device=self.device, + imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), + audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + on_step=lambda step, video, audio: progress.update(), + ) target_video = video_rows[branch.update_mask_dev] video_latent = minimax_h3_unpatchify_video_tokens( @@ -1320,14 +1401,16 @@ class MiniMaxH3Pipeline( height: int, width: int, ) -> tuple[torch.Tensor, torch.Tensor]: - with current_omni_platform.create_autocast_context( - device_type=self.device.type, - dtype=torch.float16, - enabled=True, - ): - video = self.video_vae.decode_latent(video_latent) + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + video = self.video_vae.decode_latent(video_latent) video = video[..., :height, :width].contiguous() - audio = self.audio_vae.decode_latent(audio_latent) + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) return video, audio @torch.no_grad() diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index afdd5541f..5ee194e42 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -19,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor DistributedVaeMixin, ) from vllm_omni.diffusion.distributed.parallel_state import get_dit_group +from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from .packed_tokens import minimax_h3_patchify_video_latent @@ -102,8 +103,10 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): component_path: str, *, device: torch.device, + load_device: torch.device | None = None, ) -> None: super().__init__() + self._device_target = device self.config_dict = _load_component_config(component_path) self.remote = _load_remote_component( component_path, @@ -112,13 +115,34 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): # Match the reference loader contract: video VAE weights stay FP32. # Keyframe encoding is numerically sensitive to first casting the # checkpoint through FP16; decode still runs under FP16 autocast. - self.remote.eval().to(device=device, dtype=torch.float32) + initial_device = load_device or device + self.remote.eval().to(device=initial_device, dtype=torch.float32) + self._stager = None + if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): + self._stager = PinnedModuleStager( + self.remote, + device, + pin_memory=True, + ) self.model = self.remote.model self.use_tiling = True self.use_slicing = False self.parallel_size = 1 self.device_module = torch.get_device_module() + def load_to_device(self) -> None: + if self._stager is not None: + self._stager.load() + else: + self.remote.to(self._device_target) + + def offload_to_cpu(self) -> None: + if self._stager is not None: + self._stager.offload() + else: + self.remote.to("cpu") + torch.accelerator.empty_cache() + def set_parallel_size( self, parallel_size: int, @@ -277,8 +301,10 @@ class MiniMaxH3AudioVAE(nn.Module): component_path: str, *, device: torch.device, + load_device: torch.device | None = None, ) -> None: super().__init__() + self._device_target = device self.config_dict = _load_component_config(component_path) self.remote = _load_remote_component( component_path, @@ -286,10 +312,31 @@ class MiniMaxH3AudioVAE(nn.Module): ) # The checkpoint's audio VAE contract is FP32 for both reference # encoding and waveform decoding. - self.remote.eval().to(device=device, dtype=torch.float32) + initial_device = load_device or device + self.remote.eval().to(device=initial_device, dtype=torch.float32) + self._stager = None + if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): + self._stager = PinnedModuleStager( + self.remote, + device, + pin_memory=True, + ) self.model = self.remote.model self.sample_rate = int(self.config_dict["sample_rate"]) + def load_to_device(self) -> None: + if self._stager is not None: + self._stager.load() + else: + self.remote.to(self._device_target) + + def offload_to_cpu(self) -> None: + if self._stager is not None: + self._stager.offload() + else: + self.remote.to("cpu") + torch.accelerator.empty_cache() + @torch.inference_mode() def encode_waveform( self, diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 807fad7ed..9ff464494 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -16,6 +16,7 @@ from .distributed_layerwise_backend import ( remove_distributed_block_hook, ) from .layerwise_backend import LayerWiseOffloadBackend +from .module_residency import PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading from .sequential_backend import ( ModelLevelOffloadBackend, @@ -41,6 +42,7 @@ __all__ = [ "DistributedLayerwiseOffloadBackend", "DistributedLayerwiseOffloadHook", "ModelLevelOffloadBackend", + "PinnedModuleStager", "apply_sequential_offload", "remove_sequential_offload", "apply_distributed_block_hook", diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 20b2f5e14..2db80923d 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -27,7 +27,10 @@ class OffloadConfig: 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 + # True: add DP sharding + AllGather. False: stream the standard loader's + # rank-local tensors (including TP-local shards) with H2D only. + dlo_use_allgather: bool = True + dlo_resident_layers: int = 0 # leading DiT layers kept on device model_path: str | None = None # checkpoint path for mmap weight loading @classmethod @@ -94,17 +97,26 @@ 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. + # With dlo_use_allgather=False, do not add another DP shard. Each rank + # streams the tensors produced by the standard loader, which may + # already be TP-local shards. This avoids AllGather synchronization + # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) + dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) + if dlo_resident_layers < 0: + raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}") + if dlo_resident_layers and dlo_use_allgather: + raise ValueError( + "dlo_resident_layers currently requires --dlo-no-use-allgather so " + "resident blocks use weights prepared by the standard TP-aware loader" + ) # 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)" + "streaming standard-loader rank-local weights (no DP shard or AllGather)" ) # HSDP already shards parameters into DTensors. Running distributed @@ -114,8 +126,8 @@ class OffloadConfig: 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." + "would double-shard them. Use --dlo-no-use-allgather (standard-loader " + "rank-local weights) or disable HSDP." ) return cls( @@ -124,6 +136,7 @@ class OffloadConfig: use_hsdp=use_hsdp, dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, + dlo_resident_layers=dlo_resident_layers, model_path=getattr(od_config, "model", None), ) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 2748f448e..7da7c8d37 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -219,6 +219,15 @@ class DistributedLayerwiseOffloadHook(ModelHook): weights_with_local = [] for name, t in name2weights.items(): local_t = t.to_local() if hasattr(t, "to_local") else t + mmap_transform = getattr(t, "mmap_weight_transform", None) + if callable(mmap_transform) and getattr(t, "mmap_weight_transform_pending", False): + # Some checkpoints use a layout that is converted by the + # regular weight loader (for example MiniMax-H3 grouped + # QKV). Apply that conversion one block at a time while + # copying the rank-local CPU shard. Keeping the raw + # parameter as an mmap view avoids a private full-model + # copy in every worker. + local_t = mmap_transform(local_t) weights_with_local.append((name, t, local_t)) total_numel = sum(local.numel() for _, _, local in weights_with_local) @@ -331,9 +340,9 @@ class DistributedLayerwiseOffloadHook(ModelHook): with current_omni_platform.stream(self.comm_stream): for dtype, local_shard in gpu_shards.items(): # Slice the shared (max-sized) output buffer down to this - # block's actual AllGather output size. The buffers are - # sized to the *largest* block across all groups, so for - # any smaller block the full buffer would violate the + # block's actual AllGather output size. The buffers are + # sized to the *largest* block across all groups, so for any + # smaller block the full buffer would violate the # all_gather_into_tensor contract # (output.numel() == world_size * input.numel()). # Repoint offsets are relative to the block's flattened @@ -485,6 +494,111 @@ def remove_distributed_block_hook(module: nn.Module) -> None: logger.debug("Removed distributed offload hook from %s", module.__class__.__name__) +class PinnedResidentLayerGroup: + """Keep selected layers in pinned host memory between requests. + + + TODO(offload): Extract this alongside PinnedModuleStager after the + distributed shard-and-pin operation becomes a shared storage primitive. + It currently remains here because it depends on DLO's local-shard layout. + Unlike ``module.to(device)``/``module.to("cpu")``, this group retains a + pinned CPU master copy and never copies generated device weights back to + host. Entering the denoise stage performs one asynchronous H2D pass; + leaving it only restores zero-sized placeholders and releases the device + buffers. This lets the following VAE stage reuse the same HBM. + + The resident path is intentionally local-shard only. With tensor + parallelism, the regular model loader has already produced each rank's TP + shard, so no DP AllGather is required or desirable here. + """ + + def __init__( + self, + blocks: list[nn.Module], + device: torch.device, + copy_stream: Any, + pin_memory: bool, + ) -> None: + self.device = device + self.copy_stream = copy_stream + self.loaded = False + self._states: list[dict[str, Any]] = [] + self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + + for block in blocks: + params = dict(block.named_parameters()) + bufs = dict(block.named_buffers()) + targets: dict[str, torch.Tensor] = {**params, **bufs} + cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( + params, + bufs, + dp_size=1, + rank=0, + pin_memory=pin_memory, + ) + self._states.append( + { + "targets": targets, + "cpu_shards": cpu_shards, + "metadata": metadata, + } + ) + + def load(self) -> None: + if self.loaded: + return + + # Allocate on the compute stream so the caching allocator can reuse + # blocks released by the encoder stage. Allocating on the copy stream + # would create a separate stream-local pool and inflate peak HBM. + gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + for state in self._states: + block_buffers: dict[torch.dtype, torch.Tensor] = {} + for dtype, cpu_shard in state["cpu_shards"].items(): + block_buffers[dtype] = torch.empty( + cpu_shard.shape, + dtype=dtype, + device=self.device, + ) + gpu_buffers.append(block_buffers) + + self.copy_stream.wait_stream(current_omni_platform.current_stream()) + ready = current_omni_platform.Event() + with current_omni_platform.stream(self.copy_stream): + for state, block_buffers in zip(self._states, gpu_buffers): + for dtype, cpu_shard in state["cpu_shards"].items(): + block_buffers[dtype].copy_(cpu_shard, non_blocking=True) + ready.record(self.copy_stream) + + for state, block_buffers in zip(self._states, gpu_buffers): + targets = state["targets"] + for dtype, metas in state["metadata"].items(): + gpu_buffer = block_buffers[dtype] + for meta in metas: + set_tensor_storage( + targets[meta["name"]], + gpu_buffer[meta["offset"] : meta["offset"] + meta["numel"]].view(meta["shape"]), + ) + + current_omni_platform.current_stream().wait_event(ready) + self._gpu_buffers = gpu_buffers + self.loaded = True + + def offload(self) -> None: + if not self.loaded: + return + + # Denoise kernels consume these buffers on the current compute stream. + # Synchronize once at the stage boundary; no D2H copy is necessary + # because the pinned CPU master weights were never overwritten. + current_omni_platform.synchronize() + for state in self._states: + for target in state["targets"].values(): + set_tensor_storage(target, make_offload_placeholder(target)) + self._gpu_buffers.clear() + self.loaded = False + + # ---------------------------------------------------------------------- # # Backend # # ---------------------------------------------------------------------- # @@ -502,6 +616,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): overlapped with computation. """ + _MMAP_PARAM_ATTRS = ("weight_loader", "mmap_weight_transform") + def __init__(self, config: OffloadConfig, device: torch.device): super().__init__(config, device) @@ -512,6 +628,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.rank = 0 self._blocks: list[list[nn.Module]] = [] self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] + self._resident_blocks: list[nn.Module] = [] + self._resident_layer_group: PinnedResidentLayerGroup | None = None + + def load_resident_layers(self) -> None: + """Load the model-declared leading blocks for the denoise stage.""" + if self._resident_layer_group is not None: + self._resident_layer_group.load() + + def offload_resident_layers(self) -> None: + """Release leading blocks before VAE decode to bound peak HBM.""" + if self._resident_layer_group is not None: + self._resident_layer_group.offload() + + def _remember_mmap_param_attrs(self, pipeline: nn.Module) -> None: + """Save loader metadata before ``to_empty`` replaces Parameters.""" + self._mmap_param_attrs = { + name: {attr: value for attr in self._MMAP_PARAM_ATTRS if (value := getattr(param, attr, None)) is not None} + for name, param in pipeline.named_parameters() + } + + def _attach_mmap_param_attrs( + self, + name: str, + replacement: nn.Parameter, + source: nn.Parameter | None = None, + ) -> None: + attrs = getattr(self, "_mmap_param_attrs", {}).get(name, {}) + for attr in self._MMAP_PARAM_ATTRS: + value = attrs.get(attr) + if value is None and source is not None: + value = getattr(source, attr, None) + if value is not None: + setattr(replacement, attr, value) + if attr == "mmap_weight_transform": + replacement.mmap_weight_transform_pending = True def _load_weights_via_mmap(self, pipeline: nn.Module, modules) -> None: """Load DiT weights from safetensors via mmap views (no RSS). @@ -614,6 +765,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): "instead of TP, or disable distributed layerwise offload." ) + # ``Module.to_empty`` creates new Parameter objects and drops custom + # attributes installed by vLLM weight loaders. Preserve them by full + # pipeline name so mmap replacements can reproduce checkpoint layout + # transforms such as MiniMax-H3 grouped QKV. + self._remember_mmap_param_attrs(pipeline) + # # Important: when DiT modules are nested (e.g. transformer contains # transformer.language_model), to_empty("meta") on the parent @@ -738,7 +895,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + self._attach_mmap_param_attrs(name, replacement, param) + parent._parameters[parts[-1]] = replacement loaded += 1 loaded_names.add(name) @@ -793,7 +952,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[bparts[-1]] = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=bparam.requires_grad) + self._attach_mmap_param_attrs(full_param_name, replacement, bparam) + parent._parameters[bparts[-1]] = replacement block_loaded += 1 loaded_names.add(full_param_name) @@ -904,7 +1065,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): parent = parent[int(part)] else: parent = getattr(parent, part) - parent._parameters[parts[-1]] = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + replacement = torch.nn.Parameter(tensor, requires_grad=param.requires_grad) + self._attach_mmap_param_attrs(full_name, replacement, param) + parent._parameters[parts[-1]] = replacement loaded += 1 if loaded > 0: @@ -962,24 +1125,74 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): coord.ranks, ) - def _register_on_demand_hook(self, module: nn.Module, label: str) -> None: - """Move *module* to GPU permanently. + def _register_on_demand_hook(self, module: nn.Module, label: str, *, stage_on_demand: bool = False) -> None: + """Prepare a pipeline-managed stage component or keep it resident. - 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). + Components that expose an explicit stage lifecycle are initially + offloaded and loaded by their pipeline only around encode/decode. + Other models retain the conservative resident behavior because a + generic post-forward hook can disrupt the DiT prefetch streams. """ + offload_to_cpu = getattr(module, "offload_to_cpu", None) + if stage_on_demand and callable(offload_to_cpu): + offload_to_cpu() + logger.info("Prepared %s (%s) for pipeline-managed staged offload", label, module.__class__.__name__) + return module.to(self.device) logger.info("Moved %s (%s) to GPU (resident)", label, module.__class__.__name__) + def _try_layerwise_offload_encoder(self, module: nn.Module, name: str, plan: OffloadPlan | None) -> bool: + """Stream plan-declared encoder blocks on each rank without AllGather.""" + if plan is None or name not in plan.encoder_block_attrs: + return False + if getattr(module, "_omni_layerwise_enabled", False): + return True + + from operator import attrgetter + + from vllm_omni.diffusion.offloader.layerwise_backend import apply_block_hook + + hooks = [] + block_groups = [] + copy_stream = current_omni_platform.Stream() + for block_path in plan.encoder_block_attrs[name]: + try: + blocks = attrgetter(block_path)(module) + except AttributeError: + logger.warning("Encoder offload path %s.%s was not found", name, block_path) + continue + if not isinstance(blocks, nn.ModuleList) or len(blocks) <= 1: + logger.warning("Encoder offload path %s.%s is not a streamable block list", name, block_path) + continue + group_hooks = [ + apply_block_hook(blocks[-1], blocks[0], self.device, copy_stream, self.config.pin_cpu_memory) + ] + group_hooks.extend( + apply_block_hook(block, blocks[index + 1], self.device, copy_stream, self.config.pin_cpu_memory) + for index, block in enumerate(blocks[:-1]) + ) + for index, hook in enumerate(group_hooks): + hook._prev_hook = group_hooks[index - 1] + hooks.extend(group_hooks) + block_groups.append(blocks) + + if not hooks: + return False + # The component lifecycle uses these generic attributes to keep only + # non-block encoder state resident during the encode phase. + module._omni_layerwise_hooks = hooks + module._omni_layerwise_block_groups = block_groups + module._omni_layerwise_enabled = True + logger.info( + "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", + name, + sum(len(blocks) for blocks in block_groups), + len(block_groups), + ) + return True + 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 @@ -1089,6 +1302,78 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._blocks.append(blocks) return True + def _prepare_dit_non_block_modules( + self, + dit_module: nn.Module, + dit_name: str, + blocks_attr_names: list[str], + all_dit_modules: set[int], + plan: OffloadPlan | None, + ) -> None: + """Place or hook the DiT parts that are outside its repeated blocks. + + This must run even when every repeated block is resident. Otherwise + an all-resident stage skips placement for modules such as H3's token + refiner and enters the forward pass with CPU or meta tensors. + """ + _ON_DEMAND_THRESHOLD = _ON_DEMAND_THRESHOLD_MB + for name, module in dit_module.named_children(): + if name in blocks_attr_names: + logger.debug("Skipped blocks module %s", name) + continue + + has_meta = any(getattr(param, "is_meta", False) for param in module.parameters()) + if has_meta: + self._load_module_weights_from_mmap(module, dit_name, name) + module_mb = ( + sum( + param.nelement() * param.element_size() if not getattr(param, "is_meta", False) else 0 + for param in module.parameters() + ) + / 1048576 + ) + explicitly_planned = plan is not None and name in plan.offload_submodules + if explicitly_planned or module_mb > _ON_DEMAND_THRESHOLD: + if id(module) in all_dit_modules: + logger.info("Submodule '%s' is already a DiT module, skipping layerwise offload", name) + elif self._try_layerwise_offload_submodule(module, name, plan): + pass + else: + self._register_on_demand_hook(module, name) + continue + + try: + module.to(self.device) + except (NotImplementedError, RuntimeError): + self._load_module_weights_from_mmap(module, dit_name, name) + # Non-persistent buffers such as RoPE frequencies do not + # exist in the checkpoint and must be reconstructed. + has_meta_buffer = any(getattr(buffer, "is_meta", False) for buffer in module.buffers(recurse=True)) + if has_meta_buffer: + saved_params = { + param_name: param.data.clone() + for param_name, param in module.named_parameters() + if not getattr(param, "is_meta", False) + } + module.to_empty(device=self.device) + for submodule in module.modules(): + if hasattr(submodule, "reset_parameters"): + submodule.reset_parameters() + for param_name, param in module.named_parameters(): + if param_name in saved_params: + param.data.copy_(saved_params[param_name]) + try: + module.to(self.device) + except Exception: + logger.warning("Module %s still has meta params after mmap load", name) + + for param in dit_module._parameters.values(): + if param is not None and not getattr(param, "is_meta", False): + 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) + def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("DistributedLayerwiseOffloadBackend already enabled") @@ -1110,7 +1395,19 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # When present, replaces heuristic block discovery. plan = get_offload_plan(pipeline) + if self.config.dlo_resident_layers and (plan is None or not plan.resident_dit_paths): + logger.warning( + "dlo_resident_layers=%d was requested, but this model declares no " + "resident_dit_paths; all blocks will be streamed.", + self.config.dlo_resident_layers, + ) + # Load weights via mmap for DLO+AllGather. + # TODO(offload): Add a rank-local mmap path for dlo_no_use_allgather. + # It must apply declarative checkpoint-to-runtime adapters before the + # standard TP loader shards each block, and must cover staged encoders + # and VAEs as well as DiT blocks. The current AllGather mmap path is + # deliberately not reused because it changes the weight layout. # 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 @@ -1137,10 +1434,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # 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") + for enc, enc_name in zip(modules.encoders, modules.encoder_names): + self._try_layerwise_offload_encoder(enc, enc_name, plan) + self._register_on_demand_hook( + enc, "encoder", stage_on_demand=plan is not None and enc_name in plan.on_demand_component_paths + ) + for vae, vae_name in zip(modules.vaes, modules.vae_names): + self._register_on_demand_hook( + vae, + "vae", + stage_on_demand=(plan is not None and vae_name in plan.on_demand_component_paths), + ) # Move resident modules to GPU (small modules needed every forward) for name, module in zip(modules.resident_names, modules.resident_modules): @@ -1171,93 +1475,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): dit_module.to(self.device) continue + resident_count = 0 + if plan is not None and dit_name in plan.resident_dit_paths: + resident_count = min(self.config.dlo_resident_layers, len(blocks)) + if resident_count: + resident_blocks = blocks[:resident_count] + self._resident_blocks.extend(resident_blocks) + blocks = blocks[resident_count:] + logger.info( + "Keeping %d leading blocks resident on %s; streaming %d tail blocks", + resident_count, + dit_name, + len(blocks), + ) + + self._prepare_dit_non_block_modules( + dit_module, + dit_name, + blocks_attr_names, + all_dit_modules, + plan, + ) + num_blocks = len(blocks) + if num_blocks == 0: + logger.info("All blocks for %s are resident; no streaming hooks required", dit_name) + continue if num_blocks <= 1: logger.warning( - "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", dit_name, dit_module.__class__.__name__, ) - dit_module.to(self.device) + self._resident_blocks.extend(blocks) 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") @@ -1314,7 +1566,17 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.append(block_hooks) self._blocks.append(blocks) + if self._resident_blocks: + self._resident_layer_group = PinnedResidentLayerGroup( + self._resident_blocks, + self.device, + self.copy_stream, + self.config.pin_cpu_memory, + ) + pipeline._dlo_residency_controller = self + if not self._all_hook_groups: + self.enabled = bool(self._resident_blocks) return # Unified allocation: 2 shared output buffers + 2 shared shard buffers @@ -1424,8 +1686,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._on_demand_handles = [] self._on_demand_shard_infos = [] + self.offload_resident_layers() self._blocks.clear() self._all_hook_groups.clear() + self._resident_blocks.clear() + self._resident_layer_group = None self.enabled = False logger.info("Distributed layer-wise offloading disabled") diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index d6fcb2569..b93c9ed73 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -19,6 +19,7 @@ class PipelineModules: encoders: list[nn.Module] encoder_names: list[str] vaes: list[nn.Module] + vae_names: list[str] = field(default_factory=list) resident_modules: list[nn.Module] = field(default_factory=list) resident_names: list[str] = field(default_factory=list) @@ -153,7 +154,7 @@ class ModuleDiscovery: dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared) encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared) - vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) + vaes, vae_names = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared) return PipelineModules( @@ -163,5 +164,6 @@ class ModuleDiscovery: encoder_names=encoder_names, vaes=vaes, resident_modules=residents, + vae_names=vae_names, resident_names=resident_names, ) diff --git a/vllm_omni/diffusion/offloader/module_residency.py b/vllm_omni/diffusion/offloader/module_residency.py new file mode 100644 index 000000000..da32ab520 --- /dev/null +++ b/vllm_omni/diffusion/offloader/module_residency.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""On-demand module staging backed by an immutable pinned CPU snapshot.""" + +from __future__ import annotations + +from itertools import chain +from typing import Any + +import torch +from torch import nn + +from vllm_omni.platforms import current_omni_platform + +from .tensor_utils import set_tensor_storage + + +class PinnedModuleStager: + """Stage an immutable module without copying device weights back to CPU. + + ``nn.Module.to("cpu")`` performs a device-to-host copy for every parameter + after each forward. In inference the weights are immutable, so retain one + pinned CPU master instead. ``load`` materializes fresh device tensors from + that master; ``offload`` only rebinds the Parameters/Buffers to it and + releases the device allocations. + """ + + def __init__( + self, + module: nn.Module, + device: torch.device, + *, + pin_memory: bool = True, + copy_stream: Any | None = None, + ) -> None: + self.device = device + self.copy_stream = copy_stream or current_omni_platform.Stream() + self.loaded = False + self._entries: list[tuple[torch.Tensor, torch.Tensor]] = [] + self._device_tensors: list[torch.Tensor] = [] + + # parameters()/buffers() remove duplicate objects by default, which + # preserves tied weights and buffer aliases when their storage moves. + for target in chain(module.parameters(), module.buffers()): + master = target.detach() + if master.device.type != "cpu": + master = master.to("cpu") + if pin_memory and not master.is_pinned(): + master = master.pin_memory() + set_tensor_storage(target, master) + self._entries.append((target, master)) + + def load(self) -> None: + if self.loaded: + return + + # Allocate on the compute stream so the caching allocator can reuse + # memory released by the preceding encoder/DiT stage. + device_tensors = [torch.empty_like(master, device=self.device) for _, master in self._entries] + compute_stream = current_omni_platform.current_stream() + self.copy_stream.wait_stream(compute_stream) + ready = current_omni_platform.Event() + with current_omni_platform.stream(self.copy_stream): + for device_tensor, (_, master) in zip(device_tensors, self._entries): + device_tensor.copy_(master, non_blocking=master.is_pinned()) + ready.record(self.copy_stream) + + for device_tensor, (target, _) in zip(device_tensors, self._entries): + set_tensor_storage(target, device_tensor) + compute_stream.wait_event(ready) + self._device_tensors = device_tensors + self.loaded = True + + def offload(self) -> None: + if not self.loaded: + return + + # The module has completed on the compute stream. Synchronize once at + # the stage boundary, then release HBM by rebinding to the unchanged + # host masters. No device-to-host transfer is necessary. + current_omni_platform.synchronize() + for target, master in self._entries: + set_tensor_storage(target, master) + self._device_tensors.clear() + self.loaded = False + current_omni_platform.empty_cache() + + +__all__ = ["PinnedModuleStager"] diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 575194ea0..f8943757e 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -35,10 +35,21 @@ class OffloadPlan: for large non-DiT submodules within a DiT that should be independently offloaded with their own hooks. e.g. ``{"context_encoder": "layers"}`` + resident_dit_paths: DiT paths whose leading blocks may be kept on the + device when ``dlo_resident_layers`` is nonzero. Keeping this + model-declared avoids applying a consumer-GPU tuning knob to + auxiliary or dual DiTs unintentionally. + encoder_block_attrs: Maps encoder paths to rank-local block-list paths. + These blocks are streamed with ordinary layerwise hooks, never + with the DiT AllGather group. """ + on_demand_component_paths: frozenset[str] = field(default_factory=frozenset) + block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) offload_submodules: dict[str, str] = field(default_factory=dict) + resident_dit_paths: frozenset[str] = field(default_factory=frozenset) + encoder_block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict) def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: @@ -58,4 +69,6 @@ def supports_mmap_loading(pipeline: nn.Module) -> bool: (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()) + return bool(getattr(pipeline, "_supports_mmap_loading", True)) and any( + callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules() + ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fbc2bec3a..8f2fc6ef5 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -495,6 +495,7 @@ class OrchestratorArgs: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_resident_layers: int = 0 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 b3307c3f5..2856674db 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1032,6 +1032,7 @@ class AsyncOmniEngine: "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), + "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), "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 df3d8aee4..3671eff2b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -646,15 +646,25 @@ class OmniServeCommand(CLISubcommand): 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.", + "When disabled (--dlo-no-use-allgather), each rank streams the " + "standard loader's rank-local tensors via H2D only — no additional " + "DP sharding, no AllGather, and no concurrent-request requirement.", ) 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.", + help=( + "Disable AllGather and stream standard-loader rank-local weights " + "independently (including existing TP shards)." + ), + ) + omni_config_group.add_argument( + "--dlo-resident-layers", + type=int, + default=0, + help="Keep this many leading main-DiT blocks resident on the device " + "while distributed layerwise offload streams the remaining blocks.", ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument(