[XiaodaNV]

What does this PR do?

With overlap_moe_expert_parallel_comm=True, MFSDP v2's combined 1F1B scheduler bypasses FsdpModule.pre_backward. That hook is what makes the gradient-reduction stream wait on the current stream. Without that dependency, capturing a full-iteration CUDA graph fails: the first gradient reduction runs on a stream outside the capture and raises cudaErrorStreamCaptureIsolation.

_module_post_backward_hook now makes inner_reduction_stream wait on the current stream before reducing gradients.

Test

test_mcore_moe.py::test_ep_overlap_full_iteration_cuda_graph_matches_default_schedule builds a 2-layer MoE GPTModel (EP=2, drop-padding alltoall) wrapped in FullyShardedDataParallel v2. It captures a full-iteration CUDA graph with and without overlap_moe_expert_parallel_comm, then requires bitwise-equal losses and parameters.

• Fails on main with cudaErrorStreamCaptureIsolation.
• Passes with this fix.

[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.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[svcnvidia-nemo-ci]

The regression test silently skips optimizer updates on both replay-only iterations, weakening its training-equivalence check. Its reset helper also violates the repository's explicit string-quote convention.

I inspected both changed files and the affected stream, capture and optimizer paths. I didn't execute GPU tests.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[svcnvidia-nemo-ci]

[P2] Preserve gradient bindings across graph replays

On iterations three and four, optimizer.zero_grad() clears the optimizer-facing .grad attributes to None. Graph replay runs GPU work without rerunning the Python schedule, so it doesn't repeat _install_sharded_grads. SGD therefore skips every parameter even though captured kernels still write the persistent gradient buffers. Both schedule variants freeze identically, allowing their final parameter comparisons to pass without checking continued training. FullyShardedOptimizer already handles this exact contract, but this test uses raw SGD.

Replace this line with:

Suggested change

	optimizer.zero_grad()
	optimizer.zero_grad(set_to_none=False)

This preserves the gradient views installed during warmup and zeros their persistent storage before capture and each replay. Retain the eager warmup, both schedule variants, successful-capture assertion, and bitwise loss/parameter comparisons.

Proposed acceptance checks: assert optimizer-facing gradients are non-None immediately before every optimizer.step(), including both replay-only iterations; snapshot local parameter shards and verify at least one nonempty shard changes on each replay-only update. Keep the existing rtol=0, atol=0 comparisons. Removing the production stream fix should still make the overlap variant fail during capture. These checks are proposed; I haven't executed them.

[svcnvidia-nemo-ci]

[P3] Use double quotes in the new reset helper

All four new reset assignments use single-quoted keys, contrary to the protected style guide's explicit double-quote requirement. This is a review-only convention: Black preserves string quotes, so formatting won't correct it.

Replace the four assignments in this helper with:

    FullCudaGraphWrapper.curr_iteration = {"training": 0, "validation": 0}
    FullCudaGraphWrapper.cuda_graph = {"training": None, "validation": None}
    FullCudaGraphWrapper.result = {"training": None, "validation": None}
    StaticBufferLoader.static_buffers = {"training": [], "validation": []}

Preserve both stage keys, their values and the assignment order. Acceptance is static inspection: all eight key literals use double quotes, with reset behavior unchanged.

[balasaajay]

/ok to test 3505954

[XiaodaNV]

The regression test silently skips optimizer updates on both replay-only iterations, weakening its training-equivalence check. Its reset helper also violates the repository's explicit string-quote convention.

I inspected both changed files and the affected stream, capture and optimizer paths. I didn't execute GPU tests.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

Changed

[balasaajay]

/ok to test 46c5375

[shjwudp]

My understanding is that this fix makes the RS stream wait on the compute stream via inner_reduction_stream.wait_stream before _reduce_gradient_groups allocates the gradient buffer, allowing CUDA graph capture to include inner_reduction_stream. Is that correct? 🤔

If so, I’d suggest moving this wait to a dedicated point near the start of the operation, so that it enables CUDA Graph capture without being issued too frequently. For example, could we place it immediately after context.ensure_finalized() here?

[XiaodaNV]

Done. Thanks for this suggestion. I moved the wait into microbatch() right after context.ensure_finalized(). It now runs once per microbatch scope

[wujingyue]

Should we fix the 1f1b hooks instead?

1. microbatch (being renamed to finalize_gradients in [MFSDP] Rename the microbatch scope to finalize_gradients #9000) having to deal with pre-backward violates SRP.
2. The non-1f1b path shouldn't have to doubly wait.
3. The scope is optional technically speaking.

In addition, we could wrap

Megatron-LM/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py

Lines 421 to 432 in 45eec61

	# Fork the reduce-scatter stream from the current stream once, at the
	# start of backward, so every module's post-backward reduce-scatter is
	# part of any active CUDA-graph capture. A stream only joins the
	# capture via this wait_stream edge; without it the first allocation on
	# the reduce-scatter stream falls back to a raw cudaMalloc, which is
	# illegal during capture. Later modules are covered by the post-copy
	# fork each preceding module issues before its collective.
	context.inner_reduction_stream.wait_stream(current_stream)
	# The DP-outer reduction stream needs no fork here: it never allocates (the
	# finalize reuses the persistent main_grad and its view), and it joins any
	# capture via its wait on the reduce-scatter stream before each finalize in
	# _reduce_gradient_groups.

into context.pre_forward, symmetric to context.post_backward, so the code stays SLAP (single level of abstraction).

In addition, could this disable any overlapping? Currently, the wait is done after forward; with this PR, before.

[XiaodaNV]

Reverted the microbatch() change. The root block of FsdpModule.pre_backward now goes through FsdpContext.pre_backward(), and the combined-1F1B backward pre-hook calls it for the root. The default path waits once again, and the fork is back at backward start, so overlap is unchanged

[wujingyue]

Optional: could we get rid of the is_root check?

is_root is needed in default mode because FsdpModule.pre_backward can operate on a nested FsdpModule.

For 1f1b, I expect the scheduler knows exactly the schedule and the model structure so it only needs to FsdpContext.pre_backward once at the right place.

[wujingyue]

Could we move the test to tests/unit_tests/a2a_overlap?

This may be a philosophical distinction, but I view combined_1f1b + MFSDP support as adding MFSDP support to combined_1f1b, rather than adding combined_1f1b support to MFSDP. In that model, combined_1f1b is a user of MFSDP with register_hooks=False, so its specific tests should live alongside the other 1f1b tests.