From fa2f5b45719157c9ffc24286b003590a4278ffe9 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sat, 15 Aug 2026 23:57:03 +0800 Subject: [PATCH] [Diffusion] Preserve non-DiT loading with DLO plans Signed-off-by: Hongsheng Liu --- .../distributed_layerwise_offload.md | 8 +++ .../distributed_layerwise_offload.md | 6 ++ .../model_loader/test_diffusers_loader.py | 62 +++++++++++++++++++ .../test_distributed_layerwise_backend.py | 26 ++++++++ .../model_loader/diffusers_loader.py | 33 ++++++++-- .../model_loader/host_weight_plan.py | 26 +++++++- 6 files changed, 155 insertions(+), 6 deletions(-) diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index b9bb199a2..6c6407e58 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -63,6 +63,14 @@ 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. +The plan owns only dedicated DiT component sources. If a pipeline also exposes +ordinary sources for a text encoder or another non-DiT component, the loader +still consumes those sources and includes their loaded names in its strict +coverage check. Only the source prefixes covered by the plan skip ordinary +materialization. A source that mixes DiT and non-DiT weights fails closed to +the complete ordinary-loader path because it cannot be skipped safely as a +unit. + 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 diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 7d188f4f3..6d79407eb 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -72,6 +72,12 @@ 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. +The mmap plan skips only dedicated DiT weight sources. Other component sources, +such as a text encoder loaded through the shared diffusion loader, continue to +use their ordinary component loader. A checkpoint source that mixes DiT and +non-DiT weights falls back completely rather than leaving an unplanned +component uninitialized. + With direct checkpoint mmap, the loader: 1. saves non-persistent buffers such as RoPE frequencies; diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index a6a724c7e..aa7651bb6 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -19,6 +19,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, HostWeightPlanResult, + TensorBinding, ) from vllm_omni.diffusion.models.helios import HeliosPipeline from vllm_omni.diffusion.registry import initialize_model @@ -301,6 +302,67 @@ def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch assert loader.take_host_weight_plan() is None +def test_dlo_plan_loads_component_sources_outside_planned_dit(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) + + class MixedSourceModel(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(2, 2, bias=False) + self.text_encoder = nn.Linear(2, 2, bias=False) + self.weights_sources = ( + DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."), + DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."), + ) + self.loaded_weight_names: list[str] = [] + + def load_weights(self, weights): + self.loaded_weight_names = [name for name, _ in weights] + return set(self.loaded_weight_names) + + model = MixedSourceModel() + plan = HostWeightPlan( + backing_kind="checkpoint_mmap", + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="weight", + file_path="unused", + ) + }, + planned_source_prefixes=frozenset({"transformer."}), + ) + requested_prefixes: list[str] = [] + + def get_weights(source, model=None): + del model + requested_prefixes.append(source.prefix) + yield source.prefix + "weight", torch.ones(2, 2) + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader._get_weights_iterator = get_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 requested_prefixes == ["text_encoder."] + assert model.loaded_weight_names == ["text_encoder.weight"] + assert loader.take_host_weight_plan() is plan + + def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 70741fe94..6c8248391 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -1094,6 +1094,32 @@ class TestMmapValidation: checkpoint_key="blocks.0.weight", file_path=str(checkpoint_file), ) + assert result.plan.planned_source_prefixes == frozenset({"transformer."}) + + def test_non_dedicated_component_source_falls_back_before_discovery(self): + pipeline = nn.Module() + pipeline.transformer = nn.Linear(2, 2, bias=False) + sources = ( + SimpleNamespace( + model_or_path="unused", + subfolder=None, + revision=None, + prefix="", + ), + ) + + result = build_checkpoint_mmap_plan( + pipeline, + dit_modules=(("transformer", pipeline.transformer),), + sources=sources, + model_path=None, + tensor_parallel_size=1, + use_hsdp=False, + online_quantization=False, + ) + + assert result.plan is None + assert "dedicated component weight source" in result.fallback_reason def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path): source_root = tmp_path / "partition" diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 8a88ee66c..5e05bd547 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -6,7 +6,7 @@ import glob import os import re import time -from collections.abc import Generator, Iterable +from collections.abc import Generator, Iterable, Sequence from pathlib import Path from typing import cast @@ -305,8 +305,10 @@ class DiffusersPipelineLoader: def get_all_weights( self, model: nn.Module, + sources: Sequence["ComponentSource"] | None = None, ) -> Generator[tuple[str, torch.Tensor], None, None]: - sources = self._get_weight_sources(model) + if sources is None: + sources = self._get_weight_sources(model) for source in sources: yield from self._get_weights_iterator(source, model=model) @@ -440,12 +442,13 @@ class DiffusersPipelineLoader: _dlo_group_size = _dp_size if _dp_size > 1 else _sp_size plan_result = None + weight_sources = self._get_weight_sources(model) 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), + sources=weight_sources, model_path=str(getattr(self.od_config, "model", "")) or None, tensor_parallel_size=_tp_size, use_hsdp=_use_hsdp, @@ -457,10 +460,26 @@ class DiffusersPipelineLoader: if _skip_load: logger.info( - "DLO host-weight plan active (%s, %s): skipping ordinary weight materialization", + "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s", "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local", self.host_weight_plan.backing_kind, + sorted(self.host_weight_plan.planned_source_prefixes) or "legacy DiT sources", ) + ordinary_sources = tuple( + source + for source in weight_sources + if source.prefix not in self.host_weight_plan.planned_source_prefixes + ) + if ordinary_sources: + logger.info( + "Loading %d component weight source(s) outside the DLO host-weight plan", + len(ordinary_sources), + ) + self.load_weights( + model, + sources=ordinary_sources, + planned_weights=self.host_weight_plan.bindings, + ) else: if _dist_offload and _use_ag and _has_online_quant: raise ValueError( @@ -630,12 +649,16 @@ class DiffusersPipelineLoader: model: nn.Module, *, stream_online_quant_to_cpu: bool = False, + sources: Sequence["ComponentSource"] | None = None, + planned_weights: Iterable[str] = (), ) -> None: weights_to_load = self._get_expected_parameter_names(model) - weights = self.get_all_weights(model) + weights = self.get_all_weights(model) if sources is None else self.get_all_weights(model, sources=sources) if stream_online_quant_to_cpu: weights = self._stream_online_quant_weights_to_cpu(model, weights) loaded_weights = model.load_weights(weights) + if loaded_weights is not None: + loaded_weights = set(loaded_weights).union(planned_weights) self.counter_after_loading_weights = time.perf_counter() logger.info_once( diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index ac44b0028..a61d60cfb 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -39,6 +39,7 @@ class HostWeightPlan: backing_kind: str bindings: dict[str, TensorBinding] + planned_source_prefixes: frozenset[str] = frozenset() @dataclass(frozen=True) @@ -225,6 +226,24 @@ def _collect_required_dit_tensors( return required +def _planned_source_prefixes( + sources: Sequence[object], + dit_modules: Sequence[tuple[str, nn.Module]], +) -> frozenset[str]: + """Identify component sources that the DiT-only plan owns completely.""" + if not sources: + return frozenset() + + expected = frozenset(f"{dit_name}." for dit_name, _ in dit_modules) + available = {getattr(source, "prefix", "") for source in sources} + missing = sorted(expected - available) + if missing: + raise _PlanIncompatibleError( + f"direct mmap requires a dedicated component weight source for each DiT; missing prefixes: {missing}" + ) + return expected + + def _validate_source_metadata( required: dict[str, torch.Tensor], bindings: dict[str, TensorBinding], @@ -280,8 +299,12 @@ def build_checkpoint_mmap_plan( adapter = get_direct_mmap_adapter(pipeline) try: + planned_source_prefixes = _planned_source_prefixes(sources, dit_modules) + planned_sources = tuple( + source for source in sources if getattr(source, "prefix", "") in planned_source_prefixes + ) model_to_ckpt = _build_source_map( - sources=sources, + sources=planned_sources, model_path=model_path, remap_fn=remap_fn, ) @@ -318,6 +341,7 @@ def build_checkpoint_mmap_plan( HostWeightPlan( backing_kind="checkpoint_mmap", bindings=bindings, + planned_source_prefixes=planned_source_prefixes, ) )