feat(fsdp): use TE general_gemm for mixed-precision wgrad in FSDP path by Victarry · Pull Request #3822 · NVIDIA/Megatron-LM

Victarry · 2026-03-12T07:27:48Z

Summary

Replace torch.matmul with TE general_gemm in the FSDP gradient accumulation fusion path to support mixed-precision output (e.g. bf16 input → fp32 main_grad buffer)
torch.matmul requires the out tensor dtype to match the compute dtype, which forces FSDP to use grad_reduce_in_bf16 when gradient_accumulation_fusion is enabled
TE's general_gemm supports arbitrary out_dtype via cuBLAS, removing this hard dependency
Uses inline lazy import to avoid circular dependency with megatron.core.extensions.transformer_engine

Test plan

Verified on DeepSeek-V3-proxy (MLA+MoE+MTP) with FSDP (fsdp_double_buffer=True, gradient_accumulation_fusion=True, fp32 grad buffer, TP1/PP1/EP4/DP2, 8×H100)
Training converges normally (~103 TFLOP/s/GPU, lm_loss decreases from 10.35 to 0.30)

🤖 Generated with Claude Code

copy-pr-bot · 2026-03-12T07:27:52Z

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 · 2026-03-12T07:27:57Z

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:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

shjwudp · 2026-03-12T08:15:43Z

/ok to test eb769df

Replace torch.matmul with TE general_gemm in the FSDP gradient accumulation fusion path to support mixed-precision output (e.g. bf16 input -> fp32 main_grad buffer). torch.matmul requires the `out` tensor dtype to match the compute dtype, which forces FSDP to use grad_reduce_in_bf16 when gradient_accumulation_fusion is enabled. TE's general_gemm supports arbitrary out_dtype via cuBLAS, removing this hard dependency. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Victarry · 2026-03-12T08:41:53Z

/ok to test f2235d6

shjwudp · 2026-03-24T10:05:45Z

/ok to test b28bae0

cspades

Yup, this will mean that users can use (megatron_fsdp_)main_grads_dtype=torch.float32 and (megatron_fsdp_)grad_comm_dtype=torch.float32 even with FP8 params / BF16 gradients, thanks to out_dtype, thanks @Victarry !

cspades · 2026-03-25T01:40:24Z

Rebasing on main to hopefully get rid of the CodeCov failure.

shjwudp · 2026-03-25T02:53:29Z

/ok to test 02a87e2

svcnvidia-nemo-ci · 2026-04-01T23:32:18Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-04-01T23:45:07Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-04-02T00:12:56Z

🔄 Merge queue validation started!

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

NVIDIA#3822) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Jianbin Chang <shjwudp@gmail.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>

NVIDIA#3822) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Jianbin Chang <shjwudp@gmail.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

NVIDIA#3822) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Jianbin Chang <shjwudp@gmail.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

NVIDIA#3822) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Jianbin Chang <shjwudp@gmail.com> Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>

Victarry requested review from a team as code owners March 12, 2026 07:27

svcnvidia-nemo-ci marked this pull request as draft March 12, 2026 07:27

Victarry mentioned this pull request Mar 12, 2026

[Dev] feat(fsdp): use TE general_gemm for mixed-precision wgrad in FSDP path #3824

Merged

1 of 2 tasks

Victarry marked this pull request as ready for review March 12, 2026 08:01

svcnvidia-nemo-ci requested a review from a team March 12, 2026 08:01

Victarry requested a review from shjwudp March 12, 2026 08:01

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage complexity: low labels Mar 12, 2026

shjwudp approved these changes Mar 12, 2026

View reviewed changes

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 12, 2026

Victarry force-pushed the denliu/fsdp-te-general-gemm branch from eb769df to f2235d6 Compare March 12, 2026 08:41

copy-pr-bot Bot temporarily deployed to test March 12, 2026 08:42 Inactive

jaredcasper approved these changes Mar 18, 2026

View reviewed changes

Merge branch 'main' into denliu/fsdp-te-general-gemm

b28bae0

copy-pr-bot Bot temporarily deployed to test March 24, 2026 10:06 Inactive

shjwudp requested review from a team and cspades March 24, 2026 15:21

cspades approved these changes Mar 25, 2026

View reviewed changes

Merge branch 'main' into denliu/fsdp-te-general-gemm

02a87e2

copy-pr-bot Bot temporarily deployed to test March 25, 2026 02:54 Inactive

Victarry enabled auto-merge April 1, 2026 06:55

shjwudp mentioned this pull request Apr 1, 2026

[QUESTION] How to use Megatron-FSDP with mixed precision? #4013

Closed

ericharper approved these changes Apr 1, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Apr 1, 2026

Victarry added this pull request to the merge queue Apr 1, 2026

Merged via the queue into NVIDIA:main with commit 8d7a3f8 Apr 2, 2026
66 of 70 checks passed

Victarry deleted the denliu/fsdp-te-general-gemm branch April 2, 2026 04:42

Navigation Menu

Repository navigation

feat(fsdp): use TE general_gemm for mixed-precision wgrad in FSDP path - #3822

feat(fsdp): use TE general_gemm for mixed-precision wgrad in FSDP path#3822
Victarry merged 3 commits into
NVIDIA:mainfrom
Victarry:denliu/fsdp-te-general-gemm

Victarry commented Mar 12, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Mar 12, 2026

Uh oh!

github-actions Bot commented Mar 12, 2026

Uh oh!

shjwudp commented Mar 12, 2026

Uh oh!

Victarry commented Mar 12, 2026

Uh oh!

shjwudp commented Mar 24, 2026

Uh oh!

cspades left a comment •

edited

Loading

Uh oh!

cspades commented Mar 25, 2026

Uh oh!

shjwudp commented Mar 25, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 2, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

Victarry commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

copy-pr-bot Bot commented Mar 12, 2026

Uh oh!

github-actions Bot commented Mar 12, 2026

Uh oh!

shjwudp commented Mar 12, 2026

Uh oh!

Victarry commented Mar 12, 2026

Uh oh!

shjwudp commented Mar 24, 2026

Uh oh!

cspades left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cspades commented Mar 25, 2026

Uh oh!

shjwudp commented Mar 25, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 1, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 2, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Victarry commented Mar 12, 2026 •

edited

Loading

cspades left a comment •

edited

Loading