Skip to content

[dev] fix no_shard training convergency and add unittest for no_shard - #3835

Merged
wplf merged 2 commits into
NVIDIA:devfrom
wplf:jinliang/fix-fsdp-no-shard-dev
May 27, 2026
Merged

[dev] fix no_shard training convergency and add unittest for no_shard#3835
wplf merged 2 commits into
NVIDIA:devfrom
wplf:jinliang/fix-fsdp-no-shard-dev

Conversation

@wplf

@wplf wplf commented Mar 12, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

main PR #3754

  • Fixes no_shard convergence by correcting grad_norm and all_reduce logic.
  • fix no_shard grad_norm calculation
  • fix no_shard all_reduce grad continue bug
  • Make no_shard to be compatiable without overlap_param_gather and overlap_grad_reduce
  • add unittest for no_shard
  • add assert False for no_shard and init_model_with_meta_device follow [fully_shard.py:326]

@wplf
wplf requested review from a team as code owners March 12, 2026 15:29
@copy-pr-bot

copy-pr-bot Bot commented Mar 12, 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.

@wplf wplf changed the title fix no_shard training convergency and add unittest for no_shard [dev] fix no_shard training convergency and add unittest for no_shard Mar 12, 2026
@wplf
wplf requested review from cspades and yaoyu-33 March 12, 2026 15:31

@FDecaYed FDecaYed left a comment

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.

left comments

Comment thread megatron/core/optimizer/__init__.py Outdated
wplf added 2 commits May 21, 2026 00:09
add unittest for no_shard and add empty cache to avoid OOM

add meta_device_check for no_shard following fully_shard.py:326

Signed-off-by: jinliangl <jinliangl@nvidia.com>
@wplf
wplf force-pushed the jinliang/fix-fsdp-no-shard-dev branch from 127851a to c4a4463 Compare May 21, 2026 07:11
@wplf

wplf commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c4a4463

@wplf
wplf enabled auto-merge May 21, 2026 07:51
@yaox12
yaox12 requested a review from FDecaYed May 22, 2026 01:15
@wplf
wplf dismissed FDecaYed’s stale review May 27, 2026 06:30

I've resolved this changed.

@wplf
wplf added this pull request to the merge queue May 27, 2026
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:dev with commit 2ee3bfb May 27, 2026
68 checks passed
@wplf
wplf deleted the jinliang/fix-fsdp-no-shard-dev branch May 27, 2026 07:25
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.

6 participants