[Diffusion] Register HWR mmap for direct DLO H2D (#6591)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com> Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
Hongsheng Liuandhsliuustc0106 authored and GitHub committed
1 parent
5a8b81ebeb
commit
c608910d20
28 files changed
+1869
-153
No files matched your search
@@ -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]
|
||||
@@ -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"]
|
||||
|
||||
|
||||
Reference in new issue
Block a user