From 16642dc2e8191fc915698c69abee5d456698e4ef Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 18 Aug 2026 00:21:18 +0800 Subject: [PATCH] [Diffusion] Make DLO host weight cache opt-in Signed-off-by: Hongsheng Liu --- .../distributed_layerwise_offload.md | 26 ++++++++++++------- .../distributed_layerwise_offload.md | 21 ++++++++++++--- tests/config/test_omni_config.py | 16 ++++++++++-- .../model_loader/test_diffusers_loader.py | 11 ++++---- .../test_async_omni_diffusion_config.py | 3 +++ vllm_omni/config/omni_config.py | 3 +++ vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 6 +++++ .../model_loader/diffusers_loader.py | 8 ++++++ vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 8 ++++++ 12 files changed, 85 insertions(+), 20 deletions(-) diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 61ea1aaa9..48c72c95a 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -44,10 +44,12 @@ 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, 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 +weight cache can explicitly normalize supported final ordinary-loader DiT +tensors into an immutable node-local mmap entry. Ordinary pinned weights remain +the default after checkpoint mmap fallback because bounded cache staging can +reduce throughput. Replicas can share either proven-compatible checkpoint pages +or, when selected, 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. @@ -130,9 +132,12 @@ flowchart TB The configuration fields are routed to the component that can enforce their contract. The cache root and lock timeout are consumed by the loader. +`dlo_use_host_weight_cache` explicitly selects final-layout cache publication +after direct-checkpoint mmap fallback. It defaults to false so the ordinary +pinned-loader path remains available without requiring an opt-out. `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. +positive value also 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 @@ -425,8 +430,10 @@ 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 then enter the host weight cache -automatically. Unsupported layouts keep one private runtime copy per process. +parameters. Those ordinary pinned tensors remain the default. When the host +weight cache or a positive registration budget is explicitly selected, +supported CPU layouts enter the cache; unsupported layouts keep one private +runtime copy per process. This mode means: @@ -447,7 +454,8 @@ The Phase B cache preserves the same ownership boundary. DLO neither derives 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: +strict validation, and calibration. When cache use or registration is +explicitly selected, the loader: 1. discovers final CPU DiT parameters and persistent buffers; 2. rejects unsupported non-contiguous, aliased, quantized, distributed, or diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 986a5c7af..2aad95f00 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -38,6 +38,12 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --dlo-no-use-allgather +# Share final host weights with bounded pinned staging +vllm serve /path/to/model --omni \ + --enable-distributed-layerwise-offload \ + --dlo-no-use-allgather \ + --dlo-use-host-weight-cache + # Share final host weights and register up to 80 GiB per worker for direct H2D vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ @@ -52,12 +58,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-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; zero uses bounded host staging | `0` | +| `--dlo-use-host-weight-cache` | Use the final-layout cache after checkpoint mmap fallback | `false` | +| `--dlo-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; positive values also select the cache | `0` | | `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` | -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 +No-AllGather DLO automatically uses a compatible checkpoint mmap plan. When +checkpoint mmap is unavailable, it preserves ordinary pinned loader weights by +default. `--dlo-use-host-weight-cache` explicitly selects the final-layout +cache with bounded staging; a positive registration budget also selects the +cache and attempts direct H2D. 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. @@ -68,6 +77,10 @@ controls are intentionally not separate CLI flags. directory. Do not use tmpfs or a cross-node filesystem for this version. - Cache entries are immutable and validated before use. A cache or registration failure keeps the ordinary loader weights or falls back to bounded staging. +- The bounded-staging cache is a memory-first mode: recurrent pageable-mmap to + pinned-slot copies can reduce throughput, especially across concurrent TP + engines. Keep ordinary pinned weights unless the host-memory saving is worth + that tradeoff, or use a validated registered path. - A positive registration budget must cover the complete page-aligned mapping reported in the worker log. Registration is all-or-nothing. - CUDA is the first registration backend. Platforms without an equivalent diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 45d1852ea..083ac90b0 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -912,12 +912,24 @@ def test_structured_diffusion_config_rejects_invalid_compile_granularity(): ("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): +@pytest.mark.parametrize( + ("field_name", "field_value"), + [ + ("dlo_use_host_weight_cache", True), + ("dlo_host_weight_cache_pin_limit_gib", 1.0), + ], +) +def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather_dlo( + enable_dlo, + use_allgather, + field_name, + field_value, +): 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, + **{field_name: field_value}, ) diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 2b1ac3024..029c95c61 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -375,7 +375,7 @@ 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_then_cache(monkeypatch): +def test_dlo_plan_fallback_retains_ordinary_loader_weights_by_default(monkeypatch): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod @@ -402,14 +402,13 @@ def test_dlo_plan_fallback_runs_ordinary_loader_then_cache(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") + def build_cache(*_args, **_kwargs): + pytest.fail("host weight cache must require explicit selection") monkeypatch.setattr(cache_mod, "build_host_weight_cache_plan", build_cache) assert loader.load_model(load_device="cpu") is model - assert calls == ["load", "process", "cache"] + assert calls == ["load", "process"] assert loader.take_host_weight_plan() is None @@ -434,6 +433,7 @@ 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_use_host_weight_cache=True, dlo_host_weight_cache_dir=str(tmp_path), dlo_host_weight_cache_lock_timeout=1.0, model="unused", @@ -567,6 +567,7 @@ 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_use_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 c52997aa5..f514e6d0f 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -328,6 +328,7 @@ def test_serve_cli_forwards_distributed_offload_residency(): "--omni", "--enable-distributed-layerwise-offload", "--dlo-no-use-allgather", + "--dlo-use-host-weight-cache", "--dlo-host-weight-cache-pin-limit-gib", "80", "--dlo-resident-layers", @@ -341,10 +342,12 @@ 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_use_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_use_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 84111d883..48ca8ace5 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -688,6 +688,7 @@ class _DiffusionConfigProjection: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_use_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) @@ -762,6 +763,8 @@ class _DiffusionConfigProjection: 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.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather): + raise ValueError("dlo_use_host_weight_cache 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 71e60f287..fb2d8178a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -382,6 +382,7 @@ class StageDeployConfig: enable_distributed_layerwise_offload: bool | None = None dlo_use_allgather: bool | None = None + dlo_use_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 e7799c89a..28558f026 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -708,6 +708,10 @@ 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 + # Opt into publishing/joining the final-layout host weight cache when a + # compatible direct-checkpoint mmap plan is unavailable. A positive pin + # limit also selects the cache because registration requires its layout. + dlo_use_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. @@ -951,6 +955,8 @@ class OmniDiffusionConfig: 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.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather): + raise ValueError("dlo_use_host_weight_cache 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 3d0fc438b..874da9b2c 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -550,6 +550,14 @@ class DiffusersPipelineLoader: return if getattr(self.od_config, "dlo_use_allgather", True): return + if not ( + getattr(self.od_config, "dlo_use_host_weight_cache", False) + or float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0 + ): + logger.info( + "DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback" + ) + return parallel_config = self.parallel_config tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1)) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 7fb731c9d..bb67b802b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -545,6 +545,7 @@ class OrchestratorArgs: enable_layerwise_offload: bool = False enable_distributed_layerwise_offload: bool = False dlo_use_allgather: bool = True + dlo_use_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 144455b37..edd880cfa 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1040,6 +1040,7 @@ 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_use_host_weight_cache": kwargs.get("dlo_use_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 7538466f5..41ba2fa82 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -722,6 +722,14 @@ class OmniServeCommand(CLISubcommand): "independently (including existing TP shards)." ), ) + omni_config_group.add_argument( + "--dlo-use-host-weight-cache", + action="store_true", + help=( + "Use the node-local final-layout host weight cache when direct checkpoint mmap is unavailable. " + "This trades recurrent host staging work for shared host pages unless direct registration succeeds." + ), + ) omni_config_group.add_argument( "--dlo-host-weight-cache-pin-limit-gib", type=float,