[Sunt-ing]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Local same-rank refit copies can silently read stale model weights when the source update is still queued on the caller's CUDA stream. NCCL, Gloo, NIXL, and NVSHMEM all run the local copy on a dedicated stream, but only order the caller stream after that copy; they do not order the copy stream after the producer before it reads the source tensor.

This change makes each local copy stream wait for the caller's current stream before copying. It preserves the existing copy-to-caller join and does not change the public API or remote-rank transport.

Tests

python -m torch.distributed.run --nproc-per-node 8 -m pytest -q \
  tests/unit_tests/resharding/test_copy_services.py::test_local_copy_waits_for_current_stream \
  tests/unit_tests/resharding/test_copy_services.py::test_optional_backend_local_copy_waits_for_current_stream

With the fix, every rank passes all four backend cases (4 passed per rank). Removing only the corresponding wait_stream calls makes every affected case fail with the destination retaining the stale value (maximum absolute difference 1.0). The NIXL/NVSHMEM test uses real CUDA streams and tensors while mocking only the optional remote transport, so it does not require those optional runtime dependencies in the unit-test environment.

As supplemental validation, I also exercised the public swap_model_weights() path after fused AdamW updates on a 20-layer, hidden-size 4096 GPT model. The source/destination logits were the oracle:

Backend	Current code	Fixed	Exact revert
NCCL	max diff 3.0010	0	max diff 3.3286
Gloo	max diff 2.8388	0	max diff 2.8373
NIXL	max diff 0.9583	0	max diff 7.2628
NVSHMEM	max diff 0.9849	0	max diff 0.9629

NCCL/Gloo refit latency

I measured the cost of the added dependency with the repository's examples/rl/benchmark_refit.py on one H200. Both variants used collocated TP1/PP1/EP1, BF16, a 20-layer/4096-hidden/32-head GPT model (4.23B parameters per model), sequence length 8, batch size 1, five warmup iterations, and 30 timed iterations. Each timed iteration synchronizes CUDA before refit, so this comparison measures the dependency overhead when source work is already complete.

The fixed tree differed from the current tree only by this PR. The same temporary current-main API compatibility adjustment and statistics output were applied to both trees and are not part of the PR. All values below are milliseconds; mean / median / stddev / min-max summarizes all 30 samples.

Backend	Current	Fixed	Median delta	Mean delta
NCCL	7.39 / 7.38 / 0.05 / 7.33-7.52	7.51 / 7.45 / 0.24 / 7.40-8.67	+0.07 (+0.95%)	+0.12 (+1.62%)
Gloo	16.34 / 8.28 / 44.07 / 8.24-249.67	16.21 / 8.27 / 43.47 / 8.24-246.39	-0.01 (-0.12%)	-0.13 (-0.80%)

Each Gloo arm contained one approximately 250 ms outlier, which inflates its mean and standard deviation; the median shows no measurable regression at the reported precision. If source work is still pending, the new dependency intentionally includes the time required before those weights are valid rather than allowing a stale copy to return early.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

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:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[YangFei1990]

/claude fix

[YangFei1990]

/claude assign

[svcnvidia-nemo-ci]

