Skip to content

Mark non-final microbatches for MFSDP v2 hybrid data parallelism - #7186

Merged
wujingyue merged 8 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-microbatch-no-sync
Sep 12, 2026
Merged

Mark non-final microbatches for MFSDP v2 hybrid data parallelism#7186
wujingyue merged 8 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-microbatch-no-sync

Conversation

@wujingyue

@wujingyue wujingyue commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Problem

MFSDP v2's hybrid data parallelism drops accumulated gradients: with multiple
microbatches, only the last microbatch's gradient reaches the optimizer.
reduce_partial_gradients uses context.is_last_microbatch to defer outer-DP
reduction, but the adapter inherited a no-op no_sync, leaving that flag true
for every backward. Finalizing the accumulation buffer on every backward caused
the next microbatch to clear it.

Changes

  • Enter microbatch(is_last=False) from FullyShardedDataParallelV2.no_sync.
    Install the v2 no_sync_func regardless of overlap_grad_reduce, because v2
    reduces gradients during backward in either mode.
  • Require the is_last_microbatch keyword at the gradient-reduction call site.
  • Add accumulation regression coverage and three functional cases with new
    H100 dev golden values for NGC 26.08: baseline, HSDP, and HFSDP.
  • Explicitly set expert_outer_dp_sharding_strategy='optim' in the HFSDP
    recipe and no_shard in the baseline and HSDP recipes. This PR preserves
    the current config and CLI defaults and does not depend on a default-inheritance change.

Validation

Rebased onto main at f6c33bde4; no merge commit.

  • The eight-H100 hybrid unit suite passed before the default-inheritance
    change was split into a separate PR. This PR retains the
    gradient-accumulation regression coverage.
  • Black, isort, and Ruff checks passed for the changed Core/test files.
  • All three functional cases passed in independent fresh 50-iteration runs on
    8 x H100 with NGC 26.08: deepseek_proxy_mfsdp_v2_ep2_mb2,
    deepseek_proxy_mfsdp_v2_ep2_hsdp_mb2, and deepseek_proxy_mfsdp_v2_ep2_hfsdp_mb2.
    Each passed exact loss and approximate memory checks against its new goldens.
    Iteration-1 gradient norm was 10.681 for 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 functional
harness does not compare the three configurations directly.

@copy-pr-bot

copy-pr-bot Bot commented Sep 9, 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.

return has_any_grad

def reduce_partial_gradients(
self, partial_grad: DBuffer, is_last_microbatch: bool = True

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is_last_microbatch is important enough to be required.

@wujingyue
wujingyue marked this pull request as ready for review September 10, 2026 00:46
@wujingyue
wujingyue requested review from a team as code owners September 10, 2026 00:46
MCore's schedules wrap every microbatch but the last in ``no_sync_func``.
"""
with microbatch(self.module.context, is_last=False):
yield

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Neat!

@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the Approved All necessary approvals have been made label Sep 12, 2026
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@nemo-automation-bot

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:main with commit 845d9be Sep 12, 2026
205 of 206 checks passed
@wujingyue
wujingyue deleted the agent/mfsdp-v2-microbatch-no-sync branch September 12, 2026 19:25
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: medium Run functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants