diff --git a/docs/design/feature/distributed_layerwise_offload.md b/docs/design/feature/distributed_layerwise_offload.md index 9a7a7e00b..0c0cca0f8 100644 --- a/docs/design/feature/distributed_layerwise_offload.md +++ b/docs/design/feature/distributed_layerwise_offload.md @@ -53,6 +53,10 @@ block's shard is copied to a device buffer and reconstructed with `all_gather_into_tensor` on a communication stream while the current block is executing. +Multi-rank AllGather behavior is unchanged. When the effective DLO group size +is one, the backend now uses rank-local mmap staging; this closes a previous +loader/backend mismatch that could leave parameters on the `meta` device. + ```text Compute: [Block N] [Block N+1] [Block N+2] H2D: [shard N+1] [shard N+2] diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 4e303cea7..ceec52a3c 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -33,6 +33,7 @@ from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, + has_online_quantization, should_select_dlo_mmap, supports_mmap_loading, ) @@ -909,6 +910,13 @@ class TestOffloadPlan: assert not supports_mmap_loading(Pipeline()) + def test_online_quantization_detection(self): + model = nn.Sequential(nn.Linear(1, 1)) + assert not has_online_quantization(model) + + model[0].quant_method = SimpleNamespace(uses_meta_device=True) + assert has_online_quantization(model) + @pytest.mark.parametrize( ("kwargs", "expected"), [ @@ -1015,6 +1023,55 @@ class TestMmapValidation: ) } + def test_explicit_mmap_opt_in_rejects_missing_model_tensors(self, tmp_path, patched_offload_runtime, monkeypatch): + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size", + lambda: 1, + ) + + 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): + _supports_mmap_loading = True + + 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() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + ), + torch.device("cpu"), + ) + modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"]) + + with pytest.raises(RuntimeError, match="Explicit mmap loading left 1"): + backend._load_weights_via_mmap(pipeline, modules) + 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 diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index bf1f1fb29..079a8b31e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -458,11 +458,11 @@ 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 + from vllm_omni.diffusion.offloader.offload_plan import ( + has_online_quantization, + ) + + 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/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 2fd3a8ab2..ea29257b9 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -36,6 +36,7 @@ from .module_collector import ModuleDiscovery from .offload_plan import ( OffloadPlan, get_offload_plan, + has_online_quantization, should_select_dlo_mmap, supports_mmap_loading, ) @@ -1229,6 +1230,14 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): dit_param_names.add(f"{dit_name}.{pname}") truly_missing = [n for n in remaining_meta if n in dit_param_names] if truly_missing: + if getattr(pipeline, "_supports_mmap_loading", None) is True: + raise RuntimeError( + "Explicit mmap loading left " + f"{len(truly_missing)} DiT tensors on the meta device " + f"(first 5: {truly_missing[:5]}). Check that every " + "runtime tensor has a matching weights_sources prefix " + "and checkpoint key." + ) logger.warning( "Mmap loading: %d params still on meta device after " "loading (first 5: %s). These may be loaded later by " @@ -1672,9 +1681,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # decision stays aligned with whether load_weights() was skipped. # When selection is False, the loader has already loaded weights via # regular load_weights() + _process_weights_after_loading(). - _has_online_quant = any( - getattr(getattr(module, "quant_method", None), "uses_meta_device", False) for module in pipeline.modules() - ) + _has_online_quant = has_online_quantization(pipeline) if self.config.dlo_use_allgather and _has_online_quant: raise ValueError( "Online quantization is incompatible with DLO+AllGather: " diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 904fed650..342c03cbe 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.py @@ -88,6 +88,11 @@ def supports_mmap_loading(pipeline: nn.Module) -> bool: return declared_support is True and bool(getattr(pipeline, "weights_sources", ())) +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 should_select_dlo_mmap( *, mmap_supported: bool,