[Diffusion] Preserve non-DiT loading with DLO plans

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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

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;

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
 

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"

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(

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,
         )
     )