🛠️ Claude fix commit a27ddc8d7c71 (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Resolved the single merge conflict in tests／unit＿tests／resharding／test＿copy＿services.py by taking the union of both additive sides: kept the base's NVSHMEMCopyService import and the PR's Utils import, and kept all three tests. Removed markers, fixed seam blank lines, staged the file, no commit. Verified py＿compile and zero deletions vs either parent. Could not run pytest／isort／black (not installed), so checks were static only.

Files changed by Claude

• tests/unit_tests/resharding/test_copy_services.py

Why
The merge left one unmerged file with two hunks, both add／add collisions at the same spots rather than semantic disagreements: the base added an NVSHMEM import and two end-of-file tests, while the PR added dist／Utils imports and a stream-ordering test there. Neither side changed the other's lines, so keeping both was correct. The CI log dir and steering file were empty, so no lint or test failures to clear.

DCO
@Sunt-ing, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test a27ddc8

[svcnvidia-nemo-ci]

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

[YangFei1990]

/claude strict-review

[claude]

[CRITICAL Correctness] The fix is right, but it is applied to only 2 of the 4 backends that have the identical defect.

NixlCopyService._do_local_copies() (nixl_copy_service.py:148) and NVSHMEMCopyService.run() (nvshmem_copy_service.py:138) use exactly the same side-stream pattern as the two files fixed here — a _copy_stream / _local_copy_stream created in __init__, local copies issued inside with torch.cuda.stream(...), and a current_stream().wait_stream(...) join at the end of run() — with no wait_stream(current_stream()) before the copy:

# nixl_copy_service.py:147-150
pairs = match_local_ops_by_task_id(local_sends, local_recvs, "NixlCopyService", self.rank)
with torch.no_grad(), torch.cuda.stream(self._copy_stream):   # <-- reads src with no producer ordering
    for send_op, recv_op in pairs:
        recv_op.tensor.copy_(send_op.tensor)

# nvshmem_copy_service.py:137-148
with torch.no_grad():
    with torch.cuda.stream(self._local_copy_stream):          # <-- same hazard
        for task_id, dst in self._local_recv_ops.items():
            ...
            dst.copy_(src, non_blocking=True)

Why it matters: this is the same silent stale-weight read the PR describes, on the same collocated refit path, reachable through the same public swap_model_weights() entry point — just with refit_method="nixl" or "nvshmem". NIXL's later torch.cuda.current_stream().synchronize() at line 231 is guarded by if remote_sends: and happens after the local copies are already enqueued, so it does not order them. NVSHMEM's torch_copy_stream_wrapper.synchronize() inside pipeline_executor.py:351 is a different stream on the RemoteCopyService self-move path and does not cover _local_copy_stream either. Leaving these two unfixed means a user who switches backends silently reacquires the bug the PR just closed, and the new unit test won't catch it because it only parametrizes over NCCLCopyService and GlooCopyService.

Suggested fix: add the same one-liner to both backends and extend the new test's parametrize list to cover them (guarded by the existing import/availability skips those backends already need). Given that four independent call sites now need the identical two-line dance, consider instead putting it in a small shared helper in base.py — e.g. local_copy_stream_scope(stream) — so a fifth backend can't reintroduce the hazard.

[claude]

[IMPORTANT Correctness] The hardcoded cycle count makes this regression test hardware-dependent, and it fails in the silent direction.

torch.cuda._sleep takes device cycles, not time. 100_000_000 cycles is ~55 ms on an H100 at ~1.8 GHz, but the wall-clock delay scales inversely with the SM clock — it shrinks under boost and on future/faster parts. If the gate ever elapses before the unordered copy would have run, the test passes with or without the wait_stream calls, and the guard against this regression quietly stops working with no failure to alert anyone.

This repo already has the conversion helper for exactly this reason — tests/unit_tests/determinism/utils.py:16-27 imports get_cycles_per_ms from torch.testing._internal.common_utils with a conservative fallback, with a comment noting PyTorch's own test_cuda.py uses it to express _sleep durations in milliseconds.

Suggested fix: express the gate in milliseconds so the delay is stable across clocks:

from torch.testing._internal.common_utils import get_cycles_per_ms

GATE_MS = 50
...
with torch.cuda.stream(release_stream):
    torch.cuda._sleep(int(GATE_MS * get_cycles_per_ms()))
    producer_gate.record()

[claude]

[IMPORTANT Correctness] Two test-hygiene issues that will surface as confusing CI failures rather than as this test failing on its own merits.

1. GlooCopyService() leaks a process group per invocation. With group=None, GlooCopyService.__init__ calls dist.new_group(backend="gloo") (gloo_copy_service.py:28). dist.new_group is collective — every rank must call it in the same order — and nothing here calls service.close() or destroys the group. This test file is otherwise entirely rank-agnostic (pure-Python assertions and monkeypatched stubs), so it is the first test here to require lockstep participation from all 8 ranks in the CI unit-test bucket. Since pytest collection order is stable this will likely work today, but it is fragile: any -k filter, -x early exit, or rank-divergent skip leaves ranks disagreeing on how many groups exist, which manifests as an NCCL/Gloo hang rather than a clear failure. Wrap the service in a try/finally that calls service.close(), and consider constructing the Gloo PG once at module scope in a fixture.

2. service_cls() reads the default PG, so this only exercises rank-0-vs-itself semantics correctly by accident. CopyService.__init__ takes dist.get_rank() when group is None, and the test then submits with dist.get_rank() — consistent, so the local-pair matching works on every rank. That part is fine. But every rank independently runs the full body including the _sleep gate, so the bucket pays the gate latency 8× in parallel with no added coverage; the hazard is purely rank-local. A @pytest.mark.skipif(dist.get_rank() != 0)-style guard is not safe here because of the collective new_group above — which is another reason to hoist the PG creation out of the test body first.

3. torch.testing.assert_close(destination, source) cannot fail as intended. Both tensors are read after torch.cuda.synchronize(), so if the copy read a stale source the destination would hold 0 while source holds 1 — that works. But if the fill itself were dropped, both would be 0 and the assertion would still pass. Comparing against a literal expected value makes the oracle explicit:

Suggested change

	def test_local_copy_waits_for_current_stream(service_cls):
	"""A local copy must observe writes already queued on the caller's stream."""
	Utils.initialize_distributed()
	service = service_cls()
	source = torch.zeros(1, device="cuda")
	destination = torch.zeros_like(source)
	torch.cuda.synchronize()
	producer_gate = torch.cuda.Event()
	release_stream = torch.cuda.Stream()
	with torch.cuda.stream(release_stream):
	torch.cuda._sleep(100_000_000)
	producer_gate.record()
	torch.cuda.current_stream().wait_event(producer_gate)
	source.fill_(1)
	service.submit_send(source, dist.get_rank(), task_id=0)
	service.submit_recv(destination, dist.get_rank(), task_id=0)
	service.run()
	torch.cuda.synchronize()
	torch.testing.assert_close(destination, source)
	@pytest.mark.parametrize("service_cls", [NCCLCopyService, GlooCopyService])
	def test_local_copy_waits_for_current_stream(service_cls):
	"""A local copy must observe writes already queued on the caller's stream."""
	Utils.initialize_distributed()
	service = service_cls()
	try:
	source = torch.zeros(1, device="cuda")
	destination = torch.zeros_like(source)
	torch.cuda.synchronize()
	producer_gate = torch.cuda.Event()
	release_stream = torch.cuda.Stream()
	with torch.cuda.stream(release_stream):
	torch.cuda._sleep(int(50 * get_cycles_per_ms()))
	producer_gate.record()
	torch.cuda.current_stream().wait_event(producer_gate)
	source.fill_(1)
	service.submit_send(source, dist.get_rank(), task_id=0)
	service.submit_recv(destination, dist.get_rank(), task_id=0)
	service.run()
	torch.cuda.synchronize()
	torch.testing.assert_close(destination, torch.ones_like(destination))
	finally:
	service.close()

[claude]

Strict Review Summary

CRITICAL: 1 · IMPORTANT: 2 · SUGGESTION: 1

Assessment

The core diagnosis is correct and the fix is the right one. Both NCCLCopyService.run() and GlooCopyService.run() established only the downstream half of the stream contract (current_stream().wait_stream(self._copy_stream) at the end of run()) while omitting the upstream half, so the copy stream could read source weights before a producer update queued on the caller's stream had landed. wait_stream is a device-side dependency with no host sync, so it is the correct primitive here — the latency data in the PR description (median +0.07 ms NCCL, no measurable Gloo change) matches that expectation. The CopyService docstring at base.py:48-51 already states that supports_multiple_runs_per_plan requires run() to "establish current-stream ordering for all submitted sends and recvs"; this PR brings the two implementations in line with the contract they already advertise.

Most impactful finding

The same defect exists in two backends this PR does not touch. NixlCopyService._do_local_copies() (nixl_copy_service.py:148) and NVSHMEMCopyService.run() (nvshmem_copy_service.py:138) use an identical side-stream-plus-trailing-join pattern with no upstream ordering, and both are reachable through the same public swap_model_weights() path via refit_method="nixl" / "nvshmem". I checked whether either is incidentally protected: NIXL's current_stream().synchronize() at line 231 is behind if remote_sends: and runs after the local copies are enqueued; NVSHMEM's torch_copy_stream_wrapper.synchronize() in pipeline_executor.py:351 is a different stream on the RemoteCopyService self-move path. Neither orders the local-copy stream. A user switching backends silently reacquires the stale-weight bug, and the new test won't catch it because it parametrizes only the two fixed classes.

With four call sites now needing the identical two-line dance, a shared helper in base.py would be worth more than a fourth copy of the one-liner.

Other findings

• Test gate is clock-dependent and fails silently open — torch.cuda._sleep(100_000_000) is in device cycles (~55 ms on H100 at ~1.8 GHz) and shrinks under boost or on faster parts. If the gate ever elapses too early the test passes with or without the fix, so the regression guard stops working with no visible failure. tests/unit_tests/determinism/utils.py:16-27 already vendors get_cycles_per_ms for exactly this conversion.
• GlooCopyService() leaks a collective process group per invocation — with group=None the constructor calls dist.new_group(backend="gloo") (gloo_copy_service.py:28) and nothing calls close(). This is the first test in an otherwise fully rank-agnostic file to require lockstep participation from all ranks; any -k filter or -x early exit leaves ranks disagreeing on group count, surfacing as a hang rather than a failure. Hoist the PG into a module-scoped fixture and add try/finally: service.close().
• Assertion oracle is weaker than intended — assert_close(destination, source) also passes if the fill_(1) were dropped entirely (both sides 0). Compare against torch.ones_like(destination).
• Comments at gloo_copy_service.py:67-68 and nccl_copy_service.py:27-28 document only the overlap rationale, not the ordering invariant the new line enforces — worth naming explicitly given the silent failure mode.

Risk

Low for what it changes; the residual risk is what it leaves unchanged. No public API, checkpoint-format, argument, or default changes; no backward-compatibility surface touched; the added dependency is device-side only. The two unfixed backends are the reason this shouldn't merge as-is — the fix is correct but incomplete, and partial coverage of a silent-corruption bug is the worst state to leave it in, since the presence of a targeted regression test implies coverage that isn't there.

Recommendation: extend the fix and the test parametrization to NIXL and NVSHMEM (ideally via a shared helper), then this is good to go.

[YangFei1990]

Hi @Sunt-ing, could you check the comments from agent?

[Sunt-ing]

@YangFei1990 Sure.

First, thanks for flagging NIXL and NVSHMEM. I reproduced the same bug in both backends through the public swap_model_weights() path after fused AdamW updates on a 4.23B GPT model. At the current PR head, the destination logits were stale, with maximum differences of 0.9583 for NIXL and 0.9849 for NVSHMEM. Adding one producer-to-copy stream dependency to each backend made all three iterations exact. Removing just those two lines brought the stale results back.

I extended the patch and CUDA ordering test to both optional backends. The test uses real CUDA streams and tensors and mocks only the remote transport, so it does not require the optional NIXL or NVSHMEM runtime. Under the repository's 8-GPU torchrun topology, all four backend cases pass on every rank. Reverting only the NIXL and NVSHMEM waits makes both cases fail on every rank, with the destination retaining 0 and a maximum difference of 1.0. The gate now uses get_cycles_per_ms(), and the expected all-ones tensor is independent of the source.

Each wait remains at its backend's local-copy site, one line next to the stream it orders. A shared helper would add indirection while all four backends would still need separate calls. I did not add a Gloo cleanup fixture because GlooCopyService inherits the no-op CopyService.close(), so the fixture would not release the process group. The test creates the group collectively on every rank in the same torchrun invocation.

[kvareddy]

@wdykas can you please take a look at this MR? Also, let's create a review group for refit related code.

[wdykas]

Thanks for the contribution! Nice find, thanks for updating all the colocated backends. @Sunt-ing what are you using this for?

[wdykas]

/ok to test aee431b

[Sunt-ing]

@wdykas Thanks for the approval! I'm evaluating it for collocated RL training, where the rollout model needs the latest trainer weights after each optimizer step.

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34521483023