diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 7e1fbcc1b58..26ea7774c47 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -145,6 +145,9 @@ class TEQuantizationRecipe: If no FP8 or FP4 quantization is configured, the recipe is execution in high-precision (BF16). """ + + recipe_attrs: Optional[Dict[str, Any]] = None + """Attributes to set in the TE quantization recipe class.""" custom_recipe_factory: Optional[str] = None """The path to a custom recipe factory if a custom Fp4 or Fp8 recipe is configured""" fp8_format: str = "e4m3" @@ -293,6 +296,11 @@ def _get_fp8_model_init_for_quant_recipe(qrecipe: TEQuantizationRecipe): else: raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp4_quantization_recipe}") + # Set recipe attrs + if quant_recipe is not None and qrecipe.recipe_attrs is not None: + for key, val in qrecipe.recipe_attrs.items(): + setattr(quant_recipe, key, val) + return fp8_model_init( enabled=enabled, recipe=quant_recipe, @@ -355,7 +363,12 @@ def _get_fp8_autocast_for_quant_recipe(qrecipe: TEQuantizationRecipe): if qrecipe.fp4_quantization_recipe == Fp4Recipe.nvfp4: quant_recipe = te.common.recipe.NVFP4BlockScaling() else: - raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp8_quantization_recipe}") + raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp4_quantization_recipe}") + + # Set recipe attrs + if quant_recipe is not None and qrecipe.recipe_attrs is not None: + for key, val in qrecipe.recipe_attrs.items(): + setattr(quant_recipe, key, val) return fp8_autocast(enabled=True, fp8_recipe=quant_recipe, fp8_group=amax_group) @@ -3294,9 +3307,39 @@ def forward(self, hidden_states: torch.Tensor, **kwargs) -> Tuple[Tensor, Option # Build fused impl and cache recipe lazily on first forward pass. # Both are created once and reused — avoids object creation every call. + # This recipe cache is a hack (TEFusedMLP does not need + # it). Keep it for now for expediency, but future + # developers are warned that this should be refactored to + # use the model's recipe config like the other TE modules. if not hasattr(self, '_recipe'): if os.getenv("FP4_RECIPE", "") == "nvfp4": self._recipe = te.common.recipe.NVFP4BlockScaling() + elif os.getenv("FP4_RECIPE", "") == "nvfp4_ue5m3": + + def _make_nvfp4_ue5m3_quantizer( + role: te.pytorch.QuantizerRole, + ) -> te.pytorch.NVFP4Quantizer: + """Construct NVFP4 quantizer with UE5M3 scales and no per-tensor + scale for activations.""" + tensor_type = role.tensor_type if role is not None else None + if not tensor_type: + tensor_type = "input" + with_rht = tensor_type in ("input", "grad_output") + return te.pytorch.NVFP4Quantizer( + scale_dtype=te.pytorch.DType.kFloat8UE5M3, + with_rht=with_rht, + with_post_rht_amax=with_rht, + with_2d_quantization=tensor_type == "weight", + stochastic_rounding=False, + with_random_sign_mask=False, + disable_second_level_scale=tensor_type == "input", + ) + + # Construct custom recipe for NVFP4 with UE5M3 scales + self._recipe = te.common.recipe.CustomRecipe( + qfactory=_make_nvfp4_ue5m3_quantizer + ) + self._recipe.enable_cutedsl_fused_grouped_mlp = True else: self._recipe = te.common.recipe.MXFP8BlockScaling() recipe = self._recipe diff --git a/megatron/core/fp4_utils.py b/megatron/core/fp4_utils.py index a31ba7630c8..d337a993f30 100644 --- a/megatron/core/fp4_utils.py +++ b/megatron/core/fp4_utils.py @@ -225,11 +225,9 @@ def get_fp4_recipe(config: TransformerConfig): fp8_dpa=config.fp8_dot_product_attention ) except AttributeError: - raise ValueError( - """NVFP4BlockScaling recipe is not available in this version of + raise ValueError("""NVFP4BlockScaling recipe is not available in this version of Transformer Engine. Please make sure you are using TE version - >= 2.7.0.dev0.""" - ) + >= 2.7.0.dev0.""") elif config.fp4_recipe == Fp4Recipe.custom: fp4_recipe = _get_custom_recipe(config.fp4_quantizer_factory) else: @@ -238,10 +236,14 @@ def get_fp4_recipe(config: TransformerConfig): "Please make sure you are using a compatible TE version >= 2.7.0.dev0." ) else: - raise ValueError( - """FP4 support requires TransformerEngine version >= 2.7.0.dev0 - for NVFP4BlockScaling.""" - ) + raise ValueError("""FP4 support requires TransformerEngine version >= 2.7.0.dev0 + for NVFP4BlockScaling.""") + + # Set recipe attrs + if fp4_recipe is not None and config.fp4_recipe_attrs is not None: + for key, val in config.fp4_recipe_attrs.items(): + setattr(fp4_recipe, key, val) + return fp4_recipe def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False): diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 80fcd41f3ba..f2b2c65080a 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -807,6 +807,12 @@ def get_fp8_recipe(config: TransformerConfig): fp8_format=fp8_format, override_linear_precision=(False, False, not config.fp8_wgrad), ) + + # Set recipe attrs + if fp8_recipe is not None and config.fp8_recipe_attrs is not None: + for key, val in config.fp8_recipe_attrs.items(): + setattr(fp8_recipe, key, val) + return fp8_recipe def get_fp8_recipe_for_a2a(a2a_dtype: str): diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 6eef5dee9ff..d7501325eb9 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -11,6 +11,8 @@ from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.fp4_utils import get_fp4_recipe +from megatron.core.fp8_utils import get_fp8_recipe from megatron.core.fusions.fused_bias_geglu import bias_geglu_impl from megatron.core.fusions.fused_bias_gelu import bias_gelu_impl from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl @@ -451,18 +453,13 @@ def _validate_fused_grouped_swiglu(self) -> None: def _get_fused_grouped_swiglu_recipe(self): """Create the TE recipe used to select the fused grouped MLP kernel.""" if self._fused_grouped_swiglu_recipe is None: - fp4_recipe = getattr(self.config.fp4_recipe, "value", self.config.fp4_recipe) - fp8_recipe = getattr(self.config.fp8_recipe, "value", self.config.fp8_recipe) - if self.config.fp4 and fp4_recipe == "nvfp4": - self._fused_grouped_swiglu_recipe = te.common.recipe.NVFP4BlockScaling() - elif self.config.fp8 and fp8_recipe == "mxfp8": - self._fused_grouped_swiglu_recipe = te.common.recipe.MXFP8BlockScaling() + if self.config.fp8: + self._fused_grouped_swiglu_recipe = get_fp8_recipe(self.config) + elif self.config.fp4: + self._fused_grouped_swiglu_recipe = get_fp4_recipe(self.config) else: raise ValueError( - f"{self.__class__.__name__} requires fp4_recipe='nvfp4' or " - f"fp8_recipe='mxfp8', but got fp4={self.config.fp4}, " - f"fp4_recipe={self.config.fp4_recipe}, fp8={self.config.fp8}, " - f"fp8_recipe={self.config.fp8_recipe}." + f"{self.__class__.__name__} requires either FP8 or FP4 to be enabled." ) return self._fused_grouped_swiglu_recipe diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4c39148ebd2..1811d2a3eee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -4,7 +4,7 @@ import math import warnings from dataclasses import dataclass, field -from typing import Callable, List, Literal, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -638,6 +638,9 @@ class TransformerConfig(ModelParallelConfig): under the MXFP8 autocast context. Only active when fp8=True and fp8_recipe='mxfp8'.""" + fp8_recipe_attrs: Optional[Dict[str, Any]] = None + """Attributes to set in the FP8 recipe.""" + fp8_dot_product_attention: bool = False """When set to True, use the FP8 implementation of Dot Product Attention.""" @@ -678,8 +681,9 @@ class TransformerConfig(ModelParallelConfig): supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0).""" fp4_recipe: Optional[Literal['nvfp4', 'custom']] = "nvfp4" - """If set, enables the use of FP4 precision through Transformer Engine. Currently only - 'nvfp4' is supported which uses NVFP4BlockScaling recipe for Blackwell+ architecture.""" + """If set, enables the use of FP4 precision through Transformer Engine.""" + fp4_recipe_attrs: Optional[Dict[str, Any]] = None + """Attributes to set in the FP4 recipe.""" fp4_param: bool = field( default=False, metadata={"argparse_meta": {"arg_names": ["--fp4-param-gather"]}} diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c797d648161..541f0864e1e 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2369,6 +2369,8 @@ def _add_network_size_args(parser): "heterogeneous_block_specs", "hetereogenous_dist_checkpoint", "quant_recipe", + "fp8_recipe_attrs", + "fp4_recipe_attrs", # deprecated and no CLI arg exists "tp_comm_atomic_ag", "tp_comm_atomic_rs", diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 2fc55962ae8..c462c40851f 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -1,5 +1,6 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +from contextlib import nullcontext from unittest.mock import Mock, patch import pytest @@ -7,6 +8,7 @@ import torch.nn as nn from megatron.core import fp8_utils +from megatron.core.transformer import TransformerConfig from megatron.training.utils import get_device_arch_version from tests.unit_tests.test_utilities import Utils @@ -48,6 +50,44 @@ def test_get_fp8_disabled_context_uses_disabled_te_context(is_init, config_value te_context.assert_called_once_with(enabled=False) +def _dummy_quantizer_factory(*args, **kwargs): + return None + + +@pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine required") +def test_fp8_model_init_uses_custom_recipe_attrs(monkeypatch): + if not hasattr(fp8_utils.transformer_engine.common.recipe, "CustomRecipe"): + pytest.skip("CustomRecipe requires newer Transformer Engine") + + captured = {} + + def fake_fp8_model_init(enabled=False, recipe=None, **kwargs): + captured["enabled"] = enabled + captured["recipe"] = recipe + captured["kwargs"] = kwargs + return nullcontext() + + monkeypatch.setattr(fp8_utils.transformer_engine.pytorch, "fp8_model_init", fake_fp8_model_init) + + config = TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + fp8="e4m3", + fp8_param=True, + fp8_recipe="custom", + fp8_quantizer_factory="tests.unit_tests.test_fp8_utils._dummy_quantizer_factory", + fp8_recipe_attrs={"_test_recipe_flag": "set"}, + ) + + with fp8_utils.get_fp8_context(config, is_init=True): + pass + + assert captured["enabled"] is True + assert isinstance(captured["recipe"], fp8_utils.transformer_engine.common.recipe.CustomRecipe) + assert captured["recipe"]._test_recipe_flag == "set" + + class MockTELinear(nn.Module): """Mock TE Linear module for testing.""" diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index c71f665259f..97af246d0c7 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -116,6 +116,16 @@ def _patch_fake_shared_expert_te(monkeypatch, linear_cls=_FakeTELinear): fake_te = _fake_te_module(linear_cls) monkeypatch.setattr(shared_experts_module, "HAVE_TE", True) monkeypatch.setattr(shared_experts_module, "te", fake_te) + monkeypatch.setattr( + shared_experts_module, + "get_fp8_recipe", + lambda _config: fake_te.common.recipe.MXFP8BlockScaling(), + ) + monkeypatch.setattr( + shared_experts_module, + "get_fp4_recipe", + lambda _config: fake_te.common.recipe.NVFP4BlockScaling(), + ) monkeypatch.setattr(shared_experts_module, "is_te_min_version", lambda *args, **kwargs: True) monkeypatch.setattr(shared_experts_module, "get_pg_size", lambda group: 1) monkeypatch.setattr( @@ -137,9 +147,9 @@ def _fake_shared_expert(**config_kwargs): moe_shared_expert_glu_interleave_size=32, delay_wgrad_compute=False, sequence_parallel=False, - fp4=False, + fp4=None, fp4_recipe="nvfp4", - fp8=True, + fp8="e4m3", fp8_recipe="mxfp8", ) for key, value in config_kwargs.items():