From e2796b616847b2be2587e1c2dc0cff13b657f6d0 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sat, 15 Aug 2026 16:43:56 +0800 Subject: [PATCH] [Diffusion] Add loader-owned host-weight plans for DLO Signed-off-by: Hongsheng Liu --- .../distributed_layerwise_offload.md | 101 +- .../distributed_layerwise_offload.md | 48 +- .../minimax_h3/dlo_dp2_lifecycle.py | 15 +- .../model_loader/test_diffusers_loader.py | 75 ++ .../test_minimax_h3_quantization.py | 43 + .../test_distributed_layerwise_backend.py | 570 ++++++++++-- .../diffusion/test_diffusion_model_runner.py | 9 +- .../checkpoint_adapters/__init__.py | 8 + .../checkpoint_adapters/direct_mmap.py | 108 +++ .../model_loader/diffusers_loader.py | 59 +- .../model_loader/host_weight_plan.py | 333 +++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 7 - vllm_omni/diffusion/offloader/__init__.py | 28 +- vllm_omni/diffusion/offloader/base.py | 10 +- .../distributed_layerwise_backend.py | 868 +++++++++--------- vllm_omni/diffusion/offloader/offload_plan.py | 17 - .../worker/diffusion_model_runner.py | 6 +- 17 files changed, 1717 insertions(+), 588 deletions(-) create mode 100644 vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py create mode 100644 vllm_omni/diffusion/model_loader/host_weight_plan.py diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index ade39fa27..477d8d0f1 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -4,7 +4,7 @@ This document describes distributed layerwise offload (DLO) for diffusion models. DLO keeps only a small number of DiT blocks on the accelerator and streams the remaining blocks from host memory. The distributed backend can either shard those host-side weights across an existing parallel group or keep -the standard loader's rank-local weights and avoid an additional collective. +complete rank-local block sources and avoid an additional collective. For user-facing commands, see the [distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) @@ -14,10 +14,15 @@ and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/rec DLO is implemented for multi-device diffusion execution. The default AllGather path is the primary path for DP and SP deployments. The -`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful -for standard-loader sharding, workstation bring-up, and systems where an -additional DLO collective is undesirable, but it does not reduce host weight -storage across ranks. +`--dlo-no-use-allgather` path streams complete blocks independently and adds no +DLO weight collective. + +Host storage is selected separately from the transfer protocol. The loader can +produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; +otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on +the same node can share immutable checkpoint pages when direct mmap is +selected, while the ordinary-loader fallback still keeps a private runtime +copy per process. The compatibility matrix below describes the current implementation. The unit-level guards are covered, but not every parallelism combination has a @@ -41,6 +46,28 @@ TP is deliberately not used as DLO's AllGather group. HSDP has its own parameter-sharding lifecycle and is not allowed to be sharded a second time by DLO's AllGather path. +### The loader owns host-weight planning + +Before it decides whether ordinary weight materialization can be skipped, the +diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is +accepted only when preflight proves all of the following: + +- every required DiT parameter and persistent buffer has exactly one source; +- runtime names, checkpoint keys, shapes, and dtypes match; +- the runtime topology is TP1 without HSDP or online quantization; and +- every custom loader operation is represented by a loader-owned checkpoint + adapter. + +The exact plan object is handed to DLO. The backend does not rescan checkpoint +files, repeat the capability decision, or reconstruct names from its block +topology. If preflight fails, the loader materializes weights normally and DLO +consumes those runtime tensors. + +This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline +flags such as `_supports_mmap_loading` or parameter attributes for mmap-only +transforms. Model-specific direct-layout knowledge, when required, lives in a +checkpoint adapter beside the ordinary loader. + ### AllGather path With the default `dlo_use_allgather=True`, each rank stores approximately @@ -61,6 +88,15 @@ Buffers: [current slot] [prefetch slot] [current slot] The backend uses two shared device buffers, so accelerator weight residency is bounded by the largest streamed blocks rather than the complete model. +When direct checkpoint mmap is selected, the checkpoint mappings are only the +source used to prepare each rank's persistent shard. They can be closed after +shard preparation. Across the AllGather group, those private shards total +approximately one runtime model copy. + +An effective DLO group size of one performs no collective, even when +`dlo_use_allgather=True`; it follows the rank-local transfer path described +below. + When DP is greater than one, the engine can process one request per DP rank in the same denoising wave. Because AllGather is a collective, all participating requests must take the same execution path at every denoising step. @@ -68,9 +104,20 @@ requests must take the same execution path at every denoising step. ### Rank-local path without DLO AllGather With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to -one. The regular model loader remains responsible for preparing each rank's -weights, including TP-local tensors or HSDP-managed parameters. DLO then -streams those rank-local tensors block by block using H2D copies only. +one and streams complete blocks using H2D copies only. The host backing may be +either a loader-approved checkpoint mapping or ordinary runtime tensors. + +For direct mmap, each process retains immutable safetensors views and uses two +bounded pinned host staging slots. Processes on the same node that map the same +files share physical checkpoint pages through the OS page cache. This removes +the persistent private full-model copy per pure-DP process, but each process +still packs and transfers every complete block. Sharing is node-local; each +node has its own page cache. + +When direct mmap preflight fails, the regular model loader remains responsible +for preparing each rank's weights, including TP-local tensors or HSDP-managed +parameters. In that fallback, each pure-DP process keeps a private full runtime +copy. This mode means: @@ -80,21 +127,18 @@ This mode means: not shard weights across SP ranks. - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only DLO's additional weight AllGather is disabled. -- Pure DP deployments keep a full host-side model copy per rank, subject to - shared OS page-cache behavior. +- Pure DP deployments share one checkpoint-backed copy per node when direct + mmap is selected; the ordinary-loader fallback keeps one private runtime + copy per rank. - The scheduler does not require a synchronized DP request wave for DLO. -This path is intentionally not implemented by reusing the AllGather mmap -loader. It relies on the standard loader so model-specific TP/HSDP transforms -remain intact. - ## Parallelism compatibility | Parallelism | DLO + AllGather | DLO without AllGather | |---|---|---| -| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight sharding or cross-DP DLO collective. | +| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. | | **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. | -| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. | +| **TP > 1** | Direct checkpoint mmap falls back before mutation. The ordinary loader preserves TP-local layouts before DLO applies any DP/SP host sharding. | Direct checkpoint mmap falls back to ordinary TP-aware loading; the resulting TP-local tensors are streamed without a DLO collective. Broader model and hardware validation is still needed. | | **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. | ### Combined dimensions @@ -122,11 +166,15 @@ AllGather DP multi-concurrency requires: The no-AllGather path does not impose these DLO-specific synchronized-wave requirements. -The mmap loader is used only by the supported DLO+AllGather path when the -model has a compatible checkpoint layout. Online quantization is incompatible -with that sharded mmap path; use `--dlo-no-use-allgather` or disable online -quantization. Models that do not meet the mmap requirements use the regular -loader path. +Direct checkpoint mmap can back either transfer path. It is currently limited +to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the +ordinary loader. Online quantization remains incompatible with DLO AllGather; +use `--dlo-no-use-allgather` or disable online quantization. + +A normalized runtime mmap cache, built through the ordinary loader, is the +proposed general mechanism for sharing transformed TP or quantized layouts. +That cache and its publication/lifecycle protocol are intentionally outside +this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). ## Validation coverage @@ -134,7 +182,11 @@ Current source-level validation includes: - HSDP + DLO + AllGather rejection; - HSDP + DLO without AllGather acceptance at configuration level; -- TP rejection in the DLO+AllGather mmap path; +- loader preflight fallback for TP, HSDP, online quantization, unknown custom + loaders, missing keys, and shape/dtype mismatches; +- exact loader-to-backend plan transfer and ordinary-loader fallback; +- rank-local mmap source retention, bounded two-slot staging, and adapter + transforms without parameter-side flags; - resident-layer requests requiring no-AllGather; - DP request-wave validation for denoising-step compatibility; - sharding, double-buffer, AllGather-size, and heterogeneous-block regression @@ -150,7 +202,8 @@ on the target CUDA/NCCL or CANN/HCCL hardware. scaling path. - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is not the goal. -- Use **no-AllGather** to bring up TP or workstation PCIe configurations, with - the expectation of higher host-memory use and lower validation confidence. +- Use **no-AllGather** when independent replica execution is required. TP1 + direct-mmap deployments can share checkpoint pages per node; other layouts + retain the ordinary loader's private host memory behavior. - Prefer **HSDP alone** for production HSDP deployments until the combined HSDP + DLO no-AllGather path has broader end-to-end coverage. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 15298c897..7d188f4f3 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -3,7 +3,9 @@ Distributed layerwise offloading (DLO) extends block streaming to multi-device deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of the host weights and reconstructs each layer at runtime. Without AllGather, -each rank streams its standard-loader rank-local weights independently. +each rank streams a complete block independently. Compatible TP1 deployments +can share checkpoint-backed host pages among processes on the same node; +otherwise DLO streams the ordinary loader's rank-local tensors. See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) for the implementation contract and compatibility matrix. @@ -59,21 +61,36 @@ omni = Omni( | `--enable-distributed-layerwise-offload` | Enable DLO | `false` | | `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` | | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` | -| `--dlo-no-use-allgather` | Keep standard-loader rank-local weights | `false` | +| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` | | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | -## mmap weight loading +## Host-weight loading -The DLO plus AllGather path: +The diffusion loader chooses host storage before DLO is enabled. It first +attempts to build a complete, validated direct-checkpoint mmap plan. If names, +coverage, shape, dtype, topology, or loader-callback compatibility cannot be +proven, it runs the ordinary model loader instead. DLO consumes that result and +does not make a second checkpoint-compatibility decision. + +With direct checkpoint mmap, the loader: 1. saves non-persistent buffers such as RoPE frequencies; 2. moves the normally created transformer to the meta device; 3. loads checkpoint tensors as mmap views backed by the shared OS page cache; -4. calls model-specific `post_load_weights()` conversions; and -5. restores the saved non-persistent buffers. +4. applies any loader-owned bounded layout adapters while packing blocks; +5. restores saved non-persistent buffers; and +6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle + hooks. -This avoids one full checkpoint RSS copy per rank and does not require -model-specific loading code. +For AllGather with a group larger than one, each process copies only its +persistent shard and then releases the source mapping. For no-AllGather, each +process keeps the mapping open and packs complete blocks through two bounded +pinned staging slots. Processes mapping the same files on one node share the +immutable pages; no-AllGather still performs a complete-block H2D copy in each +process. + +When the effective DLO group size is one, `dlo_use_allgather=True` does not +perform a collective and uses the same rank-local transfer behavior. ## Declarative topology @@ -103,16 +120,21 @@ must enter each collective. ## Limitations -- Online FP8 quantization is rejected with the DLO plus AllGather mmap path. - Use `--dlo-no-use-allgather` or disable online quantization. -- Tensor parallel size greater than one is rejected in the mmap path because - it bypasses TP-aware loader callbacks. The no-AllGather path retains the - standard loader and remains experimental with TP. +- Direct checkpoint mmap currently requires TP1. TP greater than one falls + back before model mutation to the ordinary TP-aware loader; DLO can then + stream that runtime layout. Broader TP combinations remain experimental. - HSDP plus AllGather is rejected to avoid double sharding. HSDP without AllGather has limited end-to-end validation. +- Online quantization uses the ordinary loader with no-AllGather. It remains + incompatible with DLO AllGather. - Resident leading layers require `--dlo-no-use-allgather` and a model `OffloadPlan` that declares eligible `resident_dit_paths`. - DP concurrency requires an explicit, identical inference-step count. +Sharing transformed TP or quantized runtime layouts through a normalized mmap +cache is a follow-up design in +[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of +the direct-checkpoint path. + See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for an end-to-end example. diff --git a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py b/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py index f8117b0f9..2cac7a666 100644 --- a/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py +++ b/examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py @@ -11,6 +11,13 @@ Examples: python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2 + # Rank-local mmap storage, with no DLO collective or DP wave scheduling. + VLLM_WORKER_MULTIPROC_METHOD=spawn \ + VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ + CUDA_VISIBLE_DEVICES=0,1 \ + python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \ + --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2-no-allgather + VLLM_WORKER_MULTIPROC_METHOD=spawn \ VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \ CUDA_VISIBLE_DEVICES=0,1 \ @@ -45,9 +52,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA") parser.add_argument( "--mode", - choices=("dlo-dp2", "request"), + choices=("dlo-dp2", "dlo-dp2-no-allgather", "request"), required=True, - help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode", + help="DLO DP2 (with or without AllGather) or ordinary non-DLO TP2 request mode", ) parser.add_argument("--steps", type=int, default=2) parser.add_argument("--duration", type=float, default=5.0) @@ -74,14 +81,14 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]: "stage_init_timeout": args.init_timeout, "init_timeout": args.init_timeout, } - if args.mode == "dlo-dp2": + if args.mode in {"dlo-dp2", "dlo-dp2-no-allgather"}: common.update( tensor_parallel_size=1, data_parallel_size=2, text_encoder_tp_size=1, vae_patch_parallel_size=1, enable_distributed_layerwise_offload=True, - dlo_use_allgather=True, + dlo_use_allgather=args.mode == "dlo-dp2", dlo_resident_layers=0, ) else: diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 69600d811..ff1049be5 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -16,6 +16,10 @@ from vllm.config.load import LoadConfig from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + HostWeightPlanResult, +) from vllm_omni.diffusion.models.helios import HeliosPipeline from vllm_omni.diffusion.registry import initialize_model @@ -258,6 +262,77 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): assert get_current_diffusion_config_or_none() is None +def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={}, + runtime_layout_key="test", + ) + loaded_ordinary_weights = False + + def load_weights(_model): + nonlocal loaded_ordinary_weights + loaded_ordinary_weights = True + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = load_weights # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(plan), + ) + + assert loader.load_model(load_device="cpu") is model + assert not loaded_ordinary_weights + assert loader.take_host_weight_plan() is plan + assert loader.take_host_weight_plan() is None + + +def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + calls: list[str] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + monkeypatch.setattr( + loader_mod, + "build_checkpoint_mmap_plan", + lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"), + ) + + assert loader.load_model(load_device="cpu") is model + assert calls == ["load", "process"] + assert loader.take_host_weight_plan() is None + + def test_hsdp_processes_quantized_weights_before_sharding(mocker): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod from vllm_omni.diffusion.offloader.module_collector import PipelineModules diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py index 5f220d6ab..cc14d92e3 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py @@ -191,6 +191,49 @@ def test_model_load_weights_transforms_before_calling_vllm_loader(): ] +def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch): + from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + get_direct_mmap_adapter, + ) + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear) + monkeypatch.setattr(h3, "Attention", _FakeAttention) + + arch = h3.MiniMaxH3DiTArchConfig( + hidden_size=1, + num_attention_heads=2, + attention_head_dim=1, + rope_inv_freq_len=1, + ) + attention = h3.MiniMaxH3Attention( + arch, + quant_config=None, + prefix="blocks.0.attn", + ) + transformer = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(transformer) + transformer.blocks = nn.ModuleList([nn.Module()]) + transformer.blocks[0].attn = attention + pipeline = nn.Module() + pipeline.transformer = transformer + + adapter = get_direct_mmap_adapter(pipeline) + assert adapter is not None + policy = adapter.policy_for( + "transformer.blocks.0.attn.qkv_proj.weight", + attention.qkv_proj.weight, + ) + assert policy is not None + assert policy.allow_custom_loader + assert policy.transform is not None + checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1) + + assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5] + assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform") + + def test_pipeline_resolves_transformer_component_quant_config(): from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( _resolve_component_quant_config, diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 4c0cceb68..b29de143f 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -18,6 +18,11 @@ from torch import nn from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + TensorBinding, + build_checkpoint_mmap_plan, +) from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy from vllm_omni.diffusion.offloader.block_discovery import ( get_blocks_attr_names, @@ -30,7 +35,10 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( 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.diffusion.offloader.offload_plan import ( + OffloadPlan, + get_offload_plan, +) from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -48,6 +56,9 @@ class DummyEvent: def record(self, _stream) -> None: return None + def synchronize(self) -> None: + return None + @contextmanager def dummy_stream(_stream): @@ -351,8 +362,9 @@ class TestDistributedLayerwiseOffloadHook: 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 + + def transform(tensor): + return tensor.flip(0) hook = DistributedLayerwiseOffloadHook( next_block=next_block, @@ -361,17 +373,17 @@ class TestDistributedLayerwiseOffloadHook: dp_size=2, rank=0, pin_memory=False, + tensor_transforms={id(next_block.weight): transform}, ) 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): + def test_transform_is_not_applied_without_a_loader_plan(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, @@ -385,6 +397,81 @@ class TestDistributedLayerwiseOffloadHook: assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0])) + def test_rank_local_mmap_retains_source_and_uses_staging(self, dist_group, patched_offload_runtime): + current_block = nn.Linear(2, 2, bias=False) + next_block = nn.Linear(2, 2, bias=False) + next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2)) + source_storage = next_block.weight.untyped_storage().data_ptr() + + def transform(tensor): + return tensor.flip(0) + + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + rank_local_mmap=True, + tensor_transforms={id(next_block.weight): transform}, + ) + hook.initialize_hook(current_block) + + source = hook.cpu_sources[torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage + assert hook.cpu_shards == {} + assert next_block.weight.numel() == 0 + + hook.cpu_staging_buffers = [ + {torch.float32: torch.empty(4)}, + {torch.float32: torch.empty(4)}, + ] + hook.prefetch_layer(slot=0, non_blocking=False) + + assert torch.equal( + next_block.weight, + torch.tensor([[2.0, 3.0], [0.0, 1.0]]), + ) + assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2)) + + def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime): + hooks = [] + for size in (4, 9): + current_block = nn.Linear(1, 1, bias=False) + next_block = nn.Module() + next_block.weight = nn.Parameter(torch.arange(size, dtype=torch.float32)) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + shared_buffers=[None, None], + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + hooks.append(hook) + + resident_block = nn.Module() + resident_block.weight = nn.Parameter(torch.arange(12, dtype=torch.float32)) + resident_group = PinnedResidentLayerGroup( + [resident_block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + staging = DistributedLayerwiseOffloadBackend._allocate_shared_cpu_staging_buffers( + hooks, + resident_group, + ) + + assert len(staging) == 2 + assert staging[0][torch.float32].numel() == 12 + assert staging[1][torch.float32].numel() == 12 + class TestPinnedResidentLayerGroup: def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime): @@ -439,23 +526,41 @@ class TestPinnedResidentLayerGroup: assert block.weight.stride() == expected_stride group.offload() - def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self): - module = nn.Linear(2, 2, bias=False) + def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime): + blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)] + expected = [] + source_storage = [] + for index, block in enumerate(blocks): + block.weight.data.fill_(index + 1) + expected.append(block.weight.detach().clone()) + source_storage.append(block.weight.untyped_storage().data_ptr()) - def transform(tensor): - return tensor.flip(0) + group = PinnedResidentLayerGroup( + blocks, + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) - module.weight.mmap_weight_transform = transform + for index, state in enumerate(group._states): + source = state["cpu_sources"][torch.float32][0]["tensor"] + assert source.untyped_storage().data_ptr() == source_storage[index] + assert state["cpu_shards"] == {} + assert len(group._cpu_staging_buffers) == 2 - backend = object.__new__(DistributedLayerwiseOffloadBackend) - backend._remember_mmap_param_attrs(module) - module.to_empty(device="meta") + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) - replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float()) - backend._attach_mmap_param_attrs("weight", replacement, module.weight) + blocks[0].weight.data.zero_() + group.offload() + group.load() + for block, weight in zip(blocks, expected): + assert torch.equal(block.weight, weight) - assert replacement.mmap_weight_transform is transform - assert replacement.mmap_weight_transform_pending is True + group.offload() + assert all(block.weight.numel() == 0 for block in blocks) def test_shared_allgather_output_is_narrowed_to_current_block( self, dist_group, patched_offload_runtime, monkeypatch @@ -598,7 +703,6 @@ class TestMmapWeightLoading: OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, - model_path=str(tmp_path), ), torch.device("cpu"), ) @@ -606,13 +710,82 @@ class TestMmapWeightLoading: dits=[pipeline.transformer], dit_names=["transformer"], ) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={ + name: TensorBinding( + checkpoint_key=name, + file_path=str(weight_file), + ) + for name in weights + }, + runtime_layout_key="test", + ) - backend._load_weights_via_mmap(pipeline, modules) + backend._load_weights_via_mmap(pipeline, modules, plan) assert pipeline.transformer.post_load_calls == 1 assert pipeline.transformer.time_embedder.weight.dtype == torch.float32 assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16 + def test_allgather_flag_with_group_size_one_uses_rank_local_mmap( + self, + tmp_path, + patched_offload_runtime, + ): + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList( + [ + nn.Linear(2, 2, bias=False), + nn.Linear(2, 2, bias=False), + ] + ) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()} + checkpoint_file = tmp_path / "model.safetensors" + save_file(weights, str(checkpoint_file)) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + assert result.plan is not None + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=1, + dlo_use_allgather=True, + ), + torch.device("cpu"), + host_weight_plan=result.plan, + ) + + backend.enable(pipeline) + + assert backend._using_rank_local_mmap + assert backend.dp_group is None + assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group) + backend.disable() + class TestGetBlocksFromDit: def test_get_blocks_from_dit_single_block_attr(self): @@ -878,23 +1051,97 @@ class TestOffloadPlan: class TestMmapValidation: - """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights.""" + """Tests for loader preflight and backend plan realization.""" - def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch): - """DLO+AllGather should reject when TP world size > 1.""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy - - # Mock TP world size = 2 - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", - lambda: 2, + def test_component_source_prefix_builds_model_namespace(self, tmp_path): + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + checkpoint_file = checkpoint_dir / "model.safetensors" + save_file( + {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)}, + str(checkpoint_file), ) + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Module() + self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)]) + self.weights_sources = [ + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + ] + + pipeline = Pipeline() + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=tuple(pipeline.weights_sources), + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.fallback_reason is None + assert result.plan is not None + assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding( + checkpoint_key="blocks.0.weight", + file_path=str(checkpoint_file), + ) + + def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path): + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + save_file( + {"blocks.0.weight": torch.ones((2, 2), dtype=torch.float32)}, + str(checkpoint_dir / "model.safetensors"), + ) + + class Transformer(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)]) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + self.weights_sources = [ + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ) + ] + + pipeline = Pipeline() + original_weight = pipeline.transformer.blocks[0].weight + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=tuple(pipeline.weights_sources), + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert "1 required DiT tensors" in result.fallback_reason + assert pipeline.transformer.blocks[0].weight is original_weight + assert not original_weight.is_meta + + def test_tp_falls_back_before_checkpoint_discovery(self): class SimpleModel(nn.Module): _layerwise_offload_blocks_attrs = ["blocks"] @@ -912,41 +1159,118 @@ class TestMmapValidation: return key pipeline = SimplePipeline() - weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta} - save_file(weights, str(tmp_path / "model.safetensors")) - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path="unused", + tensor_parallel_size=2, + use_hsdp=False, + online_quantization=False, ) - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - model_path=str(tmp_path), - ), - torch.device("cpu"), - ) - modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + assert result.plan is None + assert result.fallback_reason == "TP=2 requires the ordinary loader" - with pytest.raises(ValueError, match="Tensor Parallel"): - backend._load_weights_via_mmap(pipeline, modules) + @pytest.mark.parametrize( + ("use_hsdp", "online_quantization", "expected_reason"), + [ + (True, False, "HSDP requires the ordinary loader"), + (False, True, "online quantization requires the ordinary loader"), + ], + ) + def test_runtime_layouts_that_need_loader_callbacks_fall_back( + self, + use_hsdp, + online_quantization, + expected_reason, + ): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) - def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch): - """DLO+AllGather should work when TP world size = 1, even if params - have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1).""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy - - # Mock TP world size = 1 (default) - monkeypatch.setattr( - "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", - lambda: 1, + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path="unused", + tensor_parallel_size=1, + use_hsdp=use_hsdp, + online_quantization=online_quantization, ) + assert result.plan is None + assert result.fallback_reason == expected_reason + + @pytest.mark.parametrize( + ("checkpoint_tensor", "expected_reason"), + [ + (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"), + (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"), + ], + ) + def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason): + checkpoint_file = tmp_path / "model.safetensors" + save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file)) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(2, 2, bias=False) + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert expected_reason in result.fallback_reason + + def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path): + class Transformer(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(2, 2)) + self.register_buffer("derived", torch.ones(2), persistent=False) + + checkpoint_file = tmp_path / "model.safetensors" + save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file)) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = Transformer() + + @staticmethod + def _remap_ckpt_key(key): + return key + + pipeline = Pipeline() + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is not None + assert set(result.plan.bindings) == {"transformer.weight"} + + def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path): def custom_weight_loader(param, weight): param.data.copy_(weight) @@ -976,27 +1300,67 @@ class TestMmapValidation: json.dumps({"weight_map": {k: "model.safetensors" for k in weights}}) ) - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - model_path=str(tmp_path), - ), - torch.device("cpu"), + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, ) - modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) - # Should NOT raise — TP=1 is allowed even with custom weight_loader - backend._load_weights_via_mmap(pipeline, modules) + assert result.plan is None + assert "custom weight_loader" in result.fallback_reason + + def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path): + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + transformer = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(transformer) + transformer.proj_in = nn.Linear(2, 2, bias=False) + transformer.proj_in.weight.weight_loader = lambda *_args: None + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = transformer + + source_root = tmp_path / "partition" + checkpoint_dir = source_root / "transformer" + checkpoint_dir.mkdir(parents=True) + checkpoint_file = checkpoint_dir / "model.safetensors" + save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file)) + sources = ( + SimpleNamespace( + model_or_path=str(source_root), + subfolder="transformer", + revision=None, + prefix="transformer.", + ), + ) + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", transformer),), + sources=sources, + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is not None + assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding( + checkpoint_key="proj_in.weight", + file_path=str(checkpoint_file), + ) def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime): """validate_loaded_weights should be called after mmap loading.""" - import json - from types import SimpleNamespace - - from safetensors.torch import save_file - - from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy class ModelWithValidate(nn.Module): _layerwise_offload_blocks_attrs = ["blocks"] @@ -1038,19 +1402,65 @@ class TestMmapValidation: backend = DistributedLayerwiseOffloadBackend( OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path) + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, ), torch.device("cpu"), ) modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=(), + model_path=str(tmp_path), + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) - backend._load_weights_via_mmap(pipeline, modules) + assert result.plan is not None + backend._load_weights_via_mmap(pipeline, modules, result.plan) assert pipeline.transformer.validate_called, "validate_loaded_weights should be called" class TestConfigValidation: """Tests for configuration validation in OffloadConfig / execute_request.""" + def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch): + import vllm_omni.diffusion.offloader as offloader_module + + od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_resident_layers=0, + pin_cpu_memory=False, + model="unused", + parallel_config=SimpleNamespace( + use_hsdp=False, + data_parallel_size=1, + sequence_parallel_size=1, + ), + ) + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={}, + runtime_layout_key="test", + ) + monkeypatch.setattr( + offloader_module.current_omni_platform, + "supports_cpu_offload", + lambda: False, + ) + + with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"): + offloader_module.get_offload_backend( + od_config, + device=torch.device("cpu"), + host_weight_plan=plan, + ) + def test_hsdp_with_allgather_rejected(self): """HSDP + DLO + AllGather should raise ValueError (double sharding).""" from vllm_omni.diffusion.offloader.base import OffloadConfig diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 665196ad9..b60cec89c 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -740,6 +740,9 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): del kwargs return SimpleNamespace(transformer=torch.nn.Identity()) + def take_host_weight_plan(self): + return None + class _DummyMemoryProfiler: consumed_memory = 0 @@ -779,7 +782,11 @@ def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler) - monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None) + monkeypatch.setattr( + model_runner_module, + "get_offload_backend", + lambda od_config, device, host_weight_plan: None, + ) monkeypatch.setattr( model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend ) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py index 584562193..804699bc7 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py @@ -9,6 +9,13 @@ from .modelopt import ( ) +def get_direct_mmap_adapter(model: nn.Module): + """Resolve direct-mmap adapters lazily to avoid model-loader cycles.""" + from .direct_mmap import get_direct_mmap_adapter as _get_direct_mmap_adapter + + return _get_direct_mmap_adapter(model) + + def get_checkpoint_adapter( model: nn.Module, source: object, @@ -29,4 +36,5 @@ __all__ = [ "ModelOptMixedPrecisionCheckpointAdapter", "ModelOptNvFp4CheckpointAdapter", "get_checkpoint_adapter", + "get_direct_mmap_adapter", ] diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py new file mode 100644 index 000000000..86b23429b --- /dev/null +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Loader-side adapters for proven direct-checkpoint mmap layouts.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from functools import partial +from typing import Protocol + +import torch +from torch import nn + + +@dataclass(frozen=True) +class DirectMmapTensorPolicy: + """Loader-owned proof for one tensor's direct mmap behavior.""" + + allow_custom_loader: bool = False + transform: Callable[[torch.Tensor], torch.Tensor] | None = None + + +class DirectMmapAdapter(Protocol): + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy | None: ... + + +class _MiniMaxH3DirectMmapAdapter: + """TP1 runtime-layout contract already enforced by MiniMax-H3's loader.""" + + def __init__(self, pipeline: nn.Module) -> None: + self.pipeline = pipeline + + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy: + del target + transform = None + suffix = ".qkv_proj.weight" + if runtime_name.endswith(suffix): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + _reorder_grouped_qkv_to_qkv, + ) + + attention_path = runtime_name[: -len(suffix)] + attention = self.pipeline.get_submodule(attention_path) + if not isinstance(attention, MiniMaxH3Attention): + raise ValueError( + f"MiniMax-H3 direct-mmap adapter expected attention at {attention_path!r}, " + f"got {type(attention).__name__}" + ) + transform = partial( + _reorder_grouped_qkv_to_qkv, + num_query_groups=attention.total_num_heads, + heads_per_group=1, + head_dim=attention.head_dim, + ) + return DirectMmapTensorPolicy( + allow_custom_loader=True, + transform=transform, + ) + + +class _Cosmos3DirectMmapAdapter: + """Preserve Cosmos3's existing TP1 direct-mmap load contract.""" + + def policy_for( + self, + runtime_name: str, + target: torch.Tensor, + ) -> DirectMmapTensorPolicy: + del runtime_name, target + # Cosmos3 has no packed module mapping or TP1 semantic transform after + # its checkpoint-key remap. Its custom vLLM loaders reduce to exact + # copies at TP1; topology and metadata are checked by plan preflight. + return DirectMmapTensorPolicy(allow_custom_loader=True) + + +def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None: + """Return a loader-owned adapter without requiring a pipeline flag.""" + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTModel, + ) + + if any(isinstance(module, MiniMaxH3DiTModel) for module in pipeline.modules()): + return _MiniMaxH3DirectMmapAdapter(pipeline) + + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + Cosmos3VFMTransformer, + ) + + if any(isinstance(module, Cosmos3VFMTransformer) for module in pipeline.modules()): + return _Cosmos3DirectMmapAdapter() + return None + + +__all__ = [ + "DirectMmapAdapter", + "DirectMmapTensorPolicy", + "get_direct_mmap_adapter", +] diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index bde9a1543..8a88ee66c 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -34,6 +34,11 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, + build_checkpoint_mmap_plan, + has_online_quantization, +) from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model @@ -135,6 +140,13 @@ class DiffusersPipelineLoader: self.od_config = od_config self.quant_config = od_config.quantization_config self.parallel_config = od_config.parallel_config + self.host_weight_plan: HostWeightPlan | None = None + + def take_host_weight_plan(self) -> HostWeightPlan | None: + """Transfer the loader-produced plan to the offload backend.""" + plan = self.host_weight_plan + self.host_weight_plan = None + return plan def _prepare_weights( self, @@ -386,6 +398,7 @@ class DiffusersPipelineLoader: device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" + self.host_weight_plan = None if load_format is None: load_format = "default" # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), @@ -417,22 +430,37 @@ class DiffusersPipelineLoader: else: model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False) - # Skip load_weights only for DLO+AllGather when the model - # supports mmap loading and no online quantization is active. - # This condition MUST match the gate in - # DistributedLayerwiseOffloadBackend.enable() so that the - # loader skips ⟺ enable() uses mmap. - from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading - _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False) _use_ag = getattr(self.od_config, "dlo_use_allgather", True) _has_online_quant = self._has_online_quant(model) - _supports_mmap = supports_mmap_loading(model) + _tp_size = int(getattr(self.parallel_config, "tensor_parallel_size", 1)) + _use_hsdp = bool(getattr(self.parallel_config, "use_hsdp", False)) + _dp_size = int(getattr(self.parallel_config, "data_parallel_size", 1)) + _sp_size = int(getattr(self.parallel_config, "sequence_parallel_size", 1)) + _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size - _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant + plan_result = None + if _dist_offload: + modules = ModuleDiscovery.discover(model) + plan_result = build_checkpoint_mmap_plan( + model, + dit_modules=tuple(zip(modules.dit_names, modules.dits)), + sources=self._get_weight_sources(model), + model_path=str(getattr(self.od_config, "model", "")) or None, + tensor_parallel_size=_tp_size, + use_hsdp=_use_hsdp, + online_quantization=_has_online_quant, + ) + self.host_weight_plan = plan_result.plan + + _skip_load = self.host_weight_plan is not None if _skip_load: - logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())") + logger.info( + "DLO host-weight plan active (%s, %s): skipping ordinary weight materialization", + "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local", + self.host_weight_plan.backing_kind, + ) else: if _dist_offload and _use_ag and _has_online_quant: raise ValueError( @@ -442,9 +470,10 @@ class DiffusersPipelineLoader: "Please use --dlo-no-use-allgather or disable online " "quantization." ) - if _dist_offload and _use_ag and not _supports_mmap: + if _dist_offload and plan_result is not None: logger.info( - "DLO+AllGather active but model does not support mmap: loading weights via regular loader." + "DLO direct checkpoint mmap unavailable; using ordinary loader: %s", + plan_result.fallback_reason, ) logger.debug("Loading weights on %s ...", load_device) if offload_after_quant: @@ -498,11 +527,7 @@ class DiffusersPipelineLoader: """Whether any layer uses an online-quant method that defers weight materialization onto the ``meta`` device (upstream vLLM ``uses_meta_device=True``, e.g. online FP8).""" - for module in model.modules(): - quant_method = getattr(module, "quant_method", None) - if getattr(quant_method, "uses_meta_device", False): - return True - return False + return has_online_quantization(model) def _apply_skip_softmax_calibration(self, model: nn.Module) -> None: from vllm_omni.diffusion.attention.backends.trtllm_calibration import ( diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py new file mode 100644 index 000000000..0342c9295 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Loader-owned host-weight plans shared with diffusion offload backends.""" + +from __future__ import annotations + +import glob +import json +import os +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch +from safetensors import safe_open +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + get_direct_mmap_adapter, +) + +logger = init_logger(__name__) + +TensorTransform = Callable[[torch.Tensor], torch.Tensor] + + +@dataclass(frozen=True) +class TensorBinding: + """One runtime tensor backed by one safetensors entry.""" + + checkpoint_key: str + file_path: str + transform: TensorTransform | None = None + + +@dataclass(frozen=True) +class HostWeightPlan: + """Complete, prevalidated host backing consumed by an offload backend.""" + + backing_kind: str + bindings: dict[str, TensorBinding] + runtime_layout_key: str + + +@dataclass(frozen=True) +class HostWeightPlanResult: + """A complete plan or a fail-closed reason to use the ordinary loader.""" + + plan: HostWeightPlan | None + fallback_reason: str | None = None + + +class _PlanIncompatibleError(RuntimeError): + pass + + +_SAFETENSORS_DTYPES: dict[str, torch.dtype] = { + "BOOL": torch.bool, + "U8": torch.uint8, + "I8": torch.int8, + "I16": torch.int16, + "I32": torch.int32, + "I64": torch.int64, + "F16": torch.float16, + "BF16": torch.bfloat16, + "F32": torch.float32, + "F64": torch.float64, +} +for _safe_name, _torch_name in ( + ("F8_E4M3", "float8_e4m3fn"), + ("F8_E5M2", "float8_e5m2"), +): + if (_dtype := getattr(torch, _torch_name, None)) is not None: + _SAFETENSORS_DTYPES[_safe_name] = _dtype + + +def has_online_quantization(model: nn.Module) -> bool: + """Whether any module defers online-quantized weights on ``meta``.""" + return any(getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in model.modules()) + + +def _resolve_source_root(source: object) -> str: + source_root = os.fspath(getattr(source, "model_or_path")) + subfolder = getattr(source, "subfolder", None) + if not os.path.isdir(source_root): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + source_root = download_weights_from_hf( + model_name_or_path=source_root, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + revision=getattr(source, "revision", None), + subfolder=subfolder, + ) + return os.path.join(source_root, subfolder) if subfolder is not None else source_root + + +def _record_binding( + model_to_ckpt: dict[str, tuple[str, str]], + model_name: str, + checkpoint_key: str, + file_path: str, +) -> None: + existing = model_to_ckpt.get(model_name) + candidate = (checkpoint_key, file_path) + if existing is not None and existing != candidate: + raise _PlanIncompatibleError( + f"multiple checkpoint entries map to runtime tensor {model_name!r}: {existing!r} and {candidate!r}" + ) + model_to_ckpt[model_name] = candidate + + +def _add_safetensors_source( + checkpoint_dir: str, + prefix: str, + remap_fn: Callable[[str], str | None] | None, + model_to_ckpt: dict[str, tuple[str, str]], +) -> int: + indexed_keys = 0 + index_files = sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors.index.json"))) + if index_files: + for index_file in index_files: + with open(index_file, encoding="utf-8") as handle: + index = json.load(handle) + for checkpoint_key, filename in index.get("weight_map", {}).items(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + _record_binding( + model_to_ckpt, + model_name, + checkpoint_key, + os.path.join(checkpoint_dir, filename), + ) + indexed_keys += 1 + return indexed_keys + + for filename in sorted(glob.glob(os.path.join(checkpoint_dir, "*.safetensors"))): + with safe_open(filename, framework="pt", device="cpu") as handle: + for checkpoint_key in handle.keys(): + source_name = prefix + checkpoint_key + model_name = remap_fn(source_name) if remap_fn is not None else source_name + if model_name is not None: + _record_binding(model_to_ckpt, model_name, checkpoint_key, filename) + indexed_keys += 1 + return indexed_keys + + +def _build_source_map( + *, + sources: Sequence[object], + model_path: str | None, + remap_fn: Callable[[str], str | None] | None, +) -> dict[str, tuple[str, str]]: + model_to_ckpt: dict[str, tuple[str, str]] = {} + indexed_keys = 0 + + for source in sources: + indexed_keys += _add_safetensors_source( + _resolve_source_root(source), + getattr(source, "prefix", ""), + remap_fn, + model_to_ckpt, + ) + + if not sources: + if remap_fn is None: + raise _PlanIncompatibleError( + "the loader has no component weight sources and the pipeline has no checkpoint-key remapper" + ) + if not model_path: + raise _PlanIncompatibleError("the loader has no model path for legacy checkpoint remapping") + + resolved_model_path = model_path + if not os.path.isdir(resolved_model_path): + from vllm.model_executor.model_loader.weight_utils import ( + download_weights_from_hf, + ) + + resolved_model_path = download_weights_from_hf( + model_name_or_path=resolved_model_path, + cache_dir=None, + allow_patterns=["*.safetensors", "*.safetensors.index.json"], + ) + checkpoint_dirs = { + os.path.dirname(path) + for pattern in ("*.safetensors.index.json", "*.safetensors") + for path in glob.glob(os.path.join(resolved_model_path, "**", pattern), recursive=True) + } + for checkpoint_dir in sorted(checkpoint_dirs): + indexed_keys += _add_safetensors_source( + checkpoint_dir, + "", + remap_fn, + model_to_ckpt, + ) + + if not model_to_ckpt: + raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources") + logger.info( + "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight", + len(model_to_ckpt), + indexed_keys, + ) + return model_to_ckpt + + +def _is_persistent_buffer(module: nn.Module, local_name: str) -> bool: + parts = local_name.split(".") + owner = module.get_submodule(".".join(parts[:-1])) if len(parts) > 1 else module + return parts[-1] not in owner._non_persistent_buffers_set + + +def _collect_required_dit_tensors( + dit_modules: Sequence[tuple[str, nn.Module]], +) -> dict[str, torch.Tensor]: + required: dict[str, torch.Tensor] = {} + for dit_name, dit_module in dit_modules: + for local_name, tensor in dit_module.named_parameters(): + required[f"{dit_name}.{local_name}"] = tensor + for local_name, tensor in dit_module.named_buffers(): + if _is_persistent_buffer(dit_module, local_name): + required[f"{dit_name}.{local_name}"] = tensor + return required + + +def _validate_source_metadata( + required: dict[str, torch.Tensor], + bindings: dict[str, TensorBinding], +) -> None: + by_file: dict[str, list[tuple[str, torch.Tensor, TensorBinding]]] = {} + for runtime_name, target in required.items(): + by_file.setdefault(bindings[runtime_name].file_path, []).append((runtime_name, target, bindings[runtime_name])) + + for file_path, entries in by_file.items(): + if not os.path.isfile(file_path): + raise _PlanIncompatibleError(f"checkpoint file does not exist: {file_path}") + with safe_open(file_path, framework="pt", device="cpu") as handle: + available = set(handle.keys()) + for runtime_name, target, binding in entries: + if binding.checkpoint_key not in available: + raise _PlanIncompatibleError( + f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}" + ) + source = handle.get_slice(binding.checkpoint_key) + source_shape = tuple(source.get_shape()) + source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype()) + if source_shape != tuple(target.shape): + raise _PlanIncompatibleError( + f"shape mismatch for {runtime_name!r}: checkpoint={source_shape}, runtime={tuple(target.shape)}" + ) + if source_dtype is None or source_dtype != target.dtype: + raise _PlanIncompatibleError( + f"dtype mismatch for {runtime_name!r}: checkpoint={source.get_dtype()}, runtime={target.dtype}" + ) + + +def build_checkpoint_mmap_plan( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + sources: Sequence[object], + model_path: str | None, + tensor_parallel_size: int, + use_hsdp: bool, + online_quantization: bool, +) -> HostWeightPlanResult: + """Build a complete direct-checkpoint plan or return a fallback reason.""" + if tensor_parallel_size != 1: + return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader") + if use_hsdp: + return HostWeightPlanResult(None, "HSDP requires the ordinary loader") + if online_quantization: + return HostWeightPlanResult(None, "online quantization requires the ordinary loader") + + remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) + if not callable(remap_fn): + remap_fn = None + adapter = get_direct_mmap_adapter(pipeline) + + try: + model_to_ckpt = _build_source_map( + sources=sources, + model_path=model_path, + remap_fn=remap_fn, + ) + required = _collect_required_dit_tensors(dit_modules) + if not required: + raise _PlanIncompatibleError("no DiT parameters or persistent buffers were discovered") + + missing = sorted(set(required) - set(model_to_ckpt)) + if missing: + raise _PlanIncompatibleError( + f"{len(missing)} required DiT tensors have no checkpoint binding (first 5: {missing[:5]})" + ) + + bindings: dict[str, TensorBinding] = {} + for runtime_name, target in required.items(): + policy = adapter.policy_for(runtime_name, target) if adapter is not None else None + has_custom_loader = callable(getattr(target, "weight_loader", None)) + if has_custom_loader and not (policy and policy.allow_custom_loader): + raise _PlanIncompatibleError( + f"{runtime_name!r} requires a custom weight_loader with no direct-mmap adapter" + ) + checkpoint_key, file_path = model_to_ckpt[runtime_name] + bindings[runtime_name] = TensorBinding( + checkpoint_key=checkpoint_key, + file_path=file_path, + transform=policy.transform if policy is not None else None, + ) + + _validate_source_metadata(required, bindings) + except (OSError, ValueError, _PlanIncompatibleError) as exc: + return HostWeightPlanResult(None, str(exc)) + + return HostWeightPlanResult( + HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings=bindings, + runtime_layout_key="checkpoint-mmap:tp=1:hsdp=0:online-quant=0", + ) + ) + + +__all__ = [ + "HostWeightPlan", + "HostWeightPlanResult", + "TensorBinding", + "build_checkpoint_mmap_plan", + "has_online_quantization", +] 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 5de203721..f639f3ebb 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -540,13 +540,6 @@ class MiniMaxH3Pipeline( 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", diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 9ff464494..2502fa76b 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -5,6 +5,7 @@ import torch from vllm.logger import init_logger from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, OffloadStrategy @@ -17,7 +18,7 @@ from .distributed_layerwise_backend import ( ) from .layerwise_backend import LayerWiseOffloadBackend from .module_residency import PinnedModuleStager -from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .offload_plan import OffloadPlan, get_offload_plan from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, @@ -49,7 +50,6 @@ __all__ = [ "remove_distributed_block_hook", "get_offload_backend", "get_offload_plan", - "supports_mmap_loading", "get_blocks_attr_names", "get_blocks_from_dit", "set_blocks_attr_names", @@ -64,12 +64,15 @@ __all__ = [ def get_offload_backend( od_config: OmniDiffusionConfig, device: torch.device | None = None, + host_weight_plan: HostWeightPlan | None = None, ) -> OffloadBackend | None: """Create appropriate offload backend based on configuration. Args: od_config: OmniDiffusionConfig with offload settings device: Target device (auto-detected if None) + host_weight_plan: Exact loader-produced backing plan, if ordinary + weight materialization was skipped. Returns: OffloadBackend instance or None if offloading disabled @@ -82,12 +85,22 @@ def get_offload_backend( # Extract and validate configuration config = OffloadConfig.from_od_config(od_config) + if host_weight_plan is not None and config.strategy != OffloadStrategy.DISTRIBUTED_LAYER_WISE: + raise RuntimeError( + "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected" + ) + # Return None if no offloading requested if config.strategy == OffloadStrategy.NONE: return None # Validate platform (CUDA required for now) if not current_omni_platform.supports_cpu_offload() or current_omni_platform.get_device_count() < 1: + if host_weight_plan is not None: + raise RuntimeError( + "The loader skipped ordinary weight materialization for DLO, " + "but this platform cannot create the required offload backend" + ) logger.warning( "Current device: %s does not support CPU offloading. Skipping offloading.", current_omni_platform.get_device_name(), @@ -99,6 +112,11 @@ def get_offload_backend( try: device = current_omni_platform.get_torch_device() except (NotImplementedError, AttributeError) as exc: + if host_weight_plan is not None: + raise RuntimeError( + "The loader skipped ordinary weight materialization for DLO, " + "but the target device could not be resolved" + ) from exc logger.error("Failed to detect device: %s. Skipping offloading.", exc) return None @@ -108,7 +126,11 @@ def get_offload_backend( elif config.strategy == OffloadStrategy.LAYER_WISE: return LayerWiseOffloadBackend(config, device) elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE: - return DistributedLayerwiseOffloadBackend(config, device) + return DistributedLayerwiseOffloadBackend( + config, + device, + host_weight_plan=host_weight_plan, + ) else: logger.error("Unknown offload strategy: %s", config.strategy) return None diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 2db80923d..ea74fbfb7 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -27,11 +27,10 @@ class OffloadConfig: pin_cpu_memory: bool = True use_hsdp: bool = False dp_size: int = 1 # derived from parallel_config, not user-configurable - # True: add DP sharding + AllGather. False: stream the standard loader's - # rank-local tensors (including TP-local shards) with H2D only. + # True: add DP sharding + AllGather. False: stream complete rank-local + # blocks from the loader-selected host backing 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 def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -58,7 +57,6 @@ class OffloadConfig: parallel_config = getattr(od_config, "parallel_config", None) use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False - # Derive dp_size from parallel_config — not user-configurable. # The offload adapts to whatever DP/SP is already configured. dp_size = 1 @@ -116,7 +114,8 @@ class OffloadConfig: dp_size = 1 logger.info( "Distributed layerwise offload: dlo_use_allgather=False, " - "streaming standard-loader rank-local weights (no DP shard or AllGather)" + "streaming complete rank-local blocks (no DLO shard or AllGather); " + "the backend will select mmap or standard-loader host storage" ) # HSDP already shards parameters into DTensors. Running distributed @@ -137,7 +136,6 @@ class OffloadConfig: 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 2d6a34978..c6aa464b9 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1,22 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Distributed Layerwise Offload backend with H2D + AllGather overlap. +"""Distributed Layerwise Offload backend with double-buffered H2D. This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: -* Shards model weights across DP ranks and stores only the local shard - (1/DP_size of full model) in host pinned memory per rank. +* Optionally shards model weights across DP ranks and reconstructs each block + with AllGather, or streams a complete rank-local block without a collective. +* Can retain compatible checkpoint tensors as node-shared mmap sources instead + of creating a persistent private host copy in every rank. * Uses a fixed double-buffer scheme that keeps only two layers' worth of weights on each device at any time. -* Asynchronously pipelines both H2D transfers and AllGather communications - on dedicated streams, fully overlapping them with computation. +* Pipelines H2D transfers and, when enabled, AllGather communications on + dedicated streams, overlapping them with computation. * Is hardware-agnostic, supporting both NVIDIA GPU (CUDA) and Ascend NPU (CANN) platforms via vLLM-Omni's platform abstraction layer. """ from __future__ import annotations -import os from itertools import chain from typing import Any @@ -26,12 +27,18 @@ from torch import nn from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook +from vllm_omni.diffusion.model_loader.host_weight_plan import ( + HostWeightPlan, +) from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig from .block_discovery import get_blocks_from_dit from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading +from .offload_plan import ( + OffloadPlan, + get_offload_plan, +) from .tensor_utils import ( dtype_size as _dtype_size, ) @@ -74,6 +81,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, + tensor_transforms: dict[int, Any] | None = None, ): assert isinstance(next_block, nn.Module), "transformer block must be type `torch.nn.Module`" @@ -83,6 +92,8 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.dp_size = dp_size self.rank = rank self.pin_memory = pin_memory + self.rank_local_mmap = rank_local_mmap + self.tensor_transforms = tensor_transforms or {} self.copy_stream = copy_stream or current_omni_platform.Stream() self.comm_stream = comm_stream or current_omni_platform.Stream() @@ -98,8 +109,18 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Sharded host weights for the next block, keyed by dtype self.cpu_shards: dict[torch.dtype, torch.Tensor] = {} + # File-backed source tensors for rank-local mmap. Unlike cpu_shards, + # these remain immutable views of the checkpoint and are never pinned + # or flattened into a model-sized private allocation. + self.cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} self.metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + # Rank-local mmap uses two host staging slots shared by every hook in + # this worker. They are assigned by the backend after all block sizes + # are known, mirroring the shared device-buffer allocation. + self.cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + self.cpu_staging_events: list[Any | None] = [None, None] + # Current slot index (0 or 1). Updated dynamically by the previous # hook's prefetch_layer call via _prefetched_slot. This ensures # correct slot tracking for ALL block counts (including odd N). @@ -154,14 +175,22 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.next_block_parameters = dict(self.next_block.named_parameters()) self.next_block_buffers = dict(self.next_block.named_buffers()) - # Shard next block's weights and store local shard in pinned CPU memory - self.cpu_shards, self.metadata = self._shard_and_pin( - self.next_block_parameters, - self.next_block_buffers, - self.dp_size, - self.rank, - self.pin_memory, - ) + if self.rank_local_mmap: + self.cpu_sources, self.metadata = self._collect_mmap_sources( + self.next_block_parameters, + self.next_block_buffers, + self.tensor_transforms, + ) + else: + # Shard next block's weights and store local shard in pinned CPU memory. + self.cpu_shards, self.metadata = self._shard_and_pin( + self.next_block_parameters, + self.next_block_buffers, + self.dp_size, + self.rank, + self.pin_memory, + self.tensor_transforms, + ) # Allocate device buffers only if not using shared buffers from backend if self._owns_buffers: @@ -194,11 +223,64 @@ class DistributedLayerwiseOffloadHook(ModelHook): self._ag_output_sizes: dict[torch.dtype, int] = {} for dtype, metas in self.metadata.items(): total_numel = sum(m["numel"] for m in metas) - shard_numel = self.cpu_shards[dtype].numel() - self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel + if self.rank_local_mmap: + self._ag_output_sizes[dtype] = total_numel + else: + shard_numel = self.cpu_shards[dtype].numel() + self._ag_output_sizes[dtype] = shard_numel * self.dp_size if self.dp_size > 1 else total_numel return module + @staticmethod + def _collect_mmap_sources( + params: dict[str, nn.Parameter], + bufs: dict[str, torch.Tensor], + tensor_transforms: dict[int, Any] | None = None, + ) -> tuple[dict[torch.dtype, list[dict[str, Any]]], dict[torch.dtype, list[dict[str, Any]]]]: + """Retain file-backed tensors and replace module storage with placeholders. + + The returned sources preserve safetensors mmap storage. Runtime-layout + adapters are applied only while packing a bounded staging slot, so no + full-model anonymous CPU copy is retained by a worker. + """ + cpu_sources: dict[torch.dtype, list[dict[str, Any]]] = {} + metadata: dict[torch.dtype, list[dict[str, Any]]] = {} + offsets: dict[torch.dtype, int] = {} + + for name, target in chain(params.items(), bufs.items()): + source = target.to_local() if hasattr(target, "to_local") else target + if source.device.type != "cpu": + raise ValueError( + f"Rank-local mmap storage requires CPU checkpoint views, but {name!r} is on {source.device}." + ) + + dtype = source.dtype + offset = offsets.get(dtype, 0) + transform = (tensor_transforms or {}).get(id(target)) + + cpu_sources.setdefault(dtype, []).append( + { + "name": name, + "tensor": source.detach(), + "transform": transform, + } + ) + metadata.setdefault(dtype, []).append( + { + "name": name, + "offset": offset, + "numel": source.numel(), + "shape": source.shape, + } + ) + offsets[dtype] = offset + source.numel() + + # The detached source above keeps the mmap storage alive while the + # module parameter/buffer is rebound to the rotating device slot. + set_tensor_storage(target, make_offload_placeholder(target)) + + return cpu_sources, metadata + @staticmethod def _shard_and_pin( params: dict[str, nn.Parameter], @@ -206,6 +288,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): dp_size: int, rank: int, pin_memory: bool, + tensor_transforms: dict[int, Any] | None = None, ) -> tuple[dict[torch.dtype, torch.Tensor], dict[torch.dtype, list[dict[str, Any]]]]: """Flatten params+buffers by dtype, split into DP shards, store local shard. @@ -228,8 +311,8 @@ 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): + mmap_transform = (tensor_transforms or {}).get(id(t)) + if callable(mmap_transform): # 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 @@ -347,6 +430,52 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Prefetch: H2D + AllGather (overlapped on dedicated streams) # # ------------------------------------------------------------------ # + @staticmethod + def _pack_mmap_sources( + cpu_sources: dict[torch.dtype, list[dict[str, Any]]], + metadata: dict[torch.dtype, list[dict[str, Any]]], + slot_buffers: dict[torch.dtype, torch.Tensor], + ) -> dict[torch.dtype, torch.Tensor]: + staged: dict[torch.dtype, torch.Tensor] = {} + for dtype, metas in metadata.items(): + total_numel = sum(meta["numel"] for meta in metas) + destination = slot_buffers[dtype][:total_numel] + sources = cpu_sources[dtype] + for source_info, meta in zip(sources, metas, strict=True): + source = source_info["tensor"] + transform = source_info["transform"] + if callable(transform): + source = transform(source) + if source.dtype != dtype or source.numel() != meta["numel"]: + raise ValueError( + "mmap weight transform changed tensor layout for " + f"{source_info['name']!r}: expected dtype={dtype}, " + f"numel={meta['numel']}, got dtype={source.dtype}, " + f"numel={source.numel()}" + ) + start = meta["offset"] + destination[start : start + meta["numel"]].copy_(source.reshape(-1)) + staged[dtype] = destination + return staged + + def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]: + """Pack this block's mmap views into one bounded host staging slot.""" + previous_copy = self.cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + # Platform events normally expose synchronize(). Retain a + # correctness fallback for test and non-CUDA platform shims. + current_omni_platform.synchronize() + self.cpu_staging_events[slot] = None + + slot_buffers = self.cpu_staging_buffers[slot] + if slot_buffers is None: + raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated") + return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers) + @torch.compiler.disable def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None: """Prepare next block's weights into the shared device buffer for *slot*. @@ -362,11 +491,17 @@ class DistributedLayerwiseOffloadHook(ModelHook): assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated" if self.dp_size <= 1 or self.dp_group is None: + cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards with current_omni_platform.stream(self.copy_stream): - for dtype, cpu_shard in self.cpu_shards.items(): + for dtype, cpu_shard in cpu_weights.items(): gw = gpu_weights[dtype] - gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=non_blocking) + async_copy = non_blocking and cpu_shard.is_pinned() + gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy) evt.record(self.copy_stream) + if self.rank_local_mmap: + # The CPU slot may be overwritten only after this H2D copy has + # finished. The shared event protects reuse by another hook. + self.cpu_staging_events[slot] = evt else: gpu_shards: dict[torch.dtype, torch.Tensor] = {} shard_bufs = self.gpu_shard_buffers[slot] @@ -513,6 +648,8 @@ def apply_distributed_block_hook( comm_stream: Any | None = None, pin_memory: bool = True, shared_buffers: list[dict[torch.dtype, torch.Tensor] | None] | None = None, + rank_local_mmap: bool = False, + tensor_transforms: dict[int, Any] | None = None, ) -> DistributedLayerwiseOffloadHook: """Register a DistributedLayerwiseOffloadHook on *module*.""" registry = HookRegistry.get_or_create(module) @@ -526,6 +663,8 @@ def apply_distributed_block_hook( comm_stream=comm_stream, pin_memory=pin_memory, shared_buffers=shared_buffers, + rank_local_mmap=rank_local_mmap, + tensor_transforms=tensor_transforms, ) registry.register_hook(DistributedLayerwiseOffloadHook._HOOK_NAME, hook) return hook @@ -540,15 +679,16 @@ def remove_distributed_block_hook(module: nn.Module) -> None: class PinnedResidentLayerGroup: - """Keep selected layers in pinned host memory between requests. + """Keep selected layers available for stage-scoped device residency. 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; + pinned CPU master copy (or mmap source plus bounded staging) 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. @@ -563,10 +703,15 @@ class PinnedResidentLayerGroup: device: torch.device, copy_stream: Any, pin_memory: bool, + rank_local_mmap: bool = False, + defer_staging: bool = False, + tensor_transforms: dict[int, Any] | None = None, ) -> None: self.device = device self.copy_stream = copy_stream self.loaded = False + self.rank_local_mmap = rank_local_mmap + self.pin_memory = pin_memory self._states: list[dict[str, Any]] = [] self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] @@ -574,21 +719,49 @@ class PinnedResidentLayerGroup: 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, - ) + if rank_local_mmap: + cpu_sources, metadata = DistributedLayerwiseOffloadHook._collect_mmap_sources( + params, + bufs, + tensor_transforms, + ) + cpu_shards = {} + else: + cpu_shards, metadata = DistributedLayerwiseOffloadHook._shard_and_pin( + params, + bufs, + dp_size=1, + rank=0, + pin_memory=pin_memory, + tensor_transforms=tensor_transforms, + ) + cpu_sources = {} self._states.append( { "targets": targets, "cpu_shards": cpu_shards, + "cpu_sources": cpu_sources, "metadata": metadata, } ) + self._cpu_staging_buffers: list[dict[torch.dtype, torch.Tensor]] = [] + self._cpu_staging_events: list[Any | None] = [None, None] + if rank_local_mmap and not defer_staging: + max_sizes: dict[torch.dtype, int] = {} + for state in self._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + for _ in range(2): + buffers = {} + for dtype, total in max_sizes.items(): + buffer = torch.empty(total, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + self._cpu_staging_buffers.append(buffers) + def load(self) -> None: if self.loaded: return @@ -599,20 +772,41 @@ class PinnedResidentLayerGroup: 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, - ) + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + block_buffers[dtype] = torch.empty(total, 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) + for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)): + if self.rank_local_mmap: + slot = index % 2 + previous_copy = self._cpu_staging_events[slot] + if previous_copy is not None: + synchronize = getattr(previous_copy, "synchronize", None) + if callable(synchronize): + synchronize() + else: + current_omni_platform.synchronize() + cpu_weights = DistributedLayerwiseOffloadHook._pack_mmap_sources( + state["cpu_sources"], + state["metadata"], + self._cpu_staging_buffers[slot], + ) + else: + cpu_weights = state["cpu_shards"] + + for dtype, cpu_weight in cpu_weights.items(): + block_buffers[dtype].copy_( + cpu_weight, + non_blocking=cpu_weight.is_pinned(), + ) + if self.rank_local_mmap: + slot_ready = current_omni_platform.Event() + slot_ready.record(self.copy_stream) + self._cpu_staging_events[slot] = slot_ready ready.record(self.copy_stream) for state, block_buffers in zip(self._states, gpu_buffers): @@ -665,9 +859,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): overlapped with computation. """ - _MMAP_PARAM_ATTRS = ("weight_loader", "mmap_weight_transform") - - def __init__(self, config: OffloadConfig, device: torch.device): + def __init__( + self, + config: OffloadConfig, + device: torch.device, + host_weight_plan: HostWeightPlan | None = None, + ): super().__init__(config, device) self.copy_stream = current_omni_platform.Stream() @@ -679,6 +876,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups: list[list[DistributedLayerwiseOffloadHook]] = [] self._resident_blocks: list[nn.Module] = [] self._resident_layer_group: PinnedResidentLayerGroup | None = None + self._using_mmap = False + self._using_rank_local_mmap = False + self.host_weight_plan = host_weight_plan + self._mmap_transforms_by_tensor_id: dict[int, Any] = {} def load_resident_layers(self) -> None: """Load the model-declared leading blocks for the denoise stage.""" @@ -690,93 +891,30 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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( + def _load_weights_via_mmap( self, - name: str, - replacement: nn.Parameter, - source: nn.Parameter | None = None, + pipeline: nn.Module, + modules, + plan: HostWeightPlan, ) -> 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). + """Load DiT checkpoint tensors as file-backed safetensors views. When the transformer is created on meta device, this method replaces meta params with mmap views of the checkpoint files. - The views point to OS page cache (shared across ranks), so - no private copy is created. _shard_and_pin then copies only - the 1/dp_size shard from the mmap view to a private buffer. + The views point to OS page cache shared across ranks. AllGather mode + copies only the rank's 1/dp_size shard to a persistent private buffer; + rank-local mode retains the views and packs one block at a time into + bounded staging storage. Non-DiT modules (VAE, encoders) are NOT affected — they were created on CPU with real weights via from_pretrained. """ - import glob - import json - from safetensors import safe_open - model_path = self.config.model_path - if not model_path: - logger.warning("No model_path for mmap weight loading, skipping") - return - - # Resolve HF repo ID to local snapshot path. - # If model_path is a local directory, use it directly; otherwise - # download the safetensors index + weight files via HuggingFace Hub. - if not os.path.isdir(model_path): - from vllm.model_executor.model_loader.weight_utils import ( - download_weights_from_hf, - ) - - logger.info("model_path %s is not local, downloading from HF", model_path) - model_path = download_weights_from_hf( - model_name_or_path=model_path, - cache_dir=None, - allow_patterns=["*.safetensors", "*.safetensors.index.json"], - ) - - # Build {checkpoint_key: file_path} from safetensors index - weight_map: dict[str, str] = {} - for idx_file in glob.glob(os.path.join(model_path, "**", "*.safetensors.index.json"), recursive=True): - idx_dir = os.path.dirname(idx_file) - with open(idx_file) as f: - idx = json.load(f) - for key, filename in idx.get("weight_map", {}).items(): - weight_map[key] = os.path.join(idx_dir, filename) - - if not weight_map: - single_files = glob.glob(os.path.join(model_path, "**", "model.safetensors"), recursive=True) - for sf in single_files: - from safetensors import safe_open as _safe_open - - f = _safe_open(sf, framework="pt", device="cpu") - for key in f.keys(): - weight_map[key] = sf - logger.info("Found single-file safetensors: %s (%d keys)", sf, len(f.keys())) - - if not weight_map: - raise RuntimeError( - "Distributed layerwise offload could not find safetensors " - f"weights at {model_path}. Expected *.safetensors.index.json " - "or model.safetensors. Ensure the checkpoint path is correct." - ) - - logger.info("Loading DiT weights via mmap (meta → page cache, no RSS): %d keys", len(weight_map)) + logger.info( + "Loading DiT weights via mmap (meta -> shared page cache): %d tensors", + len(plan.bindings), + ) # --- Convert DiT modules to meta device --- # The transformer was created normally (with random weights and @@ -789,36 +927,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # __init__. Save them before meta conversion and restore after # mmap loading, so we don't need model-specific buffer rebuild code. # - # Fail closed for unsupported configurations BEFORE to_empty. - # The mmap path bypasses AutoWeightsLoader, which means TP-aware - # weight_loader callbacks (fused QKV, row-parallel, etc.) are not - # invoked. When TP > 1, these callbacks shard weights across TP - # ranks — skipping them produces incorrect weights. - # Check the actual TP world size from parallel_state (initialized - # before enable()). Note: params may have custom weight_loader - # attributes even at TP=1 (e.g. QKVParallelLinear), so we cannot - # reject based on attribute presence alone. - try: - from vllm.distributed.parallel_state import ( - get_tensor_model_parallel_world_size, - ) - - tp_world = get_tensor_model_parallel_world_size() - except Exception: - tp_world = 1 - if tp_world > 1: - raise ValueError( - "Distributed layerwise offload with mmap loading does not " - "support Tensor Parallel (TP > 1). TP-aware weight_loader " - "callbacks are bypassed by the mmap path. Use DP or SP " - "instead of TP, or disable distributed layerwise offload." - ) - - # ``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) + # The loader proved topology, coverage, source metadata, and adapter + # compatibility before it skipped ordinary weight materialization. + # This method only realizes that exact plan; it does not select or + # rediscover a checkpoint layout. # # Important: when DiT modules are nested (e.g. transformer contains @@ -861,189 +973,63 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): len(saved_buffers.get(id(dit_module), {})), ) - # Build reverse mapping {model_param_name: (ckpt_key, file_path)} using - # the pipeline's _remap_ckpt_key (if available). This handles the - # model-specific checkpoint key remapping (e.g., Cosmos3's - # layers.0.mlp_moe_gen.gate_proj → transformer.gen_layers.0.mlp.gate_proj). - remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) - model_to_ckpt: dict[str, tuple[str, str]] = {} - if remap_fn is not None: - for ckpt_key, file_path in weight_map.items(): - model_name = remap_fn(ckpt_key) - if model_name is not None: - model_to_ckpt[model_name] = (ckpt_key, file_path) - logger.info( - "Built reverse remap: %d model params from %d checkpoint keys", len(model_to_ckpt), len(weight_map) - ) - else: - # No remap function — use checkpoint keys directly - for ckpt_key, file_path in weight_map.items(): - model_to_ckpt[ckpt_key] = (ckpt_key, file_path) - - # Collect all meta params in DiT modules - dit_ids = set() - for dit_module in modules.dits: - dit_ids.update(id(m) for m in dit_module.modules()) - # Cache open file handles file_cache: dict[str, Any] = {} - loaded = 0 - skipped = 0 loaded_names: set[str] = set() - # Discover blocks first so we can distinguish block params (deferred - # to _shard_and_pin) from non-block params (loaded here as mmap views). - block_module_ids: set[int] = set() - for dit_module in modules.dits: - blocks_attr_names, blocks = get_blocks_from_dit(dit_module) - for block in blocks: - block_module_ids.update(id(m) for m in block.modules()) + # Realize the loader's exact runtime-name bindings. In particular, + # do not reconstruct names from DLO block discovery: storage planning + # and transfer topology are separate contracts. + for runtime_name, binding in plan.bindings.items(): + parent_path, _, leaf_name = runtime_name.rpartition(".") + try: + parent = pipeline.get_submodule(parent_path) + except AttributeError as exc: + raise RuntimeError(f"Host-weight plan target module {parent_path!r} no longer exists") from exc - for name, param in pipeline.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): - continue + target = parent._parameters.get(leaf_name) + is_parameter = target is not None + if target is None: + target = parent._buffers.get(leaf_name) + if target is None: + raise RuntimeError(f"Host-weight plan target tensor {runtime_name!r} no longer exists") - # Check if this param is inside a block (will be handled by - # _shard_and_pin via mmap views in section 2 below). - parent = pipeline - parts = name.split(".") - is_block_param = False - for part in parts[:-1]: - try: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - except (AttributeError, IndexError, TypeError): - break - if id(parent) in block_module_ids: - is_block_param = True - break - - if is_block_param: - # Block param — will be loaded as mmap view in section 2 - continue - - # Non-block param (merger, patch_embed, norm, proj_out, etc.) - # Look up in reverse remap mapping (handles Cosmos3 key remapping). - entry = model_to_ckpt.get(name) - if entry is None: - skipped += 1 - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view (no copy!) - parent = pipeline - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - 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) - - logger.info("Mmap weight loading: %d non-block params loaded, %d skipped", loaded, skipped) - - # Now load DiT block params from safetensors using mmap views. - # Each block param is assigned an mmap view (no RSS). - # _shard_and_pin will later copy the shard portion to a private buffer. - block_loaded = 0 - for dit_idx, dit_module in enumerate(modules.dits): - dit_name = modules.dit_names[dit_idx] - blocks_attr_names, blocks = get_blocks_from_dit(dit_module) - if not blocks: - continue - - # Determine the blocks_attr for this DiT module (e.g. "layers" or "gen_layers") - blocks_attr = None - for attr_name in blocks_attr_names: - if hasattr(dit_module, attr_name): - candidate = getattr(dit_module, attr_name) - if isinstance(candidate, nn.ModuleList) and len(candidate) > 1: - blocks_attr = attr_name - break - - for block_idx, block in enumerate(blocks): - # Build the full model param prefix for this block - # e.g. "transformer.language_model.layers.0" or "transformer.gen_layers.0" - block_full_prefix = f"{dit_name}.{blocks_attr}.{block_idx}" - - for bname, bparam in block.named_parameters(): - if not (hasattr(bparam, "is_meta") and bparam.is_meta): - continue - - # Use reverse remap to find checkpoint key - full_param_name = f"{block_full_prefix}.{bname}" - entry = model_to_ckpt.get(full_param_name) - if entry is None: - skipped += 1 - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view - parent = block - bparts = bname.split(".") - for part in bparts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - 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) - - logger.info("Mmap block loading: %d params loaded, %d skipped", block_loaded, skipped) - - # Strict validation: check that all meta params were loaded. - # This replaces the validation that AutoWeightsLoader.load_weights() - # performs in the regular load path. - remaining_meta = [ - name for name, param in pipeline.named_parameters() if hasattr(param, "is_meta") and param.is_meta - ] - if remaining_meta: - # Filter out params inside submodules that will be loaded later - # by _load_module_weights_from_mmap (non-block submodules) - dit_param_names = set() - for dit_module in modules.dits: - for pname, _ in dit_module.named_parameters(): - dit_param_names.add(pname) - truly_missing = [n for n in remaining_meta if n in dit_param_names] - if truly_missing: - logger.warning( - "Mmap loading: %d params still on meta device after " - "loading (first 5: %s). These may be loaded later by " - "_load_module_weights_from_mmap or are expected to be " - "meta (e.g. DTensor placeholders).", - len(truly_missing), - truly_missing[:5], + if binding.file_path not in file_cache: + file_cache[binding.file_path] = safe_open( + binding.file_path, + framework="pt", + device="cpu", ) + tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key) + if is_parameter: + replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad) + parent._parameters[leaf_name] = replacement + else: + replacement = tensor + parent._buffers[leaf_name] = replacement + loaded_names.add(runtime_name) + + logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names)) # Keep file handles open — _shard_and_pin will read from the mmap views. # They will be released after _shard_and_pin completes (when params are # replaced with offload placeholders). self._mmap_file_cache = file_cache - self._mmap_model_to_ckpt = model_to_ckpt - # The regular loader runs model-specific post-load transforms after # assigning checkpoint tensors. The mmap path bypasses that loader, so # preserve the same lifecycle for transforms such as Cosmos3's fp32 # timestep embedder. - for dit_module in modules.dits: + for dit_name, dit_module in zip(modules.dit_names, modules.dits): + # Restore non-persistent buffers before post-load hooks and strict + # validation. They are constructor-derived and intentionally have + # no checkpoint binding. + bufs = saved_buffers.get(id(dit_module), {}) + for name, buf in bufs.items(): + parent_path, _, leaf_name = name.rpartition(".") + parent = dit_module.get_submodule(parent_path) + restore_device = torch.device("cpu") if self._using_rank_local_mmap else self.device + parent._buffers[leaf_name] = buf.to(restore_device) + post_load_weights = getattr(dit_module, "post_load_weights", None) if callable(post_load_weights): post_load_weights() @@ -1054,73 +1040,45 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # load path (e.g. Cosmos3 checks for missing audio/action weights). validate = getattr(dit_module, "validate_loaded_weights", None) if callable(validate): - validate(loaded_names) + local_loaded_names = { + name.removeprefix(f"{dit_name}.") for name in loaded_names if name.startswith(f"{dit_name}.") + } + validate(local_loaded_names) - # Restore non-persistent buffers that were saved before meta - # conversion. These buffers (e.g. RoPE inv_freq, timestep - # freqs) are computed from formulas in __init__ and are not - # present in the checkpoint. By saving and restoring them - # generically, we avoid model-specific buffer rebuild code. - bufs = saved_buffers.get(id(dit_module), {}) - for name, buf in bufs.items(): - parent = dit_module - parts = name.split(".") - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - # Place on the same device as the module's parameters - # Place on the offloader's target device (not the module's - # current device, which may still be meta at this point). - parent._buffers[parts[-1]] = buf.to(self.device) - - def _load_module_weights_from_mmap(self, module: nn.Module, dit_name: str, child_name: str) -> None: - """Load a non-block DiT submodule's weights from safetensors via mmap.""" - from safetensors import safe_open - - if not hasattr(self, "_mmap_file_cache") or not self._mmap_file_cache: - return - - file_cache = self._mmap_file_cache - model_to_ckpt = getattr(self, "_mmap_model_to_ckpt", {}) - loaded = 0 - - for pname, param in module.named_parameters(): - if not (hasattr(param, "is_meta") and param.is_meta): + # Post-load hooks may rebind parameters (for example while casting a + # submodule), so associate deferred bounded transforms with the final + # runtime tensor objects rather than the initial mmap replacements. + self._mmap_transforms_by_tensor_id.clear() + for runtime_name, binding in plan.bindings.items(): + if binding.transform is None: continue + parent_path, _, leaf_name = runtime_name.rpartition(".") + parent = pipeline.get_submodule(parent_path) + target = parent._parameters.get(leaf_name) + if target is None: + target = parent._buffers.get(leaf_name) + if target is None: + raise RuntimeError( + f"Host-weight transform target {runtime_name!r} no longer exists after post-load processing" + ) + self._mmap_transforms_by_tensor_id[id(target)] = binding.transform - # Build full model param name and look up in reverse remap - full_name = f"{dit_name}.{child_name}.{pname}" - entry = model_to_ckpt.get(full_name) - if entry is None: - # Try without dit_name prefix - full_name = f"{child_name}.{pname}" - entry = model_to_ckpt.get(full_name) - if entry is None: - continue - - ckpt_key, file_path = entry - if file_path not in file_cache: - file_cache[file_path] = safe_open(file_path, framework="pt", device="cpu") - f = file_cache[file_path] - tensor = f.get_tensor(ckpt_key) - - # Replace meta param with mmap view - parent = module - parts = pname.split(".") - for part in parts[:-1]: - if part.isdigit(): - parent = parent[int(part)] - else: - parent = getattr(parent, part) - 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: - logger.info("Loaded %d params for submodule %s.%s via mmap", loaded, dit_name, child_name) + remaining_meta: list[str] = [] + for dit_name, dit_module in zip(modules.dit_names, modules.dits): + remaining_meta.extend( + f"{dit_name}.{name}" for name, tensor in dit_module.named_parameters() if tensor.is_meta + ) + for name, tensor in dit_module.named_buffers(): + parent_path, _, leaf_name = name.rpartition(".") + owner = dit_module.get_submodule(parent_path) + if leaf_name not in owner._non_persistent_buffers_set and tensor.is_meta: + remaining_meta.append(f"{dit_name}.{name}") + if remaining_meta: + raise RuntimeError( + "The prevalidated host-weight plan left " + f"{len(remaining_meta)} DiT tensors on the meta device " + f"(first 5: {remaining_meta[:5]})." + ) def _init_dp_group(self) -> None: """Reuse the process group initialized by parallel_state. @@ -1314,6 +1272,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) sub_hooks = [last_hook] for i, block in enumerate(blocks[:-1]): @@ -1329,6 +1289,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) sub_hooks.append(hook) @@ -1371,9 +1333,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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 @@ -1394,7 +1353,6 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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)) @@ -1437,6 +1395,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): modules = ModuleDiscovery.discover(pipeline) if not modules.dits: + if self.host_weight_plan is not None: + raise RuntimeError( + "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it" + ) logger.warning("No DiT/transformer modules found, skipping distributed layer-wise offloading") return @@ -1451,33 +1413,40 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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 - # regular load_weights() + _process_weights_after_loading(). - if self.config.dlo_use_allgather and self.dp_size > 1: - _has_online_quant = any( - getattr(getattr(module, "quant_method", None), "uses_meta_device", False) - for module in pipeline.modules() + # Storage selection belongs to the loader. DLO consumes the exact + # prevalidated plan that caused the loader to skip materialization; + # without a plan, all weights must already come from the ordinary + # loader. The transfer protocol is selected independently below. + self._using_mmap = self.host_weight_plan is not None + self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1 + if self._using_mmap: + if self.host_weight_plan.backing_kind != "checkpoint_mmap": + raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") + self._load_weights_via_mmap( + pipeline, + modules, + self.host_weight_plan, ) - if _has_online_quant: - raise ValueError( - "Online quantization is incompatible with DLO+AllGather: " - "the sharding + AllGather mechanism flattens weights by " - "dtype, which breaks quantized weight/scale layouts. " - "Please use --dlo-no-use-allgather or disable online " - "quantization." + if self._using_rank_local_mmap: + logger.info( + "DLO rank-local mmap storage enabled: checkpoint pages are " + "node-shared; each worker owns only two bounded host staging slots" ) - if supports_mmap_loading(pipeline): - self._load_weights_via_mmap(pipeline, modules) - else: - logger.info("Weights loaded via regular loader — skipping mmap (model does not support mmap)") + else: + remaining_meta = [ + name + for dit_name, dit_module in zip(modules.dit_names, modules.dits) + for name, tensor in chain( + dit_module.named_parameters(), + dit_module.named_buffers(), + ) + if getattr(tensor, "is_meta", False) + ] + if remaining_meta: + raise RuntimeError( + f"DLO received meta tensors without a loader-owned host-weight plan (first 5: {remaining_meta[:5]})" + ) + logger.info("DLO is using host tensors materialized by the ordinary loader") # Keep VAE/encoders on CPU; move to GPU on-demand via hooks. # This saves ~4.3 GB HBM per card (VAE 1.3 + encoder 1.1 + sound 1.9) @@ -1577,6 +1546,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) block_hooks: list[DistributedLayerwiseOffloadHook] = [last_hook] @@ -1593,6 +1564,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.comm_stream, self.config.pin_cpu_memory, shared_buffers=[None, None], + rank_local_mmap=self._using_rank_local_mmap, + tensor_transforms=self._mmap_transforms_by_tensor_id, ) block_hooks.append(hook) @@ -1621,11 +1594,16 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.device, self.copy_stream, self.config.pin_cpu_memory, + rank_local_mmap=self._using_rank_local_mmap, + defer_staging=bool(self._all_hook_groups), + tensor_transforms=self._mmap_transforms_by_tensor_id, ) pipeline._dlo_residency_controller = self if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) + if self._using_mmap and not self.enabled: + self._release_mmap_handles() return # Unified allocation: 2 shared output buffers + 2 shared shard buffers @@ -1639,6 +1617,21 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): unified_shard_buffers = None if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) + unified_cpu_staging = None + cpu_staging_events = None + if self._using_rank_local_mmap: + unified_cpu_staging = self._allocate_shared_cpu_staging_buffers( + all_hooks, + self._resident_layer_group, + ) + cpu_staging_events = [None, None] + if self._resident_layer_group is not None: + # Resident and streamed layers execute in the same stage and + # reuse the same host slots. Events serialize slot reuse. + self._resident_layer_group._cpu_staging_buffers = [ + buffers for buffers in unified_cpu_staging if buffers is not None + ] + self._resident_layer_group._cpu_staging_events = cpu_staging_events # Shared slot-group tracker: _shared_slot_group[slot] = group_id # that last wrote to that slot. Group-first hooks use this to @@ -1651,6 +1644,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): hook._owns_buffers = False if unified_shard_buffers is not None: hook.gpu_shard_buffers = unified_shard_buffers + if unified_cpu_staging is not None and cpu_staging_events is not None: + hook.cpu_staging_buffers = unified_cpu_staging + hook.cpu_staging_events = cpu_staging_events hook._group_id = group_idx hook._shared_slot_group = shared_slot_group @@ -1677,9 +1673,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self.enabled = True - # Release mmap file handles — _shard_and_pin has copied all shards, - # params now point to offload placeholders (not mmap views). - self._release_mmap_handles() + if self._using_mmap and not self._using_rank_local_mmap: + # AllGather mode copied each rank's persistent shard, so the source + # mappings are no longer needed. Rank-local mode retains them as + # the node-shared host master until disable(). + self._release_mmap_handles() # Assign GPU buffers to sharded on-demand modules (VAE/encoders). # Each module gets a dedicated input (shard-sized) and output @@ -1702,11 +1700,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _release_mmap_handles(self) -> None: """Release safetensors mmap file handles.""" + self._mmap_transforms_by_tensor_id.clear() if hasattr(self, "_mmap_file_cache"): self._mmap_file_cache.clear() del self._mmap_file_cache - if hasattr(self, "_mmap_model_to_ckpt"): - del self._mmap_model_to_ckpt logger.info("Released safetensors mmap file handles") def _cleanup_after_loading(self) -> None: @@ -1723,9 +1720,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): pass def disable(self) -> None: - if not self.enabled: + if not self.enabled and not hasattr(self, "_mmap_file_cache"): return + if self._using_rank_local_mmap: + current_omni_platform.synchronize() + for blocks in self._blocks: for block in blocks: remove_distributed_block_hook(block) @@ -1740,6 +1740,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.clear() self._resident_blocks.clear() self._resident_layer_group = None + self._release_mmap_handles() + self._using_mmap = False + self._using_rank_local_mmap = False self.enabled = False logger.info("Distributed layer-wise offloading disabled") @@ -1779,6 +1782,41 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) return shared_buffers + @staticmethod + def _allocate_shared_cpu_staging_buffers( + hooks: list[DistributedLayerwiseOffloadHook], + resident_group: PinnedResidentLayerGroup | None = None, + ) -> list[dict[torch.dtype, torch.Tensor] | None]: + """Allocate two bounded host slots for rank-local mmap -> device copies.""" + max_sizes: dict[torch.dtype, int] = {} + for hook in hooks: + for dtype, metas in hook.metadata.items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + if resident_group is not None: + for state in resident_group._states: + for dtype, metas in state["metadata"].items(): + total = sum(meta["numel"] for meta in metas) + max_sizes[dtype] = max(max_sizes.get(dtype, 0), total) + + pin_memory = hooks[0].pin_memory if hooks else bool(resident_group and resident_group.pin_memory) + shared_staging: list[dict[torch.dtype, torch.Tensor] | None] = [None, None] + for slot in range(2): + buffers: dict[torch.dtype, torch.Tensor] = {} + for dtype, total_numel in max_sizes.items(): + buffer = torch.empty(total_numel, dtype=dtype, device="cpu") + if pin_memory: + buffer = buffer.pin_memory() + buffers[dtype] = buffer + shared_staging[slot] = buffers + + logger.info( + "Allocated 2 shared host staging buffers for rank-local mmap (max block size: %s, pinned=%s)", + {str(k): f"{v * _dtype_size(k) / 1024 / 1024:.1f}MB" for k, v in max_sizes.items()}, + pin_memory, + ) + return shared_staging + @staticmethod def _allocate_shared_shard_buffers( hooks: list[DistributedLayerwiseOffloadHook], diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index f8943757e..2e1443df7 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -55,20 +55,3 @@ class OffloadPlan: def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None: """Retrieve the OffloadPlan declared by the pipeline, if any.""" return getattr(pipeline, "_offload_plan", None) - - -def supports_mmap_loading(pipeline: nn.Module) -> bool: - """Whether the pipeline supports mmap weight loading. - - A pipeline supports mmap if any module defines ``_remap_ckpt_key``, - which maps checkpoint keys to model parameter names. Models without - this method must use the regular ``load_weights()`` path. - - This check is shared between ``diffusers_loader.py`` (to decide whether - to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()`` - (to decide whether to call ``_load_weights_via_mmap``), ensuring both - paths use the same gate condition. - """ - return 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/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index d118a9833..4a763d651 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -292,7 +292,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): ) # Apply CPU offloading - self.offload_backend = get_offload_backend(self.od_config, device=self.device) + self.offload_backend = get_offload_backend( + self.od_config, + device=self.device, + host_weight_plan=model_loader.take_host_weight_plan(), + ) if self.offload_backend is not None: logger.info(f" Enabling offloader backend: {self.offload_backend.__class__.__name__}") self.offload_backend.enable(self.pipeline)