fix(reshard): make distributed refit reliable and faster

modelexpress_client/python/README.md

@@ -155,6 +155,7 @@ 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_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 |

modelexpress_client/python/modelexpress/envs.py

@@ -60,6 +60,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
@@ -228,6 +230,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

modelexpress_client/python/modelexpress/nixl_transfer.py

@@ -130,6 +130,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.
@@ -345,15 +349,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
 
@@ -456,10 +464,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
 
@@ -1170,6 +1180,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 = []

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,
     )