Fixed fp32 residuals by mkhona-nvidia · Pull Request #3504 · NVIDIA/Megatron-LM

mkhona-nvidia · 2026-02-19T23:15:45Z

What does this PR do ?

Fix FP32 residual stream activations in Megatron which currently silently break.

Problem

The existing fp32_residual_connection config flag was broken, it never actually kept the residual stream in fp32. The core issue was in _bias_dropout_add_func (used by every layer type), which downcast the fp32 residual to match the computation dtype (bf16/fp16):

Previously: defeated fp32 residuals by casting residual DOWN to x's dtype:

residual = residual if residual.dtype == x.dtype else residual.to(x.dtype)

This meant:
The embedding layer correctly emitted fp32 tensors when fp32_residual_connection=True

But the very first bias_dropout_add call cast the residual back to bf16. The fp32 residual stream was lost from layer 0 onwards
Additionally, the transformer layer never cast its residuals to fp32 — only the Mamba layer had its own residual_in_fp32 flag, which also got defeated by the same downcast.

Solution

Reverse the casting direction in _bias_dropout_add_func: instead of downcasting the residual, upcast x (and bias) to match the residual's dtype. This preserves the fp32 residual stream through every bias-dropout-add call.

megatron/core/fusions/fused_bias_dropout.py Upcast x/bias to residual dtype instead of downcasting residual to x dtype

megatron/core/transformer/transformer_layer.py Cast residual to fp32 at all 4 residual sites (self-attn, cross-attn, MLP, MoE router) when fp32_residual_connection=True

megatron/core/ssm/mamba_layer.py Respect config.fp32_residual_connection in addition to the existing residual_in_fp32 flag

megatron/core/transformer/transformer_config.py Auto-override pipeline_dtype to torch.float when fp32_residual_connection=True to prevent PP communication dtype mismatches

Why the old approach was wrong

The original downcast was added to prevent pipeline parallel communication hangs (the residual would upcast the output to fp32, mismatching the bf16 PP recv buffers). The correct fix is to set pipeline_dtype = torch.float when fp32 residuals are enabled, which this PR does automatically in TransformerConfig.__post_init__. The inference wrapper (abstract_model_inference_wrapper.py) already handled this correctly.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

Signed-off-by: mikail <mkhona@nvidia.com>

copy-pr-bot · 2026-02-19T23:15:49Z

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.

duncanriach

Good stuff. Thank you

janEbert · 2026-02-20T00:20:22Z

+        # When fp32 residual connections are enabled, pipeline parallel communication must
+        # use fp32 to match the dtype of the residual stream between pipeline stages.
+        if self.fp32_residual_connection and self.pipeline_dtype is not None:
+            self.pipeline_dtype = torch.float


Could you add a warnings.warn here so users aren't confused when their pipeline dtype changes?

I just realized the warning shouldn't be output when self.pipeline_dtype is already torch.float.

good point, fixed

Please only warn on a single rank if possible. We don't want output spam.

there is a warnings.warn all over that config file, not just the change I made. Shall we make another PR that fixes this for all the others?

(made a change, please let me know if that is ok, it needs log_single_rank)

janEbert · 2026-02-20T00:29:22Z

/ok to test e0e74fd

…e dtype changes Signed-off-by: mikail <mkhona@nvidia.com>

Signed-off-by: mikail <mkhona@nvidia.com>

janEbert · 2026-02-23T19:53:27Z

/ok to test c15e1dd

Signed-off-by: mikail <mkhona@nvidia.com>

mkhona-nvidia · 2026-02-24T04:44:53Z

/ok to test 5580aee

svcnvidia-nemo-ci · 2026-02-24T22:53:01Z

🔄 Merge queue validation started!

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

Signed-off-by: mikail <mkhona@nvidia.com>

…ar to unblock fp32_residual_connection When fp32_residual_connection=True the residual stream is fp32 and bias_dropout_add propagates that dtype to hidden_states. In the TE GPT layer spec input_layernorm/pre_mlp_layernorm is IdentityOp (folded into the fused TELayerNormColumnParallelLinear), so nothing casts back to params_dtype before this op. TE's set_activation_dtype strictly rejects input dtype != param dtype outside autocast, so the forward dies on the very first layer of a bf16 model with --fp32-residual-connection. Downcast input to layer_norm_weight.dtype at the top of forward. Residual-accumulation precision is preserved upstream in bias_dropout_add (fixed by NVIDIA#3504); only the input to LN+GEMM is brought back to params_dtype, which is what the op is configured to run in anyway. Add TestFp32ResidualTELayerNormLinear covering the path: SelfAttention from the TE submodules with bf16 params, fed an fp32 hidden_states. Without this fix the forward raises in TE; with the fix it returns a bf16 output of the expected shape.

