diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 7c7c508ab..3fb494478 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -232,7 +232,8 @@ 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. Runtime-cache v1 is no-AllGather only and rejects all quantization configurations, HSDP/DTensor, -CFG parallelism, PP, non-contiguous tensors, and tied/shared storage. +expert parallelism, CFG parallelism, PP, non-contiguous tensors, and +tied/shared storage. ## Validation coverage diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index f3e5d56a4..9e4632a6b 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -194,7 +194,7 @@ must enter each collective. - Online quantization uses the ordinary loader with no-AllGather. It remains incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP, - CFG-parallel, and PP layouts. + expert-parallel, CFG-parallel, and PP layouts. - 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. diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 7f48cdcb4..d80b9facf 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -12,6 +12,7 @@ import torch import torch.nn as nn from huggingface_hub import snapshot_download from vllm.config.load import LoadConfig +from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -425,6 +426,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, loader = DiffusersPipelineLoader(LoadConfig(), od_config) model = nn.Module() model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = UnquantizedLinearMethod() events: list[str] = [] runtime_plan = HostWeightPlan( backing_kind="runtime_cache", @@ -459,6 +461,7 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0)) assert kwargs["tensor_parallel_rank"] == 0 assert "data_parallel_size" not in kwargs["sequence_parallel_guard"] + assert kwargs["quantization_config"] is None return HostWeightPlanResult(runtime_plan) loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] @@ -502,6 +505,7 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat loader = DiffusersPipelineLoader(LoadConfig(), od_config) model = nn.Module() model.transformer = nn.Linear(2, 2, bias=False) + model.transformer.quant_method = object() loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] loader.load_weights = lambda _model: None # type: ignore[method-assign] loader._process_weights_after_loading = lambda *_args: None # type: ignore[method-assign] @@ -511,11 +515,12 @@ def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_pat "build_checkpoint_mmap_plan", lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"), ) - monkeypatch.setattr( - cache_mod, - "build_runtime_weight_cache_plan", - lambda *_args, **_kwargs: HostWeightPlanResult(None, "writer timed out", "lock_timeout"), - ) + + def build_cache(_pipeline, **kwargs): + assert kwargs["quantization_config"] == "module_quant_method" + return HostWeightPlanResult(None, "writer timed out", "lock_timeout") + + monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache) original_weight = model.transformer.weight.detach().clone() assert loader.load_model(load_device="cpu") is model diff --git a/tests/diffusion/model_loader/test_runtime_weight_cache.py b/tests/diffusion/model_loader/test_runtime_weight_cache.py index f50661135..6db715789 100644 --- a/tests/diffusion/model_loader/test_runtime_weight_cache.py +++ b/tests/diffusion/model_loader/test_runtime_weight_cache.py @@ -61,6 +61,7 @@ def _build(pipeline: nn.Module, cache_root: Path, **overrides): "tensor_parallel_rank": 0, "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"}, "use_hsdp": False, + "enable_expert_parallel": False, "quantization_config": None, "cfg_parallel_size": 1, "pipeline_parallel_size": 1, @@ -125,6 +126,14 @@ def test_final_runtime_content_is_authoritative(tmp_path): assert first.plan.runtime_layout_key != second.plan.runtime_layout_key +def test_expert_parallel_ownership_fails_closed(tmp_path): + result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True) + + assert result.plan is None + assert result.fallback_code == "unsupported_expert_parallel" + assert list(tmp_path.rglob("*.safetensors")) == [] + + def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path): first = _build(_Pipeline(), tmp_path) assert first.plan is not None diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7a2774d85..e77506846 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -14,6 +14,7 @@ import torch from torch import nn from vllm.config.load import LoadConfig from vllm.logger import init_logger +from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( download_safetensors_index_file_from_hf, @@ -551,6 +552,7 @@ class DiffusersPipelineLoader: modules = ModuleDiscovery.discover(model) from vllm_omni.diffusion.model_loader.runtime_weight_cache import ( + DEFAULT_LOCK_TIMEOUT_SECONDS, DEFAULT_SHARD_SIZE_BYTES, build_runtime_weight_cache_plan, ) @@ -572,12 +574,21 @@ class DiffusersPipelineLoader: "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"), } model_identity = getattr(self.od_config, "model", None) + runtime_quantization = self.quant_config + if runtime_quantization is None and any( + (quant_method := getattr(module, "quant_method", None)) is not None + and not isinstance(quant_method, UnquantizedLinearMethod) + for module in model.modules() + ): + runtime_quantization = "module_quant_method" result = build_runtime_weight_cache_plan( model, dit_modules=tuple(zip(modules.dit_names, modules.dits)), loader_type=type(self), cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None), - lock_timeout_seconds=float(getattr(self.od_config, "dlo_runtime_cache_lock_timeout", 600.0)), + lock_timeout_seconds=float( + getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS) + ), max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES, model_identity=str(model_identity) if model_identity is not None else None, revision=getattr(self.od_config, "revision", None), @@ -588,7 +599,8 @@ class DiffusersPipelineLoader: tensor_parallel_rank=tp_rank, sequence_parallel_guard=sp_guard, use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)), - quantization_config=self.quant_config, + enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)), + quantization_config=runtime_quantization, cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)), pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)), ) diff --git a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py index 537615e09..06f43104e 100644 --- a/vllm_omni/diffusion/model_loader/runtime_weight_cache.py +++ b/vllm_omni/diffusion/model_loader/runtime_weight_cache.py @@ -572,6 +572,7 @@ def build_runtime_weight_cache_plan( tensor_parallel_rank: int, sequence_parallel_guard: Mapping[str, Any], use_hsdp: bool, + enable_expert_parallel: bool, quantization_config: Any, cfg_parallel_size: int, pipeline_parallel_size: int, @@ -585,6 +586,12 @@ def build_runtime_weight_cache_plan( ) if use_hsdp: return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp") + if enable_expert_parallel: + return HostWeightPlanResult( + None, + "expert-parallel weight ownership is outside runtime-cache v1", + "unsupported_expert_parallel", + ) if quantization_config is not None: return HostWeightPlanResult( None,