Selective-FP32 runtime and native-FP32 optimizer checkpoint lifecycle

megatron/core/fp8_utils.py

@@ -831,6 +831,29 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
 
         return fp8_context
 
+    def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False):
+        """Return a context manager that disables TE quantization.
+
+        Use this around submodule construction or execution that must stay in a higher
+        precision while its enclosing module uses an FP8 or FP4 context.
+
+        Args:
+            config: Transformer configuration that controls quantization.
+            is_init: Whether to disable the parameter-initialization context instead of
+                the forward autocast context.
+
+        Returns:
+            A disabled TE quantization context when quantization is active, otherwise a
+            no-op context.
+        """
+        if is_init:
+            if not (config.fp8_param or config.fp4_param):
+                return nullcontext()
+            return transformer_engine.pytorch.fp8_model_init(enabled=False)
+        if not (config.fp8 or config.fp4):
+            return nullcontext()
+        return transformer_engine.pytorch.fp8_autocast(enabled=False)
+
 else:
 
     def get_fp8_recipe(config: TransformerConfig):
@@ -841,6 +864,10 @@ def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool
         """Returns dummy fp8 context manager since TE is not available."""
         return nullcontext()
 
+    def get_fp8_disabled_context(config: TransformerConfig, is_init: bool = False):
+        """Return a no-op context manager since TE is not available."""
+        return nullcontext()
+
 
 if HAVE_TE:
     from transformer_engine.pytorch.fp8 import FP8GlobalStateManager

megatron/core/optimizer/optimizer.py

@@ -7,7 +7,6 @@
 import math
 import warnings
 from abc import ABC, abstractmethod
-from itertools import chain
 from logging import getLogger
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -1019,25 +1018,53 @@ def sharded_state_dict(
 
         state_dict = self.state_dict()
 
+        # Optimizer state ids enumerate the inner optimizer params: the fp32 main
+        # copies of float16 params and the native fp32 params, interleaved in the
+        # original param-group order. Yield the model-side param for each inner
+        # param in that order so the ids line up even when both kinds are present.
+        def model_params_in_optimizer_order():
+            for inner_group, float16_group, fp32_group in zip(
+                self.optimizer.param_groups, self.float16_groups, self.fp32_from_fp32_groups
+            ):
+                float16_params = iter(float16_group)
+                fp32_param_ids = {id(param) for param in fp32_group}
+                for param in inner_group['params']:
+                    yield param if id(param) in fp32_param_ids else next(float16_params)
+
         id_to_sharded_param_map = get_param_id_to_sharded_param_map(
-            model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups)
+            model_sharded_state_dict, model_params_in_optimizer_order()
         )
 
         # Convert fp32_from_fp16_params
         assert len(state_dict['fp32_from_fp16_params']) == len(
             state_dict['optimizer']['param_groups']
         )
