Skip to content

[MFSDP] Join reduction stream to capture at microbatch scope entry - #9069

Draft
shjwudp wants to merge 1 commit into
NVIDIA:mainfrom
shjwudp:fix/mfsdp-finalize-scope-stream-capture
Draft

shjwudp wants to merge 1 commit into
NVIDIA:mainfrom
shjwudp:fix/mfsdp-finalize-scope-stream-capture

Conversation

@shjwudp

@shjwudp shjwudp commented Nov 6, 2026 •

Copy link
Copy Markdown
Contributor

Combined 1F1B bypasses FsdpModule.pre_backward(), so its first gradient staging allocation can occur on a reduction stream that has not joined CUDA-graph capture.

Join the inner reduction stream from the current compute stream when entering microbatch(). This reuses the existing MCore no_sync() scope without adding scheduler hooks or per-module waits.

Validation: Black 26, isort, Ruff, and diff checks passed. On NSC-LHR (4 GB300 GPUs, NeMo 26.08), both the regression and microbatch-scope tests passed on all four ranks. Removing only the added wait reproduced cudaErrorStreamCaptureIsolation on all four ranks. Slurm 642295: COMPLETED, exit 0. This is functional validation, not CI-image parity.

Related: #9000. Regression test: #8683.

@copy-pr-bot

copy-pr-bot Bot commented Nov 6, 2026

Copy link
Copy Markdown

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.

Signed-off-by: Jack Chang <jianbinc@nvidia.com>
@shjwudp
shjwudp force-pushed the fix/mfsdp-finalize-scope-stream-capture branch from 156ea5e to 0a02887 Compare November 6, 2026 06:23
@shjwudp shjwudp changed the title [MFSDP] Join reduction stream to capture at finalization scope entry [MFSDP] Join reduction stream to capture at microbatch scope entry Nov 6, 2026
@shjwudp
shjwudp changed the base branch from pull-request/9000 to main November 6, 2026 06:24

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant