diff --git a/modelexpress_client/python/README.md b/modelexpress_client/python/README.md index 398072066..defcc4933 100644 --- a/modelexpress_client/python/README.md +++ b/modelexpress_client/python/README.md @@ -219,6 +219,8 @@ register_modelexpress_loaders() | `MX_HEARTBEAT_INTERVAL_SECS` | `30` | Seconds between READY status heartbeats for published sources, including reshard rendezvous sources; keep below the server heartbeat timeout | | `MX_RESHARD_MAX_SEGMENTS_PER_COPY` | `64` | Maximum exact descriptors for one no-gather refit copy before a compatible dim-0-sharded source is pulled once into contiguous staging and sliced locally | | `MX_RESHARD_FUSED_WIRE` | `1` | Issue a refit's exact-segment, full-pull, and convert reads as one transport batch instead of draining each phase in turn. Set to `0` to restore the phased reads for an A/B comparison | +| `MX_RESHARD_BATCH_INSTALL` | `1` | Re-slice a refit's full-pulled sources with one batched `torch._foreach_copy_` instead of one `copy_()` per captured view. Issues the same copies; a per-view loop costs thousands of kernel launches whose overhead can rival the RDMA. Set to `0` to restore the per-view loop for an A/B comparison | +| `MX_RESHARD_CACHE_DESCRIPTORS` | `1` | Build NIXL read descriptors once per stable transfer plan and reuse them across refits. Set to `0` to rebuild the descriptor lists on every step for an A/B comparison | | `MX_RESHARD_REQUIRE_FULL_COVERAGE` | `0` | Fail a refit that installs less than `MX_RESHARD_COVERAGE_FLOOR` of the engine's parameter bytes. Off by default because partial and subset refit are intended; set to `1` for benchmark runs, where an incomplete refit produces timings that are the wrong magnitude | | `MX_RESHARD_COVERAGE_FLOOR` | `0.995` | Fraction of engine parameter bytes a gated refit must install. Not `1.0`: a few engine parameters, such as rotary `inv_freq`, are legitimately not refit material. Values outside `[0.0, 1.0]` are rejected | | `MX_RESHARD_HANDSHAKE_TIMEOUT_S` | `900` | Budget for the whole P2P metadata handshake, across every trainer peer and every retry. Bounds the handshake independently of the refit timeout, so one unreachable publisher cannot consume the entire refit | diff --git a/modelexpress_client/python/modelexpress/envs.py b/modelexpress_client/python/modelexpress/envs.py index 0f0ea54f7..f14be75ff 100644 --- a/modelexpress_client/python/modelexpress/envs.py +++ b/modelexpress_client/python/modelexpress/envs.py @@ -63,6 +63,8 @@ MX_MODEL_URI: Optional[str] MX_P2P_METADATA: str MX_RESHARD_FUSED_WIRE: bool + MX_RESHARD_BATCH_INSTALL: bool + MX_RESHARD_CACHE_DESCRIPTORS: bool MX_RESHARD_REQUIRE_FULL_COVERAGE: bool MX_RESHARD_COVERAGE_FLOOR: float MX_RESHARD_HANDSHAKE_TIMEOUT_S: float @@ -240,6 +242,22 @@ def _env_positive_float(name: str, default: float) -> float: "MX_MODEL_URI": lambda: os.environ.get("MX_MODEL_URI"), "MX_P2P_METADATA": lambda: os.environ.get("MX_P2P_METADATA", "1"), "MX_RESHARD_FUSED_WIRE": lambda: _env_bool("MX_RESHARD_FUSED_WIRE", True), + # Issue the per-view re-slice copies of full-pulled sources as one batched + # _foreach_copy_ instead of a copy_() per view. On by default: it is the same + # set of copies, and one launch per view means thousands of launches whose + # Python and launch overhead can rival the RDMA itself. Set to 0 to fall back + # to the per-view loop. See modelexpress.refit.reshard.receiver. + "MX_RESHARD_BATCH_INSTALL": lambda: _env_bool("MX_RESHARD_BATCH_INSTALL", True), + # Build the RDMA read descriptor lists once per plan instead of once per step. + # The descriptors are a pure function of the cached plan and the registered + # buffer addresses, neither of which changes between steps, so rebuilding them + # every refit re-derives an identical list of hundreds of thousands of objects + # in Python. On by default; set to 0 to rebuild per step for an A/B. The cache + # is dropped whenever the plan is rebuilt. See + # modelexpress.refit.reshard.receiver. + "MX_RESHARD_CACHE_DESCRIPTORS": lambda: _env_bool( + "MX_RESHARD_CACHE_DESCRIPTORS", True + ), # Refit coverage gate. The floor is a fraction of the engine's parameter # bytes; ReshardReceiver validates its range at the point of use. What a # complete refit scores is engine- and model-specific, so the default is set diff --git a/modelexpress_client/python/modelexpress/nixl_transfer.py b/modelexpress_client/python/modelexpress/nixl_transfer.py index 4cb3aeeb3..48028ca3a 100644 --- a/modelexpress_client/python/modelexpress/nixl_transfer.py +++ b/modelexpress_client/python/modelexpress/nixl_transfer.py @@ -140,6 +140,10 @@ def __init__( self._metadata: bytes = b"" self._tensor_descriptors: list[TensorDescriptor] = [] self._tensors: dict[str, torch.Tensor] = {} + # Registration descriptors must be deregistered before destroying the + # UCX-backed NIXL agent. Dropping an agent with live GPU registrations + # can abort inside ucp_worker_destroy during framework teardown. + self._registered_memory: list[Any] = [] # Remote agents this manager has loaded, so shutdown can disconnect them. # Maps agent name -> (ip, port) for agents reached over the P2P socket, or # None for agents loaded from a metadata blob. @@ -360,15 +364,19 @@ def register_tensors( alloc_tuples = [ (base, size, self._device_id, "") for base, size in allocations ] - self._agent.register_memory( - alloc_tuples, - mem_type=self._accelerator_backend.nixl_mem_type, - backends=self._backends, + self._registered_memory.append( + self._agent.register_memory( + alloc_tuples, + mem_type=self._accelerator_backend.nixl_mem_type, + backends=self._backends, + ) ) reg_count = len(allocations) else: tensor_list = list(tensors.values()) - self._agent.register_memory(tensor_list, backends=self._backends) + self._registered_memory.append( + self._agent.register_memory(tensor_list, backends=self._backends) + ) reg_count = len(tensor_list) nixl_reg_time = time.perf_counter() - nixl_reg_start @@ -507,10 +515,12 @@ def register_arena( return self.register_tensors(tensors, force_per_tensor=True) nixl_reg_start = time.perf_counter() - self._agent.register_memory( - [(base, used, self._device_id, "")], - mem_type=self._accelerator_backend.nixl_mem_type, - backends=self._backends, + self._registered_memory.append( + self._agent.register_memory( + [(base, used, self._device_id, "")], + mem_type=self._accelerator_backend.nixl_mem_type, + backends=self._backends, + ) ) nixl_reg_time = time.perf_counter() - nixl_reg_start @@ -1221,6 +1231,16 @@ def shutdown(self) -> None: atexit.unregister(self.shutdown) self._atexit_registered = False disconnected = self.disconnect_remote_agents() + if self._agent is not None: + for registered in reversed(self._registered_memory): + try: + self._agent.deregister_memory(registered) + except Exception: + logger.warning( + "Failed to deregister NIXL memory during shutdown", + exc_info=True, + ) + self._registered_memory = [] self._agent = None self._metadata = b"" self._tensor_descriptors = [] diff --git a/modelexpress_client/python/modelexpress/refit/reshard/geometry.py b/modelexpress_client/python/modelexpress/refit/reshard/geometry.py index 246568d97..3c2168ebb 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/geometry.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/geometry.py @@ -41,6 +41,7 @@ OpSpec, RecordedCopy, UnsupportedReshard, + summarize_unsupported, ) logger = logging.getLogger(__name__) @@ -71,6 +72,13 @@ torch.Tensor.flatten: "flatten", torch.Tensor.contiguous: "contiguous", torch.Tensor.chunk: "chunk", + # vLLM's fused-MoE expert loader unifies its fused and per-expert paths by + # doing `experts_shard.unbind()`, reached via `unsqueeze(0)` for a per-expert + # source. Without this entry every expert weight in a MoE model is classified + # unsupported and the refit fails closed at ~5% coverage. Like `chunk` it is a + # pure multi-return view, so the existing tuple handling in `_intercept` + # applies unchanged. + torch.Tensor.unbind: "unbind", } @@ -213,10 +221,14 @@ def _install_stamps( def make_stamp(inner, name): @functools.wraps(inner) - def stamp(p, *a, **kw): + def stamp(*a, **kw): + # Signature-transparent on purpose. vLLM's fused-MoE loader is + # invoked entirely by keyword (`param=`, `loaded_weight=`, + # `shard_id=`, `expert_id=`), so a stamp that named its first + # parameter positionally raised TypeError for every expert. recorder.current = name try: - return inner(p, *a, **kw) + return inner(*a, **kw) finally: recorder.current = None @@ -260,6 +272,7 @@ def capture_geometry( recorder = _BakeRecorder() saved = _install_stamps(model, recorder, default_weight_loader) unsupported: list[str] = [] + unsupported_reasons: dict[str, str] = {} try: # One source at a time: a single unsupported loader is attributed only # to that tensor, never the whole bake. Fused params @@ -269,8 +282,9 @@ def capture_geometry( lazy = LazyWeight(name, torch.Size(shape), dtype, "meta", recorder=recorder) try: model.load_weights([(name, lazy)]) - except UnsupportedReshard: + except UnsupportedReshard as exc: unsupported.append(name) + unsupported_reasons[name] = str(exc) finally: _restore_stamps(saved) @@ -280,8 +294,16 @@ def capture_geometry( len(unsupported), recorder.unattributed, ) + # A whole model class can fail for one reason (every fused expert source, say), + # so report the distinct causes with counts rather than a list of names. The + # names alone say which tensors are missing but never why. + for reason, count in summarize_unsupported(unsupported_reasons): + logger.warning( + "reshard capture: %d source(s) unsupported, cause: %s", count, reason + ) return CaptureResult( copies=recorder.copies, unsupported=unsupported, unattributed=recorder.unattributed, + unsupported_reasons=unsupported_reasons, ) diff --git a/modelexpress_client/python/modelexpress/refit/reshard/receiver.py b/modelexpress_client/python/modelexpress/refit/reshard/receiver.py index 7feeed15e..900d91b42 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/receiver.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/receiver.py @@ -49,6 +49,7 @@ CaptureResult, IncompleteRefit, UnsupportedReshard, + summarize_unsupported, ) logger = logging.getLogger("modelexpress.refit.reshard.receiver") @@ -267,6 +268,59 @@ def _fused_wire_enabled() -> bool: return envs.MX_RESHARD_FUSED_WIRE +def _batch_install_enabled() -> bool: + """Whether to re-slice full-pulled sources with one batched copy. + + Read at call time so an A/B can toggle it without re-importing. Set + ``MX_RESHARD_BATCH_INSTALL=0`` to fall back to one ``copy_()`` per view. + """ + return envs.MX_RESHARD_BATCH_INSTALL + + +def _cache_descriptors_enabled() -> bool: + """Whether to reuse the read descriptor lists across steps. + + Read at call time so an A/B can toggle it without re-importing. Set + ``MX_RESHARD_CACHE_DESCRIPTORS=0`` to rebuild them on every step. + """ + return envs.MX_RESHARD_CACHE_DESCRIPTORS + + +def _destinations_are_disjoint(destinations: list[torch.Tensor]) -> bool: + """Conservatively check whether destination views occupy separate storage. + + A strided view can contain holes, so its min/max storage span may overlap + another view even when their logical elements do not. Treating that case as + overlapping only loses batching; it cannot change copy order or bytes. + """ + spans_by_storage: dict[int, list[tuple[int, int]]] = {} + for tensor in destinations: + if tensor.numel() == 0: + continue + min_element = int(tensor.storage_offset()) + max_element = min_element + for size, stride in zip(tensor.shape, tensor.stride(), strict=True): + extent = (int(size) - 1) * int(stride) + min_element += min(0, extent) + max_element += max(0, extent) + element_size = int(tensor.element_size()) + spans_by_storage.setdefault( + int(tensor.untyped_storage().data_ptr()), [] + ).append( + ( + min_element * element_size, + (max_element + 1) * element_size, + ) + ) + for spans in spans_by_storage.values(): + previous_end = -1 + for start, end in sorted(spans): + if start < previous_end: + return False + previous_end = max(previous_end, end) + return True + + def _replay_ops(tensor: torch.Tensor, op_chain: tuple) -> torch.Tensor: """Replay a captured loader view chain on a staged full-source tensor.""" value = tensor @@ -339,6 +393,10 @@ def __init__( self._mx_client = MxClient(server_url=mx_server) self._plan = None # built lazily on the first refit + # Read descriptors for the cached plan, reused across steps. Tied to the + # plan's lifetime: whatever rebuilds the plan must drop this too, or the + # next refit would RDMA into the previous plan's addresses. + self._cached_descriptors: tuple | None = None self._transport: NixlReshardTransport | None = None self._recv_buffers: dict[ str, torch.Tensor @@ -444,10 +502,22 @@ def _prepare(self, timeout: float) -> None: # installed, so they would silently keep their initial (base-model) # weights for the entire run. Until the full-pull/loader path exists # (TODO), fail loudly rather than serve stale weights. + # + # Carry the capture causes, not just the names. A rejection that + # reports only which tensors are missing forces the reader back onto + # the cluster to find out which op defeated capture. + causes = summarize_unsupported( + getattr(capture, "unsupported_reasons", {}) or {} + ) + cause_text = ( + "; ".join(f"{count} x {cause}" for cause, count in causes) + if causes + else "cause not recorded at capture" + ) raise UnsupportedReshard( f"[reshard] {len(plan.fallback)} source(s) need the unimplemented " f"full-pull path (unsupported reshard ops); refusing to serve stale " - f"weights. Params: {plan.fallback[:10]}" + f"weights. Causes: {cause_text}. Params: {plan.fallback[:10]}" ) # P2P memory handshake (mirrors MX's vLLM RDMA path): fetch each trainer's # NIXL metadata (incl. its memory registrations) via its listen thread, so @@ -481,6 +551,9 @@ def _prepare(self, timeout: float) -> None: self._manager, session_to_agent, session_to_device, timeout_seconds=timeout ) self._plan = plan + # New plan, so any descriptors cached for the old one are stale by + # construction: they carry the previous plan's source addresses. + self._cached_descriptors = 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 @@ -637,6 +710,14 @@ def _log_coverage(self, capture, param_layout, all_params, plan) -> None: "copies_captured": len(capture.copies), "unsupported": len(unsupported), "unsupported_sample": [str(u)[:120] for u in unsupported[:10]], + # Grouped causes, so the harvested record explains an incomplete + # refit without needing the run's console output alongside it. + "unsupported_causes": [ + {"cause": cause[:200], "sources": count} + for cause, count in summarize_unsupported( + getattr(capture, "unsupported_reasons", {}) or {} + ) + ], "planned_wire_bytes": plan.bytes_planned(), "extra_wire_bytes": plan.extra_wire_bytes(), "descriptors": plan.descriptor_count(), @@ -703,33 +784,78 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: # dtype cast) runs after all reads complete. So the phases carry no # ordering dependency and are issued as one batch by default. Phased mode # drains each in turn and is kept for the A/B. - full_descriptors = [ - ReadDescriptor( - session=segment.session, - src_addr=segment.src_addr, - dst_addr=( - self._full_staging_ptr[full_pull.src_name] + segment.dst_byte - ), - nbytes=segment.nbytes, + # Descriptor construction is timed and cached. Timed because it is real + # per-step work that used to fall outside every stage, so it surfaced only + # as unattributed time and pushed the record below the attribution floor a + # breakdown has to clear to be worth reporting. Cached because a descriptor + # is a (session, src_addr, dst_addr, nbytes) tuple derived from the plan and + # the registered buffer addresses: the plan is built once and reused, and + # the buffers are registered once, so every step was re-deriving an + # identical list of hundreds of thousands of objects in Python. + _t = time.perf_counter() + cached = _cache_descriptors_enabled() + fused = _fused_wire_enabled() + # The fused flag is part of the key, not just the payload: the phased arm + # does not build the exact descriptors, so a cache filled under one arm + # cannot serve the other. An A/B that toggles it mid-process is exactly the + # case this is for. + reusable = ( + self._cached_descriptors is not None + and self._cached_descriptors[0] == fused + ) + if not cached or not reusable: + full_descriptors = [ + ReadDescriptor( + session=segment.session, + src_addr=segment.src_addr, + dst_addr=( + self._full_staging_ptr[full_pull.src_name] + segment.dst_byte + ), + nbytes=segment.nbytes, + ) + for full_pull in self._plan.full_pulls + for segment in full_pull.segments + ] + convert_descriptors = [ + ReadDescriptor( + session=segment.session, + src_addr=segment.src_addr, + dst_addr=self._staging_ptr[convert.param_name] + segment.dst_byte, + nbytes=segment.nbytes, + ) + for convert in self._plan.converts + for segment in convert.segments + ] + # Only the fused path issues the exact segments as descriptors; the + # phased path hands the plan to execute_transfer instead, so building + # them here would be wasted work for that arm. + exact = ( + exact_descriptors(self._plan, lambda name: self._param_ptr[name]) + if fused + else None ) - for full_pull in self._plan.full_pulls - for segment in full_pull.segments - ] - convert_descriptors = [ - ReadDescriptor( - session=segment.session, - src_addr=segment.src_addr, - dst_addr=self._staging_ptr[convert.param_name] + segment.dst_byte, - nbytes=segment.nbytes, + # nbytes of the auxiliary descriptors, summed once for the same reason. + aux_bytes = sum( + descriptor.nbytes + for descriptor in (*full_descriptors, *convert_descriptors) ) - for convert in self._plan.converts - for segment in convert.segments - ] - - if _fused_wire_enabled(): - descriptors = exact_descriptors( - self._plan, lambda name: self._param_ptr[name] + if cached: + self._cached_descriptors = ( + fused, + full_descriptors, + convert_descriptors, + exact, + aux_bytes, + ) + else: + _, full_descriptors, convert_descriptors, exact, aux_bytes = ( + self._cached_descriptors ) + stages["descriptor_build_s"] = time.perf_counter() - _t + + if fused: + assert exact is not None + descriptors = exact stats = { "segments": len(descriptors), "bytes": sum(descriptor.nbytes for descriptor in descriptors), @@ -774,10 +900,7 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: stages["wire_convert_s"] = time.perf_counter() - _t stats["segments"] += len(full_descriptors) + len(convert_descriptors) - stats["bytes"] += sum( - descriptor.nbytes - for descriptor in (*full_descriptors, *convert_descriptors) - ) + stats["bytes"] += aux_bytes # Before anything reads the receive buffers, and well before _install # commits them to live parameters. An impossible rate means the transport @@ -794,7 +917,15 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: # block. Without that, launch-bound stages read as free and whichever # stage syncs first absorbs the whole queue. if self._plan.full_pulls: + # Local re-slice of every full-pulled source. One copy_() per captured + # view means thousands of individual kernel launches, whose Python and + # launch overhead can rival the RDMA itself; _foreach_copy_ issues the + # same copies as a single batched op. _t = time.perf_counter() + batched = _batch_install_enabled() + destinations: list[torch.Tensor] = [] + source_views: list[torch.Tensor] = [] + copies_done = 0 for full_pull in self._plan.full_pulls: full_tensor = self._full_staging[full_pull.src_name] for copy in full_pull.copies: @@ -805,9 +936,28 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: copy.dest_stride, receive_buffer.storage_offset() + copy.dest_offset, ) - destination.copy_(source_view) + if batched: + destinations.append(destination) + source_views.append(source_view) + else: + destination.copy_(source_view) + copies_done += 1 + reslice_copies = len(destinations) if batched else copies_done + if batched and destinations: + if _destinations_are_disjoint(destinations): + torch._foreach_copy_(destinations, source_views) + else: + # ``_foreach_copy_`` does not define ordering for overlapping + # destinations. Preserve the captured loader order instead. + for destination, source_view in zip( + destinations, source_views, strict=True + ): + destination.copy_(source_view) torch.cuda.synchronize(self._device) stages["reslice_s"] = time.perf_counter() - _t + # Views, not sources: the per-view launch count is what batching + # removes, and the source count is already reported separately. + stages["reslice_copies"] = float(reslice_copies) # Cast the served bf16 staging into the (fp32) receive buffer - a torch # op, so the RDMA never crosses dtypes. _install writes the buffer. @@ -881,6 +1031,10 @@ def update_weights(self, step: int, *, timeout: float | None = None) -> dict: "converts": len(self._plan.converts), "fallback": len(stats["fallback"]), "fused_wire": _fused_wire_enabled(), + # The install arm this record was measured under. Without it a + # captured record cannot be attributed to a batched or per-view + # re-slice, which is the whole point of an A/B. + "batch_install": _batch_install_enabled(), **{k: round(v, 6) for k, v in stages.items()}, } # WARNING so a benchmark harness captures it without turning on INFO diff --git a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py b/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py index d2b61ecd8..389b82591 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/rendezvous.py @@ -139,8 +139,19 @@ def decode_shard_table(blob: bytes) -> list: schema = payload.get("schema") if schema != _SCHEMA: raise ValueError(f"unexpected shard-table schema {schema!r} (want {_SCHEMA!r})") + return decode_shard_entries(payload["tensors"]) + + +def decode_shard_entries(entries: list) -> list: + """Build ``PublishedTensor``s from already-parsed shard-table entries. + + Split out so a caller holding a parsed blob does not have to re-serialize the + table just to hand it back to a decoder that immediately re-parses it. On + Qwen3-30B-A3B that round-trip cost ~0.47 s per refit across 16 trainer + sources, for nothing. + """ tensors = [] - for t in payload["tensors"]: + for t in entries: shards = [ PublishedShard( agent_name=s["agent_name"], @@ -209,22 +220,48 @@ def build_sources(tensors: list) -> tuple: def merge_shard_tables(tables: list) -> list: """Merge per-rank ``list[PublishedTensor]`` into one, concatenating shards for the same source across ranks (reshard fans in cross-rank). full_shape / - dtype / elsize must agree across ranks for a given tensor name.""" + dtype / elsize must agree across ranks for a given tensor name. + + Replica publishers advertise the same geometric shard through DP/EDP + replication, so exactly one representative is retained per exact + offset/shape. Retaining all of them is what the merge used to do, and it + costs a read per replica: a replicated tensor under DP8 was pulled eight + times into the same destination bytes, and each extra owner added a P2P + handshake. It also defeats the full-pull optimization, whose dim-0 + partitioner fails closed on overlapping shards and falls back to the exact + plan that carries the duplicates. + + Selecting among candidates assumes they really are replicas. They are + byte-identical by construction when they come from DP/EDP replication, but a + publisher emitting parallelism-local names makes two *different* tensors + collide under one name, and then retaining the first installs bytes that + belong to the other. Publishers must therefore use globally unique names for + parallelism-local tensors. + """ merged: dict = {} + # name -> geometry -> first shard, insertion-ordered so the retained geometry + # sequence is deterministic. + candidates: dict = {} for table in tables: for t in table: cur = merged.get(t.name) if cur is None: merged[t.name] = PublishedTensor( - t.name, t.dtype, t.elsize, t.full_shape, list(t.shards) + t.name, t.dtype, t.elsize, t.full_shape, [] ) - continue - if cur.full_shape != t.full_shape or cur.dtype != t.dtype: + candidates[t.name] = {} + elif cur.full_shape != t.full_shape or cur.dtype != t.dtype: raise ValueError( f"tensor {t.name!r} published with inconsistent shape/dtype across ranks: " f"{cur.full_shape}/{cur.dtype} vs {t.full_shape}/{t.dtype}" ) - cur.shards.extend(t.shards) + per_geometry = candidates[t.name] + for shard in t.shards: + geometry = (tuple(shard.shard_offset), tuple(shard.shape)) + per_geometry.setdefault(geometry, shard) + + for name, tensor in merged.items(): + tensor.shards.extend(candidates[name].values()) return list(merged.values()) @@ -285,7 +322,9 @@ class RendezvousPayload(NamedTuple): permanently behind, and a consumer that excuses lagging publishers would then excuse all of its shards, going quiet instead of strict. It carries a default so that reading the stamp is a new field rather than a second unwrap function, - which is also why unpacking must now name five values. + which is also why unpacking must now name six values. ``tensor_count`` keeps + the published table size available when ``with_tensors=False`` deliberately + leaves ``tensors`` empty; ``None`` preserves direct and older callers. """ agent_metadata: bytes @@ -293,23 +332,46 @@ class RendezvousPayload(NamedTuple): metadata_endpoint: str tensors: list publisher_step: int | None = None + tensor_count: int | None = None + + def entry_count(self) -> int: + """How many shard-table entries this rank published. + + Distinct from ``len(self.tensors)`` because a payload unwrapped for a + quorum check carries the count without paying to build the table. Asking + "did this rank publish anything" must stay answerable on both. + """ + return len(self.tensors) if self.tensor_count is None else self.tensor_count + +def unwrap_rendezvous_blob( + blob: bytes, *, with_tensors: bool = True +) -> RendezvousPayload: + """Inverse of ``wrap_rendezvous_blob``. -def unwrap_rendezvous_blob(blob: bytes) -> RendezvousPayload: - """Inverse of ``wrap_rendezvous_blob``.""" + ``with_tensors=False`` decodes everything except the shard table, for callers + that only need the version stamp. Building the table dominates this function + (~97 ms vs ~9 ms per source on Qwen3-30B-A3B's 4922 entries), and a per-step + quorum check re-does it for a table that never changes between steps. + ``tensor_count`` is still populated, so an empty publisher stays detectable. + """ payload = json.loads(blob.decode("utf-8")) if payload.get("schema") != _SCHEMA: raise ValueError(f"unexpected rendezvous blob schema {payload.get('schema')!r}") agent_metadata = base64.b64decode(payload["agent_meta_b64"]) agent_name = payload["agent_name"] metadata_endpoint = payload.get("metadata_endpoint", "") - tensors = decode_shard_table( - json.dumps({"schema": _SCHEMA, "tensors": payload["tensors"]}).encode("utf-8") - ) + entries = payload["tensors"] + tensors = decode_shard_entries(entries) if with_tensors else [] raw_step = payload.get("publisher_step") publisher_step = None if raw_step is None else int(raw_step) return RendezvousPayload( - agent_metadata, agent_name, metadata_endpoint, tensors, publisher_step + agent_metadata, + agent_name, + metadata_endpoint, + tensors, + publisher_step, + len(entries), ) @@ -400,17 +462,55 @@ def close(self) -> None: self._publisher = None self._mx_source_id = None + def _fetch_metadata(self, instances: list) -> list: + """Fetch each instance's metadata, in order. + + Returns one entry per instance, ``None`` where the fetch raised. A single + rank's transport error must not abort the sweep: the poll loop's job is to + report how many ranks are actually readable, and one raised exception would + instead surface as a discovery failure with no count at all. + + Serial on purpose, which is the opposite of what the shape of the problem + suggests. Issuing these round-trips from a thread pool was tried and made + it **worse**: the fetch went from 4.02 s to ~6.9 s median on 16 sources. + The contention is not in this process. Every receiver rank runs this loop, + so 16 ranks x 16 sources put 256 requests and ~358 MB of shard table + in flight against one metadata server, and queueing there costs more than + the serial round-trips saved. Making this concurrent requires reducing what + the server has to serve first - see the note in ``discover_trainers``. + """ + return [self._get_metadata_or_none(inst) for inst in instances] + + def _get_metadata_or_none(self, inst): + try: + return self.client.get_metadata(inst.mx_source_id, inst.worker_id) + except Exception as exc: # noqa: BLE001 - one bad rank must not end the sweep + logger.warning( + "[reshard] get_metadata failed for source %s worker %s: %s", + inst.mx_source_id, + inst.worker_id, + exc, + ) + return None + def discover_trainers( self, expected_trainers: int, timeout: float = 1200.0, poll_interval: float = 1.0, + with_tensors: bool = True, ) -> list: """Block until ``expected_trainers`` trainer ranks are visible **with a non-empty shard table**, then return them. Returns ``list[RendezvousPayload]``, one per trainer rank. + ``with_tensors=False`` skips building the shard tables, for a caller that + only needs each rank's version stamp. The quorum semantics are unchanged - + emptiness is still judged on the published entry count - but a per-step + check no longer rebuilds tables that are identical every step. Anything + that reads shard geometry must leave this ``True``. + A rank counts toward the quorum only once its published table names at least one tensor. A rank that advertises READY with nothing to read has registered no memory, so satisfying the quorum with it makes the receiver @@ -424,28 +524,53 @@ def discover_trainers( """ trainer_id = self._identity("trainer") deadline = time.monotonic() + timeout + # Cost split, kept because this call dominates a MoE refit: measured at + # 51% of a 10.25 s refit on Qwen3-30B-A3B, 2.7x the weight transfer it + # precedes. On the cluster the fetch is ~4.0 s of that and the parse ~0.8 s. + # The round-trips stay serial; see ``_fetch_metadata`` for the measurement + # that ruled out a thread pool. ``fetch_s`` includes both ``list_sources`` + # and the metadata sweep. + fetch_s = 0.0 + parse_s = 0.0 + polls = 0 while True: + list_t0 = time.perf_counter() resp = self.client.list_sources( trainer_id, status_filter=p2p_pb2.SOURCE_STATUS_READY, ) instances = list(resp.instances) + fetch_s += time.perf_counter() - list_t0 payloads, empty = [], 0 # Every visible READY source is inspected, whatever the count: with # fewer sources than expected the shard-table state is exactly what a # timeout here has to report, and reaching the quorum on instances # alone says nothing about whether any of them published bytes. - for inst in instances: - meta = self.client.get_metadata(inst.mx_source_id, inst.worker_id) - if not meta.found: + # + # The fetch is the dominant term (~4.0 s of ~4.8 s here) and it is + # server-bound, not client-bound: see ``_fetch_metadata`` for the + # measurement that ruled out concurrency. Cutting it further means + # sending less, which needs ``publisher_step`` carried in the + # ``list_sources`` instance record so the quorum costs one small call + # instead of one full shard table per rank. That is a protocol change. + fetch_t0 = time.perf_counter() + metas = self._fetch_metadata(instances) + fetch_s += time.perf_counter() - fetch_t0 + for meta in metas: + if meta is None or not meta.found: continue - payload = unwrap_rendezvous_blob(meta.worker.nixl_metadata) - if not payload.tensors: + parse_t0 = time.perf_counter() + payload = unwrap_rendezvous_blob( + meta.worker.nixl_metadata, with_tensors=with_tensors + ) + parse_s += time.perf_counter() - parse_t0 + if payload.entry_count() == 0: empty += 1 continue payloads.append(payload) if len(payloads) >= expected_trainers: break + polls += 1 if len(payloads) >= expected_trainers: break if time.monotonic() >= deadline: @@ -456,6 +581,24 @@ def discover_trainers( ) time.sleep(poll_interval) + logger.warning( + "MX_DISCOVER_COST %s", + json.dumps( + { + "schema": "mx-discover-cost-v1", + "rank": self.rank, + "sources": len(payloads), + # entry_count, not len(tensors): the quorum path does not build + # the tables, and this figure is what shows the cost tracks + # source count rather than bytes moved. + "tensors": sum(p.entry_count() for p in payloads), + "tables_built": with_tensors, + "grpc_fetch_s": round(fetch_s, 6), + "blob_parse_s": round(parse_s, 6), + "polls": polls, + } + ), + ) logger.info( "[reshard] discovered %d trainer rank(s)%s: %s", len(payloads), diff --git a/modelexpress_client/python/modelexpress/refit/reshard/types.py b/modelexpress_client/python/modelexpress/refit/reshard/types.py index d5050190f..c8672ff68 100644 --- a/modelexpress_client/python/modelexpress/refit/reshard/types.py +++ b/modelexpress_client/python/modelexpress/refit/reshard/types.py @@ -35,6 +35,25 @@ class IncompleteRefit(RuntimeError): ``RuntimeError``, so a caller that catches only that is unaffected.""" +def summarize_unsupported( + reasons: dict, limit: int | None = 3 +) -> list[tuple[str, int]]: + """Group per-source capture failures by cause, most frequent first. + + Every message embeds the offending source's name and op-chain, so thousands + of sources failing for one shared reason produce thousands of textually + distinct strings. Cutting each message at its source-specific tail collapses + them, which is what makes "every expert in the model" legible as one cause + instead of 18432 unique ones. + """ + counts: dict[str, int] = {} + for message in reasons.values(): + cause = str(message).split(" on lazy ", 1)[0].strip() + counts[cause] = counts.get(cause, 0) + 1 + ranked = sorted(counts.items(), key=lambda item: (-item[1], item[0])) + return ranked[:limit] if limit is not None else ranked + + @dataclass class RecordedCopy: """One recorded scatter: read ``src_name`` sliced by ``op_chain`` and write it @@ -57,6 +76,10 @@ class CaptureResult: """Output of a bake: the recorded copies plus what could not be attributed. ``unsupported`` = source names whose loader used an unsupported op. + ``unsupported_reasons`` = that source name -> the op that defeated capture. + Without it a rejected refit reports only how many sources failed, which is + not enough to tell an unexpressible fused layout apart from a loader that + merely touched one op outside the allowlist. ``unattributed`` = copy_ calls fired with no active loader stamp. Either condition makes the update fail closed in the current receiver; there is no fallback path that serves those tensors by another route.""" @@ -64,3 +87,4 @@ class CaptureResult: copies: list = field(default_factory=list) unsupported: list = field(default_factory=list) unattributed: int = 0 + unsupported_reasons: dict = field(default_factory=dict) diff --git a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py index 2ff688a06..bdbdc1b4b 100644 --- a/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py +++ b/modelexpress_client/python/modelexpress_rl/train/engines/megatron/aliases.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any @@ -172,6 +173,187 @@ def _build_qkv_aliases( ) -> list[PublishedTensor]: if len(item.hf_names) != 3 or item.tensor.ndim != 2: raise ValueError(f"{item.name}: QKV aliasing requires 2D q/k/v weights") + has_global_q = "num_heads" in item.extras + has_global_kv = "num_kv_heads" in item.extras + if has_global_q != has_global_kv: + raise ValueError( + f"{item.name}: global QKV metadata requires both num_heads and num_kv_heads" + ) + if has_global_q: + return _build_global_qkv_aliases(item, agent_name) + return _build_legacy_qkv_aliases(item, agent_name) + + +def _qkv_source_interval(item: MegatronTensorSpec) -> tuple[int, int]: + """Return this rank's raw row interval in the global fused QKV tensor.""" + local_rows = int(item.tensor.shape[0]) + global_rows = int(item.global_shape[0]) + if item.placement_kind != "SHARD": + if local_rows != global_rows: + raise ValueError(f"{item.name}: replicated QKV shape mismatch") + return 0, global_rows + if item.shard_axis != 0 or item.local_shard_range is None: + raise ValueError(f"{item.name}: QKV shards must carry a row range") + lo, hi = (int(value) for value in item.local_shard_range) + if not 0 <= lo < hi <= global_rows or hi - lo != local_rows: + raise ValueError(f"{item.name}: inconsistent QKV source row interval") + return lo, hi + + +_Q, _K, _V = 0, 1, 2 + + +@dataclass(frozen=True) +class _QkvBand: + """One run of fused rows that belongs to a single projection. + + ``destination_start`` is where the run lands in that projection's own + tensor, which is not where it sits in the fused tensor. + """ + + projection: int + start: int + rows: int + destination_start: int + + +@dataclass(frozen=True) +class _QkvLayout: + """Global row layout of a fused QKV tensor, derived from head counts. + + Megatron lays the tensor out as one block per KV group: that group's query + rows, then its single K head, then its single V head. Blocks repeat for + every KV group, so Q, K and V rows interleave rather than forming three + contiguous regions. + """ + + head_dim: int + q_heads: int + kv_heads: int + + @property + def q_rows_per_group(self) -> int: + return (self.q_heads // self.kv_heads) * self.head_dim + + @property + def group_rows(self) -> int: + return self.q_rows_per_group + 2 * self.head_dim + + @property + def total_rows(self) -> int: + return self.kv_heads * self.group_rows + + @property + def destination_rows(self) -> tuple[int, int, int]: + """Row count of the whole q, k and v tensors this layout unpacks into.""" + kv_rows = self.kv_heads * self.head_dim + return self.q_heads * self.head_dim, kv_rows, kv_rows + + def bands(self) -> Iterator[_QkvBand]: + """Yield every projection run, in global row order.""" + for group in range(self.kv_heads): + group_start = group * self.group_rows + k_start = group_start + self.q_rows_per_group + yield _QkvBand( + _Q, group_start, self.q_rows_per_group, group * self.q_rows_per_group + ) + yield _QkvBand(_K, k_start, self.head_dim, group * self.head_dim) + yield _QkvBand( + _V, k_start + self.head_dim, self.head_dim, group * self.head_dim + ) + + +def _read_qkv_layout(item: MegatronTensorSpec) -> _QkvLayout: + """Validate the published global head metadata and derive the row layout.""" + if item.extras.get("qkv_interleave") != "by_head": + raise ValueError( + f"{item.name}: global QKV aliasing requires qkv_interleave='by_head'" + ) + if "head_dim" not in item.extras: + raise ValueError( + f"{item.name}: global QKV aliasing requires extras['head_dim']" + ) + layout = _QkvLayout( + head_dim=int(item.extras["head_dim"]), + q_heads=int(item.extras["num_heads"]), + kv_heads=int(item.extras["num_kv_heads"]), + ) + if ( + layout.head_dim < 1 + or layout.q_heads < 1 + or layout.kv_heads < 1 + or layout.q_heads % layout.kv_heads + ): + raise ValueError(f"{item.name}: invalid global Q/KV head geometry") + return layout + + +def _build_global_qkv_aliases( + item: MegatronTensorSpec, agent_name: str +) -> list[PublishedTensor]: + """Map one raw TP row interval through Megatron's global QKV interleave. + + This rank owns a single contiguous interval of fused rows. Intersecting it + with each band says which projection those rows belong to and where they + land, so a rank that happens to own no K or V rows simply matches no K or V + band. + """ + layout = _read_qkv_layout(item) + hidden = int(item.tensor.shape[1]) + if len(item.global_shape) != 2 or int(item.global_shape[1]) != hidden: + raise ValueError(f"{item.name}: QKV hidden dimension mismatch") + if int(item.global_shape[0]) != layout.total_rows: + raise ValueError(f"{item.name}: global QKV rows disagree with head metadata") + + source_lo, source_hi = _qkv_source_interval(item) + shards: tuple[list[PublishedShard], ...] = ([], [], []) + mapped_rows = 0 + for band in layout.bands(): + overlap_lo = max(source_lo, band.start) + overlap_hi = min(source_hi, band.start + band.rows) + if overlap_lo >= overlap_hi: + continue + tensor = item.tensor.narrow(0, overlap_lo - source_lo, overlap_hi - overlap_lo) + shards[band.projection].append( + PublishedShard( + agent_name=agent_name, + device_id=int(tensor.device.index or 0), + addr=int(tensor.data_ptr()), + shard_offset=(band.destination_start + overlap_lo - band.start, 0), + shape=tuple(int(dim) for dim in tensor.shape), + digest=published_digest(tensor), + ) + ) + mapped_rows += overlap_hi - overlap_lo + + if mapped_rows != int(item.tensor.shape[0]): + raise ValueError( + f"{item.name}: QKV interval mapping covered {mapped_rows} of " + f"{int(item.tensor.shape[0])} source rows" + ) + + # When KV heads are fewer than TP ranks, most publishers legitimately own no + # K or V rows. Other ranks contribute those destination intervals when the + # per-rank tables are merged. + return [ + PublishedTensor( + name=name, + dtype=str(item.tensor.dtype), + elsize=int(item.tensor.element_size()), + full_shape=(rows, hidden), + shards=projection_shards, + ) + for name, rows, projection_shards in zip( + item.hf_names, layout.destination_rows, shards, strict=True + ) + if projection_shards + ] + + +def _build_legacy_qkv_aliases( + item: MegatronTensorSpec, agent_name: str +) -> list[PublishedTensor]: + """Compatibility path for descriptors whose head counts divide across TP.""" required = ("head_dim", "num_heads_local", "num_kv_heads_local") missing = [key for key in required if key not in item.extras] if missing: diff --git a/modelexpress_client/python/tests/test_envs.py b/modelexpress_client/python/tests/test_envs.py index c8778913a..e3294f195 100644 --- a/modelexpress_client/python/tests/test_envs.py +++ b/modelexpress_client/python/tests/test_envs.py @@ -28,6 +28,8 @@ def test_defaults_when_unset(monkeypatch): "MX_GDS_TIMEOUT", "MX_HEARTBEAT_INTERVAL_SECS", "MX_RESHARD_FUSED_WIRE", + "MX_RESHARD_BATCH_INSTALL", + "MX_RESHARD_CACHE_DESCRIPTORS", ): monkeypatch.delenv(name, raising=False) @@ -49,6 +51,8 @@ def test_defaults_when_unset(monkeypatch): assert envs.MX_GDS_TIMEOUT == pytest.approx(120.0) assert envs.MX_HEARTBEAT_INTERVAL_SECS == 30 assert envs.MX_RESHARD_FUSED_WIRE is True + assert envs.MX_RESHARD_BATCH_INSTALL is True + assert envs.MX_RESHARD_CACHE_DESCRIPTORS is True def test_int_and_float_parsing(monkeypatch): diff --git a/modelexpress_client/python/tests/test_nixl_peer_lifecycle.py b/modelexpress_client/python/tests/test_nixl_peer_lifecycle.py index 350117397..16ad47817 100644 --- a/modelexpress_client/python/tests/test_nixl_peer_lifecycle.py +++ b/modelexpress_client/python/tests/test_nixl_peer_lifecycle.py @@ -32,6 +32,7 @@ class FakeAgent: def __init__(self, fail_remove: bool = False): self.removed: list[str] = [] + self.deregistered: list[object] = [] self.fail_remove = fail_remove self._fetched: set[str] = set() @@ -49,6 +50,9 @@ def remove_remote_agent(self, name: str): raise RuntimeError(f"remote metadata for agent '{name}' not found") self.removed.append(name) + def deregister_memory(self, registered): + self.deregistered.append(registered) + def _manager(agent=None, metadata=b"md", accelerator=None): mgr = NixlTransferManager( @@ -157,6 +161,26 @@ def spy(name): assert seen["agent_alive"] is True assert mgr._agent is None + def test_registered_memory_is_released_before_agent_is_dropped(self): + agent = FakeAgent() + mgr = _manager(agent=agent) + first, second = object(), object() + mgr._registered_memory = [first, second] + seen = [] + real_deregister = agent.deregister_memory + + def spy(registered): + seen.append((registered, mgr._agent is agent)) + real_deregister(registered) + + agent.deregister_memory = spy + mgr.shutdown() + + assert seen == [(second, True), (first, True)] + assert agent.deregistered == [second, first] + assert mgr._registered_memory == [] + assert mgr._agent is None + def test_removing_a_peer_twice_is_harmless(self): """The peer may already be gone, e.g. it sent us NIXLCOMM:INVL on exit.""" agent = FakeAgent() diff --git a/modelexpress_client/python/tests/test_reshard_megatron_gqa.py b/modelexpress_client/python/tests/test_reshard_megatron_gqa.py new file mode 100644 index 000000000..6799e3b40 --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_megatron_gqa.py @@ -0,0 +1,390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from modelexpress.refit.reshard.rendezvous import build_sources, merge_shard_tables +from modelexpress.refit.reshard.transfer_plan import plan_transfer +from modelexpress.refit.reshard.types import CaptureResult, RecordedCopy +from modelexpress_rl.train.engines.megatron import ( + MegatronTensorSpec, + build_hf_aliases, +) + +HF_NAMES = ("q_proj.weight", "k_proj.weight", "v_proj.weight") + + +def _global_fused(q_heads: int, kv_heads: int, head_dim: int, hidden: int): + rows = (q_heads + 2 * kv_heads) * head_dim + return torch.arange(rows * hidden, dtype=torch.float32).reshape(rows, hidden) + + +def _global_extras(q_heads: int, kv_heads: int, head_dim: int): + return { + "qkv_interleave": "by_head", + "num_heads": str(q_heads), + "num_kv_heads": str(kv_heads), + "head_dim": str(head_dim), + } + + +def _publish_all_ranks( + q_heads: int, kv_heads: int, tp_size: int, head_dim: int, hidden: int = 3 +): + fused = _global_fused(q_heads, kv_heads, head_dim, hidden) + assert fused.shape[0] % tp_size == 0 + local_rows = fused.shape[0] // tp_size + published = [] + locals_by_agent = {} + for rank in range(tp_size): + lo = rank * local_rows + hi = lo + local_rows + local = fused[lo:hi].clone() + agent = f"tp{rank}" + locals_by_agent[agent] = local + published.extend( + build_hf_aliases( + [ + MegatronTensorSpec( + name="linear_qkv.weight", + tensor=local, + role="qkv_column", + hf_names=HF_NAMES, + global_shape=tuple(fused.shape), + placement_kind="SHARD" if tp_size > 1 else "REPLICATE", + shard_axis=0 if tp_size > 1 else None, + local_shard_range=(lo, hi) if tp_size > 1 else None, + extras=_global_extras(q_heads, kv_heads, head_dim), + ) + ], + agent_name=agent, + ) + ) + return fused, locals_by_agent, published + + +def _expected_hf(fused: torch.Tensor, q_heads: int, kv_heads: int, head_dim: int): + q_per_group = q_heads // kv_heads + q_rows = q_per_group * head_dim + group_rows = q_rows + 2 * head_dim + q, k, v = [], [], [] + for group in range(kv_heads): + group_lo = group * group_rows + q.append(fused[group_lo : group_lo + q_rows]) + k.append(fused[group_lo + q_rows : group_lo + q_rows + head_dim]) + v.append(fused[group_lo + q_rows + head_dim : group_lo + group_rows]) + return tuple(torch.cat(parts) for parts in (q, k, v)) + + +def _reconstruct(published, locals_by_agent): + by_name = {name: [] for name in HF_NAMES} + full_shapes = {} + for tensor in published: + by_name[tensor.name].extend(tensor.shards) + full_shapes[tensor.name] = tensor.full_shape + + actual = [] + for name in HF_NAMES: + assert name in full_shapes + sample = next(iter(locals_by_agent.values())) + destination = torch.empty( + full_shapes[name], dtype=sample.dtype, device=sample.device + ) + coverage = torch.zeros(full_shapes[name][0], dtype=torch.int32) + for shard in by_name[name]: + local = locals_by_agent[shard.agent_name] + row_bytes = local.shape[1] * local.element_size() + byte_offset = shard.addr - local.data_ptr() + assert byte_offset % row_bytes == 0 + source_lo = byte_offset // row_bytes + rows = shard.shape[0] + destination_lo = shard.shard_offset[0] + destination[destination_lo : destination_lo + rows].copy_( + local[source_lo : source_lo + rows] + ) + coverage[destination_lo : destination_lo + rows] += 1 + assert torch.all(coverage == 1), (name, coverage) + actual.append(destination) + + for agent, local in locals_by_agent.items(): + source_coverage = torch.zeros(local.shape[0], dtype=torch.int32) + for shards in by_name.values(): + for shard in shards: + if shard.agent_name != agent: + continue + row_bytes = local.shape[1] * local.element_size() + source_lo = (shard.addr - local.data_ptr()) // row_bytes + source_coverage[source_lo : source_lo + shard.shape[0]] += 1 + assert torch.all(source_coverage == 1), (agent, source_coverage) + return tuple(actual) + + +def _tables_by_agent(published): + tables = {} + for tensor in published: + assert tensor.shards + agent = tensor.shards[0].agent_name + assert all(shard.agent_name == agent for shard in tensor.shards) + tables.setdefault(agent, []).append(tensor) + return [tables[name] for name in sorted(tables)] + + +def _full_copy(name: str, shape: tuple[int, int]) -> RecordedCopy: + return RecordedCopy( + src_name=name, + op_chain=(), + param_name=name, + dest_offset=0, + dest_shape=shape, + dest_stride=(shape[1], 1), + dest_dtype=torch.float32, + ) + + +@pytest.mark.parametrize( + ("q_heads", "kv_heads", "tp_size", "head_dim"), + [ + (32, 4, 2, 4), + (32, 4, 1, 4), + (64, 2, 8, 128), + (24, 6, 4, 2), + ], +) +def test_global_interval_aliases_cover_qkv_without_gaps_or_overlaps( + q_heads: int, kv_heads: int, tp_size: int, head_dim: int +): + fused, locals_by_agent, published = _publish_all_ranks( + q_heads, kv_heads, tp_size, head_dim + ) + + actual = _reconstruct(published, locals_by_agent) + expected = _expected_hf(fused, q_heads, kv_heads, head_dim) + + assert all( + torch.equal(got, want) for got, want in zip(actual, expected, strict=True) + ) + + +def test_kv_below_tp_only_advertises_kv_on_ranks_that_own_it(): + _, _, published = _publish_all_ranks(64, 2, 8, 128) + names_by_agent = {f"tp{rank}": set() for rank in range(8)} + for tensor in published: + for shard in tensor.shards: + names_by_agent[shard.agent_name].add(tensor.name) + + assert names_by_agent["tp0"] == {"q_proj.weight"} + assert names_by_agent["tp3"] == set(HF_NAMES) + assert names_by_agent["tp4"] == {"q_proj.weight"} + assert names_by_agent["tp7"] == set(HF_NAMES) + + +def test_two_layers_may_use_different_qkv_geometry(): + fixtures = [(64, 2, 8, 128), (32, 8, 8, 64)] + for q_heads, kv_heads, tp_size, head_dim in fixtures: + fused, locals_by_agent, published = _publish_all_ranks( + q_heads, kv_heads, tp_size, head_dim + ) + actual = _reconstruct(published, locals_by_agent) + expected = _expected_hf(fused, q_heads, kv_heads, head_dim) + assert all( + torch.equal(got, want) for got, want in zip(actual, expected, strict=True) + ) + + +def test_divisible_global_descriptors_match_legacy_aliases_byte_for_byte(): + q_heads, kv_heads, tp_size, head_dim, hidden = 32, 4, 2, 4, 3 + fused = _global_fused(q_heads, kv_heads, head_dim, hidden) + local_rows = fused.shape[0] // tp_size + for rank in range(tp_size): + lo = rank * local_rows + hi = lo + local_rows + local = fused[lo:hi].clone() + common = { + "name": "linear_qkv.weight", + "tensor": local, + "role": "qkv_column", + "hf_names": HF_NAMES, + "global_shape": tuple(fused.shape), + "placement_kind": "SHARD", + "shard_axis": 0, + "local_shard_range": (lo, hi), + } + legacy = build_hf_aliases( + [ + MegatronTensorSpec( + **common, + extras={ + "num_heads_local": str(q_heads // tp_size), + "num_kv_heads_local": str(kv_heads // tp_size), + "head_dim": str(head_dim), + }, + ) + ], + agent_name=f"tp{rank}", + ) + global_aliases = build_hf_aliases( + [ + MegatronTensorSpec( + **common, + extras=_global_extras(q_heads, kv_heads, head_dim), + ) + ], + agent_name=f"tp{rank}", + ) + assert global_aliases == legacy + + +def test_sparse_kv_tables_merge_into_a_complete_bounded_plan(): + _, _, published = _publish_all_ranks(64, 2, 8, 128) + merged = merge_shard_tables(_tables_by_agent(published)) + sources, _, _ = build_sources(merged) + capture = CaptureResult( + copies=[ + _full_copy(name, tuple(sources[name].global_shape)) for name in HF_NAMES + ] + ) + + plan = plan_transfer(capture, sources, max_segments_per_copy=1) + + assert plan.fallback == [] + assert {pull.src_name for pull in plan.full_pulls} == set(HF_NAMES) + assert plan.bytes_planned() == sum( + tensor.numel() * tensor.element_size() + for tensor in _expected_hf(_global_fused(64, 2, 128, 3), 64, 2, 128) + ) + + +def test_missing_kv_publishers_fail_closed_instead_of_silently_falling_back(): + _, _, published = _publish_all_ranks(64, 2, 8, 128) + q_only = [tensor for tensor in published if tensor.name == HF_NAMES[0]] + sources, _, _ = build_sources(merge_shard_tables(_tables_by_agent(q_only))) + capture = CaptureResult( + copies=[ + _full_copy(HF_NAMES[0], (64 * 128, 3)), + _full_copy(HF_NAMES[1], (2 * 128, 3)), + _full_copy(HF_NAMES[2], (2 * 128, 3)), + ] + ) + + plan = plan_transfer(capture, sources) + + assert set(plan.fallback) == {HF_NAMES[1], HF_NAMES[2]} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA tensors") +def test_cuda_64q_2kv_tp8_noop_and_moving_weight_refits_match_logits(): + device = torch.device("cuda", 0) + q_heads, kv_heads, tp_size, head_dim, hidden = 64, 2, 8, 128, 64 + rows = (q_heads + 2 * kv_heads) * head_dim + generator = torch.Generator(device=device).manual_seed(20260818) + base = torch.randn( + (rows, hidden), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + activation = torch.randn( + (4, hidden), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + + def refit(fused): + local_rows = fused.shape[0] // tp_size + published = [] + locals_by_agent = {} + for rank in range(tp_size): + lo = rank * local_rows + hi = lo + local_rows + agent = f"tp{rank}" + local = fused[lo:hi].clone() + locals_by_agent[agent] = local + published.extend( + build_hf_aliases( + [ + MegatronTensorSpec( + name="linear_qkv.weight", + tensor=local, + role="qkv_column", + hf_names=HF_NAMES, + global_shape=tuple(fused.shape), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(lo, hi), + extras=_global_extras(q_heads, kv_heads, head_dim), + ) + ], + agent_name=agent, + ) + ) + actual = _reconstruct(published, locals_by_agent) + expected = _expected_hf(fused, q_heads, kv_heads, head_dim) + assert all( + torch.equal(got, want) + for got, want in zip(actual, expected, strict=True) + ) + actual_logits = tuple(activation @ weight.T for weight in actual) + expected_logits = tuple(activation @ weight.T for weight in expected) + assert all( + torch.equal(got, want) + for got, want in zip(actual_logits, expected_logits, strict=True) + ) + return actual + + no_op = refit(base) + moving = refit(base + torch.tensor(0.125, dtype=base.dtype, device=device)) + + assert any( + not torch.equal(before, after) + for before, after in zip(no_op, moving, strict=True) + ) + + +@pytest.mark.parametrize( + ("extras", "global_rows", "match"), + [ + ( + {"num_heads": "64", "head_dim": "128", "qkv_interleave": "by_head"}, + 8704, + "requires both", + ), + ( + { + "num_heads": "64", + "num_kv_heads": "2", + "qkv_interleave": "by_head", + }, + 8704, + "head_dim", + ), + (_global_extras(64, 2, 128), 8192, "rows disagree"), + ( + {**_global_extras(64, 2, 128), "qkv_interleave": "unsupported"}, + 8704, + "qkv_interleave", + ), + (_global_extras(63, 2, 128), 8576, "invalid global"), + ], +) +def test_unrecoverable_global_geometry_fails_closed(extras, global_rows, match): + assert global_rows % 8 == 0 + local = torch.zeros(global_rows // 8, 3) + with pytest.raises(ValueError, match=match): + build_hf_aliases( + [ + MegatronTensorSpec( + name="linear_qkv.weight", + tensor=local, + role="qkv_column", + hf_names=HF_NAMES, + global_shape=(global_rows, 3), + placement_kind="SHARD", + shard_axis=0, + local_shard_range=(0, local.shape[0]), + extras=extras, + ) + ], + agent_name="tp0", + ) diff --git a/modelexpress_client/python/tests/test_reshard_refit_batch_install.py b/modelexpress_client/python/tests/test_reshard_refit_batch_install.py new file mode 100644 index 000000000..ab4e28059 --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_refit_batch_install.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Batched vs per-view re-slice in ReshardReceiver.update_weights - CPU, no NIXL. + +A full-pulled source is staged whole and then re-sliced locally into the receive +buffers, one copy per view the loader recorded. On a real model that is thousands +of views, and thousands of individual ``copy_()`` launches cost enough Python and +launch overhead to rival the RDMA they follow. The copies are now collected and +issued as a single ``torch._foreach_copy_``. + +The batching is only safe because the destinations are disjoint, so the tests +that matter are the ones that would catch a wrong batch: + + * every receive buffer is byte-identical between the batched and per-view + paths, over a plan with several sources and several views each; + * the values are the ones the op chains actually describe, so two paths that + are wrong in the same way still fail; + * overlapping views, where batching would be order-dependent, are not silently + reordered. + +Run: pytest tests/test_reshard_refit_batch_install.py +""" + +import torch + +from modelexpress.refit.reshard.receiver import ReshardReceiver +from modelexpress.refit.reshard.slice_plan import PullSegment +from modelexpress.refit.reshard.transfer_plan import FullPullSource, TransferPlan +from modelexpress.refit.reshard.transport import InMemoryReferenceTransport + +EL = 4 # float32 element size + + +class _Copy: + """The subset of ``RecordedCopy`` the re-slice reads.""" + + def __init__(self, *, param_name, op_chain, dest_shape, dest_stride, dest_offset): + self.param_name = param_name + self.op_chain = op_chain + self.dest_shape = dest_shape + self.dest_stride = dest_stride + self.dest_offset = dest_offset + + +class _Harness(ReshardReceiver): + """A receiver with the plan and buffers set directly. + + ``ReshardReceiver.__init__`` builds a NIXL agent and a metadata client, so + this bypasses it. Only the state ``update_weights`` touches is populated. + """ + + def __init__(self, transport) -> None: + self._device = torch.device("cpu") + self._timeout = 1.0 + self._transport = transport + self._global_rank = 0 + self._cached_descriptors = None + + def _install(self, recv_buffers) -> None: + pass + + +def _build(transport): + """Two full-pulled sources, four views in total. + + ``qkv`` is a 4x6 source re-sliced into three column blocks, the shape a + fused QKV projection produces under tensor parallelism. ``rows`` is a 4x4 + source re-sliced into two row blocks. Four views is enough that the batched + path passes a real list to ``_foreach_copy_`` and a mis-paired + destination/source shows up as wrong bytes. Returns the harness plus the + source tensors, which the caller must keep alive: the plan holds their raw + addresses. + """ + harness = _Harness(transport) + + qkv_src = torch.arange(24, dtype=torch.float32).reshape(4, 6) + rows_src = torch.arange(100, 116, dtype=torch.float32).reshape(4, 4) + + recv_q = torch.zeros(4, 2, dtype=torch.float32) + recv_k = torch.zeros(4, 2, dtype=torch.float32) + recv_v = torch.zeros(4, 2, dtype=torch.float32) + recv_rows = torch.zeros(4, 4, dtype=torch.float32) + qkv_staging = torch.zeros(4, 6, dtype=torch.float32) + rows_staging = torch.zeros(4, 4, dtype=torch.float32) + + harness._recv_buffers = { + "q": recv_q, + "k": recv_k, + "v": recv_v, + "rows": recv_rows, + } + harness._param_ptr = {} + harness._full_staging = {"qkv": qkv_staging, "rows": rows_staging} + harness._full_staging_ptr = { + "qkv": qkv_staging.data_ptr(), + "rows": rows_staging.data_ptr(), + } + harness._staging = {} + harness._staging_ptr = {} + + plan = TransferPlan( + segments=[], + full_pulls=[ + FullPullSource( + src_name="qkv", + global_shape=(4, 6), + dtype=torch.float32, + elsize=EL, + segments=[ + PullSegment( + session="s0", + src_addr=qkv_src.data_ptr(), + dst_byte=0, + nbytes=24 * EL, + param_name="qkv", + ) + ], + copies=[ + _Copy( + param_name=name, + op_chain=(("narrow", (1, start, 2), ()),), + dest_shape=(4, 2), + dest_stride=(2, 1), + dest_offset=0, + ) + for name, start in (("q", 0), ("k", 2), ("v", 4)) + ], + ), + FullPullSource( + src_name="rows", + global_shape=(4, 4), + dtype=torch.float32, + elsize=EL, + segments=[ + PullSegment( + session="s1", + src_addr=rows_src.data_ptr(), + dst_byte=0, + nbytes=16 * EL, + param_name="rows", + ) + ], + # Two row blocks into one buffer, so the destination offset is + # exercised as well as the shape. + copies=[ + _Copy( + param_name="rows", + op_chain=(("narrow", (0, start, 2), ()),), + dest_shape=(2, 4), + dest_stride=(4, 1), + dest_offset=start * 4, + ) + for start in (0, 2) + ], + ), + ], + converts=[], + exact_bytes=0, + exact_descriptor_count=0, + ) + harness._plan = plan + keepalive = (qkv_src, rows_src) + return harness, keepalive + + +def _run(monkeypatch, *, batched: bool): + monkeypatch.setenv("MX_RESHARD_BATCH_INSTALL", "1" if batched else "0") + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + transport = InMemoryReferenceTransport() + harness, keepalive = _build(transport) + metrics = harness.update_weights(step=1) + return harness, metrics, keepalive + + +def test_batched_and_per_view_produce_identical_buffers(monkeypatch): + """The correctness gate on batching: same bytes, same destinations.""" + batched, batched_metrics, _k = _run(monkeypatch, batched=True) + per_view, per_view_metrics, _k = _run(monkeypatch, batched=False) + + assert batched._recv_buffers.keys() == per_view._recv_buffers.keys() + for name, buffer in batched._recv_buffers.items(): + assert torch.equal(buffer, per_view._recv_buffers[name]), name + + # And the accounting agrees, so a dashboard cannot tell the paths apart. + assert batched_metrics["bytes_received"] == per_view_metrics["bytes_received"] + assert batched_metrics["segments"] == per_view_metrics["segments"] + + +def test_reconstructs_the_expected_values(monkeypatch): + """Buffer parity between two wrong paths would still pass, so check values.""" + qkv = torch.arange(24, dtype=torch.float32).reshape(4, 6) + rows = torch.arange(100, 116, dtype=torch.float32).reshape(4, 4) + + for batched in (True, False): + harness, _m, _k = _run(monkeypatch, batched=batched) + for name, start in (("q", 0), ("k", 2), ("v", 4)): + assert torch.equal( + harness._recv_buffers[name], qkv[:, start : start + 2] + ), f"{name} batched={batched}" + assert torch.equal(harness._recv_buffers["rows"], rows), f"batched={batched}" + + +def test_disjoint_destinations_use_the_foreach_batch(monkeypatch): + """The normal plan is pairwise disjoint and stays on the batched path.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + monkeypatch.setenv("MX_RESHARD_BATCH_INSTALL", "1") + calls = [] + real_foreach = torch._foreach_copy_ + + def record(destinations, sources): + calls.append((destinations, sources)) + return real_foreach(destinations, sources) + + monkeypatch.setattr(torch, "_foreach_copy_", record) + harness, keepalive = _build(InMemoryReferenceTransport()) + harness.update_weights(step=1) + + assert len(calls) == 1 + assert len(calls[0][0]) == 5 + assert all(t.data_ptr() for t in keepalive) + + +def test_overlapping_destinations_fall_back_to_plan_order(monkeypatch): + """Overlapping views use sequential copies because foreach order is undefined.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + monkeypatch.setenv("MX_RESHARD_BATCH_INSTALL", "1") + monkeypatch.setattr( + torch, + "_foreach_copy_", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("overlapping destinations must not be batched") + ), + ) + harness, keepalive = _build(InMemoryReferenceTransport()) + for copy in harness._plan.full_pulls[1].copies: + copy.dest_offset = 0 + + harness.update_weights(step=1) + + expected = torch.arange(108, 116, dtype=torch.float32).reshape(2, 4) + assert torch.equal(harness._recv_buffers["rows"][:2], expected) + assert all(t.data_ptr() for t in keepalive) + + +def test_stage_record_reports_the_install_arm(monkeypatch, caplog): + """A captured record must say which arm produced it, or it is unattributable.""" + import json + import logging + + monkeypatch.setenv("MX_REFIT_STAGE_RECORD", "1") + for batched in (True, False): + caplog.clear() + with caplog.at_level(logging.WARNING): + _run(monkeypatch, batched=batched) + records = [ + json.loads(message.split("MX_REFIT_STAGE ", 1)[1]) + for message in caplog.messages + if "MX_REFIT_STAGE " in message + ] + assert records, f"no stage record emitted (batched={batched})" + assert records[-1]["batch_install"] is batched + # The re-slice is attributed either way, so an A/B can be compared. + assert "reslice_s" in records[-1] + # Views, not sources: five views over two sources. Both arms must agree, + # or the launch count batching removes cannot be read off the records. + assert records[-1]["reslice_copies"] == 5 + assert records[-1]["full_pull_sources"] == 2 + + +def test_empty_full_pulls_is_a_no_op(monkeypatch): + """No full pulls means no batched copy, and no crash on an empty list.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + + for batched in (True, False): + monkeypatch.setenv("MX_RESHARD_BATCH_INSTALL", "1" if batched else "0") + transport = InMemoryReferenceTransport() + harness, keepalive = _build(transport) + harness._plan.full_pulls = [] + + metrics = harness.update_weights(step=1) + + assert metrics["full_pull_sources"] == 0 + assert all(t.data_ptr() for t in keepalive) # keep alive diff --git a/modelexpress_client/python/tests/test_reshard_refit_descriptor_cache.py b/modelexpress_client/python/tests/test_reshard_refit_descriptor_cache.py new file mode 100644 index 000000000..e5976058e --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_refit_descriptor_cache.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Reusing the read descriptors across steps - CPU only, no NIXL. + +A read descriptor is a (session, src_addr, dst_addr, nbytes) tuple derived from +the transfer plan and the registered buffer addresses. Neither changes between +steps, so building the lists once per plan rather than once per step removes work +that scales with the descriptor count: a real MoE refit issues hundreds of +thousands of them, and re-deriving that list in Python cost more than the local +re-slice it precedes. + +Caching addresses is only safe while those addresses are still the plan's, so the +tests that matter are the ones that would catch a stale cache: + + * a second refit through the cache moves the same bytes as the first, and the + same bytes an uncached refit moves; + * rebuilding the plan drops the cache, so a refit never RDMAs into the previous + plan's addresses; + * switching the fused/phased arm rebuilds rather than reusing, because the + phased arm does not build the exact descriptors at all; + * the build is reported as its own stage, so it cannot silently return to being + unattributed time. + +Run: pytest tests/test_reshard_refit_descriptor_cache.py +""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from modelexpress.refit.reshard import receiver as receiver_mod +from modelexpress.refit.reshard.transfer_plan import TransferPlan +from modelexpress.refit.reshard.types import CaptureResult +from tests.test_reshard_refit_fused_wire import _build, _RecordingTransport + + +@pytest.fixture(autouse=True) +def _cpu_only(monkeypatch): + """These refits run on CPU, where the stage syncs have nothing to wait on.""" + monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + + +def _refit(harness, step): + """One refit, returning the metrics and a copy of every receive buffer.""" + metrics = harness.update_weights(step=step) + buffers = { + name: buffer.detach().clone() + for name, buffer in harness._recv_buffers.items() + } + return metrics, buffers + + +def _zero(harness): + for buffer in harness._recv_buffers.values(): + buffer.zero_() + + +def test_cached_and_uncached_refits_move_identical_bytes(monkeypatch): + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "0") + uncached_harness, uncached_keepalive = _build(_RecordingTransport()) + _, uncached = _refit(uncached_harness, 1) + + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "1") + cached_harness, cached_keepalive = _build(_RecordingTransport()) + _, cached = _refit(cached_harness, 1) + + assert set(cached) == set(uncached) + for name in cached: + assert torch.equal(cached[name], uncached[name]), name + assert uncached_keepalive and cached_keepalive + + +def test_second_refit_through_the_cache_repeats_the_first(monkeypatch): + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "1") + harness, keepalive = _build(_RecordingTransport()) + + first_metrics, first = _refit(harness, 1) + # Zeroed so a cache that quietly moved nothing the second time cannot pass by + # leaving the first refit's bytes in place. + _zero(harness) + second_metrics, second = _refit(harness, 2) + + for name in first: + assert torch.equal(second[name], first[name]), name + assert second_metrics["bytes_received"] == first_metrics["bytes_received"] + assert second_metrics["segments"] == first_metrics["segments"] + assert keepalive + + +def test_rebuilding_the_plan_drops_the_cache(monkeypatch): + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "1") + harness, keepalive = _build(_RecordingTransport()) + _refit(harness, 1) + assert harness._cached_descriptors is not None + + new_plan = TransferPlan() + monkeypatch.setattr( + receiver_mod, + "gather_sources", + lambda *_args, **_kwargs: ({}, {}, {}, {}), + ) + monkeypatch.setattr(receiver_mod, "plan_transfer", lambda *_args: new_plan) + monkeypatch.setattr( + receiver_mod, "handshake_endpoints_for_plan", lambda *_args: {} + ) + monkeypatch.setattr(receiver_mod, "handshake_with_peers", lambda *_args: None) + monkeypatch.setattr( + receiver_mod, "NixlReshardTransport", lambda *_args, **_kwargs: object() + ) + monkeypatch.setattr(receiver_mod, "classic_cuda_alloc", nullcontext) + harness._num_trainer_sources = 0 + harness._mx_client = object() + harness._model_name = "model" + harness._global_rank = 0 + harness._manager = SimpleNamespace(register_tensors=lambda *_args: None) + harness._capture = lambda _manifest: (CaptureResult(), {}) + harness._log_coverage = lambda *_args: None + + harness._prepare(timeout=1.0) + + assert harness._plan is new_plan + assert harness._cached_descriptors is None + assert keepalive + + +def test_switching_the_wire_arm_rebuilds_the_cache(monkeypatch): + """A cache filled under one wire arm must not be served to the other. + + The phased arm hands the plan to execute_transfer instead of building exact + descriptors, so its cache entry holds None for them. Serving that to the fused + arm would read the exact segments not at all rather than into the wrong place, + which is the quiet kind of wrong: fewer bytes, no error. + """ + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "1") + monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "0") + harness, keepalive = _build(_RecordingTransport()) + phased_metrics, phased = _refit(harness, 1) + assert harness._cached_descriptors[0] is False + + monkeypatch.setenv("MX_RESHARD_FUSED_WIRE", "1") + _zero(harness) + fused_metrics, fused = _refit(harness, 2) + assert harness._cached_descriptors[0] is True + + for name in phased: + assert torch.equal(fused[name], phased[name]), name + assert fused_metrics["bytes_received"] == phased_metrics["bytes_received"] + assert keepalive + + +def test_descriptor_build_is_reported_as_a_stage(monkeypatch): + monkeypatch.setenv("MX_RESHARD_CACHE_DESCRIPTORS", "1") + harness, keepalive = _build(_RecordingTransport()) + metrics, _ = _refit(harness, 1) + # Present and non-negative rather than above a threshold: this asserts the + # stage is accounted for, and a timing floor would be a flake on shared CI. + assert "descriptor_build_s" in metrics + assert metrics["descriptor_build_s"] >= 0.0 + assert keepalive 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..c754571e4 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py +++ b/modelexpress_client/python/tests/test_reshard_refit_fused_wire.py @@ -61,6 +61,7 @@ def __init__(self, transport) -> None: # noqa: D107 - see class docstring self._timeout = 1.0 self._transport = transport self._global_rank = 0 + self._cached_descriptors = None self._install_order: list[str] = [] def _install(self, recv_buffers) -> None: diff --git a/modelexpress_client/python/tests/test_reshard_refit_geometry.py b/modelexpress_client/python/tests/test_reshard_refit_geometry.py index 684e4a65a..f93b543e8 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_geometry.py +++ b/modelexpress_client/python/tests/test_reshard_refit_geometry.py @@ -124,6 +124,55 @@ def test_unsupported_op_falls_back_per_source(): assert by_src["q"].dest_offset == 0 and by_src["norm"].op_chain == () +def test_unsupported_source_records_the_op_that_defeated_capture(): + """The count alone cannot distinguish an unexpressible fused layout from a + loader that merely touched one op outside the allowlist, so keep the cause.""" + with torch.device("meta"): + model = ToyModel(with_bad=True) + result = capture_geometry(model, _manifest(with_bad=True)) + + assert set(result.unsupported_reasons) == {"bad"} + reason = result.unsupported_reasons["bad"] + assert "unsupported op" in reason + assert "aten.mul" in reason + # The offending source and its op-chain stay in the message, which is what + # makes a single failure actionable without a re-run. + assert "'bad'" in reason + + +def test_summarize_unsupported_groups_one_cause_across_many_sources(): + """Thousands of sources failing for one reason must read as one cause. Each + message embeds its own source name, so grouping has to ignore that tail.""" + from modelexpress.refit.reshard.types import summarize_unsupported + + reasons = { + f"model.layers.0.mlp.experts.{i}.gate_proj.weight": ( + f"unsupported op aten.index_copy_ on lazy " + f"'model.layers.0.mlp.experts.{i}.gate_proj.weight' (chain=());" + ) + for i in range(128) + } + reasons["odd"] = "unsupported op aten.mul on lazy 'odd' (chain=());" + + assert summarize_unsupported(reasons) == [ + ("unsupported op aten.index_copy_", 128), + ("unsupported op aten.mul", 1), + ] + + +def test_summarize_unsupported_accepts_none_for_all_causes(): + from modelexpress.refit.reshard.types import summarize_unsupported + + reasons = { + f"source-{index}": f"cause-{index} on lazy 'source-{index}'" + for index in range(4) + } + + assert summarize_unsupported(reasons, limit=None) == [ + (f"cause-{index}", 1) for index in range(4) + ] + + def test_capture_feeds_slice_plan(): """Compose capture -> slice-plan: real captured copies drive plan_pull. The row-parallel source (full [4,8], need cols [0:4]) is strided -> 4 runs diff --git a/modelexpress_client/python/tests/test_reshard_refit_moe_experts.py b/modelexpress_client/python/tests/test_reshard_refit_moe_experts.py new file mode 100644 index 000000000..898b14701 --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_refit_moe_experts.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Fused grouped-expert (MoE) capture, mirroring vLLM's real expert loader. + +A trainer publishes MoE experts in per-expert HF form +(``...experts..gate_proj.weight``) while vLLM's destination is a fused +grouped-expert buffer (``...experts.w13_weight`` of shape +``[local_experts, 2 * inter_per_tp, hidden]``). vLLM unifies its fused and +per-expert loader paths with:: + + experts_shard = loaded_weight.unsqueeze(0) # per-expert source + loaded_experts = experts_shard.unbind() + for expert_id, loaded_expert in enumerate(loaded_experts, start=start): + param.weight_loader(param, loaded_expert, ..., expert_id=expert_id) + +``unbind`` is a pure multi-return view, but until it was allowlisted every +expert source in the model was classified unsupported and the refit failed +closed at ~5% coverage (48 layers x 128 experts x 3 projections = 18432 sources +on Qwen3-30B-A3B). Note the ``unsqueeze(0)`` / ``unbind()`` pair cancels, so the +resolved view is rank-preserving and the plan stays an axis-aligned box. + +Runs on CPU/meta in any torch env: no GPU, no vLLM. +""" + +import torch + +from modelexpress.refit.reshard.geometry import capture_geometry +from modelexpress.refit.reshard.slice_plan import Shard, plan_pull + +HIDDEN = 4 +INTER = 8 # per-expert intermediate size, full (unsharded) +TP_SIZE = 2 +TP_RANK = 0 +INTER_PER_TP = INTER // TP_SIZE +LOCAL_EXPERTS = 2 + + +class FusedMoEModel(torch.nn.Module): + """Destination holds experts fused: ``w13`` stacks gate (w1) then up (w3).""" + + def __init__(self): + super().__init__() + self.w13_weight = torch.nn.Parameter( + torch.empty(LOCAL_EXPERTS, 2 * INTER_PER_TP, HIDDEN) + ) + self.w13_weight.weight_loader = self._expert_loader + + def _expert_loader(self, *, param, loaded_weight, shard_id, expert_id): + """Mirrors vLLM's ``_load_w13``: pick this rank's slice of the source, then + the member half (w1 low, w3 high) of the stacked destination slot.""" + expert_data = param.data[expert_id] + member = 0 if shard_id == "w1" else 1 + dest = expert_data.narrow(0, member * INTER_PER_TP, INTER_PER_TP) + src = loaded_weight.narrow(0, TP_RANK * INTER_PER_TP, INTER_PER_TP) + dest.copy_(src) + + def load_weights(self, weights): + for name, loaded in weights: + expert_id = int(name.split(".experts.")[1].split(".")[0]) + shard_id = "w1" if "gate_proj" in name else "w3" + # The vLLM path under test: add a dummy expert dim, then unbind it. + experts_shard = loaded.unsqueeze(0) + for local_id, loaded_expert in enumerate( + experts_shard.unbind(), start=expert_id + ): + # Invoked entirely by keyword, exactly as vLLM's RoutedExperts + # loader does. A capture stamp that names its first parameter + # positionally raises TypeError here. + self.w13_weight.weight_loader( + param=self.w13_weight, + loaded_weight=loaded_expert, + shard_id=shard_id, + expert_id=local_id, + ) + + +def _name(expert: int, proj: str) -> str: + return f"model.layers.0.mlp.experts.{expert}.{proj}.weight" + + +def _manifest(): + return [ + (_name(e, proj), torch.float32, [INTER, HIDDEN]) + for e in range(LOCAL_EXPERTS) + for proj in ("gate_proj", "up_proj") + ] + + +def test_per_expert_sources_capture_into_fused_destination(): + """Every expert source is captured; none lands in ``unsupported``.""" + with torch.device("meta"): + model = FusedMoEModel() + result = capture_geometry(model, _manifest()) + + assert result.unsupported == [] + assert result.unsupported_reasons == {} + assert result.unattributed == 0 + # One copy per (expert, projection): both halves of every expert's w13 slot. + assert len(result.copies) == 2 * LOCAL_EXPERTS + + by_src = {c.src_name: c for c in result.copies} + assert set(by_src) == {n for n, _, _ in _manifest()} + for copy in result.copies: + assert copy.param_name == "w13_weight" + assert copy.dest_shape == (INTER_PER_TP, HIDDEN) + assert ("unbind", (), ()) in copy.op_chain + + +def test_unbind_is_recorded_and_does_not_change_rank(): + """``unsqueeze(0)`` then ``unbind()`` cancel, so the recorded chain still + resolves to a rank-preserving box. This is why an allowlist entry suffices + and no rank-collapse support is needed in the slice arithmetic.""" + with torch.device("meta"): + model = FusedMoEModel() + result = capture_geometry(model, _manifest()) + + copy = next(c for c in result.copies if c.src_name == _name(0, "gate_proj")) + ops = [name for name, _args, _kw in copy.op_chain] + assert ops[:3] == ["unsqueeze", "unbind", "__getitem__"] + assert "narrow" in ops + + +def test_expert_copies_plan_to_the_right_source_bytes(): + """The captured geometry must resolve to this TP rank's half of each expert, + proving the chain is a real box and not merely accepted by the allowlist.""" + with torch.device("meta"): + model = FusedMoEModel() + result = capture_geometry(model, _manifest()) + copy = next(c for c in result.copies if c.src_name == _name(1, "up_proj")) + + # One publisher offering the whole per-expert source, contiguous row-major. + shard = Shard( + shard_offset=(0, 0), + shape=(INTER, HIDDEN), + session="pub0", + addr=0, + elsize=4, + ) + segments = plan_pull( + copy, + global_shape=(INTER, HIDDEN), + src_dtype=torch.float32, + elsize=4, + shards=[shard], + ) + + assert segments, "expert copy produced no pull segments" + pulled = sum(s.nbytes for s in segments) + # Exactly this rank's slice: INTER_PER_TP rows of HIDDEN float32 elements. + assert pulled == INTER_PER_TP * HIDDEN * 4 + # Rank 0 reads from the start of the source. + assert min(s.src_addr for s in segments) == 0 + # It lands in expert 1's w3 (upper) half of the fused destination. + expected_dst = ( + 1 * (2 * INTER_PER_TP) * HIDDEN + INTER_PER_TP * HIDDEN + ) * 4 + assert min(s.dst_byte for s in segments) == expected_dst diff --git a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py b/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py index 8b330f7af..cde26a790 100644 --- a/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py +++ b/modelexpress_client/python/tests/test_reshard_refit_rendezvous.py @@ -291,3 +291,152 @@ def publish_metadata(self, *_args, **_kwargs): with pytest.raises(ValueError, match="must be positive"): rendezvous.publish(b"registered") + + +def test_quorum_can_skip_shard_tables_without_losing_emptiness(): + """The per-step quorum check needs each rank's version stamp, not its shard + table, and rebuilding the table dominated the call. Skipping it must not make + a rank that published nothing look like a valid member of the quorum, which is + the one thing the emptiness rule exists to prevent.""" + client = _DiscoveryClient( + [_blob("empty-rank", []), _blob("real-rank", _one_tensor())] + ) + + with pytest.raises(TimeoutError) as excinfo: + _rendezvous(client).discover_trainers( + expected_trainers=2, timeout=0, with_tensors=False + ) + + message = str(excinfo.value) + assert "1 with a non-empty shard table" in message + assert "1 empty" in message + + +def test_skipping_shard_tables_still_reports_the_entry_count(caplog): + """``entry_count`` is what keeps emptiness decidable, and it is also the figure + that showed this cost scales with source count rather than bytes moved.""" + import json + import logging + + client = _DiscoveryClient([_blob("rank-0", _one_tensor() * 3)]) + + with caplog.at_level(logging.WARNING): + (payload,) = _rendezvous(client).discover_trainers( + expected_trainers=1, timeout=0, with_tensors=False + ) + + assert payload.tensors == [] + assert payload.entry_count() == 3 + assert payload.agent_name == "rank-0" + record = next( + json.loads(message.split("MX_DISCOVER_COST ", 1)[1]) + for message in caplog.messages + if "MX_DISCOVER_COST " in message + ) + assert record["rank"] == 0 + assert record["tensors"] == 3 + assert record["tables_built"] is False + + +def test_entry_count_falls_back_to_the_table_when_not_recorded(): + """Payloads built directly, as tests and older callers do, record no count.""" + from modelexpress.refit.reshard.rendezvous import RendezvousPayload + + payload = RendezvousPayload(b"", "a", "", _one_tensor()) + assert payload.tensor_count is None + assert payload.entry_count() == 1 + + +def test_decoding_parsed_entries_matches_decoding_the_blob(): + """The shard table used to be re-serialized only to be parsed again. Removing + that round-trip must not change a single decoded field.""" + import json + + from modelexpress.refit.reshard.rendezvous import ( + _SCHEMA, + decode_shard_entries, + decode_shard_table, + encode_shard_table, + ) + + blob = encode_shard_table(_one_tensor()) + entries = json.loads(blob.decode("utf-8"))["tensors"] + + assert decode_shard_entries(entries) == decode_shard_table(blob) + assert decode_shard_table(blob) == _one_tensor() + assert json.loads(blob.decode("utf-8"))["schema"] == _SCHEMA + + +def test_metadata_fetches_stay_serial(): + """Concurrency here was measured to be slower, not faster: from a thread pool + the fetch went 4.02 s -> ~6.9 s median on 16 sources, because every receiver + rank runs this loop and the single metadata server, not this process, is the + contended resource. Pinned with a barrier that a concurrent implementation + would satisfy and a serial one cannot, so the decision cannot be quietly + reversed without this failing.""" + from threading import Barrier, BrokenBarrierError + + ranks = 4 + barrier = Barrier(ranks, timeout=0.5) + overlapped = [] + + class Barriered(_DiscoveryClient): + def get_metadata(self, source_id, worker_id): + try: + barrier.wait() + overlapped.append(source_id) + except BrokenBarrierError: + pass + return super().get_metadata(source_id, worker_id) + + client = Barriered([_blob(f"rank-{i}", _one_tensor()) for i in range(ranks)]) + + discovered = _rendezvous(client).discover_trainers( + expected_trainers=ranks, timeout=0 + ) + + assert len(discovered) == ranks + assert overlapped == [], "fetches overlapped; this path is deliberately serial" + + +def test_quorum_membership_does_not_depend_on_completion_order(): + """Concurrency must not make which ranks satisfy the quorum depend on who + answers first. With more READY sources than needed, the same prefix has to win + every time, or two receivers can disagree about the source set they read.""" + import time as _time + + class Reordering(_DiscoveryClient): + def get_metadata(self, source_id, worker_id): + # Earlier ranks answer last, inverting completion order. + index = int(source_id.rsplit("-", 1)[1]) + _time.sleep(0.02 * (4 - index)) + return super().get_metadata(source_id, worker_id) + + client = Reordering([_blob(f"rank-{i}", _one_tensor()) for i in range(4)]) + + discovered = _rendezvous(client).discover_trainers(expected_trainers=2, timeout=0) + + assert [p.agent_name for p in discovered] == ["rank-0", "rank-1"] + + +def test_one_unreadable_rank_does_not_abort_the_sweep(): + """The poll loop's value is reporting how many ranks are readable. A single + rank's transport error must therefore be counted, not raised: otherwise a + transient failure on one rank surfaces as a discovery crash with no count.""" + + class OneBadRank(_DiscoveryClient): + def get_metadata(self, source_id, worker_id): + if source_id.endswith("-1"): + raise RuntimeError("transport blip on rank 1") + return super().get_metadata(source_id, worker_id) + + client = OneBadRank([_blob(f"rank-{i}", _one_tensor()) for i in range(3)]) + + with pytest.raises(TimeoutError) as excinfo: + _rendezvous(client).discover_trainers(expected_trainers=3, timeout=0) + assert "3 READY source(s)" in str(excinfo.value) + assert "2 with a non-empty shard table" in str(excinfo.value) + + # The readable ranks are still returned when the quorum only needs them. + discovered = _rendezvous(client).discover_trainers(expected_trainers=2, timeout=0) + assert [p.agent_name for p in discovered] == ["rank-0", "rank-2"] diff --git a/modelexpress_client/python/tests/test_reshard_refit_replica_merge.py b/modelexpress_client/python/tests/test_reshard_refit_replica_merge.py new file mode 100644 index 000000000..c46e2a7f4 --- /dev/null +++ b/modelexpress_client/python/tests/test_reshard_refit_replica_merge.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Replica dedup in merge_shard_tables - pure data, no torch. + +Under DP/EDP the same geometric shard is published by several trainer ranks, at +distinct addresses on distinct NICs, holding byte-identical bytes. The merge used +to retain every one of those offers, which is a read per replica into the same +destination: correct bytes, but wire and descriptor count scaling with the DP +degree, plus a P2P handshake per redundant owner. It now retains exactly one +representative per ``(shard_offset, shape)``. + +The tests that matter are the ones that would catch a wrong merge: + + * duplicates collapse, and the resulting plan reads each region once; + * genuine cross-rank fan-in (distinct geometry) is still fully retained, for + every source - collapsing that would silently drop half the model; + * the default is the previous first-writer-wins choice. + +Run: pytest tests/test_reshard_refit_replica_merge.py +""" + +import pytest + +from modelexpress.refit.reshard.rendezvous import ( + PublishedShard, + PublishedTensor, + build_sources, + merge_shard_tables, +) +from modelexpress.refit.reshard.slice_plan import plan_pull +from modelexpress.refit.reshard.types import RecordedCopy + +EL = 2 # bytes per bfloat16 element + + +def _dp_replica_tables(replicas: int, shards_per_rank: int = 2): + """``replicas`` publishers each advertising the same geometries. + + This is the DP8 / EDP2 shape: byte-identical copies of a shard on distinct + ranks, at distinct addresses. + """ + return [ + [ + PublishedTensor( + name="weight", + dtype="torch.bfloat16", + elsize=EL, + full_shape=(8, 4), + shards=[ + PublishedShard( + agent_name=f"dp{dp}", + device_id=dp, + addr=1000 * (dp + 1) + 10 * i, + shard_offset=(4 * i, 0), + shape=(4, 4), + ) + for i in range(shards_per_rank) + ], + ) + ] + for dp in range(replicas) + ] + + +def _fan_in_tables(ranks: int = 4): + """``ranks`` publishers each owning a distinct row block: real fan-in.""" + return [ + [ + PublishedTensor( + name="weight", + dtype="torch.bfloat16", + elsize=EL, + full_shape=(16, 4), + shards=[ + PublishedShard( + agent_name=f"r{rank}", + device_id=rank, + addr=1000 * (rank + 1), + shard_offset=(4 * rank, 0), + shape=(4, 4), + ) + ], + ) + ] + for rank in range(ranks) + ] + + +def test_duplicate_replicas_collapse_to_one_owner(): + """The core of the change: eight offers of a geometry become one.""" + merged = merge_shard_tables(_dp_replica_tables(replicas=8)) + + assert len(merged) == 1 + # Two geometries, not sixteen offers. + assert len(merged[0].shards) == 2 + assert [tuple(s.shard_offset) for s in merged[0].shards] == [(0, 0), (4, 0)] + + +def test_dedup_removes_the_redundant_reads_not_just_the_offers(): + """The point is wire bytes, so assert it where the reads are planned. + + Without dedup a DP8 publish plans eight reads of every region into the same + destination bytes. The segment count and byte total must instead match what a + single publisher would produce. + """ + def segments_for(replicas): + merged = merge_shard_tables(_dp_replica_tables(replicas=replicas)) + sources, _agents, _devices = build_sources(merged) + source = sources["weight"] + # build_sources resolves the published dtype string to the real dtype, and + # plan_pull refuses a cross-dtype read, so take the dest dtype from there. + copy = RecordedCopy( + src_name="weight", + op_chain=(), + param_name="weight", + dest_offset=0, + dest_shape=(8, 4), + dest_stride=(4, 1), + dest_dtype=source.dtype, + ) + return plan_pull( + copy, + global_shape=source.global_shape, + src_dtype=source.dtype, + elsize=source.elsize, + shards=source.shards, + ) + + single = segments_for(1) + assert len(single) == 2 # one contiguous run per row block + assert sum(s.nbytes for s in single) == 32 * EL # the whole 8x4 tensor, once + + for replicas in (2, 8): + spread = segments_for(replicas) + assert len(spread) == len(single), replicas + assert sum(s.nbytes for s in spread) == sum(s.nbytes for s in single), replicas + # One owner serves the tensor, so one peer needs a handshake. + assert len({s.session for s in spread}) == 1, replicas + + +def test_default_offset_takes_the_first_publisher(): + """The default must be the previous first-writer-wins choice.""" + merged = merge_shard_tables(_dp_replica_tables(replicas=4)) + assert [s.agent_name for s in merged[0].shards] == ["dp0", "dp0"] + + +def test_merge_still_fans_in_distinct_geometry_across_ranks(): + """Non-replica shards are real fan-in: collapsing them drops model bytes.""" + tables = _fan_in_tables(ranks=4) + + merged = merge_shard_tables(tables) + assert len(merged[0].shards) == 4 + assert sorted(s.agent_name for s in merged[0].shards) == [ + "r0", + "r1", + "r2", + "r3", + ] + + +def test_replicated_fan_in_keeps_every_region_once(): + """DP replication on top of fan-in: four regions, one owner each.""" + tables = [] + for _dp in range(3): + tables.extend(_fan_in_tables(ranks=4)) + + merged = merge_shard_tables(tables) + + assert len(merged[0].shards) == 4 + assert sorted(tuple(s.shard_offset) for s in merged[0].shards) == [ + (0, 0), + (4, 0), + (8, 0), + (12, 0), + ] + + +def test_inconsistent_shape_or_dtype_still_raises(): + """Preserved from before the rewrite: disagreeing publishers are a hard error.""" + tables = _dp_replica_tables(replicas=1) + conflicting = _dp_replica_tables(replicas=1) + conflicting[0][0].full_shape = (16, 4) + + with pytest.raises(ValueError, match="inconsistent shape/dtype"): + merge_shard_tables(tables + conflicting) + + +def test_distinct_tensor_names_are_independent(): + """Dedup is per name: two names sharing a geometry must both survive.""" + tables = [] + for dp in range(2): + tables.append( + [ + PublishedTensor( + name=name, + dtype="torch.bfloat16", + elsize=EL, + full_shape=(4, 4), + shards=[ + PublishedShard( + agent_name=f"dp{dp}", + device_id=dp, + addr=1000 * (dp + 1), + shard_offset=(0, 0), + shape=(4, 4), + ) + ], + ) + for name in ("a", "b") + ] + ) + + merged = merge_shard_tables(tables) + + assert sorted(t.name for t in merged) == ["a", "b"] + assert all(len(t.shards) == 1 for t in merged)