Add config option to set attrs in TE quantization recipe

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

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):

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):

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
 

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"]}}

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",

tests/unit_tests/test_fp8_utils.py

@@ -1,12 +1,14 @@
 # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
 
+from contextlib import nullcontext
 from unittest.mock import Mock, patch
 
 import pytest
 import torch
 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."""
 

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():