diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index bf282d8c5..af4eee100 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -196,16 +196,15 @@ adopt it. ### Final-layout Host Weight Runtime consumer -> **Status:** PR2 implementation. The representation-independent identity, -> producer, and restoration contracts landed in -> [PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445); this consumer -> wires the final-layout BF16 path into eligible no-AllGather DLO. +> **Status:** PR2 landed in +> [PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). PR3 adds +> registered direct H2D plus generic HWR publication/source-digest +> optimizations without changing AllGather behavior. Final-layout Host Weight Runtime (HWR) backing is opt-in and currently applies only to no-AllGather DLO. Enable it with `--host-weight-runtime-mode preferred` (or `required`) and -`--host-weight-runtime-root `. Registration of shared mappings -and direct asynchronous H2D remain a separate PR3 transport change. +`--host-weight-runtime-root `. The modes express operator fallback policy, not different artifact formats: @@ -254,6 +253,44 @@ outcomes are authoritative: Checkpoint mmap remains an unchanged control path whenever final-layout HWR is not selected. +Local source identity and publication avoid repeated full-payload passes while +preserving the same semantic contract: + +- Canonical files without a trusted immutable Hub blob identity are hashed in + parallel. Their content IDs are cached per storage domain under the HWR root + and reused only when the path, inode, size, timestamps, symlink target, and + cache-record checksum still match. Corrupt entries are rebuilt; cache-lock or + cache-I/O failure falls back to hashing the canonical file directly. +- Producers that emit payload bytes in canonical storage-key order may declare + `ordered=True`. The filesystem writer then computes file and tensor SHA256 + values during the write and overlaps payload `fsync` with later producer + work. Producers without that guarantee retain unordered writes and use a + parallel readback checksum fallback. +- Manifest and `READY` publication still wait for every payload `fsync` and + checksum to complete. These optimizations do not weaken artifact validation + or change lease ownership. + +#### Registered mmap transport + +After DLO takes an HWR lease, the no-AllGather backend may register the lease's +complete immutable mapped ranges under the existing `pin_cpu_memory` policy. +Registration is transport state: it does not change artifact identity, the +store, tensor ownership, or H2D payload. On success, each tensor view copies +directly into the existing rotating HBM block buffers and the two private host +staging slots are not allocated. + +`--dlo-host-registration-limit-gib` is an optional per-worker preflight ceiling +over page-aligned registered bytes. Zero adds no ceiling. A disabled pinned +memory policy, unsupported platform/capability, over-budget mapping, or fully +rolled-back registration error selects the existing two-slot staging path. A +partial registration that cannot be rolled back aborts startup because closing +the lease would unmap memory still owned by the platform. + +Direct checkpoint mmap remains unchanged and continues to use staging. It may +require loader-owned per-block transforms, while the HWR artifact already +contains final runtime bytes. DLO AllGather never receives an HWR final-layout +lease and therefore never enters this registration path. + #### Pre-service transaction boundary The startup transaction does not end at restore commit. A lease-backed model is @@ -317,6 +354,7 @@ loader owns -> backend takes once before asynchronous work -> backend drains pending H2D work -> backend releases hooks, staging, and model references + -> backend unregisters every mapped range -> backend closes lease ``` @@ -326,13 +364,15 @@ it. Once the backend takes ownership, only backend abort or teardown may drain work and close it. The implementation keeps the final-layout lease through backend setup and -initial prefetch. The backend uses the existing two bounded rank-local staging -slots for the immutable HWR tensors, and its transactional `enable()` cleanup -removes partial hooks and closes the lease before reporting a setup failure. -Preferred mode then uses the runner's fresh canonical retry; required mode -propagates the failure. +initial prefetch. Successful registration copies directly from immutable HWR +views; otherwise the backend uses the existing two bounded rank-local staging +slots. Transactional cleanup drains device work, removes source references, +unregisters mappings, and only then closes the lease. An unregistration failure +retains both registration and lease for retry/process teardown. Preferred mode +then uses the runner's fresh canonical retry; required mode propagates the +failure. -#### PR2 promotion gates +#### PR2 and PR3 promotion gates - Warm hit performs zero ordinary DiT materialization and zero producer calls. - Shared warm finalization changes neither restored bytes nor backing pointers. @@ -344,6 +384,10 @@ propagates the failure. - Disabled mode and AllGather emit zero final-layout HWR interaction. - The lease carrier rejects duplicate take and serialization. - Backend setup or prefetch failure drains asynchronous work before lease close. +- Registration is all-or-nothing; rollback/unregistration failure never closes + a mapping still owned by the platform. +- Registered HWR transport bypasses host staging while preserving output and + H2D payload; unsupported registration retains the bounded staging path. - Compatible checkpoint mmap remains an unchanged benchmark and control path. - Prewarm uses one matching producer cohort per TP/SP topology and storage domain. @@ -394,11 +438,12 @@ reconstructs those tensors with their recorded layouts. Other online methods must use `--dlo-no-use-allgather` or disable online quantization until their runtime layouts are validated. -The Host Weight Runtime representation and publication contracts are merged; -see [RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414) and -[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445). The planned -no-AllGather consumer above is the PR2 implementation; registration and direct -asynchronous H2D remain a separate transport follow-up. +The Host Weight Runtime representation, publication, and no-AllGather consumer +contracts are merged; see +[RFC #6414](https://github.com/vllm-project/vllm-omni/issues/6414), +[PR #6445](https://github.com/vllm-project/vllm-omni/pull/6445), and +[PR #6486](https://github.com/vllm-project/vllm-omni/pull/6486). Registered +direct H2D is an optional transport layer over that merged lease contract. ## Validation coverage @@ -411,8 +456,12 @@ Current source-level validation includes: - ordinary-loader fallback for per-tensor online FP8 linears followed by DLO sharding of finalized weights and scales; - exact loader-to-backend plan transfer and ordinary-loader fallback; +- ordered publication hashing, overlapped durability, unordered parallel + checksum fallback, and node-local source-digest reuse/invalidation; - rank-local mmap source retention, bounded two-slot staging, and adapter transforms without parameter-side flags; +- read-only registration preflight, direct source-to-device copies, safe + fallback, partial rollback, retryable unregistration, and lease ordering; - resident-layer requests requiring no-AllGather; - DP request-wave validation for denoising-step compatibility; - bounded component allocator-cache retention, conservative/forced release, diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index ef6845965..b5656e63d 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -16,7 +16,7 @@ validation_paths: - tests/host_weight_runtime/** upstream_refs: - https://github.com/vllm-project/vllm-omni/issues/6414 -last_reviewed: 2026-08-21 +last_reviewed: 2026-08-25 --- # Host Weight Runtime @@ -84,6 +84,13 @@ another. `WeightArtifactIdentity` includes: - producer version, implementation fingerprint, manifest schema, and restorer schema. +Loader integrations may persist content digests for canonical source shards in +the selected storage domain. Reuse requires the same path, inode, size, +timestamps, symlink target, and cache-record checksum. Corrupt records are +rebuilt, and coordination or cache-I/O failure falls back to hashing the +canonical source directly. This is an identity-computation optimization, not a +second source of truth or an artifact substitution rule. + The identity excludes transfer policy such as registered mmap versus private pinned staging. Those paths move the same host representation and must not produce duplicate artifacts. @@ -237,6 +244,12 @@ Every publication computes full file and tensor hashes. Local lookup supports: and - `full_checksum`, which additionally reads and hashes every payload file. +An ordered producer may compute file and tensor digests while payload bytes are +written and overlap each closed payload's `fsync` with later producer work. An +unordered producer retains parallel readback hashing. In both cases manifest +and `READY` publication wait for every payload digest and `fsync`; the atomic +publication and validation contract is unchanged. + Filesystem verity is modeled but unsupported. Invalid artifacts receive an external deny marker before their shared lookup lock is released; a build owner then quarantines the immutable directory under the exclusive lock. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index fb13b0391..2c6bbacb4 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -65,6 +65,7 @@ omni = Omni( | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | | `--host-weight-runtime-mode {disabled,preferred,required}` | HWR policy: no interaction, populate on a miss, or require an exact hit | `disabled` | | `--host-weight-runtime-root PATH` | Writable node-local HWR store shared by workers in one storage domain; required for `preferred` and `required` | unset | +| `--dlo-host-registration-limit-gib N` | Optional per-worker ceiling for registering an HWR mmap; zero adds no ceiling | `0` | ## Host-weight loading @@ -178,12 +179,48 @@ returning to `required`. On a cold start, the canonical loader remains authoritative and publishes a validated final-layout artifact for later workers. A warm start restores the -DiT final tensors without ordinary DiT materialization, then DLO streams them -through the same two bounded host staging slots. The artifact identity includes -the TP rank/size and SP layout, so TP1, TP2 rank-local shards, and distinct SP -layouts do not alias one another. Publication failure remains separate from a -`preferred` serving startup; the next `required` startup provides the explicit -artifact-availability check. +DiT final tensors without ordinary DiT materialization. DLO then attempts to +register the immutable mapping for direct H2D; if registration is unavailable, +it streams through the same two bounded host staging slots. The artifact +identity includes the TP rank/size and SP layout, so TP1, TP2 rank-local shards, +and distinct SP layouts do not alias one another. Publication failure remains +separate from a `preferred` serving startup; the next `required` startup +provides the explicit artifact-availability check. + +For local canonical checkpoints, the first eligible worker may hash source +shards to establish immutable identity. HWR caches those digests in the same +node-local storage domain and validates file metadata before reuse, so later +workers normally avoid repeating that read. Cold BF16 publication also hashes +ordered payloads as they are written and overlaps payload durability work with +later shards; this changes startup work only, not artifact contents or runtime +H2D behavior. + +#### Registered direct H2D + +Registration is attempted automatically only for an eligible warm HWR hit when +`pin_cpu_memory` is enabled. A successful path is: + +```text +shared read-only HWR mmap -> existing rotating HBM block buffer -> GPU kernel +``` + +It removes the recurrent mmap-to-private-staging CPU copy and does not allocate +the two host staging slots. It does not reduce the H2D payload or make GPU +kernels access host memory directly. + +CUDA requires read-only host-registration capability for the immutable HWR +mapping. Unsupported capability, a positive registration limit smaller than +the complete page-aligned mapping, or a safely rolled-back registration error +falls back to bounded staging. Programmatic configurations can set +`pin_cpu_memory=False` to disable registration explicitly. A successful +registration locks the complete mapped range in host memory for that worker's +lifetime, so use +`--dlo-host-registration-limit-gib` when an operator-enforced ceiling is +required. + +Shutdown drains H2D work, releases hook/source references, unregisters every +range, and only then closes the HWR lease. AllGather and direct checkpoint mmap +retain their existing transfer paths. When HWR is disabled, DLO is disabled, or AllGather is enabled, the loader does not resolve HWR sources or construct its store. diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index c4d08f3a9..5d75b8c44 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlanResult, TensorBinding, ) +from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module from vllm_omni.diffusion.models.helios import HeliosPipeline from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.diffusion.registry import initialize_model @@ -173,6 +174,15 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( str(canonical_root / "model.safetensors"), ) store_root = tmp_path / "hwr-store" + hash_calls = 0 + original_sha256 = source_identity_module._sha256_file + + def counted_sha256(path: Path, state: object) -> str: + nonlocal hash_calls + hash_calls += 1 + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256) def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]: loader = DiffusersPipelineLoader(LoadConfig(), _hwr_config(canonical_root, store_root)) @@ -205,6 +215,8 @@ def test_hwr_cold_publication_and_warm_restore_skip_ordinary_dit_loading( assert warm_plan is not None assert warm_plan.lease_carrier is not None warm_plan.lease_carrier.close() + assert hash_calls == 1 + assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 def test_hwr_commit_failure_discards_model_and_reloads_without_hwr_or_mmap(tmp_path: Path, monkeypatch): diff --git a/tests/diffusion/model_loader/test_final_layout_host_weights.py b/tests/diffusion/model_loader/test_final_layout_host_weights.py index b8ebc27f5..71606febc 100644 --- a/tests/diffusion/model_loader/test_final_layout_host_weights.py +++ b/tests/diffusion/model_loader/test_final_layout_host_weights.py @@ -6,6 +6,8 @@ from __future__ import annotations import dataclasses import gc +import json +import threading from collections.abc import Sequence from pathlib import Path from types import SimpleNamespace @@ -24,10 +26,12 @@ from vllm_omni.diffusion.model_loader.host_weights import ( FinalLayoutRequest, FinalLayoutTensorRestorer, ImplementationIdentity, + NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind, build_final_layout_identity, ) +from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module from vllm_omni.diffusion.model_loader.host_weights.contracts import ( FINAL_LAYOUT_TENSOR_RESTORER_SCHEMA, FinalLayoutArtifactSpec, @@ -380,6 +384,7 @@ def _identity( *, request: FinalLayoutRequest | None = None, policy: FinalLayoutTensorPolicy = FINAL_LAYOUT_BF16_POLICY, + source_digest_cache: NodeSourceDigestCache | None = None, ) -> FinalLayoutIdentityContext: transformer = model.get_submodule("transformer") return build_final_layout_identity( @@ -388,6 +393,7 @@ def _identity( prepared_sources=(source,), request=request or _request(), policy=policy, + source_digest_cache=source_digest_cache, ) @@ -564,6 +570,131 @@ def test_source_identity_resolves_longest_prefix_and_rejects_ties(tmp_path: Path assert exc_info.value.code is FinalLayoutContractCode.SOURCE_COVERAGE_INVALID +def test_node_source_digest_cache_reuses_recovers_and_invalidates( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + cache_root = tmp_path / "store" + calls = 0 + original_sha256 = source_identity_module._sha256_file + + def counted_sha256(path: Path, state: object) -> str: + nonlocal calls + calls += 1 + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", counted_sha256) + first = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + second = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + + assert second.identity == first.identity + assert calls == 1 + + cache_entry = next((cache_root / "source-digests-v1" / "entries").glob("*.json")) + cache_entry.write_text("not-json", encoding="utf-8") + recovered = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert recovered.identity == first.identity + assert calls == 2 + + cache_document = json.loads(cache_entry.read_text(encoding="utf-8")) + cache_document["content_id"] = f"sha256:{'0' * 64}" + cache_entry.write_text(json.dumps(cache_document), encoding="utf-8") + checksum_recovered = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert checksum_recovered.identity == first.identity + assert calls == 3 + + source.weight_files[0].write_bytes(b"changed-canonical-source-with-a-new-size") + changed = _identity( + model, + source, + source_digest_cache=NodeSourceDigestCache(cache_root, timeout_seconds=1.0), + ) + assert changed.identity != first.identity + assert calls == 4 + + +def test_source_identity_hashes_checkpoint_shards_in_parallel( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + second_file = source.resolved_root / "model-00002.safetensors" + second_file.write_bytes(b"second-checkpoint-shard") + source = dataclasses.replace(source, weight_files=(*source.weight_files, second_file)) + barrier = threading.Barrier(2) + thread_ids: set[int] = set() + thread_ids_lock = threading.Lock() + original_sha256 = source_identity_module._sha256_file + + def synchronized_sha256(path: Path, state: object) -> str: + with thread_ids_lock: + thread_ids.add(threading.get_ident()) + barrier.wait(timeout=5) + return original_sha256(path, state) # type: ignore[arg-type] + + monkeypatch.setattr(source_identity_module, "_sha256_file", synchronized_sha256) + _identity(model, source) + + assert len(thread_ids) == 2 + + +def test_source_digest_cache_lock_timeout_falls_back_to_direct_hashing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01) + + def unavailable_lock(*_args: object, **_kwargs: object) -> object: + raise source_identity_module.FileLockTimeoutError("busy") + + monkeypatch.setattr(source_identity_module, "FileLock", unavailable_lock) + context = _identity(model, source, source_digest_cache=cache) + + assert context.identity.source.revision.startswith("content-") + assert not tuple(cache.entries.glob("*.json")) + + +def test_source_digest_cache_directory_failure_falls_back_to_direct_hashing( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _TinyPipeline() + source = _prepared_source(tmp_path) + original_mkdir = Path.mkdir + + def failing_cache_mkdir(path: Path, *args: object, **kwargs: object) -> None: + if "source-digests-v1" in path.parts: + raise OSError("injected cache directory failure") + original_mkdir(path, *args, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(Path, "mkdir", failing_cache_mkdir) + cache = NodeSourceDigestCache(tmp_path / "store", timeout_seconds=0.01) + context = _identity(model, source, source_digest_cache=cache) + + assert context.identity.source.revision.startswith("content-") + + def test_local_hex_named_symlink_target_is_content_hashed_across_startups(tmp_path: Path) -> None: model = _TinyPipeline() source_root = tmp_path / "local-symlink-source" diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py new file mode 100644 index 000000000..a3efbb40a --- /dev/null +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import ctypes +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module +from vllm_omni.diffusion.offloader.cuda_host_registration import ( + CudaHostRegistration, + _coalesce_ranges, + _mapped_regions, +) +from vllm_omni.diffusion.offloader.host_registration import ( + HostRegistrationBudgetError, + HostRegistrationCleanupError, + HostRegistrationError, + register_host_mappings, +) +from vllm_omni.host_weight_runtime import MappedHostRegion + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +class _FakeRuntime: + def __init__( + self, + register_results: list[int | Exception], + unregister_results: list[int | Exception] | None = None, + *, + read_only_supported: bool = True, + attribute_result: int = 0, + ) -> None: + self._register_results = iter(register_results) + self._unregister_results = iter(unregister_results or []) + self._read_only_supported = read_only_supported + self._attribute_result = attribute_result + self.attribute_queries: list[tuple[int, int]] = [] + self.registered: list[tuple[int, int, int]] = [] + self.unregistered: list[int] = [] + + def cudaDeviceGetAttribute(self, value, attribute: int, device: int) -> int: + self.attribute_queries.append((attribute, device)) + ctypes.cast(value, ctypes.POINTER(ctypes.c_int)).contents.value = int(self._read_only_supported) + return self._attribute_result + + def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int: + self.registered.append((pointer, size, flags)) + result = next(self._register_results) + if isinstance(result, Exception): + raise result + return result + + def cudaHostUnregister(self, pointer: int) -> int: + self.unregistered.append(pointer) + result = next(self._unregister_results, 0) + if isinstance(result, Exception): + raise result + return result + + @staticmethod + def cudaGetErrorString(error: int) -> str: + return f"error-{error}" + + +def _region( + file_name: str, + address: int, + length: int, + *, + read_only: bool = True, +) -> MappedHostRegion: + return MappedHostRegion(file_name=file_name, address=address, length=length, read_only=read_only) + + +@pytest.fixture(autouse=True) +def _cuda_runtime_prerequisites(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(registration_module.torch.accelerator, "current_device_index", lambda: 0) + + +def test_platform_dispatch_and_unsupported_device(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + regions = (_region("weights", 0x1000, 4096),) + + def create(actual_regions, *, max_bytes): + assert actual_regions is regions + assert max_bytes == 4096 + return sentinel + + monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create)) + assert register_host_mappings(regions, device=torch.device("cuda"), max_bytes=4096) is sentinel + with pytest.raises(HostRegistrationError, match="not supported on cpu"): + register_host_mappings(regions, device=torch.device("cpu"), max_bytes=None) + + +def test_ranges_align_per_file_without_merging_adjacent_files() -> None: + assert _coalesce_ranges([(0x1003, 4096), (0x2800, 1024)], page_size=4096) == ( + registration_module._AddressRange(0x1000, 0x3000), + ) + assert _mapped_regions( + ( + _region("first.safetensors", 0x1000, 4096), + _region("second.safetensors", 0x2000, 4096), + ) + ) == ( + registration_module._AddressRange(0x1000, 0x2000), + registration_module._AddressRange(0x2000, 0x3000), + ) + + +def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + with pytest.raises(HostRegistrationError, match="is writable"): + CudaHostRegistration.create((_region("weights", 0x1000, 1, read_only=False),), max_bytes=None) + with pytest.raises(HostRegistrationBudgetError, match="exceeding"): + CudaHostRegistration.create((_region("weights", 0x1003, 4096),), max_bytes=4096) + assert runtime.registered == [] + + +def test_registration_requires_read_only_capability(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], read_only_supported=False) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + with pytest.raises(HostRegistrationError, match="does not support read-only"): + CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None) + assert runtime.attribute_queries == [ + (registration_module._CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, 0) + ] + assert runtime.registered == [] + + +def test_capability_error_is_consumed_before_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], attribute_result=7) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + consumed: list[tuple[object, int]] = [] + monkeypatch.setattr( + registration_module, + "_consume_last_cuda_error", + lambda actual_runtime, error: consumed.append((actual_runtime, error)), + ) + + with pytest.raises(HostRegistrationError, match="error-7"): + CudaHostRegistration.create((_region("weights", 0x1000, 4096),), max_bytes=None) + assert consumed == [(runtime, 7)] + + +def test_partial_registration_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0, 7]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + + with pytest.raises(HostRegistrationError, match="error-7"): + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + assert runtime.unregistered == [0x1000] + + +def test_failed_rollback_exposes_active_registration_for_retry(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0, 7], unregister_results=[9, 0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + + with pytest.raises(HostRegistrationCleanupError, match="rollback errors") as error: + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + registration = error.value.active_registration + assert registration is not None + assert registration.close() == () + assert runtime.unregistered == [0x1000, 0x1000] + + +def test_successful_registration_retries_failed_close(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = _FakeRuntime([0], unregister_results=[9, 0]) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + registration = CudaHostRegistration.create((_region("weights", 0x1003, 1),), max_bytes=None) + + assert registration.total_bytes == 4096 + assert registration.region_count == 1 + assert registration.close() == ("cudaHostUnregister(0x1000) failed: error-9",) + assert registration.close() == () + assert runtime.unregistered == [0x1000, 0x1000] + + +@pytest.mark.parametrize("pending_error", [0, 801]) +def test_consume_last_cuda_error_accepts_cleared_or_matching_state( + monkeypatch: pytest.MonkeyPatch, + pending_error: int, +) -> None: + runtime = SimpleNamespace(cudaGetLastError=lambda: pending_error) + monkeypatch.setattr( + registration_module.ctypes, + "CDLL", + lambda _name: pytest.fail("the cudart handle should provide cudaGetLastError"), + ) + + registration_module._consume_last_cuda_error(runtime, 801) + + +def test_consume_last_cuda_error_falls_back_to_global_symbol(monkeypatch: pytest.MonkeyPatch) -> None: + global_runtime = SimpleNamespace(cudaGetLastError=lambda: 801) + monkeypatch.setattr(registration_module.ctypes, "CDLL", lambda _name: global_runtime) + + registration_module._consume_last_cuda_error(SimpleNamespace(), 801) + + +def test_clean_rollback_remains_recoverable_when_error_state_cannot_be_cleared( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runtime = _FakeRuntime([0, 7]) + monkeypatch.setattr(runtime, "cudaGetLastError", None, raising=False) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + + def missing_global_symbol(_name: object) -> None: + raise OSError("cudart symbols are local") + + monkeypatch.setattr( + registration_module.ctypes, + "CDLL", + missing_global_symbol, + ) + + with pytest.raises(HostRegistrationError, match="cannot clear") as error: + CudaHostRegistration.create( + ( + _region("first", 0x1000, 4096), + _region("second", 0x9000, 4096), + ), + max_bytes=None, + ) + + assert not isinstance(error.value, HostRegistrationCleanupError) + assert runtime.unregistered == [0x1000] diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 5a3390f89..9f042966b 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -5,6 +5,7 @@ import gc import json +import weakref from contextlib import contextmanager from types import SimpleNamespace @@ -17,6 +18,7 @@ from torch.distributed.tensor import DeviceMesh, DTensor, Replicate import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module from tests.helpers.runtime import get_distributed_init_method +from vllm_omni.diffusion.data import validate_dlo_host_registration_options from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, TensorBinding, @@ -33,12 +35,14 @@ from vllm_omni.diffusion.offloader.distributed_layerwise_backend import ( DistributedLayerwiseOffloadHook, PinnedResidentLayerGroup, ) +from vllm_omni.diffusion.offloader.host_registration import HostRegistrationCleanupError from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager from vllm_omni.diffusion.offloader.offload_plan import ( OffloadPlan, get_offload_plan, ) from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state +from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -703,11 +707,15 @@ class _HWRPipeline(nn.Module): class _FakeHostWeightLease: - def __init__(self, resolution_id: str): + def __init__(self, resolution_id: str, events: list[str] | None = None): self.closed = False self.provenance = SimpleNamespace(resolution_id=resolution_id) + self.mapped_regions = (MappedHostRegion("weights.safetensors", 0x1000, 4096),) + self._events = events def close(self): + if self._events is not None: + self._events.append("lease") self.closed = True @@ -757,6 +765,203 @@ def _hwr_backend(): return _HWRPipeline(), backend, carrier, lease +def test_registered_mmap_hook_bypasses_host_staging(patched_offload_runtime): + current_block = nn.Linear(2, 2, bias=False) + next_block = nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + next_block.weight.data.copy_(expected) + hook = DistributedLayerwiseOffloadHook( + next_block=next_block, + device=torch.device("cpu"), + dp_group=None, + dp_size=1, + rank=0, + pin_memory=False, + rank_local_mmap=True, + ) + hook.initialize_hook(current_block) + hook.registered_mmap = True + hook._stage_mmap_sources = lambda _slot: pytest.fail("registered mmap must bypass host staging") # type: ignore[method-assign] + + hook.prefetch_layer(slot=0, non_blocking=True) + + assert torch.equal(next_block.weight, expected) + + +def test_registered_mmap_resident_group_bypasses_host_staging(patched_offload_runtime): + block = nn.Linear(2, 2, bias=False) + expected = torch.arange(4, dtype=torch.float32).view(2, 2) + block.weight.data.copy_(expected) + group = PinnedResidentLayerGroup( + [block], + device=torch.device("cpu"), + copy_stream=DummyStream(), + pin_memory=False, + rank_local_mmap=True, + ) + group.registered_mmap = True + group._cpu_staging_buffers.clear() + + group.load() + + assert torch.equal(block.weight, expected) + + +def test_hwr_registration_uses_transport_budget(monkeypatch: pytest.MonkeyPatch): + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + plan, _, lease = _fake_hwr_plan() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + dlo_host_registration_limit_gib=1.5, + ), + torch.device("cuda"), + host_weight_plan=plan, + ) + backend._host_weight_lease = lease # type: ignore[assignment] + calls: list[tuple[object, int | None]] = [] + + class Registration: + total_bytes = 4096 + region_count = 1 + + @staticmethod + def close() -> tuple[str, ...]: + return () + + def register(regions, *, device, max_bytes): + assert device == torch.device("cuda") + calls.append((regions, max_bytes)) + return Registration() + + monkeypatch.setattr(dist_backend_module, "register_host_mappings", register) + pinned_source = SimpleNamespace(numel=lambda: 1, is_pinned=lambda: True) + + assert backend._try_register_hwr_mmap((pinned_source,)) # type: ignore[arg-type] + assert calls == [(lease.mapped_regions, int(1.5 * 1024**3))] + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)] + + backend._release_registered_mmap() + + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [] + + +def test_hwr_registration_budget_validation_is_transport_scoped(): + assert ( + validate_dlo_host_registration_options( + limit_gib=1.5, + enable_dlo=True, + use_allgather=False, + hwr_mode="preferred", + ) + == 1.5 + ) + invalid = ( + dict(limit_gib=-1, enable_dlo=True, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=float("inf"), enable_dlo=True, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=False, use_allgather=False, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=True, use_allgather=True, hwr_mode="preferred"), + dict(limit_gib=1, enable_dlo=True, use_allgather=False, hwr_mode="disabled"), + ) + for options in invalid: + with pytest.raises(ValueError): + validate_dlo_host_registration_options(**options) + + +def test_unregistration_precedes_lease_close_and_retries_failures( + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, +): + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + events: list[str] = [] + responses = iter([("busy",), ()]) + + class Registration: + @staticmethod + def close() -> tuple[str, ...]: + events.append("unregister") + return next(responses) + + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + lease = _FakeHostWeightLease("retry", events) + backend._host_weight_lease = lease # type: ignore[assignment] + backend._host_registration = Registration() + backend._using_rank_local_mmap = True + backend._using_registered_mmap = True + backend.enabled = True + monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None) + + with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"): + backend.disable() + assert not lease.closed + assert backend._host_registration is not None + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(backend._host_registration, lease)] + + backend.disable() + + assert lease.closed + assert backend._host_registration is None + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [] + assert events == ["unregister", "unregister", "lease"] + + +def test_failed_unregistration_retains_lease_after_backend_is_dropped( + monkeypatch: pytest.MonkeyPatch, + patched_offload_runtime, +): + events: list[str] = [] + + class Registration: + @staticmethod + def close() -> tuple[str, ...]: + events.append("unregister") + return ("busy",) + + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + backend = DistributedLayerwiseOffloadBackend( + OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=True, + dp_size=1, + dlo_use_allgather=False, + ), + torch.device("cpu"), + ) + lease = _FakeHostWeightLease("retained", events) + registration = Registration() + backend._host_weight_lease = lease # type: ignore[assignment] + backend._host_registration = registration + backend._using_rank_local_mmap = True + backend.enabled = True + monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None) + lease_ref = weakref.ref(lease) + + with pytest.raises(HostRegistrationCleanupError, match="failed to unregister"): + backend.disable() + del backend + del lease + gc.collect() + + retained_lease = dist_backend_module._ACTIVE_HWR_REGISTRATIONS[0][1] + assert dist_backend_module._ACTIVE_HWR_REGISTRATIONS == [(registration, retained_lease)] + assert lease_ref() is retained_lease + assert not retained_lease.closed + + dist_backend_module._ACTIVE_HWR_REGISTRATIONS.clear() + retained_lease.close() + + class _MultiBlockModel(nn.Module): _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 170faec80..709fc7283 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for WorkerWrapperBase class. @@ -12,6 +12,7 @@ This module tests the WorkerWrapperBase implementation: - Dynamic worker class extension """ +from types import SimpleNamespace from typing import Any import pytest @@ -204,6 +205,27 @@ class TestWorkerWrapperBaseDelegation: wrapper.worker.shutdown.assert_called_once() assert result is None + def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture): + events: list[str] = [] + offload_backend = mocker.Mock() + offload_backend.disable.side_effect = lambda: events.append("offload") + kv_manager = mocker.Mock() + kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv") + destroy = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env", + side_effect=lambda: events.append("distributed"), + ) + worker = DiffusionWorker.__new__(DiffusionWorker) + worker.model_runner = SimpleNamespace( + offload_backend=offload_backend, + kv_transfer_manager=kv_manager, + ) + + worker.shutdown() + + assert events == ["offload", "kv", "distributed"] + destroy.assert_called_once_with() + # ------------------------------------------------------------------------- # Tests: execute_method diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 37d4af8c3..452eec286 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -378,6 +378,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): "preferred", "--host-weight-runtime-root", "/var/cache/vllm-omni/hwr", + "--dlo-host-registration-limit-gib", + "80", ] ) @@ -387,8 +389,10 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): assert explicit_kwargs["host_weight_runtime_mode"] == "preferred" assert explicit_kwargs["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + assert explicit_kwargs["dlo_host_registration_limit_gib"] == 80 assert engine_args["host_weight_runtime_mode"] == "preferred" assert engine_args["host_weight_runtime_root"] == "/var/cache/vllm-omni/hwr" + assert engine_args["dlo_host_registration_limit_gib"] == 80 def test_serve_cli_accepts_diffusion_compile_controls(): diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py index f776f5778..330dd560b 100644 --- a/tests/host_weight_runtime/test_filesystem_store.py +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -57,7 +57,9 @@ from vllm_omni.host_weight_runtime.filesystem import ( detect_storage_class, ) from vllm_omni.host_weight_runtime.filesystem import store as filesystem_store_module +from vllm_omni.host_weight_runtime.filesystem import writer as filesystem_writer_module from vllm_omni.host_weight_runtime.filesystem.locks import FileLock +from vllm_omni.host_weight_runtime.filesystem.writer import ArtifactWriterError, FilesystemArtifactWriter pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -110,6 +112,150 @@ def _open_fd_targets() -> tuple[str, ...]: return tuple(targets) +def _direct_writer(tmp_path: Path) -> FilesystemArtifactWriter: + staging = tmp_path / "staging" + staging.mkdir() + return FilesystemArtifactWriter(staging, capacity_check=lambda _projected, _additional: None) + + +def test_ordered_writer_hashes_payload_during_writes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + weight = _weight() + scale = torch.tensor([1.0, 2.0]) + + def fail_rescan(_path: Path, _specs: tuple[TensorWriteSpec, ...]) -> object: + raise AssertionError("ordered payload must not be rescanned") + + monkeypatch.setattr(filesystem_writer_module, "_scan_payload", fail_rescan) + with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output: + # The writer's ordered contract follows canonical storage-key order. + output.write_tensor("layer.scale", scale) + output.write_rows("layer.weight", 0, weight[:2]) + output.write_rows("layer.weight", 2, weight[2:]) + + manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + payload = writer.staging_dir / "weights.safetensors" + assert manifest.files[0].sha256 == hashlib.sha256(payload.read_bytes()).hexdigest() + expected_tensor_digests = { + "layer.scale": hashlib.sha256(memoryview(scale.view(torch.uint8).numpy())).hexdigest(), + "layer.weight": hashlib.sha256(memoryview(weight.contiguous().view(torch.uint8).numpy())).hexdigest(), + } + assert {entry.name: entry.sha256 for entry in manifest.tensors} == expected_tensor_digests + + +def test_ordered_writer_rejects_out_of_order_payload(tmp_path: Path) -> None: + writer = _direct_writer(tmp_path) + output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) + + with pytest.raises(ArtifactWriterError, match="ordered tensor payload expected offset"): + output.write_tensor("layer.weight", _weight()) + + output.abort() + assert not list(writer.staging_dir.iterdir()) + + +def test_payload_fsync_overlaps_producer_work( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + started = threading.Event() + release = threading.Event() + closed = threading.Event() + original_fsync = filesystem_writer_module._fsync_payload + + def blocking_fsync(fd: int) -> None: + started.set() + if not release.wait(timeout=5): + raise TimeoutError("test did not release payload fsync") + original_fsync(fd) + + monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", blocking_fsync) + output = writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + output.write_tensor("layer.weight", _weight()) + close_errors: list[BaseException] = [] + + def close_output() -> None: + try: + output.close() + except BaseException as exc: + close_errors.append(exc) + finally: + closed.set() + + close_thread = threading.Thread(target=close_output) + close_thread.start() + try: + assert started.wait(timeout=2) + assert closed.wait(timeout=1), "payload close waited for fsync instead of overlapping it" + finally: + release.set() + close_thread.join(timeout=5) + + assert not close_errors + writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + +def test_payload_fsync_failure_prevents_publication( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + + def failing_fsync(fd: int) -> None: + try: + raise OSError(errno.EIO, "injected fsync failure") + finally: + os.close(fd) + + monkeypatch.setattr(filesystem_writer_module, "_fsync_payload", failing_fsync) + with writer.open_tensor_file("weights.safetensors", _specs(), ordered=True) as output: + output.write_tensor("layer.scale", torch.tensor([1.0, 2.0])) + output.write_tensor("layer.weight", _weight()) + + with pytest.raises(OSError, match="injected fsync failure"): + writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + assert not (writer.staging_dir / "manifest.json").exists() + assert not (writer.staging_dir / "READY.json").exists() + + +def test_unordered_payload_fallback_hashes_files_in_parallel( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + writer = _direct_writer(tmp_path) + specs = ( + TensorWriteSpec("a", (2,), torch.float32), + TensorWriteSpec("b", (2,), torch.float32), + ) + for name, spec in zip(("a.safetensors", "b.safetensors"), specs, strict=True): + with writer.open_tensor_file(name, (spec,)) as output: + output.write_tensor(spec.name, torch.tensor([1.0, 2.0])) + + original_scan = filesystem_writer_module._scan_payload + barrier = threading.Barrier(2) + thread_ids: set[int] = set() + thread_ids_lock = threading.Lock() + + def synchronized_scan(path: Path, file_specs: tuple[TensorWriteSpec, ...]) -> object: + with thread_ids_lock: + thread_ids.add(threading.get_ident()) + barrier.wait(timeout=5) + return original_scan(path, file_specs) + + monkeypatch.setattr(filesystem_writer_module, "_scan_payload", synchronized_scan) + manifest, _ = writer.finalize(_identity(), ProductionMetadata("test-producer-v1", "test-restorer-v1")) + + assert len(manifest.files) == 2 + assert len(thread_ids) == 2 + + class FakeProducer: def __init__( self, @@ -1026,7 +1172,9 @@ def test_store_and_free_space_capacity_limits_fail_before_publication(tmp_path: free_root = tmp_path / "free-limit" free_root.mkdir(mode=0o700) - unavailable_reserve = shutil.disk_usage(free_root).free + 1 + # Leave enough margin for filesystem accounting to move between the + # preflight measurement and the capacity check. + unavailable_reserve = shutil.disk_usage(free_root).free + (1 << 30) free_limited = _make_store(free_root, capacity=CapacityPolicy(min_free_bytes=unavailable_reserve)) free_result = free_limited.get_or_build( BuildRequest(identity), diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index e5c87dc88..81afc4b8c 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -699,6 +699,7 @@ class _DiffusionConfigProjection: dlo_resident_layers: int = Field(default=0, ge=0) host_weight_runtime_mode: Literal["disabled", "preferred", "required"] = "disabled" host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" diffusion_compile_dynamic: bool = Field(default=True, strict=True) @@ -756,6 +757,7 @@ class _DiffusionConfigProjection: TransformerConfig, build_attention_config, parse_kv_cache_skip_selector, + validate_dlo_host_registration_options, validate_host_weight_runtime_options, ) from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode @@ -838,6 +840,12 @@ class _DiffusionConfigProjection: mode=self.host_weight_runtime_mode, root=self.host_weight_runtime_root, ) + self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=self.dlo_host_registration_limit_gib, + enable_dlo=self.enable_distributed_layerwise_offload, + use_allgather=self.dlo_use_allgather, + hwr_mode=self.host_weight_runtime_mode, + ) if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 33636c561..4e9327530 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -391,6 +391,7 @@ class StageDeployConfig: dlo_resident_layers: int | None = None host_weight_runtime_mode: str | None = None host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float | None = None # Diffusion-specific debug and observability knobs. enable_diffusion_pipeline_profiler: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c64330c6e..dece344b2 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -113,6 +113,22 @@ def validate_host_weight_runtime_options(*, mode: object, root: object) -> None: raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") +def validate_dlo_host_registration_options( + *, + limit_gib: object, + enable_dlo: bool, + use_allgather: bool, + hwr_mode: object, +) -> float: + """Validate the optional transport budget without probing CUDA or HWR.""" + value = float(limit_gib) + if not math.isfinite(value) or value < 0: + raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0") + if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"): + raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime") + return value + + def parse_kv_cache_skip_selector( selector: str | list[int] | tuple[int, ...] | set[int] | None, ) -> set[int] | None: @@ -784,6 +800,9 @@ class OmniDiffusionConfig: # existing loader/storage path. host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None + # Optional per-worker ceiling for registering final-layout HWR mappings. + # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried. + dlo_host_registration_limit_gib: float = 0.0 pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -1164,6 +1183,12 @@ class OmniDiffusionConfig: mode=self.host_weight_runtime_mode, root=self.host_weight_runtime_root, ) + self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=self.dlo_host_registration_limit_gib, + enable_dlo=self.enable_distributed_layerwise_offload, + use_allgather=self.dlo_use_allgather, + hwr_mode=self.host_weight_runtime_mode, + ) if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): raise ValueError( diff --git a/vllm_omni/diffusion/model_loader/host_weight_loader.py b/vllm_omni/diffusion/model_loader/host_weight_loader.py index 3094cd899..4b9853555 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_loader.py +++ b/vllm_omni/diffusion/model_loader/host_weight_loader.py @@ -27,7 +27,10 @@ logger = init_logger(__name__) if TYPE_CHECKING: from vllm_omni.diffusion.model_loader.host_weights.identity_adapter import FinalLayoutIdentityContext - from vllm_omni.diffusion.model_loader.host_weights.source_identity import PreparedWeightSource + from vllm_omni.diffusion.model_loader.host_weights.source_identity import ( + NodeSourceDigestCache, + PreparedWeightSource, + ) class _HWRCommitError(RuntimeError): @@ -215,6 +218,7 @@ class HWRLoaderMixin: *, load_format: str, sources: Sequence[object], + source_digest_cache: NodeSourceDigestCache | None = None, ) -> FinalLayoutIdentityContext: from vllm.distributed.parallel_state import get_tensor_model_parallel_rank @@ -293,6 +297,7 @@ class HWRLoaderMixin: prepared_sources=prepared_sources, request=request, policy=FINAL_LAYOUT_BF16_POLICY, + source_digest_cache=source_digest_cache, ) def _resolve_hwr( @@ -306,7 +311,10 @@ class HWRLoaderMixin: sources: Sequence[object], ) -> dict[str, object] | None: """Resolve an eligible no-AllGather final-layout HWR transaction.""" - from vllm_omni.diffusion.model_loader.host_weights import FinalLayoutTensorRestorer + from vllm_omni.diffusion.model_loader.host_weights import ( + FinalLayoutTensorRestorer, + NodeSourceDigestCache, + ) from vllm_omni.host_weight_runtime import ( HostWeightLeaseCarrier, HostWeightRuntime, @@ -347,7 +355,6 @@ class HWRLoaderMixin: raise ValueError(f"required Host Weight Runtime path is ineligible: {message}") logger.info("Host Weight Runtime is ineligible; using the canonical DLO path: %s", message) return None - context = self._build_hwr_context(model, modules, load_format=load_format, sources=sources) root_value = getattr(self.od_config, "host_weight_runtime_root", None) if not isinstance(root_value, str) or not root_value.strip(): raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root") @@ -362,6 +369,17 @@ class HWRLoaderMixin: ), ) ) + source_digest_cache = NodeSourceDigestCache( + root, + timeout_seconds=runtime.config.wait.coordination_timeout_seconds, + ) + context = self._build_hwr_context( + model, + modules, + load_format=load_format, + sources=sources, + source_digest_cache=source_digest_cache, + ) resolution = runtime.resolve(context.identity) state: dict[str, object] = { "mode": mode, diff --git a/vllm_omni/diffusion/model_loader/host_weights/__init__.py b/vllm_omni/diffusion/model_loader/host_weights/__init__.py index a7709353d..1a1927fdb 100644 --- a/vllm_omni/diffusion/model_loader/host_weights/__init__.py +++ b/vllm_omni/diffusion/model_loader/host_weights/__init__.py @@ -19,7 +19,7 @@ from .producers import ( FinalLayoutBF16Producer, ) from .restorer import FinalLayoutTensorRestorer -from .source_identity import PreparedWeightSource, WeightSourceKind +from .source_identity import NodeSourceDigestCache, PreparedWeightSource, WeightSourceKind __all__ = [ "FINAL_LAYOUT_BF16_POLICY", @@ -32,6 +32,7 @@ __all__ = [ "FinalLayoutRequest", "FinalLayoutTensorRestorer", "ImplementationIdentity", + "NodeSourceDigestCache", "PreparedWeightSource", "WeightSourceKind", "build_final_layout_identity", diff --git a/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py index 416f03a45..30c9baac1 100644 --- a/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py +++ b/vllm_omni/diffusion/model_loader/host_weights/identity_adapter.py @@ -27,6 +27,7 @@ from .contracts import ( ) from .source_identity import ( FinalLayoutSourceContext, + NodeSourceDigestCache, PreparedWeightSource, resolve_final_layout_source_identity, ) @@ -78,6 +79,7 @@ def build_final_layout_identity( prepared_sources: Sequence[PreparedWeightSource], request: FinalLayoutRequest, policy: FinalLayoutTensorPolicy, + source_digest_cache: NodeSourceDigestCache | None = None, ) -> FinalLayoutIdentityContext: """Build exact source, representation, layout, and implementation identity.""" if not isinstance(request, FinalLayoutRequest): @@ -99,6 +101,7 @@ def build_final_layout_identity( prepared_sources, model_id=request.model_id, target_names=target_names, + digest_cache=source_digest_cache, ) dit_names = tuple(name for name, _ in dit_modules) model_contracts = tuple(binding.contract for binding in bindings) diff --git a/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py index 622256f24..04a48356e 100644 --- a/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py +++ b/vllm_omni/diffusion/model_loader/host_weights/producers/final_layout_bf16.py @@ -286,7 +286,7 @@ class FinalLayoutBF16Producer: for record in shard ) file_name = f"model-{index:05d}-of-{shard_count:05d}.safetensors" - with writer.open_tensor_file(file_name, specs) as output: + with writer.open_tensor_file(file_name, specs, ordered=True) as output: for record in shard: output.write_tensor(record.name, record.tensor.detach()) diff --git a/vllm_omni/diffusion/model_loader/host_weights/source_identity.py b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py index 800bdf288..5591928c8 100644 --- a/vllm_omni/diffusion/model_loader/host_weights/source_identity.py +++ b/vllm_omni/diffusion/model_loader/host_weights/source_identity.py @@ -5,9 +5,13 @@ from __future__ import annotations import hashlib +import json import os import stat +import time +import uuid from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from enum import Enum from pathlib import Path @@ -15,6 +19,7 @@ from pathlib import Path import regex as re from vllm_omni.host_weight_runtime import CanonicalJson, WeightSourceIdentity +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, FileLockTimeoutError from vllm_omni.host_weight_runtime.identity import canonical_json from .contracts import ( @@ -101,6 +106,140 @@ class _FileSnapshot: ) +@dataclass(frozen=True) +class _FileState: + size: int + device: int + inode: int + mtime_ns: int + ctime_ns: int + symlink_target: str | None + + def to_dict(self) -> dict[str, object]: + return { + "size": self.size, + "device": self.device, + "inode": self.inode, + "mtime_ns": self.mtime_ns, + "ctime_ns": self.ctime_ns, + "symlink_target": self.symlink_target, + } + + +def _observe(path: Path) -> _FileState: + try: + current = path.stat() + except OSError as exc: + raise ValueError(f"cannot stat canonical weight file {path}") from exc + if not stat.S_ISREG(current.st_mode): + raise ValueError(f"canonical weight source is not a regular file: {path}") + return _FileState( + size=current.st_size, + device=current.st_dev, + inode=current.st_ino, + mtime_ns=current.st_mtime_ns, + ctime_ns=current.st_ctime_ns, + symlink_target=os.readlink(path) if path.is_symlink() else None, + ) + + +def _sha256_file(path: Path, expected: _FileState) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + before = os.fstat(handle.fileno()) + while chunk := handle.read(_HASH_CHUNK_BYTES): + digest.update(chunk) + after = os.fstat(handle.fileno()) + observed = (before.st_size, before.st_dev, before.st_ino, before.st_mtime_ns, before.st_ctime_ns) + completed = (after.st_size, after.st_dev, after.st_ino, after.st_mtime_ns, after.st_ctime_ns) + expected_stat = (expected.size, expected.device, expected.inode, expected.mtime_ns, expected.ctime_ns) + if observed != expected_stat or completed != expected_stat or _observe(path) != expected: + raise FinalLayoutContractError( + FinalLayoutContractCode.SOURCE_CHANGED, + f"canonical weight file changed while hashing: {path}", + ) + return digest.hexdigest() + + +class NodeSourceDigestCache: + """Share content-backed local-source fingerprints across node workers.""" + + def __init__(self, root: Path, *, timeout_seconds: float) -> None: + self.entries = root / "source-digests-v1" / "entries" + self.locks = root / "source-digests-v1" / "locks" + self.timeout_seconds = timeout_seconds + self._available = True + try: + self.entries.mkdir(parents=True, exist_ok=True, mode=0o700) + self.locks.mkdir(parents=True, exist_ok=True, mode=0o700) + except OSError: + self._available = False + + @staticmethod + def _cached_content_id(entry: Path, path: Path, state: _FileState) -> str | None: + try: + cached = json.loads(entry.read_bytes()) + if not isinstance(cached, dict): + return None + cached_state = cached["state"] + content_id = cached["content_id"] + record_sha256 = cached["record_sha256"] + except (FileNotFoundError, OSError, json.JSONDecodeError, KeyError, TypeError): + return None + record = { + "content_id": content_id, + "schema_version": cached.get("schema_version"), + "state": cached_state, + } + if ( + record["schema_version"] != 1 + or cached_state != state.to_dict() + or not isinstance(content_id, str) + or re.fullmatch(r"sha256:[0-9a-f]{64}", content_id) is None + or not isinstance(record_sha256, str) + or hashlib.sha256(canonical_json(record)).hexdigest() != record_sha256 + or _observe(path) != state + ): + return None + return content_id + + def content_id(self, path: Path, state: _FileState) -> str: + if not self._available: + return f"sha256:{_sha256_file(path, state)}" + key = hashlib.sha256(os.fsencode(path.absolute())).hexdigest() + entry = self.entries / f"{key}.json" + content_id: str | None = None + try: + with FileLock( + self.locks / f"{key}.lock", + exclusive=True, + deadline=time.monotonic() + self.timeout_seconds, + ): + if (cached := self._cached_content_id(entry, path, state)) is not None: + return cached + content_id = f"sha256:{_sha256_file(path, state)}" + record: dict[str, object] = { + "content_id": content_id, + "schema_version": 1, + "state": state.to_dict(), + } + record["record_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest() + temporary = entry.with_name(f".{entry.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp") + try: + temporary.write_bytes(canonical_json(record)) + os.chmod(temporary, 0o600) + os.replace(temporary, entry) + finally: + temporary.unlink(missing_ok=True) + return content_id + except (FileLockTimeoutError, OSError): + # Digest sharing is an optimization. Coordination or cache I/O + # failures must not make a readable canonical source unavailable. + if content_id is not None and _observe(path) == state: + return content_id + return f"sha256:{_sha256_file(path, state)}" + + @dataclass(frozen=True) class _SourceSnapshot: semantic: CanonicalJson @@ -129,14 +268,6 @@ class FinalLayoutSourceContext: ) -def _sha256_file(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - while chunk := handle.read(_HASH_CHUNK_BYTES): - digest.update(chunk) - return digest.hexdigest() - - def _hf_blob_content_id( path: Path, *, @@ -200,7 +331,10 @@ def _logical_model_id(value: str) -> str: return value -def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot: +def _snapshot_source( + source: PreparedWeightSource, + digest_cache: NodeSourceDigestCache | None, +) -> _SourceSnapshot: root = source.resolved_root.expanduser().resolve() if not root.is_dir(): raise ValueError(f"canonical source root is not a directory: {root}") @@ -210,8 +344,7 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot: if _IMMUTABLE_REVISION_RE.fullmatch(requested) is not None: immutable_revision = requested.lower() - files: list[_FileSnapshot] = [] - for path in sorted(source.weight_files, key=lambda item: str(item)): + def snapshot_file(path: Path) -> _FileSnapshot: candidate = path.expanduser() if not candidate.is_absolute(): candidate = root / candidate @@ -220,34 +353,34 @@ def _snapshot_source(source: PreparedWeightSource) -> _SourceSnapshot: relative_name = candidate.relative_to(root).as_posix() except ValueError as exc: raise ValueError(f"weight file {candidate} is outside its canonical source root {root}") from exc - try: - current = candidate.stat() - except OSError as exc: - raise ValueError(f"cannot stat canonical weight file {candidate}") from exc - if not stat.S_ISREG(current.st_mode): - raise ValueError(f"canonical weight source is not a regular file: {candidate}") - symlink_target = os.readlink(candidate) if candidate.is_symlink() else None + current = _observe(candidate) content_id = _hf_blob_content_id( candidate, source=source, source_root=root, ) if content_id is None: - content_id = f"sha256:{_sha256_file(candidate)}" - files.append( - _FileSnapshot( - path=candidate, - relative_name=relative_name, - size=current.st_size, - device=current.st_dev, - inode=current.st_ino, - mtime_ns=current.st_mtime_ns, - ctime_ns=current.st_ctime_ns, - symlink_target=symlink_target, - content_id=content_id, + content_id = ( + digest_cache.content_id(candidate, current) + if digest_cache is not None + else f"sha256:{_sha256_file(candidate, current)}" ) + return _FileSnapshot( + path=candidate, + relative_name=relative_name, + size=current.size, + device=current.device, + inode=current.inode, + mtime_ns=current.mtime_ns, + ctime_ns=current.ctime_ns, + symlink_target=current.symlink_target, + content_id=content_id, ) + paths = sorted(source.weight_files, key=lambda item: str(item)) + with ThreadPoolExecutor(max_workers=min(8, len(paths))) as executor: + files = list(executor.map(snapshot_file, paths)) + file_semantics = [file.semantic_dict() for file in files] file_fingerprint = hashlib.sha256(canonical_json(file_semantics)).hexdigest() revision = immutable_revision or f"content-{file_fingerprint}" @@ -306,11 +439,12 @@ def resolve_final_layout_source_identity( *, model_id: str, target_names: frozenset[str], + digest_cache: NodeSourceDigestCache | None = None, ) -> FinalLayoutSourceContext: """Resolve exact source identity for files covering the owned tensors.""" bindings = _resolve_target_sources(prepared_sources, target_names) selected_sources = frozenset(bindings.values()) - snapshot_by_source = {source: _snapshot_source(source) for source in selected_sources} + snapshot_by_source = {source: _snapshot_source(source, digest_cache) for source in selected_sources} snapshots = tuple(sorted(snapshot_by_source.values(), key=lambda snapshot: snapshot.semantic.encoded)) source_semantics = [snapshot.semantic.to_value() for snapshot in snapshots] target_bindings = [ @@ -345,6 +479,7 @@ def resolve_final_layout_source_identity( __all__ = [ "FinalLayoutSourceContext", + "NodeSourceDigestCache", "PreparedWeightSource", "WeightSourceKind", "resolve_final_layout_source_identity", diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 4e0eb1b5b..c4ad93482 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from abc import ABC, abstractmethod from dataclasses import dataclass @@ -10,7 +10,7 @@ import torch from torch import nn from vllm.logger import init_logger -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig, validate_dlo_host_registration_options logger = init_logger(__name__) @@ -52,6 +52,9 @@ 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 + # Optional per-worker ceiling for registering an HWR mmap. Zero means no + # additional ceiling; pin_cpu_memory controls whether registration is tried. + dlo_host_registration_limit_gib: float = 0.0 @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -122,6 +125,12 @@ class OffloadConfig: # requirements (concurrent requests, dummy run skip). dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True) dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0)) + dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0), + enable_dlo=enable_distributed_layerwise_offload, + use_allgather=dlo_use_allgather, + hwr_mode=getattr(od_config, "host_weight_runtime_mode", "disabled"), + ) if dlo_resident_layers < 0: raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}") if dlo_resident_layers and dlo_use_allgather: @@ -157,6 +166,7 @@ class OffloadConfig: dp_size=dp_size, dlo_use_allgather=dlo_use_allgather, dlo_resident_layers=dlo_resident_layers, + dlo_host_registration_limit_gib=dlo_host_registration_limit_gib, ) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py new file mode 100644 index 000000000..418b5d9fb --- /dev/null +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA implementation of read-only host-mapping registration.""" + +from __future__ import annotations + +import ctypes +import mmap +from collections import defaultdict +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch + +from vllm_omni.host_weight_runtime import MappedHostRegion + +from .host_registration import ( + HostRegistrationBudgetError, + HostRegistrationCleanupError, + HostRegistrationError, +) + +_CUDA_HOST_REGISTER_READ_ONLY = 0x08 +_CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113 + + +class _CudaRuntime(Protocol): + def cudaHostRegister(self, address: int, size: int, flags: int) -> int: ... + + def cudaHostUnregister(self, address: int) -> int: ... + + def cudaGetErrorString(self, error: int) -> str | bytes: ... + + def cudaGetLastError(self) -> int: ... + + +@dataclass(frozen=True) +class _AddressRange: + start: int + end: int + + @property + def size(self) -> int: + return self.end - self.start + + +def _coalesce_ranges( + ranges: Sequence[tuple[int, int]], + page_size: int = mmap.PAGESIZE, +) -> tuple[_AddressRange, ...]: + """Page-align and merge overlapping ranges from one backing mapping.""" + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + + aligned: list[_AddressRange] = [] + for start, size in ranges: + if start <= 0 or size < 0: + raise ValueError(f"invalid host range start={start}, size={size}") + if size == 0: + continue + aligned_start = start - start % page_size + end = start + size + aligned_end = ((end + page_size - 1) // page_size) * page_size + aligned.append(_AddressRange(aligned_start, aligned_end)) + + merged: list[_AddressRange] = [] + for region in sorted(aligned, key=lambda item: (item.start, item.end)): + if merged and region.start <= merged[-1].end: + previous = merged[-1] + merged[-1] = _AddressRange(previous.start, max(previous.end, region.end)) + else: + merged.append(region) + return tuple(merged) + + +def _mapped_regions(regions: Sequence[MappedHostRegion]) -> tuple[_AddressRange, ...]: + """Resolve page spans without merging unrelated artifact files.""" + by_file: dict[str, list[tuple[int, int]]] = defaultdict(list) + for region in regions: + if not region.read_only: + raise HostRegistrationError( + f"CUDA direct H2D requires an immutable mapping, but {region.file_name!r} is writable" + ) + by_file[region.file_name].append((region.address, region.length)) + return tuple(span for ranges in by_file.values() for span in _coalesce_ranges(ranges)) + + +def _error_message(runtime: _CudaRuntime, error: int) -> str: + try: + message = runtime.cudaGetErrorString(error) + except Exception: + return str(error) + if isinstance(message, bytes): + return message.decode(errors="replace") + return str(message) + + +def _consume_last_cuda_error(runtime: _CudaRuntime, expected_error: int) -> None: + """Clear one handled CUDA Runtime error before returning to PyTorch.""" + try: + get_last_error = getattr(runtime, "cudaGetLastError", None) + if get_last_error is None: + get_last_error = ctypes.CDLL(None).cudaGetLastError + get_last_error.argtypes = [] + get_last_error.restype = ctypes.c_int + pending_error = int(get_last_error()) + except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc: + raise HostRegistrationError("cannot clear CUDA's pending error after host-registration failure") from exc + + if pending_error not in (0, expected_error): + raise HostRegistrationError( + f"host registration returned CUDA error {expected_error}, but cudaGetLastError reported {pending_error}" + ) + + +def _handled_error_message(runtime: _CudaRuntime, error: int) -> str: + error_code = int(error) + message = _error_message(runtime, error_code) + try: + _consume_last_cuda_error(runtime, error_code) + except HostRegistrationError as exc: + raise HostRegistrationError(f"{message}; {exc}") from exc + return message + + +def _supports_read_only_host_registration(runtime: _CudaRuntime) -> bool: + """Query support required by immutable artifact mappings.""" + try: + get_attribute = getattr(runtime, "cudaDeviceGetAttribute", None) + if get_attribute is None: + get_attribute = ctypes.CDLL(None).cudaDeviceGetAttribute + get_attribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int] + get_attribute.restype = ctypes.c_int + except (AttributeError, OSError) as exc: + raise HostRegistrationError("cannot query CUDA read-only host-registration support") from exc + + supported = ctypes.c_int() + try: + error = get_attribute( + ctypes.byref(supported), + _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED, + torch.accelerator.current_device_index(), + ) + except Exception as exc: + raise HostRegistrationError(f"cannot query CUDA read-only host-registration support: {exc}") from exc + if int(error) != 0: + raise HostRegistrationError( + f"cudaDeviceGetAttribute(read-only host registration) failed: {_handled_error_message(runtime, error)}" + ) + return bool(supported.value) + + +class CudaHostRegistration: + """Own CUDA registrations for immutable file-backed CPU regions.""" + + def __init__(self, runtime: _CudaRuntime, regions: tuple[_AddressRange, ...]) -> None: + self._runtime = runtime + self._regions = regions + + @classmethod + def create( + cls, + regions: Sequence[MappedHostRegion], + *, + max_bytes: int | None, + ) -> CudaHostRegistration: + mapped = _mapped_regions(regions) + total_bytes = sum(region.size for region in mapped) + 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" + ) + if not mapped: + raise HostRegistrationError("no non-empty host ranges were available for registration") + if not torch.cuda.is_available(): + raise HostRegistrationError("CUDA is not available") + + try: + runtime = torch.cuda.cudart() + except Exception as exc: + raise HostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc + if not _supports_read_only_host_registration(runtime): + raise HostRegistrationError( + "CUDA device does not support read-only host registration required by immutable HWR mappings" + ) + + registered: list[_AddressRange] = [] + try: + for region in mapped: + error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY) + if int(error) != 0: + raise HostRegistrationError( + "cudaHostRegister(read-only) failed for " + f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}" + ) + registered.append(region) + except Exception as exc: + rollback_errors: list[str] = [] + rollback_failed: list[_AddressRange] = [] + for region in reversed(registered): + try: + error = runtime.cudaHostUnregister(region.start) + if int(error) != 0: + rollback_errors.append( + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(runtime, error)}" + ) + rollback_failed.append(region) + except Exception as rollback_exc: + rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}") + rollback_failed.append(region) + if rollback_errors: + active_registration = cls(runtime, tuple(reversed(rollback_failed))) + raise HostRegistrationCleanupError( + f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}", + active_registration=active_registration, + ) from exc + if isinstance(exc, HostRegistrationError): + raise + raise HostRegistrationError(f"CUDA host registration raised: {exc}") from exc + + return cls(runtime, mapped) + + @property + def total_bytes(self) -> int: + return sum(region.size for region in self._regions) + + @property + def region_count(self) -> int: + return len(self._regions) + + def close(self) -> tuple[str, ...]: + """Unregister every range and retain failures for a later retry.""" + errors: list[str] = [] + failed: list[_AddressRange] = [] + for region in reversed(self._regions): + try: + error = self._runtime.cudaHostUnregister(region.start) + if int(error) != 0: + errors.append( + f"cudaHostUnregister({region.start:#x}) failed: {_handled_error_message(self._runtime, error)}" + ) + failed.append(region) + except Exception as exc: + errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}") + failed.append(region) + self._regions = tuple(reversed(failed)) + return tuple(errors) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e53b6d78e..7063c2759 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Distributed Layerwise Offload backend with double-buffered H2D. This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: @@ -18,6 +18,8 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that: from __future__ import annotations +import threading +import time from itertools import chain from typing import Any @@ -35,6 +37,12 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig from .block_discovery import get_blocks_from_dit +from .host_registration import ( + HostRegistration, + HostRegistrationCleanupError, + HostRegistrationError, + register_host_mappings, +) from .module_collector import ModuleDiscovery from .offload_plan import ( OffloadPlan, @@ -51,6 +59,38 @@ from .tensor_utils import ( logger = init_logger(__name__) +# A backend normally owns both objects below. This process-lifetime safety +# owner prevents HostWeightLease.__del__ from unmapping storage when cleanup +# failure unwinds startup and the backend itself becomes unreachable. A clean +# retry removes the pair before closing the lease. +_ACTIVE_HWR_REGISTRATIONS: list[tuple[HostRegistration, HostWeightLease]] = [] +_ACTIVE_HWR_REGISTRATIONS_LOCK = threading.Lock() + + +def _retain_active_hwr_registration( + registration: HostRegistration, + lease: HostWeightLease, +) -> None: + with _ACTIVE_HWR_REGISTRATIONS_LOCK: + if not any( + candidate is registration and candidate_lease is lease + for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS + ): + _ACTIVE_HWR_REGISTRATIONS.append((registration, lease)) + + +def _forget_active_hwr_registration( + registration: HostRegistration, + lease: HostWeightLease, +) -> None: + with _ACTIVE_HWR_REGISTRATIONS_LOCK: + _ACTIVE_HWR_REGISTRATIONS[:] = [ + (candidate, candidate_lease) + for candidate, candidate_lease in _ACTIVE_HWR_REGISTRATIONS + if candidate is not registration or candidate_lease is not lease + ] + + # Threshold (in MB) for deciding whether a non-block DiT submodule should # use layerwise offload (streaming hooks) or be moved to GPU as a resident # module. Submodules larger than this are offloaded to save HBM; smaller @@ -94,6 +134,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): self.rank = rank self.pin_memory = pin_memory self.rank_local_mmap = rank_local_mmap + self.registered_mmap = False self.tensor_transforms = tensor_transforms or {} self.copy_stream = copy_stream or current_omni_platform.Stream() @@ -445,6 +486,26 @@ class DistributedLayerwiseOffloadHook(ModelHook): # Prefetch: H2D + AllGather (overlapped on dedicated streams) # # ------------------------------------------------------------------ # + @staticmethod + def _resolve_mmap_source( + source_info: dict[str, Any], + meta: dict[str, Any], + dtype: torch.dtype, + ) -> torch.Tensor: + source = source_info["tensor"] + transform = source_info["transform"] + if callable(transform): + source = transform(source) + if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]: + raise ValueError( + "mmap weight transform changed tensor layout for " + f"{source_info['name']!r}: expected dtype={dtype}, " + f"shape={tuple(meta['shape'])}, stride={meta['stride']}; " + f"got dtype={source.dtype}, shape={tuple(source.shape)}, " + f"stride={source.stride()}" + ) + return source + @staticmethod def _pack_mmap_sources( cpu_sources: dict[torch.dtype, list[dict[str, Any]]], @@ -457,18 +518,7 @@ class DistributedLayerwiseOffloadHook(ModelHook): destination = slot_buffers[dtype][:total_numel] sources = cpu_sources[dtype] for source_info, meta in zip(sources, metas, strict=True): - source = source_info["tensor"] - transform = source_info["transform"] - if callable(transform): - source = transform(source) - if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]: - raise ValueError( - "mmap weight transform changed tensor layout for " - f"{source_info['name']!r}: expected dtype={dtype}, " - f"shape={tuple(meta['shape'])}, stride={meta['stride']}; " - f"got dtype={source.dtype}, shape={tuple(source.shape)}, " - f"stride={source.stride()}" - ) + source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype) start = meta["offset"] physical_storage = destination[start : start + meta["numel"]] if source.is_contiguous(): @@ -500,6 +550,32 @@ class DistributedLayerwiseOffloadHook(ModelHook): raise RuntimeError(f"cpu_staging_buffers[{slot}] was not allocated") return self._pack_mmap_sources(self.cpu_sources, self.metadata, slot_buffers) + @staticmethod + def _copy_mmap_sources_to_device( + cpu_sources: dict[torch.dtype, list[dict[str, Any]]], + metadata: dict[torch.dtype, list[dict[str, Any]]], + gpu_buffers: dict[torch.dtype, torch.Tensor], + *, + non_blocking: bool, + ) -> None: + """Copy registered source views directly into flattened device buffers.""" + for dtype, metas in metadata.items(): + destination = gpu_buffers[dtype] + sources = cpu_sources[dtype] + for source_info, meta in zip(sources, metas, strict=True): + source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype) + start = meta["offset"] + physical_storage = destination[start : start + meta["numel"]] + async_copy = non_blocking and source.is_pinned() + if source.is_contiguous(): + physical_storage.copy_(source.flatten(), non_blocking=async_copy) + else: + torch.as_strided( + physical_storage, + size=source.shape, + stride=source.stride(), + ).copy_(source, non_blocking=async_copy) + @torch.compiler.disable def prefetch_layer(self, slot: int, non_blocking: bool = True) -> None: """Prepare next block's weights into the shared device buffer for *slot*. @@ -515,17 +591,27 @@ class DistributedLayerwiseOffloadHook(ModelHook): assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated" if self.dp_size <= 1 or self.dp_group is None: - cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards - with current_omni_platform.stream(self.copy_stream): - for dtype, cpu_shard in cpu_weights.items(): - gw = gpu_weights[dtype] - async_copy = non_blocking and cpu_shard.is_pinned() - gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy) - evt.record(self.copy_stream) - if self.rank_local_mmap: - # The CPU slot may be overwritten only after this H2D copy has - # finished. The shared event protects reuse by another hook. - self.cpu_staging_events[slot] = evt + if self.rank_local_mmap and self.registered_mmap: + with current_omni_platform.stream(self.copy_stream): + self._copy_mmap_sources_to_device( + self.cpu_sources, + self.metadata, + gpu_weights, + non_blocking=non_blocking, + ) + evt.record(self.copy_stream) + else: + cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards + with current_omni_platform.stream(self.copy_stream): + for dtype, cpu_shard in cpu_weights.items(): + gw = gpu_weights[dtype] + async_copy = non_blocking and cpu_shard.is_pinned() + gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy) + evt.record(self.copy_stream) + if self.rank_local_mmap: + # The CPU slot may be overwritten only after this H2D copy has + # finished. The shared event protects reuse by another hook. + self.cpu_staging_events[slot] = evt else: gpu_shards: dict[torch.dtype, torch.Tensor] = {} shard_bufs = self.gpu_shard_buffers[slot] @@ -735,6 +821,7 @@ class PinnedResidentLayerGroup: self.copy_stream = copy_stream self.loaded = False self.rank_local_mmap = rank_local_mmap + self.registered_mmap = False self.pin_memory = pin_memory self._states: list[dict[str, Any]] = [] self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = [] @@ -805,6 +892,14 @@ class PinnedResidentLayerGroup: ready = current_omni_platform.Event() with current_omni_platform.stream(self.copy_stream): for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)): + if self.rank_local_mmap and self.registered_mmap: + DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device( + state["cpu_sources"], + state["metadata"], + block_buffers, + non_blocking=True, + ) + continue if self.rank_local_mmap: slot = index % 2 previous_copy = self._cpu_staging_events[slot] @@ -902,8 +997,10 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._resident_layer_group: PinnedResidentLayerGroup | None = None self._using_mmap = False self._using_rank_local_mmap = False + self._using_registered_mmap = False self.host_weight_plan = host_weight_plan self._host_weight_lease: HostWeightLease | None = None + self._host_registration: HostRegistration | None = None self._mmap_transforms_by_tensor_id: dict[int, Any] = {} def load_resident_layers(self) -> None: @@ -916,6 +1013,136 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self._resident_layer_group is not None: self._resident_layer_group.offload() + def _rank_local_source_tensors( + self, + hooks: list[DistributedLayerwiseOffloadHook], + ) -> tuple[torch.Tensor, ...]: + """Return the exact CPU sources that direct H2D would consume.""" + tensors: list[torch.Tensor] = [] + seen: set[int] = set() + + def collect(cpu_sources: dict[torch.dtype, list[dict[str, Any]]]) -> None: + for sources in cpu_sources.values(): + for source in sources: + tensor = source["tensor"] + if id(tensor) not in seen: + seen.add(id(tensor)) + tensors.append(tensor) + + for hook in hooks: + collect(hook.cpu_sources) + if self._resident_layer_group is not None: + for state in self._resident_layer_group._states: + collect(state["cpu_sources"]) + return tuple(tensors) + + def _try_register_hwr_mmap(self, source_tensors: tuple[torch.Tensor, ...]) -> bool: + """Register the complete final-layout lease under pinned-memory policy.""" + lease = self._host_weight_lease + if lease is None: + return False + if not self.config.pin_cpu_memory: + logger.info("HWR mmap registration disabled by pin_cpu_memory=False; using bounded host staging") + return False + if not lease.mapped_regions or not source_tensors: + logger.warning("HWR mmap registration found no mapped sources; using bounded host staging") + return False + + limit_gib = self.config.dlo_host_registration_limit_gib + max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None + started = time.perf_counter() + try: + registration = register_host_mappings( + lease.mapped_regions, + device=self.device, + max_bytes=max_bytes, + ) + try: + unpinned = [tensor for tensor in source_tensors if tensor.numel() and not tensor.is_pinned()] + except Exception as exc: + errors = registration.close() + if errors: + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + raise HostRegistrationCleanupError( + "CUDA registration succeeded but pinned-source verification failed, " + f"and rollback failed: {errors[:3]}" + ) from exc + raise HostRegistrationError(f"cannot verify registered HWR sources: {exc}") from exc + if unpinned: + errors = registration.close() + if errors: + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + raise HostRegistrationCleanupError( + "CUDA registration succeeded but PyTorch rejected mapped sources, " + f"and rollback failed: {errors[:3]}" + ) + raise HostRegistrationError( + "CUDA registration succeeded but PyTorch did not recognize " + f"{len(unpinned)} mapped source(s) as pinned" + ) + except HostRegistrationCleanupError as exc: + # Falling back could close a lease while the platform still owns + # one of its mappings. Fail startup and retain ownership for retry. + active_registration = exc.active_registration + if active_registration is not None: + self._host_registration = active_registration + _retain_active_hwr_registration(active_registration, lease) + logger.exception("HWR mmap registration rollback failed") + raise + except HostRegistrationError as exc: + logger.warning("HWR registered direct H2D unavailable (%s); using bounded host staging", exc) + return False + + self._host_registration = registration + _retain_active_hwr_registration(registration, lease) + logger.info( + "Registered %.2f GiB of HWR mmap in %d range(s) for direct H2D in %.3f s", + registration.total_bytes / 1024**3, + registration.region_count, + time.perf_counter() - started, + ) + return True + + def _configure_hwr_transfer(self, hooks: list[DistributedLayerwiseOffloadHook]) -> None: + """Select registered direct H2D or bounded staging once per backend.""" + plan = self.host_weight_plan + if ( + (not hooks and self._resident_layer_group is None) + or not self._using_rank_local_mmap + or plan is None + or plan.backing_kind != "host_weight_runtime" + ): + return + + source_tensors = self._rank_local_source_tensors(hooks) + self._using_registered_mmap = self._try_register_hwr_mmap(source_tensors) + for hook in hooks: + hook.registered_mmap = self._using_registered_mmap + if self._resident_layer_group is not None: + self._resident_layer_group.registered_mmap = self._using_registered_mmap + if self._using_registered_mmap: + self._resident_layer_group._cpu_staging_buffers.clear() + + def _release_registered_mmap(self) -> None: + """Release every platform registration before closing the HWR lease.""" + registration = self._host_registration + if registration is None: + return + errors = registration.close() + if errors: + lease = self._host_weight_lease + if lease is not None and not lease.closed: + _retain_active_hwr_registration(registration, lease) + logger.error("HWR mmap unregistration failed; retaining lease mappings for retry: %s", errors[:3]) + raise HostRegistrationCleanupError(f"failed to unregister {len(errors)} HWR mmap range(s)") + lease = self._host_weight_lease + if lease is not None: + _forget_active_hwr_registration(registration, lease) + self._host_registration = None + logger.info("Unregistered HWR mmap ranges") + def _load_weights_via_mmap( self, pipeline: nn.Module, @@ -1465,11 +1692,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): raise RuntimeError("DLO received a closed Host Weight Runtime lease") # The final-layout restorer has already rebound the model to # immutable host tensors. Treat those tensors as mmap-like - # sources for the bounded two-slot rank-local staging path; - # do not reinitialize or rematerialize the restored model. + # sources; transport setup below selects registered direct + # H2D or the bounded two-slot staging fallback. self._using_rank_local_mmap = True logger.info( - "DLO consuming final-layout Host Weight Runtime lease %s with bounded rank-local staging", + "DLO consuming final-layout Host Weight Runtime lease %s", self._host_weight_lease.provenance.resolution_id, ) elif self.host_weight_plan.backing_kind == "checkpoint_mmap": @@ -1483,7 +1710,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): if self._using_rank_local_mmap: logger.info( "DLO rank-local host storage enabled: source pages are " - "node-shared; each worker owns only two bounded host staging slots" + "node-shared; transfer setup will select registered direct H2D or bounded host staging" ) else: remaining_meta = [ @@ -1653,6 +1880,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): ) pipeline._dlo_residency_controller = self + all_hooks = [hook for group in self._all_hook_groups for hook in group] + self._configure_hwr_transfer(all_hooks) + if not self._all_hook_groups: self.enabled = bool(self._resident_blocks) if self._using_mmap and not self.enabled: @@ -1662,17 +1892,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): # Unified allocation: 2 shared output buffers + 2 shared shard buffers # sized to the max block across ALL module groups (gen_layers + # language_model). Groups execute sequentially, so 2 buffers suffice. - all_hooks: list[DistributedLayerwiseOffloadHook] = [] - for group in self._all_hook_groups: - all_hooks.extend(group) - unified_buffers = self._allocate_shared_buffers(all_hooks) unified_shard_buffers = None if self.dp_size > 1: unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks) unified_cpu_staging = None cpu_staging_events = None - if self._using_rank_local_mmap: + if self._using_rank_local_mmap and not self._using_registered_mmap: unified_cpu_staging = self._allocate_shared_cpu_staging_buffers( all_hooks, self._resident_layer_group, @@ -1753,6 +1979,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def _release_mmap_handles(self) -> None: """Release source handles and the transport-owned HWR lease.""" + if self._host_registration is not None: + raise HostRegistrationCleanupError("cannot close HWR mappings while host registration is still active") self._mmap_transforms_by_tensor_id.clear() if hasattr(self, "_mmap_file_cache"): self._mmap_file_cache.clear() @@ -1781,6 +2009,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): def disable(self) -> None: has_open_lease = self._host_weight_lease is not None and not self._host_weight_lease.closed + has_registration = self._host_registration is not None has_carrier = ( self.host_weight_plan is not None and self.host_weight_plan.lease_carrier is not None @@ -1796,12 +2025,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): not self.enabled and not hasattr(self, "_mmap_file_cache") and not has_open_lease + and not has_registration and not has_carrier and not has_partial_hooks ): return - if self._using_rank_local_mmap: + if self._using_rank_local_mmap or has_registration: current_omni_platform.synchronize() for blocks in self._blocks: @@ -1818,9 +2048,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): self._all_hook_groups.clear() self._resident_blocks.clear() self._resident_layer_group = None + self._release_registered_mmap() self._release_mmap_handles() self._using_mmap = False self._using_rank_local_mmap = False + self._using_registered_mmap = False self.enabled = False logger.info("Distributed layer-wise offloading disabled") diff --git a/vllm_omni/diffusion/offloader/host_registration.py b/vllm_omni/diffusion/offloader/host_registration.py new file mode 100644 index 000000000..c9f225122 --- /dev/null +++ b/vllm_omni/diffusion/offloader/host_registration.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Platform-neutral lifecycle for registering existing host mappings.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +import torch + +from vllm_omni.host_weight_runtime import MappedHostRegion + + +class HostRegistrationError(RuntimeError): + """Raised when existing host ranges cannot be registered safely.""" + + +class HostRegistrationBudgetError(HostRegistrationError): + """Raised before mutation when the complete mapping exceeds its budget.""" + + +class HostRegistrationCleanupError(HostRegistrationError): + """Raised when registered ranges cannot be released safely.""" + + def __init__(self, message: str, *, active_registration: HostRegistration | None = None) -> None: + super().__init__(message) + self.active_registration = active_registration + + +class HostRegistration(Protocol): + """Transport-owned registration for one process-local host mapping.""" + + @property + def total_bytes(self) -> int: ... + + @property + def region_count(self) -> int: ... + + def close(self) -> tuple[str, ...]: ... + + +def register_host_mappings( + regions: Sequence[MappedHostRegion], + *, + device: torch.device, + max_bytes: int | None, +) -> HostRegistration: + """Register immutable ranges with the active platform implementation.""" + if device.type == "cuda": + from .cuda_host_registration import CudaHostRegistration + + return CudaHostRegistration.create(regions, max_bytes=max_bytes) + raise HostRegistrationError(f"host-mapping registration is not supported on {device.type}") + + +__all__ = [ + "HostRegistration", + "HostRegistrationBudgetError", + "HostRegistrationCleanupError", + "HostRegistrationError", + "register_host_mappings", +] diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 25bca55e1..f71c2bf8a 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Diffusion Worker for vLLM-Omni. @@ -897,11 +897,18 @@ class DiffusionWorker: def shutdown(self) -> None: """Shutdown the worker and cleanup distributed environment.""" - if self.model_runner is not None: - mgr = getattr(self.model_runner, "kv_transfer_manager", None) - if mgr is not None: - mgr.shutdown_prefetch() - destroy_distributed_env() + try: + if self.model_runner is not None: + mgr = getattr(self.model_runner, "kv_transfer_manager", None) + try: + offload_backend = getattr(self.model_runner, "offload_backend", None) + if offload_backend is not None: + offload_backend.disable() + finally: + if mgr is not None: + mgr.shutdown_prefetch() + finally: + destroy_distributed_env() class CustomPipelineWorkerExtension: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c618d60d0..8fdda71cb 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -571,6 +571,7 @@ class OrchestratorArgs: dlo_resident_layers: int = 0 host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None + dlo_host_registration_limit_gib: float = 0.0 boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e5a0e44e9..4b44e7067 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1061,6 +1061,7 @@ class AsyncOmniEngine: "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), + "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), "diffusion_compile_granularity": ( "regional" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index c14d306dc..87f4f8834 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -760,6 +760,16 @@ class OmniServeCommand(CLISubcommand): "the same persistent path for population and serving." ), ) + omni_config_group.add_argument( + "--dlo-host-registration-limit-gib", + type=float, + default=0.0, + help=( + "Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. " + "Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration " + "under the existing pinned-memory policy and fall back to bounded staging when unavailable." + ), + ) # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", diff --git a/vllm_omni/host_weight_runtime/filesystem/writer.py b/vllm_omni/host_weight_runtime/filesystem/writer.py index 8ecd17799..283ee2a9e 100644 --- a/vllm_omni/host_weight_runtime/filesystem/writer.py +++ b/vllm_omni/host_weight_runtime/filesystem/writer.py @@ -12,17 +12,19 @@ import os import struct import uuid from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass from pathlib import Path, PurePosixPath from types import TracebackType import torch -from safetensors import safe_open from ..identity import WeightArtifactIdentity, canonical_json from ..manifest import FileManifestEntry, PublicationMarker, TensorManifestEntry, WeightManifest from ..protocols import ProductionMetadata, TensorWriteSpec _HASH_CHUNK_BYTES = 8 * 1024**2 +_MAX_HASH_WORKERS = 8 _MANIFEST_FILE = "manifest.json" _READY_FILE = "READY.json" @@ -79,21 +81,47 @@ def _fsync_directory(path: Path) -> None: os.close(fd) -def _file_sha256(path: Path) -> str: - digest = hashlib.sha256() +def _fsync_payload(fd: int) -> None: + try: + os.fsync(fd) + finally: + os.close(fd) + + +@dataclass(frozen=True) +class _PayloadDigests: + file: str + tensors: dict[str, str] + + +def _scan_payload(path: Path, specs: tuple[TensorWriteSpec, ...]) -> _PayloadDigests: + """Calculate file and tensor digests in one ordered read.""" + file_digest = hashlib.sha256() + tensor_digests: dict[str, str] = {} with path.open("rb") as handle: - while chunk := handle.read(_HASH_CHUNK_BYTES): - digest.update(chunk) - return digest.hexdigest() - - -def _tensor_sha256(tensor: torch.Tensor) -> str: - digest = hashlib.sha256() - byte_view = tensor.detach().reshape(-1).view(torch.uint8) - for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES): - chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES] - digest.update(memoryview(chunk.numpy())) - return digest.hexdigest() + prefix = handle.read(8) + if len(prefix) != 8: + raise ArtifactWriterError(f"truncated safetensors header in {path.name!r}") + header_size = struct.unpack(" None: self._owner = owner self.relative_name = _safe_relative_name(relative_name) @@ -144,6 +174,7 @@ class StreamingSafeTensorFile: self.specs = tuple(sorted(specs, key=lambda spec: spec.name)) self._specs = {spec.name: spec for spec in self.specs} self._coverage = {spec.name: _Coverage(spec.nbytes) for spec in self.specs} + self._ordered = ordered self._closed = False offsets: dict[str, int] = {} @@ -167,6 +198,7 @@ class StreamingSafeTensorFile: total_size = self._payload_offset + cursor owner._reserve_file(total_size) self._fd = os.open(self._path, os.O_CREAT | os.O_EXCL | os.O_RDWR | os.O_CLOEXEC, 0o600) + prefix = struct.pack(" None: + file_offset = self._payload_offset + self._offsets[name] + byte_offset + if self._ordered and file_offset != self._write_cursor: + raise ArtifactWriterError( + f"ordered tensor payload expected offset {self._write_cursor}, got {file_offset} for {name!r}" + ) coverage = self._coverage[name] coverage.add(byte_offset, byte_offset + len(data)) - _pwrite_all(self._fd, data, self._payload_offset + self._offsets[name] + byte_offset) + _pwrite_all(self._fd, data, file_offset) + if self._ordered: + assert self._file_digest is not None + self._file_digest.update(data) + self._tensor_digests[name].update(data) + self._write_cursor += len(data) def close(self) -> None: if self._closed: @@ -238,11 +283,19 @@ class StreamingSafeTensorFile: self.abort() raise ArtifactWriterError(f"tensor payload closed with incomplete writes: {missing}") self._closed = True + digests = None + if self._ordered: + assert self._file_digest is not None + digests = _PayloadDigests( + self._file_digest.hexdigest(), + {name: digest.hexdigest() for name, digest in self._tensor_digests.items()}, + ) try: - os.fsync(self._fd) - finally: + os.fchmod(self._fd, 0o444) + self._owner._closed(self, digests, self._fd) + except BaseException: os.close(self._fd) - self._owner._closed(self) + raise def abort(self) -> None: if self._closed: @@ -281,20 +334,25 @@ class FilesystemArtifactWriter: self._capacity_check = capacity_check self._active: set[StreamingSafeTensorFile] = set() self._files: dict[str, tuple[TensorWriteSpec, ...]] = {} + self._digests: dict[str, _PayloadDigests] = {} self._reserved_bytes = 0 self._finalized = False + self._fsync_executor: ThreadPoolExecutor | None = None + self._fsync_futures: list[Future[None]] = [] def open_tensor_file( self, relative_name: str, specs: tuple[TensorWriteSpec, ...], + *, + ordered: bool = False, ) -> StreamingSafeTensorFile: if self._finalized: raise ArtifactWriterError("cannot add payloads after artifact finalization") relative_name = _safe_relative_name(relative_name) if relative_name in self._files or any(item.relative_name == relative_name for item in self._active): raise ArtifactWriterError(f"artifact payload {relative_name!r} was declared more than once") - return StreamingSafeTensorFile(self, relative_name, specs) + return StreamingSafeTensorFile(self, relative_name, specs, ordered=ordered) def _reserve_file(self, size: int) -> None: self._capacity_check(self._reserved_bytes + size, size) @@ -306,9 +364,14 @@ class FilesystemArtifactWriter: def _opened(self, writer: StreamingSafeTensorFile) -> None: self._active.add(writer) - def _closed(self, writer: StreamingSafeTensorFile) -> None: + def _closed(self, writer: StreamingSafeTensorFile, digests: _PayloadDigests | None, fd: int) -> None: self._active.remove(writer) self._files[writer.relative_name] = writer.specs + if digests is not None: + self._digests[writer.relative_name] = digests + if self._fsync_executor is None: + self._fsync_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="hwr-fsync") + self._fsync_futures.append(self._fsync_executor.submit(_fsync_payload, fd)) def _aborted(self, writer: StreamingSafeTensorFile, size: int) -> None: self._active.discard(writer) @@ -318,6 +381,19 @@ class FilesystemArtifactWriter: for writer in tuple(self._active): with contextlib.suppress(Exception): writer.abort() + with contextlib.suppress(Exception): + self._wait_for_payload_fsync() + + def _wait_for_payload_fsync(self) -> None: + if self._fsync_executor is None: + return + executor = self._fsync_executor + futures = tuple(self._fsync_futures) + self._fsync_executor = None + self._fsync_futures.clear() + executor.shutdown(wait=True) + for future in futures: + future.result() def finalize( self, @@ -335,46 +411,50 @@ class FilesystemArtifactWriter: if metadata.restorer_schema != identity.producer.restorer_schema: raise ArtifactWriterError("production metadata does not match the identity's restorer schema") + self._wait_for_payload_fsync() + tensor_entries: list[TensorManifestEntry] = [] file_entries: list[FileManifestEntry] = [] seen_tensors: set[str] = set() - for file_name in sorted(self._files): + file_names = sorted(self._files) + unhashed = [file_name for file_name in file_names if file_name not in self._digests] + if unhashed: + with ThreadPoolExecutor(max_workers=min(_MAX_HASH_WORKERS, len(unhashed))) as executor: + scans = executor.map( + lambda file_name: _scan_payload(self.staging_dir / file_name, self._files[file_name]), + unhashed, + ) + self._digests.update(zip(unhashed, scans, strict=True)) + + for file_name in file_names: path = self.staging_dir / file_name - os.chmod(path, 0o444) - with path.open("rb") as handle: - os.fsync(handle.fileno()) + specs = self._files[file_name] + digests = self._digests[file_name] file_entries.append( FileManifestEntry( file_name=file_name, size=path.stat().st_size, - sha256=_file_sha256(path), + sha256=digests.file, ) ) - expected = {spec.name: spec for spec in self._files[file_name]} - with safe_open(str(path), framework="pt", device="cpu") as handle: - if set(handle.keys()) != set(expected): - raise ArtifactWriterError(f"published tensor keys differ in {file_name!r}") - for name in sorted(expected): - if name in seen_tensors: - raise ArtifactWriterError(f"artifact tensor {name!r} occurs in multiple payload files") - seen_tensors.add(name) - spec = expected[name] - tensor = handle.get_tensor(name) - if tuple(tensor.shape) != spec.shape or tensor.dtype != spec.dtype: - raise ArtifactWriterError(f"published tensor {name!r} differs from its declared metadata") - entry = TensorManifestEntry( - name=name, + for spec in specs: + if spec.name in seen_tensors: + raise ArtifactWriterError(f"artifact tensor {spec.name!r} occurs in multiple payload files") + seen_tensors.add(spec.name) + tensor_entries.append( + TensorManifestEntry( + name=spec.name, kind=spec.kind, role=spec.role, file_name=file_name, - storage_key=name, - shape=tuple(tensor.shape), - dtype=str(tensor.dtype), - stride=tuple(tensor.stride()), - nbytes=tensor.numel() * tensor.element_size(), - sha256=_tensor_sha256(tensor), + storage_key=spec.name, + shape=spec.shape, + dtype=str(spec.dtype), + stride=tuple(torch.empty(spec.shape, dtype=spec.dtype, device="meta").stride()), + nbytes=spec.nbytes, + sha256=digests.tensors[spec.name], ) - tensor_entries.append(entry) + ) ordered_tensor_entries = tuple(sorted(tensor_entries, key=lambda entry: entry.name)) content_digest = hashlib.sha256() diff --git a/vllm_omni/host_weight_runtime/protocols.py b/vllm_omni/host_weight_runtime/protocols.py index c5572c9fb..f3738c1f3 100644 --- a/vllm_omni/host_weight_runtime/protocols.py +++ b/vllm_omni/host_weight_runtime/protocols.py @@ -88,6 +88,8 @@ class ArtifactWriter(Protocol): self, relative_name: str, specs: tuple[TensorWriteSpec, ...], + *, + ordered: bool = False, ) -> TensorFileWriter: ...