From e0e74fd2636fade5f7962c879ceb6247a02476a7 Mon Sep 17 00:00:00 2001 From: mikail Date: Thu, 19 Feb 2026 15:11:34 -0800 Subject: [PATCH 1/6] Fixed fp32 residuals Signed-off-by: mikail --- megatron/core/fusions/fused_bias_dropout.py | 14 +- megatron/core/ssm/mamba_block.py | 5 - megatron/core/ssm/mamba_layer.py | 6 +- .../core/transformer/transformer_config.py | 6 + .../core/transformer/transformer_layer.py | 8 + .../fusions/test_bias_dropout_fusion.py | 294 +++++++++++++++++- 6 files changed, 317 insertions(+), 16 deletions(-) diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py index 336452562b3..2eb4007f75c 100644 --- a/megatron/core/fusions/fused_bias_dropout.py +++ b/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 diff --git a/megatron/core/ssm/mamba_block.py b/megatron/core/ssm/mamba_block.py index 1a5cfdc6f71..64c283fe355 100644 --- a/megatron/core/ssm/mamba_block.py +++ b/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, diff --git a/megatron/core/ssm/mamba_layer.py b/megatron/core/ssm/mamba_layer.py index 8c2f59369a5..9b6d39a581a 100644 --- a/megatron/core/ssm/mamba_layer.py +++ b/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) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1aad3c4b89f..c37ae0ff231 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -930,6 +930,12 @@ 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: + 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." diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index c9cf57a4eb0..435fe26db8c 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/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( diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index 6303a87ba26..3b1304a0f6d 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -1,8 +1,17 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + import pytest import torch -from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.fusions.fused_bias_dropout import ( + _bias_dropout_add_func, + get_bias_dropout_add, +) + +# --------------------------------------------------------------------------- +# Existing test: fused vs. unfused parity (same dtype) +# --------------------------------------------------------------------------- @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("training", [True, False]) @@ -46,3 +55,286 @@ def test_bias_dropout_add(dtype, training): # In‑place check for inference assert out_fused.data_ptr() == x_fused.data_ptr() assert torch.allclose(out_fused, x_fused, **tols) + + +# ============================================================================ +# Tests for fp32 residual connection fix +# ============================================================================ +# +# The fix reverses the casting direction in _bias_dropout_add_func: when +# x is bf16/fp16 and residual is fp32, x (and bias) are upcast to fp32 +# so that the residual stream stays in fp32. The OLD (broken) code did the +# opposite: it downcast the fp32 residual to match x's dtype, destroying +# the fp32 residual stream from the very first layer. +# ============================================================================ + + +class TestFp32ResidualPreservation: + """Tests that _bias_dropout_add_func preserves fp32 residual dtype.""" + + device = "cuda" + B, H = 16, 64 + + # -- helpers -------------------------------------------------------- + + @staticmethod + def _reference_bias_dropout_add(x, bias, residual, prob, training): + """Manual reference computation in fp32 (no dropout for determinism).""" + x_fp32 = x.float() + r_fp32 = residual.float() + if bias is not None: + b_fp32 = bias.float() + return r_fp32 + torch.nn.functional.dropout( + x_fp32 + b_fp32, p=prob, training=training + ) + return r_fp32 + torch.nn.functional.dropout( + x_fp32, p=prob, training=training + ) + + # -- core: output dtype must follow residual, not x ----------------- + + @pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("training", [True, False]) + @pytest.mark.parametrize("has_bias", [True, False]) + def test_output_dtype_is_residual_dtype(self, x_dtype, training, has_bias): + """The output tensor must have the same dtype as the residual (fp32), + NOT x's dtype. This is the central invariant of the fix.""" + x = torch.randn(self.B, self.H, dtype=x_dtype, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) + bias = torch.randn(self.H, dtype=x_dtype, device=self.device) if has_bias else None + + out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=training) + + assert out.dtype == torch.float32, ( + f"Output dtype {out.dtype} must be fp32 (residual dtype), not {x_dtype}" + ) + + # -- numerical correctness of the upcast path ---------------------- + + @pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("has_bias", [True, False]) + def test_numerical_correctness_with_fp32_residual(self, x_dtype, has_bias): + """Forward result should match a reference computed entirely in fp32.""" + torch.manual_seed(7) + x = torch.randn(self.B, self.H, dtype=x_dtype, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) + bias = torch.randn(self.H, dtype=x_dtype, device=self.device) if has_bias else None + + out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=True) + + ref = self._reference_bias_dropout_add(x, bias, residual, prob=0.0, training=True) + # fp32 tolerance – the only imprecision is the upcast from bf16/fp16 + assert torch.allclose(out, ref, rtol=1e-5, atol=1e-5), ( + f"Max diff = {(out - ref).abs().max().item()}" + ) + + # -- backward: gradients flow correctly through the upcast --------- + + @pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("has_bias", [True, False]) + def test_backward_with_fp32_residual(self, x_dtype, has_bias): + """Gradients should be computed for x and residual when dtypes differ.""" + x = torch.randn( + self.B, self.H, dtype=x_dtype, device=self.device, requires_grad=True + ) + residual = torch.randn( + self.B, self.H, dtype=torch.float32, device=self.device, requires_grad=True + ) + bias = ( + torch.randn(self.H, dtype=x_dtype, device=self.device, requires_grad=True) + if has_bias + else None + ) + + out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=True) + grad_out = torch.randn_like(out) + out.backward(grad_out) + + assert x.grad is not None, "x.grad must not be None" + assert residual.grad is not None, "residual.grad must not be None" + assert x.grad.dtype == x_dtype, f"x.grad dtype should be {x_dtype}" + assert residual.grad.dtype == torch.float32, "residual.grad must stay fp32" + if has_bias: + assert bias.grad is not None, "bias.grad must not be None" + + # -- same dtype: no regression (both fp32 or both bf16) ------------ + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("has_bias", [True, False]) + def test_same_dtype_no_regression(self, dtype, has_bias): + """When x and residual share the same dtype, behaviour is unchanged.""" + torch.manual_seed(99) + x = torch.randn(self.B, self.H, dtype=dtype, device=self.device) + residual = torch.randn(self.B, self.H, dtype=dtype, device=self.device) + bias = torch.randn(self.H, dtype=dtype, device=self.device) if has_bias else None + + out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=True) + + assert out.dtype == dtype + ref = self._reference_bias_dropout_add(x, bias, residual, prob=0.0, training=True) + ref = ref.to(dtype) + tols = dict(rtol=1e-5, atol=1e-5) if dtype == torch.float32 else dict(rtol=2e-2, atol=1e-2) + assert torch.allclose(out, ref, **tols) + + # -- inference in-place optimisation still fires for same dtype ---- + + def test_inplace_inference_same_dtype(self): + """In eval mode with no grad, the in-place path should still work.""" + x = torch.randn(self.B, self.H, dtype=torch.bfloat16, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.bfloat16, device=self.device) + bias = torch.randn(self.H, dtype=torch.bfloat16, device=self.device) + + x_ptr = x.data_ptr() + out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=False) + # In-place: output should reuse x's storage + assert out.data_ptr() == x_ptr + + # -- inference with mixed dtypes should NOT be in-place ------------ + + @pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16]) + def test_no_inplace_when_dtypes_differ(self, x_dtype): + """When x is low-precision but residual is fp32, the upcast creates + a new tensor so the in-place optimisation must NOT fire.""" + x = torch.randn(self.B, self.H, dtype=x_dtype, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) + + x_ptr = x.data_ptr() + out = _bias_dropout_add_func((x, None), residual, prob=0.0, training=False) + # The upcast `x = x.to(residual.dtype)` creates a new tensor, + # so the output must NOT alias the original x buffer. + assert out.dtype == torch.float32 + assert out.data_ptr() != x_ptr + + # -- dropout prob > 0 still works with mixed dtypes ---------------- + + @pytest.mark.parametrize("has_bias", [True, False]) + def test_dropout_with_fp32_residual(self, has_bias): + """Smoke test: non-zero dropout with mixed dtypes doesn't crash.""" + x = torch.randn(self.B, self.H, dtype=torch.bfloat16, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) + bias = ( + torch.randn(self.H, dtype=torch.bfloat16, device=self.device) + if has_bias + else None + ) + + out = _bias_dropout_add_func((x, bias), residual, prob=0.5, training=True) + assert out.dtype == torch.float32 + # With dropout, some elements should be zeroed (before residual add) + # so the output shouldn't be identical to the no-dropout case. + out_no_drop = _bias_dropout_add_func( + (x, bias), residual, prob=0.0, training=True + ) + # They *can* be equal with very low probability; just check dtypes + assert out_no_drop.dtype == torch.float32 + + # -- get_bias_dropout_add wrappers preserve fp32 residual ---------- + + @pytest.mark.parametrize("training", [True, False]) + @pytest.mark.parametrize("fused", [True, False]) + def test_get_bias_dropout_add_fp32_residual(self, training, fused): + """All four (training × fused) wrappers returned by get_bias_dropout_add + should preserve fp32 residual dtype.""" + fn = get_bias_dropout_add(training=training, fused=fused) + x = torch.randn(self.B, self.H, dtype=torch.bfloat16, device=self.device) + residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) + bias = torch.randn(self.H, dtype=torch.bfloat16, device=self.device) + + out = fn((x, bias), residual, prob=0.0) + assert out.dtype == torch.float32, ( + f"get_bias_dropout_add(training={training}, fused={fused}) " + f"returned {out.dtype}, expected fp32" + ) + + +# ============================================================================ +# Tests simulating the multi-layer residual stream scenario +# ============================================================================ + + +class TestFp32ResidualStreamAcrossLayers: + """Simulates what happens across multiple transformer layers to ensure + the fp32 residual stream is not degraded.""" + + device = "cuda" + + def test_residual_stays_fp32_across_simulated_layers(self): + """Simulate N layers of bias-dropout-add with bf16 x and fp32 residual. + The residual should remain fp32 throughout — this is the scenario that + was broken before the fix.""" + B, H = 4, 32 + num_layers = 8 + + # Start with an fp32 residual (as the embedding layer would emit) + residual = torch.randn(B, H, dtype=torch.float32, device=self.device) + + for layer_idx in range(num_layers): + # Each layer produces bf16 output (as the attention/MLP would) + x = torch.randn(B, H, dtype=torch.bfloat16, device=self.device) + bias = torch.randn(H, dtype=torch.bfloat16, device=self.device) + + residual = _bias_dropout_add_func( + (x, bias), residual, prob=0.0, training=True + ) + + assert residual.dtype == torch.float32, ( + f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" + ) + + def test_residual_stays_fp32_no_bias(self): + """Same multi-layer simulation but without bias tensors.""" + B, H = 4, 32 + num_layers = 8 + + residual = torch.randn(B, H, dtype=torch.float32, device=self.device) + + for layer_idx in range(num_layers): + x = torch.randn(B, H, dtype=torch.bfloat16, device=self.device) + + residual = _bias_dropout_add_func( + (x, None), residual, prob=0.0, training=True + ) + + assert residual.dtype == torch.float32, ( + f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" + ) + + def test_fp32_residual_precision_advantage(self): + """Demonstrate that fp32 residuals accumulate more accurately than + bf16 residuals over many additions — the whole point of the feature.""" + B, H = 2, 16 + num_layers = 50 + torch.manual_seed(42) + + # Ground truth: everything in fp64 + residual_fp64 = torch.randn(B, H, dtype=torch.float64, device=self.device) + + # Track fp32 and bf16 residual streams + residual_fp32 = residual_fp64.float() + residual_bf16 = residual_fp64.to(torch.bfloat16) + + for _ in range(num_layers): + x_fp64 = torch.randn(B, H, dtype=torch.float64, device=self.device) + x_bf16 = x_fp64.to(torch.bfloat16) + + # fp64 reference + residual_fp64 = residual_fp64 + x_fp64 + + # fp32 residual path (the fix) + residual_fp32 = _bias_dropout_add_func( + (x_bf16, None), residual_fp32, prob=0.0, training=True + ) + + # bf16 residual path (the old broken behaviour) + residual_bf16 = _bias_dropout_add_func( + (x_bf16, None), residual_bf16.to(torch.bfloat16), prob=0.0, training=True + ) + + err_fp32 = (residual_fp32.double() - residual_fp64).abs().mean().item() + err_bf16 = (residual_bf16.double() - residual_fp64).abs().mean().item() + + # fp32 residual should be meaningfully more precise + assert err_fp32 < err_bf16, ( + f"fp32 residual error ({err_fp32:.6e}) should be less than " + f"bf16 residual error ({err_bf16:.6e})" + ) From cafae2ccfa3d36d41685bcb9eba292a4193549aa Mon Sep 17 00:00:00 2001 From: mikail Date: Thu, 19 Feb 2026 16:34:37 -0800 Subject: [PATCH 2/6] add a warnings.warn here so users aren't confused when their pipeline dtype changes Signed-off-by: mikail --- megatron/core/transformer/transformer_config.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c37ae0ff231..20f01227dda 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -934,6 +934,11 @@ def __post_init__(self): # 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: + warnings.warn( + 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: From 9aef970450bd6918cbb41514ca49df7193141037 Mon Sep 17 00:00:00 2001 From: mikail Date: Thu, 19 Feb 2026 16:46:12 -0800 Subject: [PATCH 3/6] don't warn when dtype is already fp32 --- megatron/core/transformer/transformer_config.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 20f01227dda..bf6730f9ce9 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -934,11 +934,12 @@ def __post_init__(self): # 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: - warnings.warn( - 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." - ) + if self.pipeline_dtype != torch.float: + warnings.warn( + 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: From 72aa2c771174334da76f8189b957bec032ccb41d Mon Sep 17 00:00:00 2001 From: mikail Date: Mon, 23 Feb 2026 11:28:16 -0800 Subject: [PATCH 4/6] added linting for unit test Signed-off-by: mikail --- .../fusions/test_bias_dropout_fusion.py | 60 +++++++------------ 1 file changed, 21 insertions(+), 39 deletions(-) diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index 3b1304a0f6d..e62648de574 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -3,16 +3,14 @@ import pytest import torch -from megatron.core.fusions.fused_bias_dropout import ( - _bias_dropout_add_func, - get_bias_dropout_add, -) +from megatron.core.fusions.fused_bias_dropout import _bias_dropout_add_func, get_bias_dropout_add # --------------------------------------------------------------------------- # Existing test: fused vs. unfused parity (same dtype) # --------------------------------------------------------------------------- + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("training", [True, False]) def test_bias_dropout_add(dtype, training): @@ -84,12 +82,8 @@ def _reference_bias_dropout_add(x, bias, residual, prob, training): r_fp32 = residual.float() if bias is not None: b_fp32 = bias.float() - return r_fp32 + torch.nn.functional.dropout( - x_fp32 + b_fp32, p=prob, training=training - ) - return r_fp32 + torch.nn.functional.dropout( - x_fp32, p=prob, training=training - ) + return r_fp32 + torch.nn.functional.dropout(x_fp32 + b_fp32, p=prob, training=training) + return r_fp32 + torch.nn.functional.dropout(x_fp32, p=prob, training=training) # -- core: output dtype must follow residual, not x ----------------- @@ -105,9 +99,9 @@ def test_output_dtype_is_residual_dtype(self, x_dtype, training, has_bias): out = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=training) - assert out.dtype == torch.float32, ( - f"Output dtype {out.dtype} must be fp32 (residual dtype), not {x_dtype}" - ) + assert ( + out.dtype == torch.float32 + ), f"Output dtype {out.dtype} must be fp32 (residual dtype), not {x_dtype}" # -- numerical correctness of the upcast path ---------------------- @@ -124,9 +118,9 @@ def test_numerical_correctness_with_fp32_residual(self, x_dtype, has_bias): ref = self._reference_bias_dropout_add(x, bias, residual, prob=0.0, training=True) # fp32 tolerance – the only imprecision is the upcast from bf16/fp16 - assert torch.allclose(out, ref, rtol=1e-5, atol=1e-5), ( - f"Max diff = {(out - ref).abs().max().item()}" - ) + assert torch.allclose( + out, ref, rtol=1e-5, atol=1e-5 + ), f"Max diff = {(out - ref).abs().max().item()}" # -- backward: gradients flow correctly through the upcast --------- @@ -134,9 +128,7 @@ def test_numerical_correctness_with_fp32_residual(self, x_dtype, has_bias): @pytest.mark.parametrize("has_bias", [True, False]) def test_backward_with_fp32_residual(self, x_dtype, has_bias): """Gradients should be computed for x and residual when dtypes differ.""" - x = torch.randn( - self.B, self.H, dtype=x_dtype, device=self.device, requires_grad=True - ) + x = torch.randn(self.B, self.H, dtype=x_dtype, device=self.device, requires_grad=True) residual = torch.randn( self.B, self.H, dtype=torch.float32, device=self.device, requires_grad=True ) @@ -212,19 +204,13 @@ def test_dropout_with_fp32_residual(self, has_bias): """Smoke test: non-zero dropout with mixed dtypes doesn't crash.""" x = torch.randn(self.B, self.H, dtype=torch.bfloat16, device=self.device) residual = torch.randn(self.B, self.H, dtype=torch.float32, device=self.device) - bias = ( - torch.randn(self.H, dtype=torch.bfloat16, device=self.device) - if has_bias - else None - ) + bias = torch.randn(self.H, dtype=torch.bfloat16, device=self.device) if has_bias else None out = _bias_dropout_add_func((x, bias), residual, prob=0.5, training=True) assert out.dtype == torch.float32 # With dropout, some elements should be zeroed (before residual add) # so the output shouldn't be identical to the no-dropout case. - out_no_drop = _bias_dropout_add_func( - (x, bias), residual, prob=0.0, training=True - ) + out_no_drop = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=True) # They *can* be equal with very low probability; just check dtypes assert out_no_drop.dtype == torch.float32 @@ -273,13 +259,11 @@ def test_residual_stays_fp32_across_simulated_layers(self): x = torch.randn(B, H, dtype=torch.bfloat16, device=self.device) bias = torch.randn(H, dtype=torch.bfloat16, device=self.device) - residual = _bias_dropout_add_func( - (x, bias), residual, prob=0.0, training=True - ) + residual = _bias_dropout_add_func((x, bias), residual, prob=0.0, training=True) - assert residual.dtype == torch.float32, ( - f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" - ) + assert ( + residual.dtype == torch.float32 + ), f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" def test_residual_stays_fp32_no_bias(self): """Same multi-layer simulation but without bias tensors.""" @@ -291,13 +275,11 @@ def test_residual_stays_fp32_no_bias(self): for layer_idx in range(num_layers): x = torch.randn(B, H, dtype=torch.bfloat16, device=self.device) - residual = _bias_dropout_add_func( - (x, None), residual, prob=0.0, training=True - ) + residual = _bias_dropout_add_func((x, None), residual, prob=0.0, training=True) - assert residual.dtype == torch.float32, ( - f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" - ) + assert ( + residual.dtype == torch.float32 + ), f"Layer {layer_idx}: residual dtype degraded to {residual.dtype}" def test_fp32_residual_precision_advantage(self): """Demonstrate that fp32 residuals accumulate more accurately than From c15e1dd78cdb2036331318a2c8af15e20a01d779 Mon Sep 17 00:00:00 2001 From: mikail Date: Mon, 23 Feb 2026 11:28:41 -0800 Subject: [PATCH 5/6] log only on 1 rank when changing pipeline dtype Signed-off-by: mikail --- megatron/core/transformer/transformer_config.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bf6730f9ce9..f811b09fd95 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/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 @@ -935,10 +939,12 @@ def __post_init__(self): # 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: - warnings.warn( + 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." + f"residual stream dtype between pipeline stages.", ) self.pipeline_dtype = torch.float From 5580aee2baa3b5fc7a81f4e67a32b3932c3f9dcb Mon Sep 17 00:00:00 2001 From: mikail Date: Mon, 23 Feb 2026 20:44:26 -0800 Subject: [PATCH 6/6] used isosort for linting Signed-off-by: mikail --- tests/unit_tests/fusions/test_bias_dropout_fusion.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit_tests/fusions/test_bias_dropout_fusion.py b/tests/unit_tests/fusions/test_bias_dropout_fusion.py index e62648de574..f8b23900543 100644 --- a/tests/unit_tests/fusions/test_bias_dropout_fusion.py +++ b/tests/unit_tests/fusions/test_bias_dropout_fusion.py @@ -5,7 +5,6 @@ from megatron.core.fusions.fused_bias_dropout import _bias_dropout_add_func, get_bias_dropout_add - # --------------------------------------------------------------------------- # Existing test: fused vs. unfused parity (same dtype) # ---------------------------------------------------------------------------