diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 730d55472..315a1defc 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -289,6 +289,13 @@ 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-mode` defaults to `auto`, preserving this behavior. +Select `disabled` to bypass HWR mapping registration and use bounded host staging +regardless of mapping size or budget. This leaves the existing pin-memory policy +for staging buffers unchanged. The equivalent offline/stage configuration field +is `dlo_host_registration_mode`; an explicit disable requires enabled HWR with +no-AllGather DLO. + `--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 @@ -296,6 +303,14 @@ 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. +Automatic registration logs its entry with device and budget, then each CUDA +region's ordinal/total and byte size before calling `cudaHostRegister`. A +completion message follows each successful call. These identify the last +entered boundary if progress stops; they do not impose a timeout. A synchronous +CUDA call cannot be safely cancelled by a Python timer, and a stuck process +still requires external supervision. The explicit disable option avoids this +registration path; it does not repair the underlying driver hang. + 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 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index f91e6af42..ca8bf1a28 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -193,6 +193,39 @@ def test_resolve_execution_mode_rejects_unknown_execution_type(): omni_config_module._resolve_execution_mode("unknown_execution_type") +@pytest.mark.parametrize("registration_mode", ["auto", "disabled"]) +def test_hwr_registration_policy_reaches_offloader(tmp_path: Path, registration_mode: str): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.offloader.base import OffloadConfig + + options = { + "enable_distributed_layerwise_offload": True, + "dlo_use_allgather": False, + "host_weight_runtime_mode": "preferred", + "host_weight_runtime_root": str(tmp_path / "unused-store"), + "dlo_host_registration_mode": registration_mode, + } + config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2", cli_overrides=options) + stage = config.stage_by_id(0) + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.dlo_host_registration_mode == registration_mode + offline = OmniDiffusionConfig(**options) + transport = OffloadConfig.from_od_config(offline) + assert transport.dlo_host_registration_mode == registration_mode + assert transport.pin_cpu_memory + assert transport.dlo_host_registration_limit_gib == 0.0 + assert not (tmp_path / "unused-store").exists() + + +@pytest.mark.parametrize("registration_mode", ["unknown", None, "disabled"]) +def test_hwr_registration_policy_rejects_invalid_or_ineligible_config(registration_mode): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + for config_class in (OmniDiffusionConfig, omni_config_module._DiffusionConfigProjection): + with pytest.raises(ValueError, match="dlo_host_registration_mode"): + config_class(dlo_host_registration_mode=registration_mode) + + def test_from_pipeline_config_preserves_current_pipeline_config_object(): omni_config = _from_pipeline_key("minicpmo_4_5") pipeline = _resolve_pipeline_or_skip("minicpmo_4_5") diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py index a3efbb40a..bbed5a3e3 100644 --- a/tests/diffusion/offloader/test_cuda_host_registration.py +++ b/tests/diffusion/offloader/test_cuda_host_registration.py @@ -124,6 +124,33 @@ def test_registration_rejects_writable_or_over_budget_before_cuda(monkeypatch: p assert runtime.registered == [] +@pytest.mark.parametrize("second_result", [0, 7]) +def test_registration_logs_entry_before_each_runtime_call(monkeypatch: pytest.MonkeyPatch, second_result: int) -> None: + runtime = _FakeRuntime([0, second_result]) + messages: list[str] = [] + monkeypatch.setattr(registration_module.logger, "info", lambda message, *args: messages.append(message % args)) + original_register = runtime.cudaHostRegister + + def register(pointer: int, size: int, flags: int) -> int: + ordinal = len(runtime.registered) + 1 + assert messages[-1] == f"Registering HWR host range {ordinal}/2: {size} bytes" + return original_register(pointer, size, flags) + + monkeypatch.setattr(runtime, "cudaHostRegister", register) + monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime) + monkeypatch.setattr(registration_module, "_consume_last_cuda_error", lambda _runtime, _error: None) + regions = (_region("first", 0x1000, 4096), _region("second", 0x3000, 4096)) + if second_result: + with pytest.raises(HostRegistrationError, match="error-7"): + CudaHostRegistration.create(regions, max_bytes=None) + assert runtime.unregistered == [0x1000] + else: + registration = CudaHostRegistration.create(regions, max_bytes=None) + assert registration.close() == () + assert sum(message.startswith("Registering ") for message in messages) == 2 + assert sum(message.startswith("Registered ") for message in messages) == (1 if second_result else 2) + + 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) diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 05d7d93b5..0a838a369 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -96,7 +96,7 @@ class TestDistributedLayerwiseOffloadHook: hook.offload_layer() assert not hook.is_materialized - def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch): + def test_initialize_failure_keeps_next_block_materialized(self, monkeypatch, patched_offload_runtime): current_block = nn.Linear(2, 2) next_block = nn.Linear(2, 2) expected = {name: tensor.detach().clone() for name, tensor in next_block.state_dict().items()} @@ -1373,10 +1373,14 @@ class TestMmapWeightLoading: synchronize.assert_called() - def test_hwr_registration_failure_falls_back_to_bounded_staging( + @pytest.mark.parametrize("registration_mode", ["auto", "disabled"]) + @pytest.mark.parametrize("limit_gib", [0.0, 1.5]) + def test_hwr_registration_selection_uses_bounded_staging( self, monkeypatch, patched_offload_runtime, + registration_mode, + limit_gib, ): plan, carrier, lease = _fake_hwr_plan("hwr-registration-fallback") backend = DistributedLayerwiseOffloadBackend( @@ -1385,6 +1389,8 @@ class TestMmapWeightLoading: pin_cpu_memory=True, dp_size=1, dlo_use_allgather=False, + dlo_host_registration_mode=registration_mode, + dlo_host_registration_limit_gib=limit_gib, ), torch.device("cpu"), host_weight_plan=plan, @@ -1394,12 +1400,14 @@ class TestMmapWeightLoading: def fail_registration(*args, **kwargs): del args, kwargs + assert registration_mode == "auto", "disabled mode entered platform registration" raise HostRegistrationError("registration unavailable in CPU test") def allocate_unpinned_staging(hooks, resident_group=None): # Registration selection still observes pin_cpu_memory=True. The # CPU-only test avoids asking the host for CUDA-pinned allocations. for hook in hooks: + assert hook.pin_memory, "registration policy disabled staging pinning" hook.pin_memory = False return original_staging_allocator(hooks, resident_group) @@ -1412,6 +1420,7 @@ class TestMmapWeightLoading: assert backend._using_rank_local_mmap assert not backend._using_registered_mmap assert backend._host_registration is None + assert backend.config.pin_cpu_memory assert all(len(hook.cpu_staging_buffers) == 2 for group in backend._all_hook_groups for hook in group) assert not lease.closed @@ -2630,7 +2639,9 @@ class TestDistributedComponentSelection: assert registry is None or registry.get_hook("distributed_layerwise_offload") is None torch.testing.assert_close(block.weight, expected) - def test_multirank_enable_failure_cleanup_skips_restore_collective(self, monkeypatch, mocker): + def test_multirank_enable_failure_cleanup_skips_restore_collective( + self, monkeypatch, mocker, patched_offload_runtime + ): backend = DistributedLayerwiseOffloadBackend( OffloadConfig( strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, @@ -2811,7 +2822,7 @@ class TestDistributedComponentSelection: with pytest.raises(ValueError, match="not declared replicated"): backend.enable(pipeline) - def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): + def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self, patched_offload_runtime): """Every rank must reject an unsafe encoder group, including stub ranks.""" backend = DistributedLayerwiseOffloadBackend( OffloadConfig( diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 616d5486e..d1c730dbe 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -507,6 +507,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-mode", + "disabled", "--dlo-host-registration-limit-gib", "80", ] @@ -522,6 +524,8 @@ def test_serve_cli_forwards_hwr_policy_for_no_allgather_dlo(): 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 + assert explicit_kwargs["dlo_host_registration_mode"] == "disabled" + assert engine_args["dlo_host_registration_mode"] == "disabled" def test_serve_cli_accepts_diffusion_compile_controls(): diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 1a1c0c3ae..b7bbfb2af 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -739,6 +739,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_mode: str = "auto" dlo_host_registration_limit_gib: float = Field(default=0.0, ge=0) pin_cpu_memory: bool = True diffusion_compile_granularity: Literal["regional", "full"] = "regional" @@ -895,6 +896,7 @@ class _DiffusionConfigProjection: ) self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( limit_gib=self.dlo_host_registration_limit_gib, + mode=self.dlo_host_registration_mode, enable_dlo=self.enable_distributed_layerwise_offload, use_allgather=self.dlo_use_allgather, hwr_mode=self.host_weight_runtime_mode, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 06840e55a..c538ca177 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -483,6 +483,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_mode: 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 00764ea46..298191e8d 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -119,8 +119,15 @@ def validate_dlo_host_registration_options( enable_dlo: bool, use_allgather: bool, hwr_mode: object, + mode: object = "auto", ) -> float: """Validate the optional transport budget without probing CUDA or HWR.""" + if mode not in ("auto", "disabled"): + raise ValueError("dlo_host_registration_mode must be auto or disabled") + if mode == "disabled" and (not enable_dlo or use_allgather or hwr_mode == "disabled"): + raise ValueError( + "disabled dlo_host_registration_mode requires enabled no-AllGather DLO and Host Weight Runtime" + ) if not isinstance(limit_gib, (int, float, str)): raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}") value = float(limit_gib) @@ -833,6 +840,7 @@ class OmniDiffusionConfig: 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_mode: str = "auto" dlo_host_registration_limit_gib: float = 0.0 pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading @@ -1244,6 +1252,7 @@ class OmniDiffusionConfig: ) self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options( limit_gib=self.dlo_host_registration_limit_gib, + mode=self.dlo_host_registration_mode, enable_dlo=self.enable_distributed_layerwise_offload, use_allgather=self.dlo_use_allgather, hwr_mode=self.host_weight_runtime_mode, diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index bb047a7eb..c1afb2c17 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -73,12 +73,15 @@ class OffloadConfig: # 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 + dlo_host_registration_mode: str = "auto" # ``None`` preserves the model's legacy plan-driven component topology; # a frozenset is an explicit compact-API selection. components: frozenset[str] | None = None dlo_transfers: dict[str, DLOTransfer] | None = None def __post_init__(self) -> None: + if self.dlo_host_registration_mode not in ("auto", "disabled"): + raise ValueError("dlo_host_registration_mode must be auto or disabled") if self.components is not None: self.components = parse_offload_components(self.components) if self.dlo_transfers is None: @@ -174,7 +177,9 @@ class OffloadConfig: dlo_transfers = dict(resolved.transfers) dlo_resident_layers = resolved.resident_layers dit_uses_allgather = resolved.uses_allgather(DIT_COMPONENT) + registration_mode = getattr(od_config, "dlo_host_registration_mode", "auto") dlo_host_registration_limit_gib = validate_dlo_host_registration_options( + mode=registration_mode, limit_gib=getattr(od_config, "dlo_host_registration_limit_gib", 0.0), enable_dlo=enable_distributed_layerwise_offload, use_allgather=dit_uses_allgather, @@ -210,6 +215,7 @@ class OffloadConfig: dlo_use_allgather=dit_uses_allgather, dlo_resident_layers=dlo_resident_layers, dlo_host_registration_limit_gib=dlo_host_registration_limit_gib, + dlo_host_registration_mode=registration_mode, components=components, dlo_transfers=dlo_transfers, ) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py index 418b5d9fb..b5347b6c4 100644 --- a/vllm_omni/diffusion/offloader/cuda_host_registration.py +++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py @@ -6,12 +6,14 @@ from __future__ import annotations import ctypes import mmap +import time from collections import defaultdict from collections.abc import Sequence from dataclasses import dataclass from typing import Protocol import torch +from vllm.logger import init_logger from vllm_omni.host_weight_runtime import MappedHostRegion @@ -21,6 +23,8 @@ from .host_registration import ( HostRegistrationError, ) +logger = init_logger(__name__) + _CUDA_HOST_REGISTER_READ_ONLY = 0x08 _CUDA_DEVICE_ATTRIBUTE_HOST_REGISTER_READ_ONLY_SUPPORTED = 113 @@ -187,7 +191,9 @@ class CudaHostRegistration: registered: list[_AddressRange] = [] try: - for region in mapped: + for index, region in enumerate(mapped, start=1): + logger.info("Registering HWR host range %d/%d: %d bytes", index, len(mapped), region.size) + started = time.perf_counter() error = runtime.cudaHostRegister(region.start, region.size, _CUDA_HOST_REGISTER_READ_ONLY) if int(error) != 0: raise HostRegistrationError( @@ -195,6 +201,9 @@ class CudaHostRegistration: f"[{region.start:#x}, {region.end:#x}): {_handled_error_message(runtime, error)}" ) registered.append(region) + logger.info( + "Registered HWR host range %d/%d in %.3f s", index, len(mapped), time.perf_counter() - started + ) except Exception as exc: rollback_errors: list[str] = [] rollback_failed: list[_AddressRange] = [] diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e626d87a5..f04060e0b 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -1080,6 +1080,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): lease = self._host_weight_lease if lease is None: return False + if self.config.dlo_host_registration_mode == "disabled": + logger.info("HWR mmap registration disabled by transport policy; using bounded host staging") + 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 @@ -1089,6 +1092,12 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): limit_gib = self.config.dlo_host_registration_limit_gib max_bytes = int(limit_gib * 1024**3) if limit_gib > 0 else None + logger.info( + "Starting HWR mmap registration on %s: %d source range(s), budget_bytes=%s", + self.device, + len(lease.mapped_regions), + max_bytes, + ) started = time.perf_counter() try: registration = register_host_mappings( diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 812f4c3e7..998ed8ef2 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -582,6 +582,7 @@ class OrchestratorArgs: dlo_resident_layers: int = 0 host_weight_runtime_mode: str = "disabled" host_weight_runtime_root: str | None = None + dlo_host_registration_mode: str = "auto" dlo_host_registration_limit_gib: float = 0.0 boundary_ratio: float | None = None flow_shift: float | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 608d605e9..2648e3093 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1145,6 +1145,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_mode": kwargs.get("dlo_host_registration_mode", "auto"), "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": ( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ea56acb1a..141372ed6 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -780,6 +780,17 @@ class OmniServeCommand(CLISubcommand): "the same persistent path for population and serving." ), ) + omni_config_group.add_argument( + "--dlo-host-registration-mode", + choices=("auto", "disabled"), + default="auto", + help=( + "HWR transport registration policy for eligible no-AllGather DLO. " + "auto attempts registered direct H2D under the existing pin-memory and budget policy; " + "disabled selects bounded staging without registering the HWR mappings. " + "This preserves the pin-memory policy for staging buffers." + ), + ) omni_config_group.add_argument( "--dlo-host-registration-limit-gib", type=float,