Mark non-final microbatches for MFSDP v2 hybrid data parallelism - #7186
Merged
wujingyue merged 8 commits intoSep 12, 2026
Conversation
wujingyue
force-pushed
the
agent/mfsdp-v2-microbatch-no-sync
branch
from
September 9, 2026 22:29
fb2a1e1 to
6988cda
Compare
wujingyue
commented
Sep 9, 2026
| return has_any_grad | ||
|
|
||
| def reduce_partial_gradients( | ||
| self, partial_grad: DBuffer, is_last_microbatch: bool = True |
Contributor
Author
There was a problem hiding this comment.
is_last_microbatch is important enough to be required.
shjwudp
approved these changes
Sep 10, 2026
| MCore's schedules wrap every microbatch but the last in ``no_sync_func``. | ||
| """ | ||
| with microbatch(self.module.context, is_last=False): | ||
| yield |
wujingyue
force-pushed
the
agent/mfsdp-v2-microbatch-no-sync
branch
from
September 10, 2026 02:12
56c420d to
feb6ee3
Compare
ko3n1g
approved these changes
Sep 11, 2026
maanug-nv
approved these changes
Sep 11, 2026
wujingyue
force-pushed
the
agent/mfsdp-v2-microbatch-no-sync
branch
from
September 11, 2026 17:08
53a7225 to
3e8f3b1
Compare
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
wujingyue
force-pushed
the
agent/mfsdp-v2-microbatch-no-sync
branch
from
September 12, 2026 15:33
c055dae to
5771c56
Compare
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34711461179 |
balasaajay
added a commit
to balasaajay/Megatron-LM
that referenced
this pull request
Sep 13, 2026
Restore the previous calculator probe, then temporarily reverse the source signature change from NVIDIA#7186. This returns parameter_group.py to the recorded Testmon baseline while preserving existing keyword callers. Keep dependency, container, workflow, and test files unchanged. Signed-off-by: Ajay Balasa <abalasa@nvidia.com>
balasaajay
added a commit
to balasaajay/Megatron-LM
that referenced
this pull request
Sep 13, 2026
Restore the preceding context-parallel validation probe, then temporarily remove the adapter no_sync override introduced in NVIDIA#7186. This returns the adapter file to the recorded baseline and intentionally tests whether selected hybrid accumulation tests catch the behavioral regression. Restore the override after recording results; keep test, dependency, container, and CI files unchanged. Signed-off-by: Ajay Balasa <abalasa@nvidia.com>
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
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.
Problem
MFSDP v2's hybrid data parallelism drops accumulated gradients: with multiple
microbatches, only the last microbatch's gradient reaches the optimizer.
reduce_partial_gradientsusescontext.is_last_microbatchto defer outer-DPreduction, but the adapter inherited a no-op
no_sync, leaving that flag truefor every backward. Finalizing the accumulation buffer on every backward caused
the next microbatch to clear it.
Changes
microbatch(is_last=False)fromFullyShardedDataParallelV2.no_sync.Install the v2
no_sync_funcregardless ofoverlap_grad_reduce, because v2reduces gradients during backward in either mode.
is_last_microbatchkeyword at the gradient-reduction call site.H100 dev golden values for NGC 26.08: baseline, HSDP, and HFSDP.
expert_outer_dp_sharding_strategy='optim'in the HFSDPrecipe and
no_shardin the baseline and HSDP recipes. This PR preservesthe current config and CLI defaults and does not depend on a default-inheritance change.
Validation
Rebased onto
mainatf6c33bde4; no merge commit.change was split into a separate PR. This PR retains the
gradient-accumulation regression coverage.
8 x H100 with NGC 26.08:
deepseek_proxy_mfsdp_v2_ep2_mb2,deepseek_proxy_mfsdp_v2_ep2_hsdp_mb2, anddeepseek_proxy_mfsdp_v2_ep2_hfsdp_mb2.Each passed exact loss and approximate memory checks against its new goldens.
Iteration-1 gradient norm was
10.681for all three.The recipes now explicitly select the same effective expert outer strategies
used in those runs; training was not rerun after that mechanical change.
Exact loss checks remain enabled. Approximate loss tolerance alone does not
reliably detect the original accumulation bug, so these cases must not adopt
NON_DETERMINSTIC_RESULTS=1. Each case checks its own golden; the functionalharness does not compare the three configurations directly.