diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index c6fc75e0f..0b1740068 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -651,7 +651,7 @@ Loading precedence: CLI args > environment variables > config file > defaults. |--------|---------| | `__init__.py` | Package init, exports `register_modelexpress_loaders()` for callers to register the `modelexpress` and `mx` loaders with vLLM | | `client.py` | `MxClient` - gRPC client wrapping `PublishMetadata`, `ListSources`, `GetMetadata`, and `UpdateStatus` RPCs | -| `accelerators/` | `AcceleratorBackend` boundary for accelerator-specific torch device control and fast-path capability gates, split into `base.py` (protocol), `cuda.py` (`CudaAcceleratorBackend`), and `xpu.py` (`XpuAcceleratorBackend`). CUDA and XPU are implemented backends; XPU keeps CUDA-only fast paths (pool registration, VMM arena, GDS) disabled and falls back to generic per-tensor NIXL registration. Further backends can be added behind the same interface | +| `accelerators/` | `AcceleratorBackend` boundary for accelerator-specific torch device control and fast-path capability gates, split into `base.py` (protocol), `cuda.py` (`CudaAcceleratorBackend`), and `xpu.py` (`XpuAcceleratorBackend`). CUDA and XPU are implemented backends; XPU keeps CUDA-only fast paths (pool registration, VMM arena, GDS) disabled, falls back to generic per-tensor NIXL registration, and requires no classic-allocator pool for registered buffers. Further backends can be added behind the same interface | | `nixl_transfer.py` | `NixlTransferManager` - NIXL agent lifecycle, tensor registration, RDMA transfers | | `refit/` | Engine-agnostic live-refit primitives. `RefitTimingRecorder` provides normalized stage timing; `reshard/` provides loader-observed geometry capture, slice/transfer planning, rendezvous, and transport abstractions | | `gds_transfer.py` | GPUDirect Storage availability check and transfer utilities | @@ -753,6 +753,35 @@ STALE. Long-lived framework integrations still need to call `close()` from their lifecycle; SIGKILL and mid-transfer failure recovery remain follow-up work. +The receiver's device-specific work goes through `AcceleratorBackend` rather than +`torch.cuda` directly: per-stage synchronization, the backend handed to +`NixlTransferManager`, and the allocation scope for its receive and staging +buffers. `registered_buffer_alloc_scope()` selects that scope from +`requires_classic_alloc_pool()`, which states a requirement rather than a +capability. CUDA requires it and XPU does not: CUDA uses the +classic-`cudaMalloc` pool in `reshard/cuda_pool.py` because its caching allocator +under `expandable_segments:True` can return VMM ranges that register successfully +but fail during RDMA WRITE when `nvidia_peermem` cannot pin the underlying pages +(that module's docstring records the exact failure signature). XPU uses its normal +allocator because no equivalent hazard is known or has been observed on XPU. +Successful XPU registration does not prove the WRITE-time hazard absent, so ruling +it out would take a real RDMA write into an XPU buffer allocated under an +expandable-segment equivalent. `torch.xpu` does expose `MemPool` and +`XPUPluggableAllocator`, so an alternate XPU pool could be implemented if that +test ever says one is needed. + +The receiver applies no publisher/target accelerator compatibility policy, and +this is a known gap rather than a decision that cross-family refit is safe. The +rendezvous identity deliberately carries no accelerator (the receiver builds it +to *discover* the trainer, before it knows anything the trainer served) and the +shard table carries no publisher family, so `accelerators_compatible` has nothing +to compare on this path. No pairing is rejected on accelerator-family grounds; +other NIXL, fabric, or model-geometry constraints may still prevent transfer. +Closing the gap means +publishing the source family in the shard table and comparing both endpoints +through that gate; a target-only check would be unable to express the constraint, +since compatibility is a property of the source-target pair. + `engines/vllm/refit/receiver.py` supplies the vLLM-specific boundaries: capture on an unquantized meta twin, then installation through vLLM's layerwise reload and `process_weights_after_loading` path. Unsupported loader operations fail diff --git a/modelexpress_client/python/modelexpress/accelerators/base.py b/modelexpress_client/python/modelexpress/accelerators/base.py index 50c40efd8..6ee11b6bf 100644 --- a/modelexpress_client/python/modelexpress/accelerators/base.py +++ b/modelexpress_client/python/modelexpress/accelerators/base.py @@ -70,3 +70,15 @@ def supports_vmm(self) -> bool: def supports_gds(self) -> bool: """Return whether GPUDirect Storage loading is supported.""" ... + + def requires_classic_alloc_pool(self) -> bool: + """Return whether NIXL-registered buffers must come from a separate + classic-allocation pool on this backend. + + A requirement, not a capability: True where this backend's default torch + allocator can hand out a range the HCA cannot pin, so registered buffers + have to be scoped into a classic-allocation pool instead (see + ``refit/reshard/cuda_pool.py`` for the hazard and the CUDA + implementation). False means the default allocator is already suitable. + """ + ... diff --git a/modelexpress_client/python/modelexpress/accelerators/cuda.py b/modelexpress_client/python/modelexpress/accelerators/cuda.py index cdc1a8ba5..454ebed8d 100644 --- a/modelexpress_client/python/modelexpress/accelerators/cuda.py +++ b/modelexpress_client/python/modelexpress/accelerators/cuda.py @@ -60,3 +60,6 @@ def supports_vmm(self) -> bool: def supports_gds(self) -> bool: return True + + def requires_classic_alloc_pool(self) -> bool: + return True diff --git a/modelexpress_client/python/modelexpress/accelerators/xpu.py b/modelexpress_client/python/modelexpress/accelerators/xpu.py index 9bc751c1a..89184963c 100644 --- a/modelexpress_client/python/modelexpress/accelerators/xpu.py +++ b/modelexpress_client/python/modelexpress/accelerators/xpu.py @@ -78,3 +78,11 @@ def supports_vmm(self) -> bool: def supports_gds(self) -> bool: return False + + def requires_classic_alloc_pool(self) -> bool: + # The classic pool works around a CUDA-specific interaction between + # expandable-segment allocations and nvidia_peermem pinning; no equivalent + # hazard is known or has been observed on XPU. Not a proof of absence - + # refit/reshard/cuda_pool.py has the failure signature and + # docs/ARCHITECTURE.md what would settle it. + return False diff --git a/modelexpress_client/python/modelexpress/refit/README.md b/modelexpress_client/python/modelexpress/refit/README.md index 2470975cd..22fb79d89 100644 --- a/modelexpress_client/python/modelexpress/refit/README.md +++ b/modelexpress_client/python/modelexpress/refit/README.md @@ -110,7 +110,9 @@ Refit has two independent optimization surfaces: [`MdlLoader`](../engines/vllm/refit/installer.py) is a separate experimental vLLM installer called Mapped Direct Load (MDL). It caches direct, fused, and expert destination views so warm updates can copy into known slots instead of repeating general loader dispatch. MDL can consume partial input batches, but the reshard transport in this package does not yet expose a selector that reduces wire bytes for partial updates. The two features must not be treated as one end-to-end partial-refit path until that selector is wired and validated. -Each receiver keeps one load-time receive buffer per captured destination, plus a source-dtype conversion staging buffer for any parameter whose served dtype differs from its load-time dtype. Both come from classic CUDA allocations and stay registered for the receiver's lifetime, because re-registering per refit is what the cached plan exists to avoid. +Each receiver keeps one load-time receive buffer per captured destination, plus a source-dtype conversion staging buffer for any parameter whose served dtype differs from its load-time dtype. Both stay registered for the receiver's lifetime, because re-registering per refit is what the cached plan exists to avoid. These buffers use a backend-selected allocation scope: CUDA uses a classic `cudaMalloc`-backed pool because its caching allocator under `expandable_segments` can return VMM ranges that register successfully but fail during RDMA WRITE when `nvidia_peermem` cannot pin the underlying pages. XPU uses its normal allocator because no equivalent hazard is known or has been observed on XPU; successful XPU registration does not prove the WRITE-time hazard absent. `torch.xpu` does expose `MemPool` and `XPUPluggableAllocator`, so an alternate XPU pool could be implemented if one is ever needed. The selection is `AcceleratorBackend.requires_classic_alloc_pool()`. + +No publisher/target accelerator compatibility is checked on this path: neither the rendezvous identity nor the shard table carries the publisher's family, so `metadata/payload.py` has nothing to compare. No pairing is rejected on accelerator-family grounds; other NIXL, fabric, or model-geometry constraints may still prevent transfer. This is documented as a gap, not as a judgement that cross-family refit is validated; closing it means publishing the source family and comparing both endpoints. That buffer shape is why the vLLM receiver installs through `process_weights_after_loading` (PWAL) rather than MDL. The receiver reconstructs *load-time* tensors, and a quantized model still needs the engine's post-load processing to derive its runtime representation from them. MDL is appropriate only when the incoming tensors already match the validated runtime representation, which is why it is a separate opt-in path rather than the default. diff --git a/modelexpress_client/python/modelexpress/refit/reshard/__init__.py b/modelexpress_client/python/modelexpress/refit/reshard/__init__.py index c69df8fe9..148244424 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/__init__.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/__init__.py @@ -45,6 +45,7 @@ ReadDescriptor, Transport, ) +from modelexpress.refit.reshard.alloc_scope import registered_buffer_alloc_scope from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc from modelexpress.refit.reshard.receiver import ReshardReceiver from modelexpress.refit.reshard.rendezvous import ( @@ -86,6 +87,7 @@ "paired_runs", "plan_pull", "plan_transfer", + "registered_buffer_alloc_scope", "shard_region", "tensor_digest", "wrap_rendezvous_blob", diff --git a/modelexpress_client/python/modelexpress/refit/reshard/alloc_scope.py b/modelexpress_client/python/modelexpress/refit/reshard/alloc_scope.py new file mode 100644 index 000000000..81c729937 --- /dev/null +++ b/modelexpress_client/python/modelexpress/refit/reshard/alloc_scope.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Where the reshard receiver's receive and staging buffers are allocated from. + +Registering a buffer with NIXL constrains how it may be allocated, and the +constraint is a property of the accelerator rather than of the refit. This module +holds the family-agnostic selection so callers do not have to name a specific +allocator; :mod:`modelexpress.refit.reshard.cuda_pool` holds the one +accelerator-specific implementation that exists. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import TYPE_CHECKING, ContextManager + +if TYPE_CHECKING: + from modelexpress.accelerators import AcceleratorBackend + + +def registered_buffer_alloc_scope( + backend: "AcceleratorBackend", +) -> ContextManager[None]: + """Return the allocation scope ``backend`` needs for NIXL-registered buffers. + + A backend that reports ``requires_classic_alloc_pool()`` gets + :func:`~modelexpress.refit.reshard.cuda_pool.classic_cuda_alloc`; every other + backend gets a no-op and allocates normally, which is the correct answer here + rather than a fallback. CUDA is the only dedicated pool implementation; fail + clearly if another backend requests one instead of silently using CUDA code. + """ + if backend.requires_classic_alloc_pool(): + if backend.name != "cuda": + raise NotImplementedError( + f"No registered-buffer allocation pool is implemented for " + f"backend {backend.name!r}" + ) + # Imported per call rather than at module scope so this selection carries + # no import-time dependency on the one accelerator-specific + # implementation. Note this does not by itself keep cuda_pool out of a + # non-CUDA process: the package __init__ re-exports classic_cuda_alloc, + # so importing anything under refit.reshard loads that module anyway. + from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc + + return classic_cuda_alloc() + return nullcontext() + + +__all__ = ["registered_buffer_alloc_scope"] diff --git a/modelexpress_client/python/modelexpress/refit/reshard/receiver.py b/modelexpress_client/python/modelexpress/refit/reshard/receiver.py index 7feeed15e..882f819d4 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/receiver.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/receiver.py @@ -32,9 +32,10 @@ import torch from modelexpress import envs +from modelexpress.accelerators import accelerator_backend_for from modelexpress.client import MxClient from modelexpress.nixl_transfer import NixlTransferManager -from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc +from modelexpress.refit.reshard.alloc_scope import registered_buffer_alloc_scope from modelexpress.refit.reshard.rendezvous import gather_sources from modelexpress.refit.reshard.transfer_plan import ( exact_descriptors, @@ -327,13 +328,24 @@ def __init__( self._num_trainer_sources = num_trainer_sources self._timeout = timeout self._global_rank = global_rank + self._backend = accelerator_backend_for(device) + + # TODO(publisher-accelerator): nothing here checks the publisher's + # accelerator family, because the rendezvous identity and the shard table + # do not carry one. Both same-family and cross-family pairings are + # unchecked by accelerator family. Publish the source family in the shard + # table and run it through + # metadata/payload.py::accelerators_compatible. # TODO(transport-agnostic): the receiver is engine-agnostic but still # transport-bound to NIXL (this manager, NixlReshardTransport, and the # fetch_remote_and_wait P2P handshake in _prepare). Abstract these behind # a transport interface so non-NIXL backends can plug in. self._manager = NixlTransferManager( - agent_name=agent_name, device_id=local_rank, listen_port=listen_port + agent_name=agent_name, + device_id=local_rank, + listen_port=listen_port, + accelerator_backend=self._backend, ) self._manager.initialize() self._mx_client = MxClient(server_url=mx_server) @@ -484,13 +496,13 @@ def _prepare(self, timeout: float) -> None: # dtype-mismatched sources (e.g. a bf16-served router for an fp32 dest): # one persistent bf16 STAGING buffer per convert param, registered as an - # RDMA target (classic cudaMalloc so the HCA can RDMA into it); each refit - # we RDMA into staging then cast staging -> the (load-time) receive buffer. + # RDMA target using the backend-selected allocation scope; each refit we + # RDMA into staging then cast staging -> the (load-time) receive buffer. # Allocation and registration happen in three places below (convert # staging, full-pull staging, receive buffers). They are accumulated into # one figure each rather than reported per buffer class, because the - # actionable question is how much of a cold start is cudaMalloc versus HCA - # registration. + # actionable question is how much of a cold start is allocation versus + # HCA registration. alloc_s = 0.0 register_s = 0.0 @@ -498,7 +510,7 @@ def _prepare(self, timeout: float) -> None: self._staging_ptr = {} if plan.converts: _t = time.perf_counter() - with classic_cuda_alloc(): + with registered_buffer_alloc_scope(self._backend): self._staging = { c.param_name: torch.empty( c.dest_shape, dtype=c.src_dtype, device=self._device @@ -520,7 +532,7 @@ def _prepare(self, timeout: float) -> None: self._full_staging_ptr = {} if plan.full_pulls: _t = time.perf_counter() - with classic_cuda_alloc(): + with registered_buffer_alloc_scope(self._backend): self._full_staging = { full_pull.src_name: torch.empty( full_pull.global_shape, @@ -543,15 +555,16 @@ def _prepare(self, timeout: float) -> None: } # Receive buffers: one per captured param at its CAPTURED (load-time) - # shape/dtype, classic cudaMalloc, registered once. The live params are - # NOT RDMA targets; _install() writes the buffers into the live params. + # shape/dtype, allocated through the backend-selected scope and registered + # once. The live params are NOT RDMA targets; _install() writes the buffers + # into the live params. # Segment params (captured == served) are the RDMA targets - register them # + point _param_ptr at them. Convert params (router) are captured fp32 -> # their bf16 staging is the RDMA target and the refit casts into the buffer. seg_params = {seg.param_name for seg in plan.segments} self._recv_buffers = {} _t = time.perf_counter() - with classic_cuda_alloc(): + with registered_buffer_alloc_scope(self._backend): for name in all_params: shape, dtype = param_layout[name] self._recv_buffers[name] = torch.empty( @@ -806,7 +819,7 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: receive_buffer.storage_offset() + copy.dest_offset, ) destination.copy_(source_view) - torch.cuda.synchronize(self._device) + self._backend.synchronize(self._device.index) stages["reslice_s"] = time.perf_counter() - _t # Cast the served bf16 staging into the (fp32) receive buffer - a torch @@ -817,12 +830,12 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: self._recv_buffers[convert.param_name].copy_( self._staging[convert.param_name] ) - torch.cuda.synchronize(self._device) + self._backend.synchronize(self._device.index) stages["convert_s"] = time.perf_counter() - _t _t = time.perf_counter() self._install(self._recv_buffers) - torch.cuda.synchronize(self._device) + self._backend.synchronize(self._device.index) stages["install_s"] = time.perf_counter() - _t metrics = { diff --git a/modelexpress_client/python/tests/conftest.py b/modelexpress_client/python/tests/conftest.py index 131eb30b1..3f306da8d 100644 --- a/modelexpress_client/python/tests/conftest.py +++ b/modelexpress_client/python/tests/conftest.py @@ -29,6 +29,7 @@ class MockAcceleratorBackend: pool_reg: bool = False vmm: bool = False gds: bool = False + classic_alloc_pool: bool = False set_device_calls: list[int] = field(default_factory=list) synchronize_calls: list[int | None] = field(default_factory=list) empty_cache_calls: int = 0 @@ -63,6 +64,9 @@ def supports_vmm(self) -> bool: def supports_gds(self) -> bool: return self.gds + def requires_classic_alloc_pool(self) -> bool: + return self.classic_alloc_pool + @pytest.fixture def mock_accelerator_backend_cls(): diff --git a/modelexpress_client/python/tests/test_accelerator_backend.py b/modelexpress_client/python/tests/test_accelerator_backend.py index 6ffa10759..ece3986d0 100644 --- a/modelexpress_client/python/tests/test_accelerator_backend.py +++ b/modelexpress_client/python/tests/test_accelerator_backend.py @@ -5,7 +5,7 @@ from __future__ import annotations -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -34,6 +34,7 @@ def test_cuda_backend_uses_nixl_vram_segment(self): assert backend.supports_pool_reg() is True assert backend.supports_vmm() is True assert backend.supports_gds() is True + assert backend.requires_classic_alloc_pool() is True def test_cuda_backend_delegates_torch_calls(self, monkeypatch): calls = [] @@ -100,6 +101,7 @@ def test_xpu_backend_uses_nixl_vram_segment_and_disables_cuda_fast_paths(self): assert backend.supports_pool_reg() is False assert backend.supports_vmm() is False assert backend.supports_gds() is False + assert backend.requires_classic_alloc_pool() is False def test_xpu_backend_delegates_torch_calls(self, monkeypatch): calls = [] @@ -315,6 +317,76 @@ def apply_weight_iter(self, result: LoadResult, weights_iter): ): assert GdsStrategy().is_available(ctx) is False + def test_registered_buffer_scope_is_a_noop_without_a_pool_requirement( + self, + mock_accelerator_backend_cls, + ): + from modelexpress.refit.reshard import cuda_pool + from modelexpress.refit.reshard.alloc_scope import ( + registered_buffer_alloc_scope, + ) + + backend = mock_accelerator_backend_cls(classic_alloc_pool=False) + + with patch.object( + cuda_pool, + "_get_pool", + side_effect=AssertionError("classic pool should not be built"), + ): + scope = registered_buffer_alloc_scope(backend) + assert isinstance(scope, type(nullcontext())) + with scope: + pass + + def test_registered_buffer_scope_uses_the_classic_pool_when_required( + self, + mock_accelerator_backend_cls, + ): + from modelexpress.refit.reshard import cuda_pool + from modelexpress.refit.reshard.alloc_scope import ( + registered_buffer_alloc_scope, + ) + + backend = mock_accelerator_backend_cls( + name="cuda", + classic_alloc_pool=True, + ) + entered = [] + + @contextmanager + def fake_use_mem_pool(pool, device=None): + entered.append(pool) + yield + + with patch.object(cuda_pool, "_get_pool", return_value="pool"): + with patch.object(torch.cuda, "use_mem_pool", fake_use_mem_pool): + with registered_buffer_alloc_scope(backend): + pass + + assert entered == ["pool"] + + def test_registered_buffer_scope_rejects_an_unimplemented_backend_pool( + self, + mock_accelerator_backend_cls, + ): + from modelexpress.refit.reshard import cuda_pool + from modelexpress.refit.reshard.alloc_scope import ( + registered_buffer_alloc_scope, + ) + + backend = mock_accelerator_backend_cls( + name="rocm", + classic_alloc_pool=True, + ) + + with patch.object( + cuda_pool, + "_get_pool", + side_effect=AssertionError("CUDA pool should not be built"), + ): + with pytest.raises(NotImplementedError, match="rocm"): + registered_buffer_alloc_scope(backend) + def test_vmm_runtime_noops_when_backend_does_not_support_arena( self, monkeypatch, diff --git a/modelexpress_client/python/tests/test_reshard_refit_accelerator_wiring.py b/modelexpress_client/python/tests/test_reshard_refit_accelerator_wiring.py new file mode 100644 index 000000000..4db53b111 --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_refit_accelerator_wiring.py @@ -0,0 +1,235 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""How a reshard receiver wires itself to its accelerator. + +The receiver's device-specific work - the allocation scope for NIXL-registered +buffers, the per-stage synchronizes, and the backend the transfer manager +registers memory through - all resolve from the constructor's ``device`` rather +than assuming CUDA. These tests pin that wiring for both implemented families. + +The transfer-manager assertion is the one that catches a silent regression: +``NixlTransferManager`` defaults to a CUDA backend when none is passed, so +forgetting to hand it the receiver's backend registers XPU memory through CUDA +device calls while every other part of the receiver looks correct. + +Nothing here asserts a publisher/target compatibility policy, because the +receiver has none to assert: the rendezvous identity and the shard table carry no +publisher accelerator, so both same-family and cross-family pairings are +unchecked by accelerator family. See the ``TODO(publisher-accelerator)`` in +``receiver.py``. + +Run: pytest tests/test_reshard_refit_accelerator_wiring.py +""" + +from contextlib import contextmanager, nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from modelexpress.accelerators import CudaAcceleratorBackend, XpuAcceleratorBackend +from modelexpress.refit.reshard.alloc_scope import registered_buffer_alloc_scope +from modelexpress.refit.reshard.receiver import ReshardReceiver + +_KWARGS = dict( + model_name="m", + mx_server="127.0.0.1:8011", + agent_name="agent0", + local_rank=0, + global_rank=0, + num_trainer_sources=1, + listen_port=9000, +) + + +@pytest.fixture +def build_receiver(): + """Construct a receiver with the NIXL agent and metadata client stubbed out.""" + + def _build(device): + with patch( + "modelexpress.refit.reshard.receiver.NixlTransferManager" + ) as manager_cls, patch("modelexpress.refit.reshard.receiver.MxClient"): + manager_cls.return_value = MagicMock() + receiver = ReshardReceiver(device=device, **_KWARGS) + return receiver, manager_cls + + return _build + + +@pytest.fixture +def xpu_available(monkeypatch): + monkeypatch.setattr( + torch, "xpu", SimpleNamespace(is_available=lambda: True), raising=False + ) + + +class TestBackendResolution: + def test_cuda_device_resolves_the_cuda_backend(self, build_receiver): + receiver, _ = build_receiver(torch.device("cuda", 0)) + + assert isinstance(receiver._backend, CudaAcceleratorBackend) + assert receiver._backend.requires_classic_alloc_pool() is True + + def test_xpu_device_resolves_the_xpu_backend(self, build_receiver, xpu_available): + receiver, _ = build_receiver(torch.device("xpu", 0)) + + assert isinstance(receiver._backend, XpuAcceleratorBackend) + assert receiver._backend.requires_classic_alloc_pool() is False + + @pytest.mark.parametrize("device_type", ["cuda", "xpu"]) + def test_backend_is_passed_to_the_transfer_manager( + self, + build_receiver, + xpu_available, + device_type, + ): + """Otherwise the manager falls back to its own CUDA default and registers + this backend's memory through CUDA device calls.""" + receiver, manager_cls = build_receiver(torch.device(device_type, 0)) + + passed = manager_cls.call_args.kwargs["accelerator_backend"] + assert passed is receiver._backend + assert passed.name == device_type + + +class TestAllocationScope: + def test_cuda_receiver_scopes_registered_buffers_into_the_classic_pool( + self, + build_receiver, + ): + receiver, _ = build_receiver(torch.device("cuda", 0)) + entered = [] + + @contextmanager + def fake_use_mem_pool(pool, device=None): + entered.append(pool) + yield + + with patch( + "modelexpress.refit.reshard.cuda_pool._get_pool", return_value="pool" + ), patch.object(torch.cuda, "use_mem_pool", fake_use_mem_pool): + with registered_buffer_alloc_scope(receiver._backend): + pass + + assert entered == ["pool"] + + def test_xpu_receiver_allocates_registered_buffers_normally( + self, + build_receiver, + xpu_available, + ): + receiver, _ = build_receiver(torch.device("xpu", 0)) + + with patch( + "modelexpress.refit.reshard.cuda_pool._get_pool", + side_effect=AssertionError("classic pool should not be built"), + ): + scope = registered_buffer_alloc_scope(receiver._backend) + assert isinstance(scope, type(nullcontext())) + with scope: + pass + + def test_prepare_scopes_every_registered_buffer_class(self, monkeypatch): + from tests.test_reshard_refit_fused_wire import _build, _RecordingTransport + + receiver, keepalive = _build(_RecordingTransport()) + plan = receiver._plan + receiver._manager = MagicMock() + receiver._mx_client = MagicMock() + receiver._model_name = "m" + receiver._num_trainer_sources = 3 + receiver._capture = MagicMock( + return_value=( + SimpleNamespace( + copies=[ + SimpleNamespace(param_name="exact"), + SimpleNamespace(param_name="strided"), + SimpleNamespace(param_name="router"), + ], + unsupported=[], + ), + { + "exact": ((8,), torch.float32), + "strided": ((4, 2), torch.float32), + "router": ((4,), torch.float32), + }, + ) + ) + receiver._log_coverage = MagicMock() + sessions = {"s0": "a0", "s1": "a1", "s2": "a2"} + sources = { + name: SimpleNamespace(dtype=torch.float32, global_shape=(1,)) + for name in ("exact", "strided", "router") + } + entered = [] + + @contextmanager + def recording_scope(_backend): + entered.append(True) + yield + + monkeypatch.setattr( + "modelexpress.refit.reshard.receiver.gather_sources", + lambda *_args, **_kwargs: ( + sources, + sessions, + {session: 0 for session in sessions}, + {agent: "host:1" for agent in sessions.values()}, + ), + ) + monkeypatch.setattr( + "modelexpress.refit.reshard.receiver.plan_transfer", + lambda *_args, **_kwargs: plan, + ) + monkeypatch.setattr( + "modelexpress.refit.reshard.receiver.handshake_with_peers", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "modelexpress.refit.reshard.receiver.NixlReshardTransport", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "modelexpress.refit.reshard.receiver.registered_buffer_alloc_scope", + recording_scope, + ) + + receiver._prepare(timeout=1.0) + + assert len(entered) == 3 + assert receiver._manager.register_tensors.call_count == 3 + assert all(tensor.data_ptr() for tensor in keepalive) + + +class TestStageSynchronization: + def test_every_stage_sync_goes_through_the_backend(self, monkeypatch): + """The receiver holds no ``torch.cuda`` call of its own, so a refit on a + non-CUDA target synchronizes through that target's backend. Asserted by + running a real refit against a stub backend: were any stage still calling + ``torch.cuda`` directly, the count here would be short.""" + from tests.test_reshard_refit_fused_wire import _RecordingTransport, _build + + monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "1") + harness, _keepalive = _build(_RecordingTransport()) + + harness.update_weights(step=1) + + # re-slice, dtype cast, install - the three stages this plan exercises. + assert harness._backend.synchronize_calls == [None, None, None] + + def test_receiver_module_holds_no_direct_torch_cuda_call(self): + """A regression fence: the point of the backend boundary is that this file + names no accelerator directly.""" + from pathlib import Path + + import modelexpress.refit.reshard.receiver as receiver_module + + source = Path(receiver_module.__file__).read_text() + code = "\n".join( + line for line in source.splitlines() if not line.strip().startswith("#") + ) + assert "torch.cuda" not in code + assert "torch.xpu" not in code diff --git a/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py b/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py index 6f8e362a3..be26bf43d 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py +++ b/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py @@ -32,6 +32,7 @@ TransferPlan, ) from modelexpress.refit.reshard.transport import InMemoryReferenceTransport +from tests.conftest import MockAcceleratorBackend EL = 4 # float32 element size @@ -62,6 +63,9 @@ def __init__(self, transport) -> None: # noqa: D107 - see class docstring self._transport = transport self._global_rank = 0 self._install_order: list[str] = [] + # The stage syncs go through the accelerator backend, so a stub backend is + # what keeps this CPU-only rather than patching a torch device module. + self._backend = MockAcceleratorBackend() def _install(self, recv_buffers) -> None: self._install_order.append("install") @@ -175,7 +179,6 @@ def __init__(self, *, param_name, op_chain, dest_shape, dest_stride, dest_offset def _run(monkeypatch, *, fused: bool): monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "1" if fused else "0") - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) transport = _RecordingTransport() harness, keepalive = _build(transport) metrics = harness.update_weights(step=1) @@ -245,7 +248,6 @@ def test_accounting_covers_all_three_groups(monkeypatch): def test_empty_full_pull_and_convert_groups_are_skipped(monkeypatch): """A plan with only exact segments must still issue exactly one batch.""" - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) for fused in (True, False): monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "1" if fused else "0") diff --git a/modelexpress_client/python/tests/test_reshard_refit_stage_record.py b/modelexpress_client/python/tests/test_reshard_refit_stage_record.py index ddc515f5d..dfd830267 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_stage_record.py +++ b/modelexpress_client/python/tests/test_reshard_refit_stage_record.py @@ -39,7 +39,6 @@ def _run(monkeypatch, caplog, *, fused=True, enabled=True): monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "1" if fused else "0") monkeypatch.setenv("MX_REFIT_STAGE_RECORD", "1" if enabled else "0") monkeypatch.delenv("MX_RESHARD_MAX_GBPS", raising=False) - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) transport = _RecordingTransport() harness, keepalive = _build(transport) with caplog.at_level(logging.WARNING): @@ -141,7 +140,6 @@ def test_setup_costs_land_on_the_step_that_paid_them(monkeypatch, caplog): # clear the ceiling itself; a low one inherited from the environment aborts the # refit before any of the record assertions below are reached. monkeypatch.delenv("MX_RESHARD_MAX_GBPS", raising=False) - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) transport = _RecordingTransport() harness, keepalive = _build(transport) @@ -193,7 +191,6 @@ def test_a_skipped_stage_is_absent_rather_than_zero(monkeypatch, caplog): # Built directly rather than through _run, so the ceiling has to be cleared here # too, for the same reason. monkeypatch.delenv("MX_RESHARD_MAX_GBPS", raising=False) - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) transport = _RecordingTransport() harness, keepalive = _build(transport) harness._plan.converts = [] @@ -226,7 +223,6 @@ def test_an_exact_phase_that_moved_nothing_is_not_timed(monkeypatch, caplog): monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "0") monkeypatch.setenv("MX_REFIT_STAGE_RECORD", "1") monkeypatch.delenv("MX_RESHARD_MAX_GBPS", raising=False) - monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) harness, keepalive = _build(_RecordingTransport()) # A bounded plan, which is the only way a real plan reaches zero segments: # plan_transfer counts every copy into exact_descriptor_count and then moves the