diff --git a/modelexpress_client/python/modelexpress_rl/train/adapter.py b/modelexpress_client/python/modelexpress_rl/train/adapter.py index dff225ca6..15f1a6971 100644 --- a/modelexpress_client/python/modelexpress_rl/train/adapter.py +++ b/modelexpress_client/python/modelexpress_rl/train/adapter.py @@ -10,11 +10,20 @@ from collections.abc import Callable from dataclasses import dataclass from enum import Enum -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + import torch class NixlMetadataProvider(Protocol): - """Narrow NIXL manager surface required to publish trainer manifests.""" + """NIXL manager surface required to publish trainer manifests. + + Exposes the agent metadata every adapter needs plus ``register_tensors``. + NIXL can only transfer registered memory, so an adapter registers whatever + source buffers it owns (staging arenas, in-place local storage) before + building its manifest, so the published ``nixl_metadata`` covers them. + """ @property def agent_name(self) -> str: @@ -31,6 +40,10 @@ def listen_port(self) -> int | None: """Return the local NIXL metadata-listener port, when enabled.""" ... + def register_tensors(self, tensors: dict[str, torch.Tensor]) -> bytes: + """Register buffers with NIXL and return the refreshed agent metadata.""" + ... + class TrainerStagingMode(str, Enum): """How a trainer adapter preserves a version's immutable source bytes.""" diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/__init__.py b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/__init__.py new file mode 100644 index 000000000..6b5f77941 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FSDP/DTensor trainer-engine adapter for RL refit (async-capable).""" + +from .adapter import FSDPTrainerAdapter +from .publisher import ( + LocalTensorShard, + build_fsdp_reshard_manifest, + capture_local_shards, +) + +__all__ = [ + "FSDPTrainerAdapter", + "LocalTensorShard", + "build_fsdp_reshard_manifest", + "capture_local_shards", +] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/adapter.py b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/adapter.py new file mode 100644 index 000000000..25ee5c271 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/adapter.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FSDP/DTensor implementation of the trainer-engine adapter contract. + +Setup is one-time; the per-step source geometry is re-read from the state_dict +the client passes each stage, so a trainer that re-materializes its state_dict +(CPU offload, gathered state dict) still publishes the latest weights: + +- COPY_TO_DEVICE (default): ``initialize`` allocates one persistent wire-dtype + arena per shard and registers them once. Each stage snapshots the live weights + into those stable arenas (cast to the wire dtype only when the source differs); + ``publish_ready`` fences the async copy. Robust to a moving source because the + registered arena never moves. +- IN_PLACE (optimization): ``initialize`` registers the DTensor local storage + directly (contiguous, so RDMA-registerable) and serves it with no copy. Its + premise is stable storage: the registered address must not change, so each + stage asserts the source still sits where it was registered and fails toward + COPY_TO_DEVICE otherwise. The source must already be the wire dtype (no + in-place cast). +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.distributed as dist + +from modelexpress.refit.reshard.cuda_pool import classic_cuda_alloc +from modelexpress_rl.train.adapter import ( + CompletionFence, + NixlMetadataProvider, + StagedWeightVersionShardData, + TrainerEngineAdapter, + TrainerStagingMode, + WeightPayloadFormat, + WeightVersionShardManifest, +) + +from .publisher import ( + WIRE_DTYPE, + LocalTensorShard, + build_fsdp_reshard_manifest, + capture_local_shards, +) + + +def _source_reuse_unsupported() -> None: + raise NotImplementedError( + "FSDP publishing requires the RL framework to retire the published " + "version before resuming training; version-retirement signaling is not wired" + ) + + +class FSDPTrainerAdapter(TrainerEngineAdapter): + """Expose FSDP/DTensor state-dict shards through the trainer contract. + + ``initialize`` fixes the shard layout and registers the source buffers once; + ``stage_shard`` re-reads the rank-local views each step and either snapshots + them into the persistent arenas (COPY) or serves them in place (IN_PLACE). + """ + + def __init__( + self, + *, + manager: NixlMetadataProvider, + nixl_metadata_endpoint: str, + ) -> None: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("FSDP distributed process group is not initialized") + self._manager = manager + self._nixl_metadata_endpoint = nixl_metadata_endpoint + self._source_slot_id = f"publisher:global-rank:{dist.get_rank()}" + self._initialized = False + self._staging_mode: TrainerStagingMode | None = None + # name -> (global_shape, shard_offset, local_shape) fixed at initialize(). + self._expected_layout: dict[ + str, tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] + ] = {} + self._arenas: dict[str, torch.Tensor] = {} # COPY: name -> registered arena + # name -> the address we registered (the arena for COPY, the live source + # for IN_PLACE). The served buffer must keep sitting here. + self._registered_addrs: dict[str, int] = {} + + @property + def source_slot_id(self) -> str: + return self._source_slot_id + + @property + def supported_staging_modes(self) -> frozenset[TrainerStagingMode]: + return frozenset({TrainerStagingMode.COPY_TO_DEVICE, TrainerStagingMode.IN_PLACE}) + + @property + def supported_payload_formats(self) -> frozenset[WeightPayloadFormat]: + # Sharding lives in the manifest; the payload is the (sharded) full tensor. + return frozenset({WeightPayloadFormat.FULL_TENSOR}) + + def initialize( + self, *, shards: list[LocalTensorShard], staging_mode: TrainerStagingMode + ) -> None: + """Fix the shard layout and register the source buffers (idempotent).""" + if self._initialized: + return + if staging_mode not in self.supported_staging_modes: + raise NotImplementedError( + f"FSDPTrainerAdapter does not support {staging_mode.value} staging" + ) + names = frozenset(s.name for s in shards) + if len(names) != len(shards): + raise ValueError("FSDP shard names are not unique within this rank") + self._expected_layout = { + s.name: (s.global_shape, s.shard_offset, s.local_shape) for s in shards + } + + if staging_mode is TrainerStagingMode.COPY_TO_DEVICE: + self._allocate_and_register_arenas(shards) + else: # IN_PLACE + self._register_sources_in_place(shards) + + self._staging_mode = staging_mode + self._initialized = True + + def _allocate_and_register_arenas(self, shards: list[LocalTensorShard]) -> None: + """Allocate one persistent bf16 arena per shard and register them once.""" + # TODO(staging-followups): + # - register via a single VmmArena + register_arena (one dmabuf MR). + # - env var to stage into CPU pinned host memory vs GPU device memory + # (host staging frees device memory when the GPU is tight). + with classic_cuda_alloc(): + self._arenas = { + s.name: torch.empty( + s.local_shape, dtype=WIRE_DTYPE, device=s.source_tensor.device + ) + for s in shards + } + self._manager.register_tensors( + {self._register_key(i, s.name): self._arenas[s.name] for i, s in enumerate(shards)} + ) + self._registered_addrs = { + name: arena.data_ptr() for name, arena in self._arenas.items() + } + + def _register_sources_in_place(self, shards: list[LocalTensorShard]) -> None: + """Register the live local storage as the served buffer (no copy).""" + for shard in shards: + self._require_in_place_servable(shard) + self._manager.register_tensors( + {self._register_key(i, s.name): s.source_tensor for i, s in enumerate(shards)} + ) + self._registered_addrs = {s.name: s.source_tensor.data_ptr() for s in shards} + + def stage_shard( + self, + *, + tensors: Any, + staging_mode: TrainerStagingMode, + payload_format: WeightPayloadFormat, + ) -> StagedWeightVersionShardData: + """Capture one immutable, rank-local FSDP version shard.""" + if staging_mode not in self.supported_staging_modes: + raise NotImplementedError( + f"FSDPTrainerAdapter does not support {staging_mode.value} staging" + ) + if payload_format not in self.supported_payload_formats: + raise NotImplementedError( + f"FSDPTrainerAdapter does not support {payload_format.value} payloads" + ) + + # Re-read the rank-local views from THIS step's state_dict so a + # re-materialized source still publishes the latest weights; these same + # shards seed the one-time setup on the first stage (single capture). + shards = self._capture(tensors) + self.initialize(shards=shards, staging_mode=staging_mode) + if staging_mode is not self._staging_mode: + raise ValueError( + f"FSDPTrainerAdapter initialized for {self._staging_mode.value} " + f"staging; cannot stage {staging_mode.value}" + ) + self._require_same_layout(shards) + + if staging_mode is TrainerStagingMode.COPY_TO_DEVICE: + publish_ready = self._snapshot_into_arenas(shards) + else: # IN_PLACE serves live storage; nothing to copy. + self._require_sources_pinned(shards) + publish_ready = CompletionFence(lambda: None) + + return self._staged(shards, publish_ready) + + def _snapshot_into_arenas(self, shards: list[LocalTensorShard]) -> CompletionFence: + """Copy each rank-local source into its persistent registered arena. + + ``copy_`` casts to bf16 only when the source dtype differs. The arena is + the served buffer, so point each shard at it. + """ + stream = torch.cuda.current_stream() if torch.cuda.is_available() else None + for shard in shards: + arena = self._arenas[shard.name] + arena.copy_(shard.source_tensor) + shard.staging_tensor = arena + if stream is not None: + done = torch.cuda.Event() + done.record(stream) + return CompletionFence(done.synchronize) + return CompletionFence(lambda: None) + + def _require_sources_pinned(self, shards: list[LocalTensorShard]) -> None: + """Fail unless every source still sits where it was registered. + + IN_PLACE publishes the registered address, so a moved source would + advertise stale (freed or reused) memory. Fail toward COPY_TO_DEVICE. + """ + for shard in shards: + self._require_in_place_servable(shard) + if shard.source_tensor.data_ptr() != self._registered_addrs[shard.name]: + raise NotImplementedError( + f"{shard.name}: source storage moved since registration; " + "IN_PLACE requires stable storage, use COPY_TO_DEVICE" + ) + + def _capture(self, tensors: Any) -> list[LocalTensorShard]: + if not isinstance(tensors, dict): + raise TypeError("tensors must be an FSDP state_dict (dict[str, Tensor])") + shards = capture_local_shards(tensors) + if not shards: + raise ValueError("no local FSDP shards to publish") + return shards + + def _require_same_layout(self, shards: list[LocalTensorShard]) -> None: + expected_names = frozenset(self._expected_layout) + names = frozenset(s.name for s in shards) + if names != expected_names: + missing = sorted(expected_names - names) + extra = sorted(names - expected_names) + raise ValueError( + "FSDP tensor set changed since initialize " + f"(missing={missing[:5]} extra={extra[:5]})" + ) + for shard in shards: + layout = (shard.global_shape, shard.shard_offset, shard.local_shape) + if layout != self._expected_layout[shard.name]: + raise ValueError( + f"{shard.name}: shard geometry changed since initialize " + f"(was {self._expected_layout[shard.name]}, now {layout}); " + "a trainer must keep a fixed shard layout across steps" + ) + + @staticmethod + def _register_key(index: int, name: str) -> str: + return f"__pub__{index}__{name}" + + @staticmethod + def _require_in_place_servable(shard: LocalTensorShard) -> None: + if shard.source_tensor.dtype != WIRE_DTYPE: + raise NotImplementedError( + f"{shard.name}: IN_PLACE serves the source dtype but wire is " + f"{WIRE_DTYPE}; use COPY_TO_DEVICE to cast" + ) + if not shard.source_tensor.is_contiguous(): + raise NotImplementedError( + f"{shard.name}: IN_PLACE requires a contiguous local shard; " + "use COPY_TO_DEVICE for this tensor" + ) + + def _staged( + self, shards: list[LocalTensorShard], publish_ready: CompletionFence + ) -> StagedWeightVersionShardData: + blob = build_fsdp_reshard_manifest( + manager=self._manager, + shards=shards, + metadata_endpoint=self._nixl_metadata_endpoint, + ) + wire_elsize = torch.empty((), dtype=WIRE_DTYPE).element_size() + total_bytes = sum(math.prod(s.local_shape) * wire_elsize for s in shards) + return StagedWeightVersionShardData( + manifest=WeightVersionShardManifest( + data=blob, + tensor_count=len({s.name for s in shards}), + total_bytes=total_bytes, + transport="NIXL", + ), + publish_ready=publish_ready, + # Mirror Megatron: the framework must retire the version before + # resuming training. Retirement is not wired to this fence, so fail + # rather than claim the source is safe to mutate. + source_reuse_ready=CompletionFence(_source_reuse_unsupported), + # Keep the served buffers alive while the version can be selected. + buffer_owner=tuple(s.served_tensor for s in shards), + ) + + +__all__ = ["FSDPTrainerAdapter"] diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/publisher.py b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/publisher.py new file mode 100644 index 000000000..5b8d0b474 --- /dev/null +++ b/modelexpress_client/python/modelexpress_rl/train/engines/fsdp/publisher.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Extract FSDP/DTensor local shards and describe them as MX reshard sources. + +Extracts the rank-local shards from an FSDP ``state_dict`` and emits MX's +engine-neutral manifest (``PublishedTensor`` / ``PublishedShard`` + +``wrap_rendezvous_blob``). HF-name conversion is deliberately NOT done here: the +receiver captures how these trainer-format sources land in the vLLM param layout +(see ``modelexpress_rl/inference/reshard/fsdp``). + +Extraction rules (per state_dict tensor, floating point only): +- unsharded (not a DTensor): every rank holds + publishes the full tensor; + the identical copies across ranks are deduped by box upstream of the + transfer planner (merge_shard_tables) +- replicated DTensor: same as unsharded +- sharded DTensor: this rank serves its per-dim local box (general: FSDP dim 0, + tensor-parallel dim 1, or 2-D meshes) via compute_local_shape_and_global_offset +- served as the wire dtype (WIRE_DTYPE), cast into the staging arena for + COPY_TO_DEVICE only when the source differs +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass + +import torch +from torch.distributed.tensor import DTensor +from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + +from modelexpress.refit.reshard.rendezvous import ( + PublishedShard, + PublishedTensor, + wrap_rendezvous_blob, +) +from modelexpress_rl.train.adapter import NixlMetadataProvider + +logger = logging.getLogger("modelexpress_rl.train.engines.fsdp.publisher") + +# The dtype weights are served on the wire as. The cast only actually happens +# when the source dtype differs (e.g. an fp32 master): a matching source copies +# as-is, and IN_PLACE can serve a matching source with no copy at all. +# TODO: make this configurable at client initialization; hardcoded to bf16 for now. +WIRE_DTYPE = torch.bfloat16 + + +@dataclass +class LocalTensorShard: + """One rank-local source shard extracted from the FSDP state_dict. + + ``source_tensor`` is the live (or detached) rank-local view. ``shard_offset`` + is the per-dim offset of this shard's box inside the global tensor (all-zero + for unsharded/replicated). ``staging_tensor`` is set only for COPY_TO_DEVICE + and is the WIRE_DTYPE registered arena the source is copied into (copy_ + converts only if the source dtype differs). + """ + + name: str + global_shape: tuple[int, ...] + shard_offset: tuple[int, ...] + local_shape: tuple[int, ...] + source_tensor: torch.Tensor + staging_tensor: torch.Tensor | None = None + + @property + def served_tensor(self) -> torch.Tensor: + """The tensor NIXL actually registers/serves (staging if copied, else source).""" + return self.staging_tensor if self.staging_tensor is not None else self.source_tensor + + +def capture_local_shards( + state_dict: dict[str, torch.Tensor], +) -> list[LocalTensorShard]: + """Extract each rank's local source shards from an FSDP state_dict. + + Every rank publishes what it holds, including the tensors it holds in full + (unsharded and replicated). Those redundant full copies are identical + across ranks; deduping them to one shard per box is handled upstream of + the transfer planner (merge_shard_tables), not here. + """ + shards: list[LocalTensorShard] = [] + skipped: list[str] = [] + for name, value in state_dict.items(): + if not value.is_floating_point(): + skipped.append(name) + continue + full_shape = tuple(value.shape) + zero_offset = tuple(0 for _ in full_shape) + + # TODO(dedup-staging): for the full-copy tensors below (unsharded + + # replicated), under COPY_TO_DEVICE every rank stages a redundant copy; + # consider staging + publishing from rank 0 only in that mode. + + # Unsharded (not a DTensor): this rank holds the full tensor; publish it. + if not isinstance(value, DTensor): + shards.append( + LocalTensorShard( + name=name, + global_shape=full_shape, + shard_offset=zero_offset, + local_shape=full_shape, + source_tensor=value.detach(), + ) + ) + continue + + placements = value.placements + local_shape, global_offset = compute_local_shape_and_global_offset( + value.shape, value.device_mesh, placements + ) + local = value.to_local().detach() + if tuple(local.shape) != tuple(local_shape): + local = local[tuple(slice(size) for size in local_shape)] + + # Replicated DTensor: this rank holds the full tensor; publish it. + if all(placement.is_replicate() for placement in placements): + shards.append( + LocalTensorShard( + name=name, + global_shape=full_shape, + shard_offset=zero_offset, + local_shape=tuple(local_shape), + source_tensor=local, + ) + ) + continue + + # Sharded DTensor: this rank's per-dim box (general — FSDP dim 0, TP dim 1, + # or 2-D meshes). compute_local_shape_and_global_offset gives the full + # per-dim offset directly, including uneven splits. + if local.numel(): + shards.append( + LocalTensorShard( + name=name, + global_shape=full_shape, + shard_offset=tuple(int(off) for off in global_offset), + local_shape=tuple(local_shape), + source_tensor=local, + ) + ) + if skipped: + logger.debug( + "capture_local_shards: skipped %d non-floating-point entries: %s", + len(skipped), + sorted(skipped), + ) + return shards + + +def build_fsdp_reshard_manifest( + *, + manager: NixlMetadataProvider, + shards: list[LocalTensorShard], + metadata_endpoint: str, +) -> bytes: + """Describe already-registered FSDP source shards as an MX manifest blob. + + Assumes ``shards`` have their ``served_tensor`` registered with ``manager`` + (addr resolved via ``data_ptr``). Groups shards by name into one + ``PublishedTensor`` each (one shard per rank; fan-in of multiple ranks' + shards happens on the receive/plan side via the side table). + """ + if not metadata_endpoint or ":" not in metadata_endpoint: + raise ValueError("metadata_endpoint must be an explicit host:port") + agent_name = str(manager.agent_name) + if not shards: + raise ValueError("no local shards to publish") + + by_name: dict[str, PublishedTensor] = {} + for shard in shards: + served = shard.served_tensor + if not served.is_contiguous(): + raise ValueError(f"{shard.name}: served tensor must be contiguous for RDMA") + addr = served.data_ptr() + if addr <= 0: + raise ValueError(f"{shard.name}: shard has invalid address") + published_shard = PublishedShard( + agent_name=agent_name, + device_id=served.device.index if served.device.type == "cuda" else 0, + addr=addr, + shard_offset=tuple(shard.shard_offset), + shape=tuple(shard.local_shape), + ) + tensor = by_name.get(shard.name) + if tensor is None: + by_name[shard.name] = PublishedTensor( + name=shard.name, + dtype=str(served.dtype), + elsize=served.element_size(), + full_shape=tuple(shard.global_shape), + shards=[published_shard], + ) + else: + tensor.shards.append(published_shard) + + published = list(by_name.values()) + return wrap_rendezvous_blob( + manager.nixl_metadata, + agent_name, + metadata_endpoint, + published, + ) + + +__all__ = ["WIRE_DTYPE", "LocalTensorShard", "capture_local_shards", "build_fsdp_reshard_manifest"] diff --git a/modelexpress_client/python/tests/test_refit_fsdp_adapter.py b/modelexpress_client/python/tests/test_refit_fsdp_adapter.py new file mode 100644 index 000000000..b995f8e20 --- /dev/null +++ b/modelexpress_client/python/tests/test_refit_fsdp_adapter.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from modelexpress_rl.train.adapter import TrainerStagingMode, WeightPayloadFormat +from modelexpress_rl.train.engines.fsdp.adapter import FSDPTrainerAdapter + +ADAPTER = "modelexpress_rl.train.engines.fsdp.adapter" + + +class _Manager: + agent_name = "trainer-r0" + nixl_metadata = b"agent-metadata" + listen_port = 19000 + + def __init__(self): + self.registered = [] + + def register_tensors(self, tensors): + self.registered.append(dict(tensors)) + return self.nixl_metadata + + +@pytest.fixture +def dist_ready(monkeypatch): + monkeypatch.setattr(f"{ADAPTER}.dist.is_available", lambda: True) + monkeypatch.setattr(f"{ADAPTER}.dist.is_initialized", lambda: True) + monkeypatch.setattr(f"{ADAPTER}.dist.get_rank", lambda: 0) + + +def _adapter(manager=None): + return FSDPTrainerAdapter( + manager=manager or _Manager(), nixl_metadata_endpoint="host:1234" + ) + + +def _stage(adapter, state_dict, mode=TrainerStagingMode.IN_PLACE): + return adapter.stage_shard( + tensors=state_dict, + staging_mode=mode, + payload_format=WeightPayloadFormat.FULL_TENSOR, + ) + + +def test_requires_initialized_distributed_engine(monkeypatch): + monkeypatch.setattr(f"{ADAPTER}.dist.is_available", lambda: True) + monkeypatch.setattr(f"{ADAPTER}.dist.is_initialized", lambda: False) + + with pytest.raises(RuntimeError, match="distributed process group"): + _adapter() + + +def test_source_slot_id_is_rank_stamped(dist_ready): + assert _adapter().source_slot_id == "publisher:global-rank:0" + + +def test_in_place_stage_registers_once_and_blocks_source_reuse(dist_ready): + manager = _Manager() + adapter = _adapter(manager) + state_dict = {"w": torch.ones(2, 4, dtype=torch.bfloat16)} + + staged = _stage(adapter, state_dict) + + assert staged.manifest.tensor_count == 1 + assert staged.manifest.total_bytes == 2 * 4 * 2 # bf16 elsize + assert staged.manifest.transport == "NIXL" + staged.publish_ready.wait() # IN_PLACE performs no copy: no-op + with pytest.raises(NotImplementedError, match="retire the published"): + staged.source_reuse_ready.wait() + + # Re-staging the same weights must not re-register (setup is one-time). + _stage(adapter, state_dict) + assert len(manager.registered) == 1 + + +def test_in_place_rejects_a_moved_source(dist_ready): + adapter = _adapter() + _stage(adapter, {"w": torch.ones(2, 4, dtype=torch.bfloat16)}) + + # Same name/shape/dtype but fresh storage: the registered address is stale. + with pytest.raises(NotImplementedError, match="source storage moved"): + _stage(adapter, {"w": torch.ones(2, 4, dtype=torch.bfloat16)}) + + +def test_stage_rejects_a_changed_tensor_set(dist_ready): + adapter = _adapter() + state_dict = {"w": torch.ones(2, 4, dtype=torch.bfloat16)} + _stage(adapter, state_dict) + + state_dict["b"] = torch.ones(4, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="tensor set changed"): + _stage(adapter, state_dict) + + +def test_stage_rejects_a_changed_shard_geometry(dist_ready): + adapter = _adapter() + _stage(adapter, {"w": torch.ones(2, 4, dtype=torch.bfloat16)}) + + # Same name, different local shape: geometry must stay fixed after initialize. + with pytest.raises(ValueError, match="shard geometry changed"): + _stage(adapter, {"w": torch.ones(4, 4, dtype=torch.bfloat16)}) + + +def test_in_place_requires_wire_dtype_source(dist_ready): + adapter = _adapter() + with pytest.raises(NotImplementedError, match="use COPY_TO_DEVICE to cast"): + _stage(adapter, {"w": torch.ones(2, 4, dtype=torch.float32)}) + + +def test_in_place_requires_contiguous_source(dist_ready): + adapter = _adapter() + strided = torch.ones(4, 4, dtype=torch.bfloat16).t() + with pytest.raises(NotImplementedError, match="contiguous"): + _stage(adapter, {"w": strided}) + + +def test_staging_mode_cannot_change_after_initialize(dist_ready): + adapter = _adapter() + _stage(adapter, {"w": torch.ones(2, 4, dtype=torch.bfloat16)}) + + with pytest.raises(ValueError, match="initialized for"): + _stage( + adapter, + {"w": torch.ones(2, 4, dtype=torch.bfloat16)}, + mode=TrainerStagingMode.COPY_TO_DEVICE, + ) + + +def test_unsupported_staging_mode_is_rejected(dist_ready): + adapter = _adapter() + with pytest.raises(NotImplementedError, match="COPY_TO_HOST"): + _stage( + adapter, + {"w": torch.ones(2, 4, dtype=torch.bfloat16)}, + mode=TrainerStagingMode.COPY_TO_HOST, + ) + + +def test_unsupported_payload_format_is_rejected(dist_ready): + adapter = _adapter() + with pytest.raises(NotImplementedError, match="XOR_DELTA"): + adapter.stage_shard( + tensors={"w": torch.ones(2, 4, dtype=torch.bfloat16)}, + staging_mode=TrainerStagingMode.IN_PLACE, + payload_format=WeightPayloadFormat.XOR_DELTA, + ) + + +def test_non_dict_tensors_is_rejected(dist_ready): + adapter = _adapter() + with pytest.raises(TypeError, match="state_dict"): + _stage(adapter, [torch.ones(2, 4, dtype=torch.bfloat16)]) diff --git a/modelexpress_client/python/tests/test_refit_fsdp_publisher.py b/modelexpress_client/python/tests/test_refit_fsdp_publisher.py new file mode 100644 index 000000000..eb03521ab --- /dev/null +++ b/modelexpress_client/python/tests/test_refit_fsdp_publisher.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from modelexpress.refit.reshard.rendezvous import unwrap_rendezvous_blob +from modelexpress_rl.train.engines.fsdp.publisher import ( + LocalTensorShard, + build_fsdp_reshard_manifest, + capture_local_shards, +) + + +class _Manager: + agent_name = "trainer-r0" + nixl_metadata = b"agent-metadata" + listen_port = 19000 + + +def test_capture_publishes_full_copies_and_skips_non_float(): + state_dict = { + "w": torch.arange(8, dtype=torch.float32).reshape(2, 4), + "b": torch.ones(4, dtype=torch.bfloat16), + "step": torch.arange(4, dtype=torch.long), # non-float: skipped + } + + shards = capture_local_shards(state_dict) + + by_name = {s.name: s for s in shards} + assert set(by_name) == {"w", "b"} + for shard in shards: + # Plain (non-DTensor) tensors are held in full by every rank. + assert shard.shard_offset == (0,) * len(shard.global_shape) + assert shard.local_shape == shard.global_shape + assert shard.staging_tensor is None + assert by_name["w"].global_shape == (2, 4) + + +def test_capture_gives_a_scalar_tensor_a_rank_zero_offset(): + state_dict = {"loss_scale": torch.tensor(1.0, dtype=torch.float32)} + + (shard,) = capture_local_shards(state_dict) + + assert shard.global_shape == () + assert shard.shard_offset == () + assert shard.local_shape == () + + +def test_build_manifest_describes_the_served_buffer(): + tensor = torch.zeros(2, 4, dtype=torch.bfloat16) + shard = LocalTensorShard( + name="w", + global_shape=(2, 4), + shard_offset=(0, 0), + local_shape=(2, 4), + source_tensor=tensor, + ) + + blob = build_fsdp_reshard_manifest( + manager=_Manager(), shards=[shard], metadata_endpoint="host:1234" + ) + + payload = unwrap_rendezvous_blob(blob) + assert payload.agent_metadata == b"agent-metadata" + (published,) = payload.tensors + assert published.name == "w" + assert published.dtype == "torch.bfloat16" + assert published.elsize == 2 + assert tuple(published.full_shape) == (2, 4) + (pshard,) = published.shards + assert pshard.agent_name == "trainer-r0" + assert pshard.addr == tensor.data_ptr() + assert tuple(pshard.shard_offset) == (0, 0) + assert tuple(pshard.shape) == (2, 4) + + +def test_build_manifest_groups_multiple_shards_under_one_tensor(): + top = torch.zeros(2, 4, dtype=torch.bfloat16) + bottom = torch.zeros(2, 4, dtype=torch.bfloat16) + shards = [ + LocalTensorShard("w", (4, 4), (0, 0), (2, 4), top), + LocalTensorShard("w", (4, 4), (2, 0), (2, 4), bottom), + ] + + payload = unwrap_rendezvous_blob( + build_fsdp_reshard_manifest( + manager=_Manager(), shards=shards, metadata_endpoint="host:1234" + ) + ) + + (published,) = payload.tensors + offsets = sorted(tuple(s.shard_offset) for s in published.shards) + assert offsets == [(0, 0), (2, 0)] + + +def test_build_manifest_rejects_non_contiguous_served_tensor(): + strided = torch.zeros(4, 4, dtype=torch.bfloat16)[:, ::2] + shard = LocalTensorShard("w", (4, 2), (0, 0), (4, 2), strided) + + with pytest.raises(ValueError, match="contiguous"): + build_fsdp_reshard_manifest( + manager=_Manager(), shards=[shard], metadata_endpoint="host:1234" + ) + + +def test_build_manifest_rejects_an_invalid_address(): + meta = torch.zeros(2, 4, dtype=torch.bfloat16, device="meta") + shard = LocalTensorShard("w", (2, 4), (0, 0), (2, 4), meta) + + with pytest.raises(ValueError, match="invalid address"): + build_fsdp_reshard_manifest( + manager=_Manager(), shards=[shard], metadata_endpoint="host:1234" + ) + + +def test_build_manifest_requires_host_port_endpoint(): + shard = LocalTensorShard( + "w", (2, 4), (0, 0), (2, 4), torch.zeros(2, 4, dtype=torch.bfloat16) + ) + + with pytest.raises(ValueError, match="host:port"): + build_fsdp_reshard_manifest( + manager=_Manager(), shards=[shard], metadata_endpoint="no-port" + )