Refit: local plan building, node-add cache key, and NIXL backend#5872
Conversation
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
|
/ok to test 14b97e0 |
svcnvidia-nemo-ci
added
Final Review
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Light review — the refactor looks correct and self-consistent (clean rename with no lingering build_centralized_reshard_plan references, valid _PlanCacheKey.world_size addition, and deterministic ordering that holds because gathered_pairs is identical across ranks and Python dicts preserve insertion order).
One non-blocking suggestion on test coverage: this PR adds a substantial NIXL backend plus a centralized→local plan-building refactor, but no new unit tests. The newly-extracted pure functions are testable without GPUs or a live process group and would directly guard the PR's core correctness claim — that sender and receiver independently derive matching task_ids. A cheap CPU-only test in tests/unit_tests/resharding/test_planner.py could construct synthetic (src_meta, dst_meta) rosters and assert:
build_plan_from_rosters(rosters, my_global_rank=A)send ops andbuild_plan_from_rosters(rosters, my_global_rank=B)recv ops share the sametask_idfor each transfer (the determinism guarantee);- adding a rank to the rosters and replaying keeps existing
task_ids stable (the node-add / membership-growth path).
Everything else LGTM.
Replace the centralized (rank-0 gather/build/scatter) reshard planner with local
plan building: every rank all-gathers parameter metadata and replays one
deterministic schedule (_iter_global_transfer_ops), keeping only its own
send/recv ops. The iteration order and monotonic task_id depend only on the
gathered metadata, so a send op and its matching recv op agree on task_id with
no central authority, dropping the rank-0 bottleneck and the scatter. The gather
and the build are split (index_metadata_rosters + build_plan_from_rosters) so the
build can also run off rosters assembled another way, and the schedule iterates
sorted roster ranks so a growing rank set rebuilds without renumbering existing
task_ids. CPU-only unit tests in test_planner.py cover the cross-rank task_id
determinism and the node-add stability.
Include world_size in the reshard plan cache key so growing the destination pool
(adding inference nodes) rebuilds the plan instead of reusing a stale one.
Add NixlCopyService: a cross-cluster, non-collocated backend over NIXL (UCX/RDMA)
that keeps refit off the torch process group. Agent metadata, registrations, and
a {task_id: (addr,len,dev)} map of the recv buffers are exchanged once in a cached
handshake; every refit after that is pure NIXL and sender-driven. A receiver
signals each source that its buffers are free, the source waits for that, syncs
its weights, and issues one WRITE per receiver; the WRITE carries a notification
so the receiver knows its weights landed. These ready/data notifications order
producer and consumer per refit with no barrier and no per-refit collective. The
service also augments a host-only UCX_TLS (e.g. tcp) with cuda_copy,cuda_ipc so
GPU tensors transfer instead of segfaulting.
Membership can also grow without a collective: local_roster_entry registers a
node's recv buffers and returns its (agent metadata, recv descriptors) entry, and
update_membership installs an orchestrator-relayed roster on every rank, adding
new peers point-to-point. Notifications are tagged with an orchestrator-provided
membership epoch so a node that joins later stays in sync.
Signed-off-by: William Dykas <wdykas@nvidia.com>
14b97e0 to
6792b29
Compare
|
|
||
| # Name by group rank, which is unique across the (possibly cross-world) | ||
| # group. dist.get_rank() would collide: separate worlds each have a rank 0. | ||
| self.agent_name = agent_name or f"refit-nixl-rank-{self.rank}" |
There was a problem hiding this comment.
this is specific to Nemo-rl which has two different worlds and the ranks can conflict
| descs = {task_id: self._mem_desc(t) for task_id, t in items} | ||
| return (self.agent_name, self.agent.get_agent_metadata(), descs) | ||
|
|
||
| def update_membership(self, roster: Dict[int, tuple], epoch: int) -> None: |
There was a problem hiding this comment.
This is for node joining which can implemented later
There was a problem hiding this comment.
Should we check-in this code for when we plan to implement elasticity, then? Feels odd to have untested code in main.
|
LGTM |
|
/claude strict-review |
Is there a way to test this? And would this elasticity currently be limited to DP in inference? |
|
/ok to test d6fb491 |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30034797201 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30049102792 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30051556418 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30052010057 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30052267431 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30052578747 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30055536833 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30116069434 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30117537409 |
6 participants
The point of these changes is that this would allow us to hopefully be able to handle arbitrary nodes coming in and out of the training and inference groups without resetting the job.
What does this PR do?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.