+        # State ids of the fp32 main copies only, skipping native fp32 params.
+        float16_param_ids_per_group = []
+        for state_group, inner_group, fp32_group in zip(
+            state_dict['optimizer']['param_groups'],
+            self.optimizer.param_groups,
+            self.fp32_from_fp32_groups,
+        ):
+            fp32_param_ids = {id(param) for param in fp32_group}
+            float16_param_ids_per_group.append(
+                [
+                    param_id
+                    for param_id, param in zip(state_group['params'], inner_group['params'])
+                    if id(param) not in fp32_param_ids
+                ]
+            )
         state_dict['fp32_from_fp16_params'] = [
             [
                 make_sharded_optimizer_tensor(
                     id_to_sharded_param_map[param_id],
                     fp32_param,
                     prefix=f'optimizer.state.fp32_param',
                 )
-                for param_id, fp32_param in zip(state_group['params'], fp32_group)
+                for param_id, fp32_param in zip(param_ids, fp32_group)
             ]
-            for fp32_group, state_group in zip(
-                state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups']
+            for fp32_group, param_ids in zip(
+                state_dict['fp32_from_fp16_params'], float16_param_ids_per_group
             )
         ]
 

megatron/core/transformer/module.py

@@ -433,6 +433,40 @@ def float_conversion(val):
     return conversion_helper(val, float_conversion)
 
 
+def mark_keep_in_fp32(tensor: torch.Tensor) -> torch.Tensor:
+    """Mark a parameter or buffer so that ``Float16Module`` keeps it in FP32.
+
+    Args:
+        tensor: The parameter or buffer to mark.
+
+    Returns:
+        The same tensor, for call-site convenience.
+    """
+    tensor.keep_in_fp32 = True
+    return tensor
+
+
+def convert_module_to_dtype_except_fp32_marked(
+    module: torch.nn.Module, dtype: torch.dtype
+) -> torch.nn.Module:
+    """Cast floating-point parameters and buffers except those marked to stay in FP32.
+
+    Args:
+        module: The module to convert in place.
+        dtype: The target floating-point dtype.
+
+    Returns:
+        The converted module.
+    """
+    return module._apply(
+        lambda tensor: (
+            tensor.to(dtype)
+            if tensor.is_floating_point() and not getattr(tensor, 'keep_in_fp32', False)
+            else tensor
+        )
+    )
+
+
 class Float16Module(MegatronModule):
     """Float 16 Module.
 
@@ -455,13 +489,17 @@ def __init__(self, config: TransformerConfig, module: torch.nn.Module):
         self.pg_collection = getattr(module, 'pg_collection', None)
 
         if self.fp16:
-            self.add_module('module', module.half())
+            self.add_module(
+                'module', convert_module_to_dtype_except_fp32_marked(module, torch.half)
+            )
 
             def float16_convertor(val):
                 return val.half()
 
         elif self.bf16:
-            self.add_module('module', module.bfloat16())
+            self.add_module(
+                'module', convert_module_to_dtype_except_fp32_marked(module, torch.bfloat16)
+            )
 
             def float16_convertor(val):
                 return val.bfloat16()

tests/unit_tests/dist_checkpointing/test_optimizer.py

@@ -79,6 +79,25 @@ def sharded_state_dict(self):
         return sharded_state_dict
 
 
+class NativeFp32Model(torch.nn.Module):
+    """Three parameters that can be converted to an interleaved BF16/FP32/BF16 group."""
+
+    def __init__(self):
+        super().__init__()
+        self.pre = torch.nn.Linear(8, 8, bias=False)
+        self.gate = torch.nn.Parameter(torch.zeros(24, dtype=torch.float32))
+        self.post = torch.nn.Linear(8, 8, bias=False)
+        self.config = TransformerConfig(
+            hidden_size=8, num_attention_heads=1, num_layers=1, bf16=True
+        )
+
+    def sharded_state_dict(self):
+        return {
+            key: ShardedTensor.from_rank_offsets(key, value)
+            for key, value in self.state_dict(keep_vars=True).items()
+        }
+
+
 class SwigluFactoryModel(torch.nn.Module):
     def __init__(self, pp_separate_model: bool = False):
         super().__init__()
@@ -238,6 +257,59 @@ def test_optimizer_params(self, tmp_path_dist_ckpt):
             ]
         )
 
+    def test_float16_optimizer_with_native_fp32_params(self):
+        """Native FP32 state ids must remain correct between two BF16 parameters."""
+        from megatron.core.optimizer import OptimizerConfig
+        from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params
+        from megatron.core.transformer.module import (
+            convert_module_to_dtype_except_fp32_marked,
+            mark_keep_in_fp32,
+        )
+
+        Utils.initialize_model_parallel(1, 1)
+        model = NativeFp32Model().cuda()
+        model.gate = mark_keep_in_fp32(model.gate)
+        convert_module_to_dtype_except_fp32_marked(model, torch.bfloat16)
+        assert model.pre.weight.dtype == torch.bfloat16
+        assert model.gate.dtype == torch.float32
+        assert model.post.weight.dtype == torch.bfloat16
+
+        # Use an explicit BF16/FP32/BF16 optimizer order. Module.parameters()
+        # would yield the root gate before parameters owned by child modules.
+        ordered_params = [model.pre.weight, model.gate, model.post.weight]
+        for param in ordered_params:
+            param.grad = torch.zeros_like(param)
+        inner_optim = Adam(ordered_params)
+        inner_optim.step()
+
+        optim = Float16OptimizerWithFloat16Params(
+            inner_optim,
+            OptimizerConfig(optimizer='adam', lr=1e-4, bf16=True),
+            None,
+            lambda opt, cfg: None,
+        )
+        sharded_state_dict = optim.sharded_state_dict(model.sharded_state_dict())
+
+        # FP32 main copies pair with the BF16 params only, in optimizer order.
+        fp32_params = sharded_state_dict['fp32_from_fp16_params'][0]
+        assert [(sharded.key, tuple(sharded.data.shape)) for sharded in fp32_params] == [
+            ('optimizer.state.fp32_param.pre.weight', (8, 8)),
+            ('optimizer.state.fp32_param.post.weight', (8, 8)),
+        ]
+
+        # Per-param state maps every param, including the native FP32 one, to the right key.
+        state = sharded_state_dict['optimizer']['state']
+        expected = {0: ('pre.weight', (8, 8)), 1: ('gate', (24,)), 2: ('post.weight', (8, 8))}
+        for param_id, (model_key, shape) in expected.items():
+            for state_key in ('exp_avg', 'exp_avg_sq'):
+                sharded = state[param_id][state_key]
+                assert sharded.key == f'optimizer.state.{state_key}.{model_key}', sharded.key
+                assert tuple(sharded.data.shape) == shape, (
+                    param_id,
+                    sharded.key,
+                    sharded.data.shape,
+                )
+
 
 def initialize_pp_agnostic_model(pre_process=True, post_process=True, seed=0, **config_kwargs):
     torch.manual_seed(seed)

tests/unit_tests/test_fp8_utils.py

@@ -10,6 +10,31 @@
 from tests.unit_tests.test_utilities import Utils
 
 
+@pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed")
+@pytest.mark.parametrize(
+    ("is_init", "config_values", "te_helper"),
+    [
+        (
+            False,
+            {"fp8": "hybrid", "fp4": None, "fp8_param": False, "fp4_param": False},
+            "fp8_autocast",
+        ),
+        (True, {"fp8": None, "fp4": None, "fp8_param": True, "fp4_param": False}, "fp8_model_init"),
+    ],
+)
+def test_get_fp8_disabled_context_uses_disabled_te_context(is_init, config_values, te_helper):
+    config = Mock(**config_values)
+    disabled_context = Mock()
+
+    with patch.object(
+        fp8_utils.transformer_engine.pytorch, te_helper, return_value=disabled_context
+    ) as te_context:
+        result = fp8_utils.get_fp8_disabled_context(config, is_init=is_init)
+
+    assert result is disabled_context
+    te_context.assert_called_once_with(enabled=False)
+
+
 class MockTELinear(nn.Module):
     """Mock TE Linear module for testing."""
 

tests/unit_tests/transformer/test_module.py

@@ -4,7 +4,7 @@
 import torch
 
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
-from megatron.core.transformer.module import Float16Module, MegatronModule
+from megatron.core.transformer.module import Float16Module, MegatronModule, mark_keep_in_fp32
 from megatron.core.transformer.transformer_config import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
@@ -163,3 +163,18 @@ def test_bf16_module(self):
         x = torch.ones((2, 2)).cuda()
         # inputs are converted to bf16 then outputs are converted to fp32
         assert bf16_module(x).dtype == torch.float32
+
+    @pytest.mark.parametrize(
+        ('precision', 'dtype'), [('fp16', torch.float16), ('bf16', torch.bfloat16)]
+    )
+    def test_keep_in_fp32_params(self, precision, dtype):
+        transformer_config = self.transformer_config
+        megatron_module = self.megatron_module
+        megatron_module.fp32_param = mark_keep_in_fp32(
+            torch.nn.Parameter(torch.zeros(4, dtype=torch.float32, device='cuda'))
+        )
+        setattr(transformer_config, precision, True)
+        float16_module = Float16Module(config=transformer_config, module=megatron_module)
+
+        assert float16_module.module.linear.weight.dtype == dtype
+        assert float16_module.module.fp32_param.dtype == torch.float32