feat(reshard-refit): add no-gather slice-resharding weight refit#496
Open
tanushriya910 wants to merge 2 commits into
Open
feat(reshard-refit): add no-gather slice-resharding weight refit#496tanushriya910 wants to merge 2 commits into
tanushriya910 wants to merge 2 commits into
Conversation
Adds modelexpress.reshard_refit: an engine-agnostic core for slice-level trainer->inference weight transfer over NIXL/RDMA. It captures which slice of each full source tensor an engine's own weight loader reads (geometry), intersects those against the published source shards (slice_plan), and pulls exactly the needed byte segments (transfer_plan) - no all-gather and no per-model conversion specs. Includes a NIXL / in-memory transport split behind a Transport protocol, a shard-geometry rendezvous over MxClient, a classic-cudaMalloc MemPool for NIXL-registrable buffers, and an engine-agnostic ReshardReceiver with a vLLM subclass (unquantized meta-twin capture + layerwise-reload / PWAL install, CUDA-graph-safe bare-attr and MLA handling). Adds NixlTransferManager fetch_remote_and_wait for the P2P memory-registration handshake. Signed-off-by: Tanushriya Singh <tanushriyas@nvidia.com>
Contributor
WalkthroughAdds a complete slice-level weight resharding pipeline: geometry capture, shard rendezvous, byte-range planning, CUDA/NIXL transport, generic receiver orchestration, and vLLM-specific installation with quantization and MLA handling. ChangesReshard refit pipeline
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (5)
modelexpress_client/python/modelexpress/nixl_transfer.py (1)
727-737: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThird near-identical poll/timeout/error loop in this file.
execute_read_batch's wait loop duplicates the samecheck_xfer_state/timeout/sleep(0.001)pattern already present inreceive_from_source(lines 624-638) andreceive_dram_into_buffer(lines 829-847). Extracting a small shared_wait_for_xfer(handle, timeout_seconds, label)helper would remove the duplication and unify the (currently slightly different) error messages across all three call sites.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/nixl_transfer.py` around lines 727 - 737, The transfer wait logic is duplicated across execute_read_batch, receive_from_source, and receive_dram_into_buffer. Add a shared _wait_for_xfer(handle, timeout_seconds, label) helper that performs polling, timeout handling, terminal-status checks, and the existing sleep interval, then replace all three inline loops with calls that provide the appropriate operation label and preserve each method’s surrounding cleanup behavior.modelexpress_client/python/modelexpress/reshard_refit/slice_plan.py (1)
227-236: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto thezip()calls over shard geometry.
intersect(here), and thezip()s inpaired_runs(L250-251) andplan_pull's shard loop (L299), assumebox_a/box_b/sh.shard_offset/sh.shapeare always rank-consistent. That data is decoded from wire/JSON shard metadata (rendezvous.py), not locally trusted, so a malformed shard entry would silently truncate instead of raising a clear error.Proposed fix (apply to all 4 sites)
- for (a0, a1), (b0, b1) in zip(box_a, box_b): + for (a0, a1), (b0, b1) in zip(box_a, box_b, strict=True):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/reshard_refit/slice_plan.py` around lines 227 - 236, Update all four shard-geometry zip calls in intersect, paired_runs, and plan_pull to use strict=True, ensuring rank-mismatched metadata raises an error instead of silently truncating. Preserve the existing intersection and planning behavior for rank-consistent inputs.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/reshard_refit/transport/nixl.py (1)
58-67: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winInconsistent fail-fast policy:
session_to_agentraises,session_to_devicesilently defaults.Missing agent mapping raises
KeyError(Line 65-66), but a missing device mapping silently falls back to device 0 (Line 67). A wrong/missing device id causes a READ from the wrong GPU memory - silent data corruption is worse than a crash.rendezvous.build_sourcescurrently populates both maps together for every shard, so this shouldn't trigger via that path today, but the transport doesn't enforce that invariant itself.Proposed fix
- device_id = self._session_to_device.get(session, 0) + if session not in self._session_to_device: + raise KeyError(f"no remote device id registered for session {session!r}") + device_id = self._session_to_device[session]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/reshard_refit/transport/nixl.py` around lines 58 - 67, Update read in the session-group handling to require a device mapping for every session, matching the existing fail-fast behavior for _session_to_agent. Remove the default device ID fallback from _session_to_device.get and raise a clear error when the mapping is absent, while preserving the existing mapped-device path.modelexpress_client/python/modelexpress/reshard_refit/rendezvous.py (1)
53-62: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the caught exception in
_mx_version.Catching bare
Exceptionmasks bugs unrelated to a missing package (e.g. a corrupt metadata read), and both sides silently falling back to different values only manifests as a confusing rendezvous timeout later.importlib.metadata.versionraisesPackageNotFoundErrorspecifically when the package isn't installed - catch that instead.Proposed fix
from importlib.metadata import version as pkg_version try: return pkg_version("modelexpress") - except Exception: + except PackageNotFoundError: return "0.0.0"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/reshard_refit/rendezvous.py` around lines 53 - 62, Update `_mx_version` to catch only `PackageNotFoundError` from `importlib.metadata` when the `modelexpress` package is unavailable, while preserving the `"0.0.0"` fallback for that case. Allow other metadata errors to propagate instead of masking them.Source: Linters/SAST tools
modelexpress_client/python/modelexpress/reshard_refit/transfer_plan.py (1)
91-112: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse keyword args for
RecordedCopyPrefer keywords here (and in the matching call ingeometry.py) so a futureRecordedCopyfield change can’t silently swap arguments.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelexpress_client/python/modelexpress/reshard_refit/transfer_plan.py` around lines 91 - 112, Update the staging-buffer construction in the reshard planning flow to instantiate RecordedCopy with keyword arguments, explicitly naming each field currently passed positionally. Apply the same keyword-argument change to the matching RecordedCopy call in geometry.py, without changing values or behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelexpress_client/python/modelexpress/engines/vllm/reshard.py`:
- Around line 172-188: Update the bare-attribute restore loop in the snapshot
restoration flow to log an error whenever the current tensor and boot tensor
have mismatched shapes or dtypes and the copy is skipped. Include the
attribute/module identity and both tensor metadata in the diagnostic, while
preserving the existing copy-and-reattach behavior for matching tensors and
reattachment behavior afterward.
In `@modelexpress_client/python/modelexpress/nixl_transfer.py`:
- Around line 684-699: The diagnostic block in execute_read_batch currently
performs metadata checks and info-level formatting on every read. Gate the
entire check_remote_metadata and sample logging block behind
logger.isEnabledFor(logging.DEBUG), switch the output to logger.debug, and
narrow the exception handling where possible or explicitly annotate the
intentional broad catch while preserving transfer execution if diagnostics fail.
In `@modelexpress_client/python/modelexpress/reshard_refit/receiver.py`:
- Around line 165-274: The receiver currently only logs plan.fallback entries,
leaving unsupported parameters uninitialized because the promised
full-pull/loader path is never invoked. Update the preparation or update flow
around _prepare and update_weights to execute the existing full-pull/loader
mechanism for every fallback source, while preserving RDMA handling for
supported segments and ensuring fallback parameters are installed before
returning from an update.
In `@modelexpress_client/python/tests/test_reshard_refit_nixl_transport.py`:
- Line 51: Rename the unused unpacked variable a_agent to _a_agent in the
by_agent["trainer-agent-A"] assignment, preserving the existing tuple unpacking
and all other values unchanged.
---
Nitpick comments:
In `@modelexpress_client/python/modelexpress/nixl_transfer.py`:
- Around line 727-737: The transfer wait logic is duplicated across
execute_read_batch, receive_from_source, and receive_dram_into_buffer. Add a
shared _wait_for_xfer(handle, timeout_seconds, label) helper that performs
polling, timeout handling, terminal-status checks, and the existing sleep
interval, then replace all three inline loops with calls that provide the
appropriate operation label and preserve each method’s surrounding cleanup
behavior.
In `@modelexpress_client/python/modelexpress/reshard_refit/rendezvous.py`:
- Around line 53-62: Update `_mx_version` to catch only `PackageNotFoundError`
from `importlib.metadata` when the `modelexpress` package is unavailable, while
preserving the `"0.0.0"` fallback for that case. Allow other metadata errors to
propagate instead of masking them.
In `@modelexpress_client/python/modelexpress/reshard_refit/slice_plan.py`:
- Around line 227-236: Update all four shard-geometry zip calls in intersect,
paired_runs, and plan_pull to use strict=True, ensuring rank-mismatched metadata
raises an error instead of silently truncating. Preserve the existing
intersection and planning behavior for rank-consistent inputs.
In `@modelexpress_client/python/modelexpress/reshard_refit/transfer_plan.py`:
- Around line 91-112: Update the staging-buffer construction in the reshard
planning flow to instantiate RecordedCopy with keyword arguments, explicitly
naming each field currently passed positionally. Apply the same keyword-argument
change to the matching RecordedCopy call in geometry.py, without changing values
or behavior.
In `@modelexpress_client/python/modelexpress/reshard_refit/transport/nixl.py`:
- Around line 58-67: Update read in the session-group handling to require a
device mapping for every session, matching the existing fail-fast behavior for
_session_to_agent. Remove the default device ID fallback from
_session_to_device.get and raise a clear error when the mapping is absent, while
preserving the existing mapped-device path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8ab1562d-568a-4704-83df-09f1656fe659
📒 Files selected for processing (17)
modelexpress_client/python/modelexpress/engines/vllm/reshard.pymodelexpress_client/python/modelexpress/nixl_transfer.pymodelexpress_client/python/modelexpress/reshard_refit/__init__.pymodelexpress_client/python/modelexpress/reshard_refit/cuda_pool.pymodelexpress_client/python/modelexpress/reshard_refit/geometry.pymodelexpress_client/python/modelexpress/reshard_refit/receiver.pymodelexpress_client/python/modelexpress/reshard_refit/rendezvous.pymodelexpress_client/python/modelexpress/reshard_refit/slice_plan.pymodelexpress_client/python/modelexpress/reshard_refit/transfer_plan.pymodelexpress_client/python/modelexpress/reshard_refit/transport/__init__.pymodelexpress_client/python/modelexpress/reshard_refit/transport/base.pymodelexpress_client/python/modelexpress/reshard_refit/transport/nixl.pymodelexpress_client/python/modelexpress/reshard_refit/types.pymodelexpress_client/python/tests/test_reshard_refit_geometry.pymodelexpress_client/python/tests/test_reshard_refit_nixl_transport.pymodelexpress_client/python/tests/test_reshard_refit_slice_plan.pymodelexpress_client/python/tests/test_reshard_refit_transfer.py
- reshard.py: log (not silently skip) bare-attr restores whose shape/dtype changed across a refit, so stale re-attachment is visible - nixl_transfer.py: gate the per-READ diagnostic behind DEBUG (debug level) so steady-state refits don't pay check_remote_metadata + formatting; annotate the intentional broad except (BLE001) - receiver.py: raise UnsupportedReshard on non-empty plan.fallback instead of only warning - fallback params are never pulled, so warning-and-continuing silently served stale weights (full-pull path is a TODO) - test: rename unused unpack a_agent -> _a_agent (RUF059)
This was referenced Jul 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds modelexpress.reshard_refit: an engine-agnostic core for slice-level trainer->inference weight transfer over NIXL/RDMA. It captures which slice of each full source tensor an engine's own weight loader reads (geometry), intersects those against the published source shards (slice_plan), and pulls exactly the needed byte segments (transfer_plan) - no all-gather and no per-model conversion specs.
Includes a NIXL / in-memory transport split behind a Transport protocol, a shard-geometry rendezvous over MxClient, a classic-cudaMalloc MemPool for NIXL-registrable buffers, and an engine-agnostic ReshardReceiver with a vLLM subclass (unquantized meta-twin capture + layerwise-reload / PWAL install, CUDA-graph-safe bare-attr and MLA handling). Adds NixlTransferManager fetch_remote_and_wait for the P2P memory-registration handshake.
##Testing
Implemented and validated bit-exact (max_abs_diff=0) on dense Qwen3-0.6B, fp8 MoE mini-glm, and Qwen3-30B-A3B (128-expert) on nscale B200
Summary by CodeRabbit
New Features
Bug Fixes
Tests