[Diffusion] Add loader-owned host-weight plans for DLO

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

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",
 ]

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",
+]

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 (

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",
+]