diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index fa134c03c..61ea1aaa9 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -43,13 +43,14 @@ DLO weight collective. Host storage is selected separately from the transfer protocol. The loader can produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; -otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B -host weight cache can normalize those final ordinary-loader DiT tensors into an -immutable node-local mmap entry. Consequently, replicas can share either -proven-compatible checkpoint pages or equivalent final runtime layouts. An -opt-in host-registration budget can make the complete final mapping directly -H2D-copyable when the platform provides a backend, avoiding the recurrent host -packing otherwise required by mmap. CUDA is the first implementation. +otherwise it uses the ordinary loader. In no-AllGather mode, the Phase B host +weight cache automatically normalizes supported final ordinary-loader DiT +tensors into an immutable node-local mmap entry. Consequently, replicas can +share either proven-compatible checkpoint pages or equivalent final runtime +layouts. An opt-in host-registration budget can make the complete final mapping +directly H2D-copyable when the platform provides a backend, avoiding the +recurrent host packing otherwise required by mmap. CUDA is the first +implementation. The Phase A shared-mmap support boundary is TP1. TP greater than one is an ordinary-loader compatibility path: DLO can consume the resulting TP-local @@ -128,14 +129,14 @@ flowchart TB | `DiffusionWorker` | Calling backend teardown before destroying the distributed environment | Managing backend-owned mappings or registrations directly | The configuration fields are routed to the component that can enforce their -contract. `dlo_enable_host_weight_cache`, its cache root, and lock timeout are -consumed by the loader. `dlo_host_weight_cache_pin_limit_gib` is also visible -while loading because a positive value deliberately selects the -transform-complete host weight cache; the backend later consumes the same -budget for registration. `dlo_use_allgather` controls loader plan routing -because the host weight cache is no-AllGather-only, and it selects the backend -transfer protocol. This is coordination through typed configuration and -`HostWeightPlan`, not through model-specific global state. +contract. The cache root and lock timeout are consumed by the loader. +`dlo_host_weight_cache_pin_limit_gib` is also visible while loading because a +positive value deliberately selects the transform-complete host weight cache; +the backend later consumes the same budget for registration. +`dlo_use_allgather` controls loader plan routing because the host weight cache +is no-AllGather-only, and it selects the backend transfer protocol. This is +coordination through typed configuration and `HostWeightPlan`, not through +model-specific global state. ### Shared startup and lifecycle @@ -424,8 +425,8 @@ or failed registration retains the bounded staging behavior. When direct mmap preflight fails, the regular model loader remains responsible for preparing each rank's weights, including TP-local tensors or HSDP-managed -parameters. Supported CPU layouts may then enter the opt-in host weight cache. -Unsupported layouts keep one private runtime copy per process. +parameters. Supported CPU layouts then enter the host weight cache +automatically. Unsupported layouts keep one private runtime copy per process. This mode means: @@ -443,9 +444,10 @@ This mode means: ### Post-load host weight cache The Phase B cache preserves the same ownership boundary. DLO neither derives -cache keys nor interprets model loaders. When enabled with no-AllGather, every -rank first completes the ordinary loader, custom callbacks, post-load -processing, strict validation, and calibration. The loader then: +cache keys nor interprets model loaders. No-AllGather DLO first attempts a +compatible direct-checkpoint mmap plan. When that plan is unavailable, every +rank completes the ordinary loader, custom callbacks, post-load processing, +strict validation, and calibration. The loader then automatically: 1. discovers final CPU DiT parameters and persistent buffers; 2. rejects unsupported non-contiguous, aliased, quantized, distributed, or @@ -481,9 +483,10 @@ model-sized private pinned storage, but adds a recurrent host-to-host copy and can amplify TP synchronization waits when several ranks compete for host bandwidth. -This version is deliberately opt-in and steady-state-oriented. It requires -roughly one local-disk model copy per runtime identity, has no eviction, and -still performs ordinary loading plus full content-hash passes in every rank. +This version is deliberately steady-state-oriented. It requires roughly one +local-disk model copy per runtime identity, has no eviction, and still performs +ordinary loading plus full content-hash passes in every rank. Cache publication +or validation failure falls back to the already-valid ordinary host tensors. Skip-load-on-hit is a separate follow-up after loader side effects can be modeled safely. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 84733c913..986a5c7af 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -43,7 +43,6 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --tensor-parallel-size 2 \ --dlo-no-use-allgather \ - --dlo-enable-host-weight-cache \ --dlo-host-weight-cache-pin-limit-gib 80 ``` @@ -53,14 +52,15 @@ vllm serve /path/to/model --omni \ | --- | --- | --- | | `--enable-distributed-layerwise-offload` | Enable DLO | `false` | | `--dlo-no-use-allgather` | Stream complete rank-local blocks independently | `false` | -| `--dlo-enable-host-weight-cache` | Share final no-AllGather DiT weights through a node-local mmap cache | `false` | | `--dlo-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; zero uses bounded host staging | `0` | | `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` | -The host weight cache uses `~/.cache/vllm-omni/dlo-host-weights`. Programmatic -configuration may override `dlo_host_weight_cache_dir` and the writer lock -timeout; these advanced storage controls are intentionally not separate CLI -flags. +No-AllGather DLO automatically uses a compatible checkpoint mmap plan, or +builds/joins the final-layout host weight cache after ordinary loading when a +checkpoint plan is unavailable. The cache uses +`~/.cache/vllm-omni/dlo-host-weights`. Programmatic configuration may override +`dlo_host_weight_cache_dir` and the writer lock timeout; these advanced storage +controls are intentionally not separate CLI flags. ## Operational notes diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 108d997a5..1912628b7 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -908,6 +908,19 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity(): omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block") +@pytest.mark.parametrize( + ("enable_dlo", "use_allgather"), + [(False, False), (True, True)], +) +def test_structured_diffusion_config_restricts_host_weight_registration_to_no_allgather_dlo(enable_dlo, use_allgather): + with pytest.raises(ValueError, match="requires no-AllGather distributed layerwise offload"): + omni_config_module._DiffusionConfigProjection( + enable_distributed_layerwise_offload=enable_dlo, + dlo_use_allgather=use_allgather, + dlo_host_weight_cache_pin_limit_gib=1.0, + ) + + def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage(): pipeline = _resolve_pipeline_or_skip("qwen3_tts") legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index e1b9d8202..2b1ac3024 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -375,8 +375,9 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch): assert loader.take_host_weight_plan() is plan -def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): +def test_dlo_plan_fallback_runs_ordinary_loader_then_cache(monkeypatch): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod od_config = SimpleNamespace( dtype=torch.float32, @@ -401,8 +402,14 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"), ) + def build_cache(_pipeline, **_kwargs): + calls.append("cache") + return HostWeightPlanResult(None, "unsupported test layout", "unsupported_layout") + + monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache) + assert loader.load_model(load_device="cpu") is model - assert calls == ["load", "process"] + assert calls == ["load", "process", "cache"] assert loader.take_host_weight_plan() is None @@ -427,7 +434,6 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat quantization_config=None, enable_distributed_layerwise_offload=True, dlo_use_allgather=False, - dlo_enable_host_weight_cache=True, dlo_host_weight_cache_dir=str(tmp_path), dlo_host_weight_cache_lock_timeout=1.0, model="unused", @@ -509,7 +515,6 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m quantization_config=None, enable_distributed_layerwise_offload=True, dlo_use_allgather=False, - dlo_enable_host_weight_cache=True, dlo_host_weight_cache_dir=str(tmp_path), dlo_host_weight_cache_pin_limit_gib=1.0, model="unused", @@ -562,7 +567,6 @@ def test_dlo_host_weight_cache_failure_retains_ordinary_weights(monkeypatch, tmp quantization_config=None, enable_distributed_layerwise_offload=True, dlo_use_allgather=False, - dlo_enable_host_weight_cache=True, dlo_host_weight_cache_dir=str(tmp_path), model="unused", ) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index a99d13344..c52997aa5 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -328,7 +328,6 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", - "--dlo-enable-host-weight-cache", "--dlo-host-weight-cache-pin-limit-gib", "80", "--dlo-resident-layers", @@ -342,12 +341,10 @@ def test_serve_cli_forwards_distributed_offload_residency(): assert args.enable_distributed_layerwise_offload is True assert args.dlo_use_allgather is False - assert args.dlo_enable_host_weight_cache is True assert args.dlo_host_weight_cache_pin_limit_gib == 80 assert args.dlo_resident_layers == 20 assert engine_args["enable_distributed_layerwise_offload"] is True assert engine_args["dlo_use_allgather"] is False - assert engine_args["dlo_enable_host_weight_cache"] is True assert engine_args["dlo_host_weight_cache_dir"] is None assert engine_args["dlo_host_weight_cache_lock_timeout"] == 600.0 assert engine_args["dlo_host_weight_cache_pin_limit_gib"] == 80 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 3641f6d59..321aa53a3 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -687,7 +687,6 @@ class _DiffusionConfigProjection: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True - dlo_enable_host_weight_cache: bool = False dlo_host_weight_cache_dir: str | None = None dlo_host_weight_cache_lock_timeout: float = Field(default=600.0, gt=0) dlo_host_weight_cache_pin_limit_gib: float = Field(default=0.0, ge=0) @@ -753,12 +752,10 @@ class _DiffusionConfigProjection: from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode from vllm_omni.quantization import build_quant_config - if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload: - raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload") - if self.dlo_enable_host_weight_cache and self.dlo_use_allgather: - raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False") - if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache: - raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True") + if self.dlo_host_weight_cache_pin_limit_gib and ( + not self.enable_distributed_layerwise_offload or self.dlo_use_allgather + ): + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload") if self.tf_model_config is None: self.tf_model_config = TransformerConfig() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index c97e1c23e..71e60f287 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -382,7 +382,6 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None - dlo_enable_host_weight_cache: bool | None = None dlo_host_weight_cache_dir: str | None = None dlo_host_weight_cache_lock_timeout: float | None = None dlo_host_weight_cache_pin_limit_gib: float | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d86f42139..e7799c89a 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -708,11 +708,6 @@ class OmniDiffusionConfig: # This avoids AllGather synchronization, while host memory follows the # loader's existing rank-local layout instead of adding a second DP shard. dlo_use_allgather: bool = True - # Build or join a node-local mmap cache after ordinary loading. This is - # used when direct checkpoint mmap is unavailable, or when direct CUDA - # registration requires a transform-complete final layout. The first - # version is no-AllGather only. - dlo_enable_host_weight_cache: bool = False # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-host-weights. dlo_host_weight_cache_dir: str | None = None # Maximum wait for the per-layout POSIX writer lock. @@ -952,12 +947,10 @@ class OmniDiffusionConfig: raise ValueError("dlo_host_weight_cache_lock_timeout must be positive") if not math.isfinite(self.dlo_host_weight_cache_pin_limit_gib) or self.dlo_host_weight_cache_pin_limit_gib < 0: raise ValueError("dlo_host_weight_cache_pin_limit_gib must be non-negative") - if self.dlo_enable_host_weight_cache and not self.enable_distributed_layerwise_offload: - raise ValueError("dlo_enable_host_weight_cache requires distributed layerwise offload") - if self.dlo_enable_host_weight_cache and self.dlo_use_allgather: - raise ValueError("dlo_enable_host_weight_cache currently requires dlo_use_allgather=False") - if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_enable_host_weight_cache: - raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_enable_host_weight_cache=True") + if self.dlo_host_weight_cache_pin_limit_gib and ( + not self.enable_distributed_layerwise_offload or self.dlo_use_allgather + ): + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload") if self.diffusion_compile_granularity not in {"regional", "full"}: raise ValueError( "diffusion_compile_granularity must be 'regional' or 'full', " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index ec0814a74..3d0fc438b 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -444,7 +444,6 @@ class DiffusersPipelineLoader: _prefer_registered_host_weight_cache = ( _dist_offload and not _use_ag - and bool(getattr(self.od_config, "dlo_enable_host_weight_cache", False)) and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0 ) @@ -549,10 +548,7 @@ class DiffusersPipelineLoader: return if not getattr(self.od_config, "enable_distributed_layerwise_offload", False): return - if not getattr(self.od_config, "dlo_enable_host_weight_cache", False): - return if getattr(self.od_config, "dlo_use_allgather", True): - logger.info("DLO host weight cache is no-AllGather only; leaving AllGather storage unchanged") return parallel_config = self.parallel_config diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fe4b1926a..7fb731c9d 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -545,7 +545,6 @@ class OrchestratorArgs: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True - dlo_enable_host_weight_cache: bool = False dlo_host_weight_cache_dir: str | None = None dlo_host_weight_cache_lock_timeout: float = 600.0 dlo_host_weight_cache_pin_limit_gib: float = 0.0 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index d5465f3c0..144455b37 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1040,7 +1040,6 @@ class AsyncOmniEngine: "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), - "dlo_enable_host_weight_cache": kwargs.get("dlo_enable_host_weight_cache", False), "dlo_host_weight_cache_dir": kwargs.get("dlo_host_weight_cache_dir", None), "dlo_host_weight_cache_lock_timeout": kwargs.get("dlo_host_weight_cache_lock_timeout", 600.0), "dlo_host_weight_cache_pin_limit_gib": kwargs.get("dlo_host_weight_cache_pin_limit_gib", 0.0), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 88aca8f6f..7538466f5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -722,15 +722,6 @@ class OmniServeCommand(CLISubcommand): "independently (including existing TP shards)." ), ) - omni_config_group.add_argument( - "--dlo-enable-host-weight-cache", - action="store_true", - help=( - "After ordinary loading, publish or join a node-local mmap cache " - "for final DiT weights. This first version requires " - "--dlo-no-use-allgather." - ), - ) omni_config_group.add_argument( "--dlo-host-weight-cache-pin-limit-gib", type=float,