diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 94684cd81..e3f3e3dcf 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -46,13 +46,14 @@ 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 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. +the default after checkpoint mmap fallback when the cache is not selected; cache +publication and validation add startup and disk lifecycle costs. Replicas can +share either proven-compatible checkpoint pages or, when selected, equivalent +final runtime layouts. Under the existing +pinned-memory policy, the backend attempts to make the complete final mapping +directly H2D-copyable when the platform provides registration, avoiding the +recurrent host packing otherwise required by mmap. An optional positive limit +adds a per-worker safety ceiling. 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 @@ -107,7 +108,7 @@ flowchart TB Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend Contract -.->|"passed once by the runner"| Backend - Backend -.->|"host weight cache + positive budget"| Registration["host_registration
platform-neutral lifecycle"] + Backend -.->|"host weight cache + pinned-memory policy"| Registration["host_registration
platform-neutral lifecycle"] Registration --> CUDA["CUDA registration backend"] Backend --> Hooks["DistributedLayerwiseOffloadHook
one circular hook per streamed block"] Forward["pipeline.forward"] --> Hooks @@ -119,7 +120,7 @@ flowchart TB | Component | Owns | Must not own | | --- | --- | --- | | CLI, engine arguments, and `OmniDiffusionConfig` | User policy, defaults, validation, and propagation to every worker | Checkpoint inspection, cache publication, or buffer allocation | -| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, and registration budget | Host-weight compatibility decisions | +| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, pinned-memory policy, and optional registration ceiling | Host-weight compatibility decisions | | `DiffusersPipelineLoader` | Ordinary loading lifecycle and the decision to return an optional `HostWeightPlan` | DLO process groups or block-stream scheduling | | Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Host weight cache identity or transfer policy | | `host_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives | @@ -133,11 +134,14 @@ 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 +instead of direct-checkpoint mmap. 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 also deliberately selects the transform-complete host weight -cache; the backend later consumes the same budget for registration. +With `pin_cpu_memory=True`, selecting the host weight cache deliberately chooses +the transform-complete layout even when direct checkpoint mmap is available, +because only that final layout can be registered without deferred transforms. +The backend consumes `dlo_host_weight_cache_pin_limit_gib` as an optional +registration ceiling; zero applies no additional ceiling and does not select +the cache by itself. `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 @@ -246,19 +250,22 @@ sequenceDiagram participant P as platform registration participant H as block hooks - L->>L: checkpoint mmap preflight - alt loader selects a compatible checkpoint plan - L-->>B: checkpoint plan via runner - else ordinary loading + alt host weight cache explicitly selected L->>L: load, mutate, validate, and calibrate final tensors - opt host weight cache enabled - L->>C: publish or join final-layout entry - C-->>L: validated host weight cache plan or fail-closed fallback + L->>C: publish or join final-layout entry + C-->>L: validated host weight cache plan or fail-closed fallback + L-->>B: cache plan or materialized fallback via runner + else automatic checkpoint selection + L->>L: checkpoint mmap preflight + alt loader selects a compatible checkpoint plan + L-->>B: checkpoint plan via runner + else ordinary loading + L->>L: load, mutate, validate, and calibrate final tensors + L-->>B: materialized tensors via runner end - L-->>B: plan or materialized tensors via runner end B->>B: consume complete rank-local block sources - opt host weight cache plan and positive registration budget + opt host weight cache plan and pinned-memory policy B->>P: register the complete immutable mapping P-->>B: direct-H2D source or bounded-staging fallback end @@ -267,7 +274,7 @@ sequenceDiagram alt registered host weight cache mapping H->>H: copy complete next block directly to device else pageable mmap mapping - H->>H: pack into a pinned staging slot, then copy to device + H->>H: pack into a bounded staging slot, then copy to device else ordinary runtime tensors H->>H: copy complete next block to device end @@ -287,8 +294,8 @@ the source of the complete-block H2D copy. | Ordinary runtime tensors | AllGather group larger than one | Each rank retains a pinned local weight shard | shard H2D → DLO AllGather → device slot | | Checkpoint mmap | AllGather group larger than one | Each rank prepares its persistent shard, then releases the source mapping | shard H2D → DLO AllGather → device slot | | Ordinary runtime tensors | Rank-local / no-AllGather | Each rank retains the complete loader-produced local layout | complete-block H2D → device slot | -| Checkpoint or host weight cache mmap, zero registration budget | Rank-local / no-AllGather | Immutable mapping plus two bounded pinned host-staging slots | mmap views → pinned staging slot → device slot | -| Host weight cache mmap, positive registration budget | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot | +| Checkpoint mmap, or host weight cache with pinned memory disabled/registration unavailable | Rank-local / no-AllGather | Immutable mapping plus two bounded host-staging slots | mmap views → staging slot → device slot | +| Host weight cache mmap with pinned memory enabled and registration available | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot | The three axes therefore remain independent: @@ -453,14 +460,15 @@ the persistent private full-model copy per pure-DP process, but each process still packs and transfers every complete block. Sharing is node-local; each node has its own page cache. -For a host weight cache mapping, a positive registration budget replaces those -host slots with direct copies from the shared final tensor views. A zero budget -or failed registration retains the bounded staging behavior. +For a host weight cache mapping, the default pinned-memory policy attempts to +replace those host slots with direct copies from the shared final tensor views. +Disabling pinned memory or a failed registration retains 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. Those ordinary pinned tensors remain the default. When the host -weight cache or a positive registration budget is explicitly selected, +When direct mmap preflight fails or the host weight cache is selected, the +regular model loader remains responsible for preparing each rank's weights, +including TP-local tensors or HSDP-managed parameters. Ordinary pinned tensors +remain the default without cache selection. When the host weight cache is +explicitly selected, supported CPU layouts enter the cache; unsupported layouts keep one private runtime copy per process. @@ -512,13 +520,13 @@ error retains the already-valid ordinary tensors. DLO maps this plan without moving the DiT to `meta` and without rerunning any post-load hook. It changes tensor storage while preserving `Parameter` objects, then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing -streaming hooks. The source mappings remain open through DLO disable. With the -default zero registration budget they feed the existing bounded two-slot host -staging path. Every streamed block is first packed from pageable mmap storage -into one of those pinned slots and is then copied to the device. This removes -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. +streaming hooks. The source mappings remain open through DLO disable. With +pinned memory disabled, or when registration is unavailable, they feed the +existing bounded two-slot host-staging path. Every streamed block is first +packed from pageable mmap storage into a staging slot and is then copied to the +device. This removes 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 steady-state-oriented. It requires roughly one local-disk model copy per runtime identity, has no eviction, and still performs @@ -529,19 +537,20 @@ modeled safely. ### Registered host weight cache transfer -`dlo_host_weight_cache_pin_limit_gib > 0` opts a worker into registration of -the complete host weight cache mapping when its platform supports it. It also -changes loader selection: no-AllGather uses the ordinary loader and -final-layout host weight cache even when a TP1 direct-checkpoint plan is -available. A checkpoint binding can carry a deferred adapter such as +Selecting the host weight cache with `pin_cpu_memory=True` opts a worker into +registration of the complete mapping when its platform supports it. It also +changes loader selection: no-AllGather uses the ordinary loader and final-layout +host weight cache even when a TP1 direct-checkpoint plan is available. A +checkpoint binding can carry a deferred adapter such as grouped-QKV reordering. Registering that raw source would not remove the adapter's recurrent CPU allocation/copy, whereas the host weight cache contains the transform-complete final tensors. After mapping and validating every final tensor, DLO groups source storages by backing file, page-aligns their address ranges, and coalesces overlapping or -adjacent ranges within that mapping. The complete byte count is checked -against the per-worker budget before any platform mutation. The CUDA backend +adjacent ranges within that mapping. When a positive per-worker ceiling is +configured, the complete byte count is checked against it before any platform +mutation. The CUDA backend first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every range with the read-only host-registration flag and verifies that PyTorch recognizes every mapped tensor as pinned. A device without that capability @@ -642,10 +651,10 @@ artifacts before making hardware or performance claims. Track that work in direct-mmap deployments can share checkpoint pages per node. Enable the host weight cache only when independent replicas have repeated final layouts and benchmark the target CPU/memory topology, especially for matching - coordinates across multiple TP engines. On CUDA, provide an explicit - registration budget large enough for the complete final layout when - recurrent staging cost is unacceptable; leave it at zero when page-locking - that layout is not operationally acceptable. + coordinates across multiple TP engines. With the default pinned-memory + policy, supported CUDA workers prefer registered direct H2D. Configure a + positive ceiling when an explicit page-locking bound is required, or disable + pinned memory when pageable staging is operationally preferable. ### HSDP diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 1c27509ad..b7a7b971c 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -38,17 +38,18 @@ vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --dlo-no-use-allgather -# Share final host weights with bounded pinned staging +# Share final host weights; prefer registered direct H2D under the pinned-memory policy 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 +# Apply an 80 GiB per-worker safety ceiling to the same registration attempt vllm serve /path/to/model --omni \ --enable-distributed-layerwise-offload \ --tensor-parallel-size 2 \ --dlo-no-use-allgather \ + --dlo-use-host-weight-cache \ --dlo-host-weight-cache-pin-limit-gib 80 ``` @@ -58,15 +59,18 @@ 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-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-use-host-weight-cache` | Use the final-layout cache and prefer registration when pinned memory is enabled | `false` | +| `--dlo-host-weight-cache-pin-limit-gib GIB` | Optional per-worker registration ceiling; zero means no additional ceiling | `0` | | `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` | 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 +default. `--dlo-use-host-weight-cache` explicitly selects the transform-complete +final-layout cache. Under the default `pin_cpu_memory=True` policy, DLO attempts +to register the complete immutable mapping for direct H2D. Setting +`pin_cpu_memory=False` through structured or programmatic configuration retains +the staged mmap path instead. A positive pin limit is only an optional safety +ceiling and does not select the cache. 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. @@ -77,15 +81,16 @@ 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 +- The bounded-staging fallback 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. + engines. Use it intentionally by disabling pinned host memory, or when the + active platform cannot register the immutable mapping. +- When a positive registration ceiling is configured, it must cover the + complete page-aligned mapping reported in the worker log. Registration is + all-or-nothing. Zero applies no additional ceiling. - CUDA is the first registration backend and requires read-only host registration support for immutable cache mappings. Platforms or devices - without that capability continue to use bounded staging; a positive budget + without that capability continue to use bounded staging; pinned-memory policy requests direct H2D but does not guarantee it. - Registration is process-local but does not duplicate the underlying file pages. Each worker must still satisfy its platform and OS page-locking limits. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 083ac90b0..ab9935791 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -933,6 +933,15 @@ def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather ) +def test_structured_diffusion_config_requires_cache_for_registration_ceiling(): + with pytest.raises(ValueError, match="requires dlo_use_host_weight_cache=True"): + omni_config_module._DiffusionConfigProjection( + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + dlo_host_weight_cache_pin_limit_gib=1.0, + ) + + @pytest.mark.parametrize( "field_name", [ diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 029c95c61..b9f5f85f9 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -498,7 +498,8 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat assert loader.take_host_weight_plan() is runtime_plan -def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(monkeypatch, tmp_path): +@pytest.mark.parametrize("pin_cpu_memory", [True, False]) +def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod @@ -515,8 +516,10 @@ 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_use_host_weight_cache=True, + pin_cpu_memory=pin_cpu_memory, dlo_host_weight_cache_dir=str(tmp_path), - dlo_host_weight_cache_pin_limit_gib=1.0, + dlo_host_weight_cache_pin_limit_gib=0.0, model="unused", ) loader = DiffusersPipelineLoader(LoadConfig(), od_config) @@ -536,7 +539,7 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] def unexpected_checkpoint_plan(*_args, **_kwargs): - pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan") + pytest.fail("selected host weight cache mode must not select the raw checkpoint mmap plan") def build_cache(_pipeline, **_kwargs): events.append("cache") diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index 3510de5cc..745628782 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -119,6 +119,20 @@ def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> No assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel +def test_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + sources = {"weights": []} + + def create(actual_sources, *, max_bytes): + assert actual_sources is sources + assert max_bytes is None + return sentinel + + monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create)) + + assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel + + def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None: assert _coalesce_ranges( [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)], @@ -143,6 +157,21 @@ def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytes assert runtime.registered == [] +def test_registration_without_additional_budget_registers_complete_mapping(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) + runtime = _FakeRuntime([0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + registration = CudaHostRegistration.create( + {"weights": [_FakeTensor(0x1003, 4096)]}, # type: ignore[list-item] + max_bytes=None, + ) + + assert registration.total_bytes == 8192 + assert len(runtime.registered) == 1 + assert registration.close() == [] + + def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) runtime = _FakeRuntime([0], read_only_registration_supported=False) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 3a735a636..c03e58c1d 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -1207,6 +1207,68 @@ class TestMmapWeightLoading: assert events == ["unregister", "mmap"] + def test_host_weight_cache_registration_defaults_to_pinned_memory_policy( + self, + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, + ): + class Registration: + total_bytes = 4096 + region_count = 1 + + @staticmethod + def close() -> list[str]: + return [] + + registration = Registration() + create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = [] + + def register(sources, *, device, max_bytes): + assert device == torch.device("cuda") + create_calls.append((sources, max_bytes)) + return registration + + monkeypatch.setattr(dist_backend_module, "register_host_mappings", register) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dlo_use_allgather=False, + dlo_host_weight_cache_pin_limit_gib=0.0, + ), + torch.device("cuda"), + ) + sources = {"runtime.safetensors": [torch.ones(1)]} + backend._host_weight_cache_mapped_sources = sources + + assert backend._try_register_host_weight_cache_mmap() + assert create_calls == [(sources, None)] + assert backend._host_registration is registration + + def test_host_weight_cache_registration_respects_disabled_pinned_memory( + self, + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, + ): + monkeypatch.setattr( + dist_backend_module, + "register_host_mappings", + lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"), + ) + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dlo_use_allgather=False, + dlo_host_weight_cache_pin_limit_gib=0.0, + ), + torch.device("cuda"), + ) + backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]} + + assert not backend._try_register_host_weight_cache_mmap() + assert backend._host_registration is None + def test_host_weight_cache_registration_falls_back_on_unsupported_platform( self, patched_offload_runtime, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 48ca8ace5..3608b1f46 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -765,6 +765,8 @@ class _DiffusionConfigProjection: 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.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache: + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True") if self.tf_model_config is None: self.tf_model_config = TransformerConfig() diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 28558f026..400de9296 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -708,16 +708,15 @@ 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. + # Opt into publishing/joining the transform-complete final-layout host + # weight cache instead of direct-checkpoint mmap. 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. dlo_host_weight_cache_lock_timeout: float = 600.0 - # Maximum file-backed memory a worker may register with CUDA. Zero keeps - # the bounded mmap-to-pinned staging path. + # Optional maximum file-backed memory a worker may register with CUDA. + # Zero applies no additional ceiling; pin_cpu_memory controls registration. dlo_host_weight_cache_pin_limit_gib: float = 0.0 # Leading main-DiT blocks kept resident by distributed layerwise offload. dlo_resident_layers: int = 0 @@ -957,6 +956,8 @@ class OmniDiffusionConfig: 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.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache: + raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True") 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 874da9b2c..02cedefa4 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -441,15 +441,13 @@ class DiffusersPipelineLoader: _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 - _prefer_registered_host_weight_cache = ( - _dist_offload - and not _use_ag - and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0 + _select_host_weight_cache = ( + _dist_offload and not _use_ag and bool(getattr(self.od_config, "dlo_use_host_weight_cache", False)) ) plan_result = None weight_sources = self._get_weight_sources(model) - if _dist_offload and not _prefer_registered_host_weight_cache: + if _dist_offload and not _select_host_weight_cache: modules = ModuleDiscovery.discover(model) plan_result = build_checkpoint_mmap_plan( model, @@ -461,14 +459,14 @@ class DiffusersPipelineLoader: online_quantization=_has_online_quant, ) self.host_weight_plan = plan_result.plan - elif _prefer_registered_host_weight_cache: + elif _select_host_weight_cache: # Direct checkpoint bindings can carry deferred transforms # (for example grouped-QKV reordering), so registering the # raw checkpoint mapping would not make every recurrent H2D # source directly copyable. Materialize once, apply every # post-load mutation, then publish/register the final layout. logger.info( - "DLO host weight cache registration requested; using the ordinary loader " + "DLO host weight cache selected; using the ordinary loader " "once to publish a transform-complete final mmap layout" ) @@ -550,10 +548,7 @@ 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 - ): + if not getattr(self.od_config, "dlo_use_host_weight_cache", False): logger.info( "DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback" ) diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 433909afc..7b02c75d2 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -31,8 +31,8 @@ class OffloadConfig: # blocks from the loader-selected host backing with H2D only. dlo_use_allgather: bool = True dlo_resident_layers: int = 0 # leading DiT layers kept on device - # Per-worker budget for registering shared host weight cache mappings with CUDA. - # Zero retains bounded host staging. + # Optional per-worker ceiling for registering shared host weight cache mappings. + # Zero applies no additional ceiling; pin_cpu_memory controls registration. dlo_host_weight_cache_pin_limit_gib: float = 0.0 @classmethod diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index 1183c9a80..a52596d72 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -182,13 +182,11 @@ class CudaHostRegistration: cls, sources_by_mapping: Mapping[str, Sequence[torch.Tensor]], *, - max_bytes: int, + max_bytes: int | None, ) -> CudaHostRegistration: - if max_bytes <= 0: - raise HostRegistrationBudgetError("CUDA host-registration budget is disabled") regions = _tensor_regions(sources_by_mapping) total_bytes = sum(region.size for region in regions) - if total_bytes > max_bytes: + if max_bytes is not None and total_bytes > max_bytes: raise HostRegistrationBudgetError( f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget" ) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index de298794b..3105496d9 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1269,20 +1269,16 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) def _try_register_host_weight_cache_mmap(self) -> bool: - """Register complete cache mappings when the operator budget permits.""" + """Register complete cache mappings under the pinned-memory policy.""" pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib - if pin_limit_gib <= 0: - return False if not self.config.pin_cpu_memory: - logger.warning( - "DLO host weight cache registration requires pin_cpu_memory=True; using bounded host staging" - ) + logger.info("DLO host weight cache registration disabled by pin_cpu_memory=False; using host staging") return False if not self._host_weight_cache_mapped_sources: logger.warning("DLO host weight cache registration found no mapped sources; using bounded host staging") return False - max_bytes = int(pin_limit_gib * 1024**3) + max_bytes = int(pin_limit_gib * 1024**3) if pin_limit_gib > 0 else None started = time.perf_counter() try: registration = register_host_mappings( diff --git a/vllm_omni/diffusion/offloader/host_registration.py b/vllm_omni/diffusion/offloader/host_registration.py index 72f44f218..51c0d05d9 100644 --- a/vllm_omni/diffusion/offloader/host_registration.py +++ b/vllm_omni/diffusion/offloader/host_registration.py @@ -38,12 +38,13 @@ def register_host_mappings( sources_by_mapping: Mapping[str, Sequence[torch.Tensor]], *, device: torch.device, - max_bytes: int, + max_bytes: int | None, ) -> HostRegistration: """Register mappings with the active platform or report unsupported use. CUDA is the first implementation. Other platforms retain DLO's bounded host-staging path until they provide an equivalent registration backend. + ``max_bytes=None`` applies no additional ceiling to the complete mapping. """ if device.type == "cuda": from .cuda_host_registration import CudaHostRegistration diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 41ba2fa82..839d85da9 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -726,8 +726,9 @@ class OmniServeCommand(CLISubcommand): "--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." + "Use the node-local final-layout host weight cache for no-AllGather DLO. " + "With the default pinned-memory policy, supported platforms register the immutable mapping " + "for direct H2D; otherwise DLO uses bounded host staging." ), ) omni_config_group.add_argument( @@ -735,9 +736,9 @@ class OmniServeCommand(CLISubcommand): type=float, default=0.0, help=( - "Maximum GiB of read-only host weight cache mappings each worker may " - "register for direct H2D when the platform supports it. Zero keeps " - "bounded pinned staging." + "Optional maximum GiB of read-only host weight cache mappings each worker may " + "register for direct H2D. Zero applies no additional registration ceiling; " + "this option requires --dlo-use-host-weight-cache." ), ) omni_config_group.add_argument(