Signed-off-by: mikail <mkhona@nvidia.com>

Fixed fp32 residuals

e0e74fd

Signed-off-by: mikail <mkhona@nvidia.com>

mkhona-nvidia requested review from a team as code owners February 19, 2026 23:15

svcnvidia-nemo-ci requested a review from a team February 19, 2026 23:15

mkhona-nvidia self-assigned this Feb 19, 2026

duncanriach approved these changes Feb 19, 2026

View reviewed changes

deepakn94 approved these changes Feb 19, 2026

View reviewed changes

janEbert reviewed Feb 20, 2026

View reviewed changes

janEbert approved these changes Feb 20, 2026

View reviewed changes

janEbert added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Feb 20, 2026

svcnvidia-nemo-ci added this to the Core 0.16 milestone Feb 20, 2026

janEbert enabled auto-merge February 20, 2026 00:30

mkhona-nvidia added 2 commits February 19, 2026 16:34

add a warnings.warn here so users aren't confused when their pipelin…

cafae2c

…e dtype changes Signed-off-by: mikail <mkhona@nvidia.com>

don't warn when dtype is already fp32

9aef970

deepakn94 disabled auto-merge February 20, 2026 02:57

mkhona-nvidia added the Final Review PR is in the "final review" stage label Feb 20, 2026

ananthsub approved these changes Feb 23, 2026

View reviewed changes

mkhona-nvidia added 2 commits February 23, 2026 11:28

added linting for unit test

72aa2c7

Signed-off-by: mikail <mkhona@nvidia.com>

log only on 1 rank when changing pipeline dtype

c15e1dd

Signed-off-by: mikail <mkhona@nvidia.com>

used isosort for linting

5580aee

Signed-off-by: mikail <mkhona@nvidia.com>

copy-pr-bot Bot temporarily deployed to test February 24, 2026 04:45 Inactive

deepakn94 added this pull request to the merge queue Feb 24, 2026

Merged via the queue into NVIDIA:main with commit 08857d9 Feb 25, 2026
77 of 81 checks passed

ko3n1g pushed a commit to ko3n1g/Megatron-LM that referenced this pull request Feb 26, 2026

Fixed fp32 residuals (NVIDIA#3504)

c70efe5

Signed-off-by: mikail <mkhona@nvidia.com>

BoxiangW pushed a commit to BoxiangW/Megatron-LM that referenced this pull request Mar 4, 2026

Fixed fp32 residuals (NVIDIA#3504)

d5d0819

Signed-off-by: mikail <mkhona@nvidia.com>

yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026

Fixed fp32 residuals (NVIDIA#3504)

2ecddd1

Signed-off-by: mikail <mkhona@nvidia.com>

yezhengmao1 mentioned this pull request May 30, 2026

[BUG] Fix/fp32 residual te layernorm linear #5073

Open

5 tasks

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

Fixed fp32 residuals (NVIDIA#3504)

1638b3f

Signed-off-by: mikail <mkhona@nvidia.com>

Uh oh!

Conversation

mkhona-nvidia commented Feb 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Problem

Solution

Why the old approach was wrong

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Feb 19, 2026

Uh oh!

duncanriach left a comment

Choose a reason for hiding this comment

Uh oh!

janEbert Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

mkhona-nvidia Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

janEbert Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

janEbert Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

mkhona-nvidia Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

deepakn94 Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

mkhona-nvidia Feb 23, 2026 • 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!

janEbert commented Feb 20, 2026

Uh oh!

janEbert commented Feb 23, 2026

Uh oh!

mkhona-nvidia commented Feb 24, 2026

Uh oh!

svcnvidia-nemo-ci commented Feb 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

mkhona-nvidia commented Feb 19, 2026 •

edited

Loading

(Step 1): Add PR label `Expert Review`

mkhona-nvidia Feb 23, 2026 •

edited

Loading