Fixed fp32 residuals

megatron/core/fusions/fused_bias_dropout.py

@@ -26,12 +26,14 @@ def _bias_dropout_add_func(x_with_bias, residual, prob, training):
         and (bias is None or not bias.requires_grad)
     )
 
-    # If we want to train mixed precision, then the output of this function
-    # should be half precision. However, in AMP O1, the input (residual) is
-    # in fp32, and it will up-cast the result to fp32, causing pipeline parallel
-    # GPU communication to hang. Therefore, we need to cast residual to the same
-    # dtype as x.
-    residual = residual if residual.dtype == x.dtype else residual.to(x.dtype)
+    # For fp32 residual connections: upcast x (and bias) to residual's dtype so that
+    # the addition and output remain in fp32, preserving numerical precision in the
+    # residual stream across layers. When fp32_residual_connection is enabled,
+    # pipeline parallel communication dtype should be set to fp32 accordingly.
+    if x.dtype != residual.dtype:
+        x = x.to(residual.dtype)
+        if bias is not None:
+            bias = bias.to(residual.dtype)
 
     # The Dropout operation, Residual Addition and the tensor returning can be
     # done generically outside the if statement, but that stops fusing of Bias

megatron/core/ssm/mamba_block.py

@@ -53,8 +53,6 @@ class MambaStack(MegatronModule):
     Args:
         config (TransformerConfig): the model configuration
         submodules (MambaStackSubmodules): the submodules for the stack
-        residual_in_fp32 (bool, optional): whether to do residual connections
-            in fp32. Defaults to False.
         pre_process (bool, optional): whether to include an embedding layer.
             Defaults to True.
         hybrid_attention_ratio (float, optional): the target ratio of attention layers to
@@ -77,7 +75,6 @@ def __init__(
         self,
         config: TransformerConfig,
         submodules: MambaStackSubmodules,
-        residual_in_fp32=False,
         pre_process: bool = True,
         hybrid_attention_ratio: float = 0.0,
         hybrid_mlp_ratio: float = 0.0,
@@ -90,7 +87,6 @@ def __init__(
         is_mtp_layer: bool = False,
     ) -> None:
         super().__init__(config=config)
-        self.residual_in_fp32 = residual_in_fp32
         self.pre_process = pre_process
         self.post_layer_norm = post_layer_norm
         self.post_process = post_process
@@ -146,7 +142,6 @@ def __init__(
                     layer = build_module(
                         submodules.mamba_layer,
                         config=self.config,
-                        residual_in_fp32=residual_in_fp32,
                         layer_number=i + 1 + pp_layer_offset,
                         pp_layer_offset=pp_layer_offset,
                         pg_collection=pg_collection,

megatron/core/ssm/mamba_layer.py

@@ -69,7 +69,6 @@ def __init__(
         config: TransformerConfig,
         submodules: MambaLayerSubmodules,
         layer_number: int = 1,
-        residual_in_fp32=False,
         pg_collection: ProcessGroupCollection = None,
         pp_layer_offset: int = 0,
     ):
@@ -80,7 +79,6 @@ def __init__(
         self.config = config
         self.submodules_config = submodules
         self.layer_number = layer_number
-        self.residual_in_fp32 = residual_in_fp32
         self.hidden_dropout = config.hidden_dropout
         self.mixer = build_module(
             submodules.mixer,
@@ -136,8 +134,8 @@ def forward(
         inference_context = deprecate_inference_params(inference_context, inference_params)
 
         residual = hidden_states
-        if self.residual_in_fp32:
-            residual = residual.to(torch.float32)
+        if self.config.fp32_residual_connection:
+            residual = residual.float()
 
         hidden_states = hidden_states.to(dtype=self.config.params_dtype)
         hidden_states = apply_module(self.norm)(hidden_states)

megatron/core/transformer/transformer_config.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import logging
 import warnings
 from dataclasses import dataclass, field
 from typing import Callable, List, Literal, Optional, Tuple, Union
@@ -12,6 +13,7 @@
 from megatron.core.transformer.enums import AttnBackend, CudaGraphScope
 from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout
 
+from .._rank_utils import log_single_rank
 from ..fusions.fused_bias_geglu import quick_gelu
 from ..model_parallel_config import ModelParallelConfig
 from ..utils import (
@@ -22,6 +24,8 @@
     scaled_init_method_normal,
 )
 
+logger = logging.getLogger(__name__)
+
 try:
     from packaging.version import Version as PkgVersion
 
@@ -930,6 +934,20 @@ def __post_init__(self):
         details.
         """
         super().__post_init__()
+
+        # 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:
+            if self.pipeline_dtype != torch.float:
+                log_single_rank(
+                    logger,
+                    logging.WARNING,
+                    f"fp32_residual_connection is enabled, overriding pipeline_dtype "
+                    f"from {self.pipeline_dtype} to torch.float to match the "
+                    f"residual stream dtype between pipeline stages.",
+                )
+            self.pipeline_dtype = torch.float
+
         if self.fp16 and self.bf16:
             raise ValueError(
                 f"Only one of self.fp16: {self.fp16} and self.bf16 {self.bf16} should be True."

megatron/core/transformer/transformer_layer.py

@@ -583,6 +583,8 @@ def _forward_attention(
 
         # Residual connection.
         residual = hidden_states
+        if self.config.fp32_residual_connection:
+            residual = residual.float()
 
         # Optional Input Layer norm
         if self.recompute_input_layernorm:
@@ -651,6 +653,8 @@ def _forward_attention(
 
         # Residual connection.
         residual = hidden_states
+        if self.config.fp32_residual_connection:
+            residual = residual.float()
 
         # Optional Layer norm after self-attention
         pre_cross_attn_layernorm_output = apply_module(self.pre_cross_attn_layernorm)(hidden_states)
@@ -715,6 +719,8 @@ def _forward_mlp(
 
         # Residual connection.
         residual = hidden_states
+        if self.config.fp32_residual_connection:
+            residual = residual.float()
 
         # Optional Layer norm post the cross-attention.
         pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)
@@ -1292,6 +1298,8 @@ def _forward_mlp_router(self, hidden_states, padding_mask=None):
         """
 
         residual = hidden_states
+        if self.config.fp32_residual_connection:
+            residual = residual.float()
         self.mlp.fwd_execution_map = "route"
         pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)
         router_outputs = self.mlp(