From 786d17ba292c8e48b24d883a472aed0bcab2eb78 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 7 Aug 2026 10:39:24 +0000 Subject: [PATCH 1/6] Add config option to set attrs in TE quantization recipe Signed-off-by: Tim Moon --- .../core/extensions/transformer_engine.py | 52 +++++++++++++++++++ megatron/core/fp4_utils.py | 20 +++---- megatron/core/fp8_utils.py | 14 +++-- .../core/transformer/moe/shared_experts.py | 17 +++--- .../core/transformer/transformer_config.py | 42 ++++++++------- megatron/training/arguments.py | 2 + 6 files changed, 104 insertions(+), 43 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 7f7424e6c02..2af57c2becb 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, @@ -357,6 +365,11 @@ def _get_fp8_autocast_for_quant_recipe(qrecipe: TEQuantizationRecipe): else: raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp8_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) @@ -3222,9 +3235,48 @@ 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.""" + + from transformer_engine.pytorch.quantization import RecipeState + + # Parse tensor role + tensor_type = role.tensor_type if role is not None else None + + # Construct NVFP4 quantizer based on default NVFP4 recipe + fp4_recipe = te.common.recipe.NVFP4BlockScaling( + fp8_format=te.common.recipe.Format.UE5M3 + ) + (quantizer,) = RecipeState.create( + fp4_recipe, + mode="forward" if tensor_type != "grad_output" else "backward", + num_quantizers=1, + roles=[role], + ).make_quantizers() + + # Disable per-tensor scale for activations + if tensor_type == "input": + quantizer.disable_second_level_scale = True + + return quantizer + + # 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..a02b8fae620 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 - Transformer Engine. Please make sure you are using TE version - >= 2.7.0.dev0.""" - ) + 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.""") 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 895d46e9b3d..3e7a9caa4ef 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -347,11 +347,9 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F try: custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory) except AttributeError: - raise ValueError( - """CustomRecipe recipe is not available in this version of - Transformer Engine. Please make sure you are using TE version - >= 2.9.0.dev0.""" - ) + raise ValueError("""CustomRecipe recipe is not available in this version of + Transformer Engine. Please make sure you are using TE version + >= 2.9.0.dev0.""") return custom_recipe @@ -809,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_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False): diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 038a162f899..629b4a98d03 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 @@ -436,18 +438,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 66cfaded213..8bee69918a1 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 @@ -74,8 +74,8 @@ class TransformerConfig(ModelParallelConfig): mtp_loss_scaling_factor: Optional[float] = 0.1 """Weighting factor of Multi-Token Prediction (MTP) loss. - We compute the average of the MTP losses across all depths, - and multiply it the scaling factor to obtain the overall MTP loss, + We compute the average of the MTP losses across all depths, + and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. """ @@ -109,8 +109,8 @@ class TransformerConfig(ModelParallelConfig): - list: e.g., [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]. - PipelineParallelLayerLayout: a PipelineParallelLayerLayout object. If given either a string or a list, it will be transferred into a PipelineParallelLayerLayout - in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers - in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., + in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers + in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., vpp(i)pp(j), vpp(i)pp(j+1), ..., vpp(-1)pp(-2), vpp(-1)pp(-1). In the inner lists of layers, 'embedding' or 'E' denotes the embedding layer, 'loss' or 'L' denotes the loss function, and 'decoder' or 't' denotes the transformer decoder layer. @@ -161,8 +161,8 @@ class TransformerConfig(ModelParallelConfig): """Softmax scale for attention scaling.""" softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla' - """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. - Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and + """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. + Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and and learnable offset.""" num_query_groups: Optional[int] = field( @@ -222,7 +222,7 @@ class TransformerConfig(ModelParallelConfig): The stored input is casted back to the original precision before backprop compuatation.""" glu_linear_offset: float = 0.0 - """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only + """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only used when gated_linear_unit is True""" activation_func_clamp_value: Optional[float] = None @@ -347,7 +347,7 @@ class TransformerConfig(ModelParallelConfig): # linear attention #################### linear_attention_freq: Optional[Union[int, List[int]]] = None - """Frequency between LA (linear attention) layers + """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. Accepts either: - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer @@ -389,13 +389,13 @@ class TransformerConfig(ModelParallelConfig): embedding_init_method: Optional[Callable] = None """ - Method to initialize weights of the embedding layer. If None, will be set as described + Method to initialize weights of the embedding layer. If None, will be set as described in init_method above. """ embedding_init_method_std: Optional[float] = None """ - Standard deviation of the zero mean normal for the default initialization method for the + Standard deviation of the zero mean normal for the default initialization method for the embedding layer. If None, will be set to init_method_std. Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying weight decay on embedding weights to avoid shrinkage towards zero. @@ -616,6 +616,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.""" @@ -652,12 +655,13 @@ class TransformerConfig(ModelParallelConfig): fp4: Optional[Literal['e2m1']] = field( default=None, metadata={"argparse_meta": {"arg_names": ["--fp4-format"]}} ) - """If set, enables the use of FP4 precision through Transformer Engine. Currently only + """If set, enables the use of FP4 precision through Transformer Engine. Currently only 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"]}} @@ -690,12 +694,12 @@ class TransformerConfig(ModelParallelConfig): in the hidden_states gradient.""" moe_shared_expert_gate: bool = False - """Enable gate for shared expert. Only effective when + """Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.""" moe_shared_expert_overlap: bool = False """Enable overlapping between shared expert computations and dispatcher communications. - Without this, the shared experts execute before the router. + Without this, the shared experts execute before the router. Only effective when moe-shared-expert-intermediate-size is set. """ @@ -810,7 +814,7 @@ class TransformerConfig(ModelParallelConfig): no memory: the bias is replaced by the latest global-batch quantile estimate each step.""" moe_router_force_load_balancing: bool = False - """[Experimental] Force load balancing with random logits for MoE router, supports naive topk + """[Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark.""" moe_router_force_biased: Optional[float] = None @@ -1112,7 +1116,7 @@ class TransformerConfig(ModelParallelConfig): batch_invariant_mode: bool = False """If true, uses batch-invariant kernels that provide deterministic forward execution regardless of batch size. This ensures bitwise identical results when the same inputs are processed - in different batch configurations. This will significantly affect speed of + in different batch configurations. This will significantly affect speed of training and inference as the kernels are not full optimized. Defaults to False.""" @@ -3012,7 +3016,7 @@ class MLATransformerConfig(TransformerConfig): cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. - This is only for the dynamic inference backend and requires that + This is only for the dynamic inference backend and requires that Flash MLA is installed.""" mla_down_proj_fusion: bool = False diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index d4f9eb9c0de..24d8313f225 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2257,6 +2257,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", From a4692bfa93f00272b39b6fa6532b7e94282c4403 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 21 Aug 2026 00:54:59 +0000 Subject: [PATCH 2/6] Hard-code grouped MLP NVFP4-UE5M3 recipe Matches fused kernel from cuDNN Frontend. Signed-off-by: Tim Moon --- .../core/extensions/transformer_engine.py | 31 +++++++------------ 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 2af57c2becb..49f9d210093 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -3249,28 +3249,19 @@ def _make_nvfp4_ue5m3_quantizer( ) -> te.pytorch.NVFP4Quantizer: """Construct NVFP4 quantizer with UE5M3 scales and no per-tensor scale for activations.""" - - from transformer_engine.pytorch.quantization import RecipeState - - # Parse tensor role tensor_type = role.tensor_type if role is not None else None - - # Construct NVFP4 quantizer based on default NVFP4 recipe - fp4_recipe = te.common.recipe.NVFP4BlockScaling( - fp8_format=te.common.recipe.Format.UE5M3 + 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", ) - (quantizer,) = RecipeState.create( - fp4_recipe, - mode="forward" if tensor_type != "grad_output" else "backward", - num_quantizers=1, - roles=[role], - ).make_quantizers() - - # Disable per-tensor scale for activations - if tensor_type == "input": - quantizer.disable_second_level_scale = True - - return quantizer # Construct custom recipe for NVFP4 with UE5M3 scales self._recipe = te.common.recipe.CustomRecipe( From 6e3e1c1809d4038193775e7e019a269a9825044e Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 27 Aug 2026 07:58:33 +0000 Subject: [PATCH 3/6] Add test for fp8_recipe_attrs config Co-authored-by: Codex Signed-off-by: Tim Moon --- tests/unit_tests/test_fp8_utils.py | 44 ++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index dc65d541455..905c56a98ab 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 @@ -23,6 +25,48 @@ reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10" +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.""" From 0f5f8479516bf280d837995a14ace1c6601e9860 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 27 Aug 2026 08:27:00 +0000 Subject: [PATCH 4/6] Fix incorrect recipe config in shared expert test Co-authored-by: Codex Signed-off-by: Tim Moon --- .../transformer/moe/test_shared_experts.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/transformer/moe/test_shared_experts.py b/tests/unit_tests/transformer/moe/test_shared_experts.py index 798313f188c..57bf7d9df71 100644 --- a/tests/unit_tests/transformer/moe/test_shared_experts.py +++ b/tests/unit_tests/transformer/moe/test_shared_experts.py @@ -106,6 +106,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( @@ -126,9 +136,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(): From 3cb0458feb62052dcd5d74b39226371671d04059 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 27 Aug 2026 08:40:04 +0000 Subject: [PATCH 5/6] Drive-by fix suggested by @claude Signed-off-by: Tim Moon --- megatron/core/extensions/transformer_engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index b874a42e179..26ea7774c47 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -363,7 +363,7 @@ 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: From ced16e142b8c1ed3ad50d0305df84713ae4e07c9 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Thu, 27 Aug 2026 09:03:31 +0000 Subject: [PATCH 6/6] Apply linter and restore autoremoved whitespace Signed-off-by: Tim Moon --- megatron/core/fp4_utils.py | 2 +- megatron/core/fp8_utils.py | 4 +-- .../core/transformer/transformer_config.py | 32 +++++++++---------- tests/unit_tests/test_fp8_utils.py | 6 +--- 4 files changed, 20 insertions(+), 24 deletions(-) diff --git a/megatron/core/fp4_utils.py b/megatron/core/fp4_utils.py index a02b8fae620..d337a993f30 100644 --- a/megatron/core/fp4_utils.py +++ b/megatron/core/fp4_utils.py @@ -226,7 +226,7 @@ def get_fp4_recipe(config: TransformerConfig): ) except AttributeError: raise ValueError("""NVFP4BlockScaling recipe is not available in this version of - Transformer Engine. Please make sure you are using TE version + Transformer Engine. Please make sure you are using TE version >= 2.7.0.dev0.""") elif config.fp4_recipe == Fp4Recipe.custom: fp4_recipe = _get_custom_recipe(config.fp4_quantizer_factory) diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 013b441a7f7..f2b2c65080a 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -347,8 +347,8 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F try: custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory) except AttributeError: - raise ValueError("""CustomRecipe recipe is not available in this version of - Transformer Engine. Please make sure you are using TE version + raise ValueError("""CustomRecipe recipe is not available in this version of + Transformer Engine. Please make sure you are using TE version >= 2.9.0.dev0.""") return custom_recipe diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 7cf6dd4c9c1..1811d2a3eee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -75,8 +75,8 @@ class TransformerConfig(ModelParallelConfig): mtp_loss_scaling_factor: Optional[float] = 0.1 """Weighting factor of Multi-Token Prediction (MTP) loss. - We compute the average of the MTP losses across all depths, - and multiply it the scaling factor to obtain the overall MTP loss, + We compute the average of the MTP losses across all depths, + and multiply it the scaling factor to obtain the overall MTP loss, which serves as an additional training objective. """ @@ -117,8 +117,8 @@ class TransformerConfig(ModelParallelConfig): - list: e.g., [['embedding', 'decoder'], ['decoder', 'decoder', 'decoder', 'loss']]. - PipelineParallelLayerLayout: a PipelineParallelLayerLayout object. If given either a string or a list, it will be transferred into a PipelineParallelLayerLayout - in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers - in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., + in post init. Let i = a * pp_size + b, then layout[i] gives a list of the layers + in the a-th vpp stage and the b-th pp stage, i.e., vpp(0)pp(0), vpp(0)pp(1), ..., vpp(i)pp(j), vpp(i)pp(j+1), ..., vpp(-1)pp(-2), vpp(-1)pp(-1). In the inner lists of layers, 'embedding' or 'E' denotes the embedding layer, 'loss' or 'L' denotes the loss function, and 'decoder' or 't' denotes the transformer decoder layer. @@ -169,8 +169,8 @@ class TransformerConfig(ModelParallelConfig): """Softmax scale for attention scaling.""" softmax_type: Literal['vanilla', 'off-by-one', 'learnable'] = 'vanilla' - """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. - Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and + """Applies modified softmax from https://www.evanmiller.org/attention-is-off-by-one.html. + Supports both TE FusedAttention and local unfused attention. Supports both a fixed offset and and learnable offset.""" num_query_groups: Optional[int] = field( @@ -230,7 +230,7 @@ class TransformerConfig(ModelParallelConfig): The stored input is casted back to the original precision before backprop compuatation.""" glu_linear_offset: float = 0.0 - """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only + """Offset term in the GLU activation function: activation_func(x[0]) * (x[1] + offset). Only used when gated_linear_unit is True""" activation_func_clamp_value: Optional[float] = None @@ -362,7 +362,7 @@ class TransformerConfig(ModelParallelConfig): # linear attention #################### linear_attention_freq: Optional[Union[int, List[int]]] = None - """Frequency between LA (linear attention) layers + """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. Accepts either: - An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer @@ -404,13 +404,13 @@ class TransformerConfig(ModelParallelConfig): embedding_init_method: Optional[Callable] = None """ - Method to initialize weights of the embedding layer. If None, will be set as described + Method to initialize weights of the embedding layer. If None, will be set as described in init_method above. """ embedding_init_method_std: Optional[float] = None """ - Standard deviation of the zero mean normal for the default initialization method for the + Standard deviation of the zero mean normal for the default initialization method for the embedding layer. If None, will be set to init_method_std. Setting this to a value around 1.0 may avoid loss spikes in training. Setting this to any value will also skip applying weight decay on embedding weights to avoid shrinkage towards zero. @@ -677,7 +677,7 @@ class TransformerConfig(ModelParallelConfig): fp4: Optional[Literal['e2m1']] = field( default=None, metadata={"argparse_meta": {"arg_names": ["--fp4-format"]}} ) - """If set, enables the use of FP4 precision through Transformer Engine. Currently only + """If set, enables the use of FP4 precision through Transformer Engine. Currently only supports 'nvfp4' which uses NVFP4BlockScaling recipe (requires TE >= 2.7.0.dev0).""" fp4_recipe: Optional[Literal['nvfp4', 'custom']] = "nvfp4" @@ -716,12 +716,12 @@ class TransformerConfig(ModelParallelConfig): in the hidden_states gradient.""" moe_shared_expert_gate: bool = False - """Enable gate for shared expert. Only effective when + """Enable gate for shared expert. Only effective when moe-shared-expert-intermediate-size is set.""" moe_shared_expert_overlap: bool = False """Enable overlapping between shared expert computations and dispatcher communications. - Without this, the shared experts execute before the router. + Without this, the shared experts execute before the router. Only effective when moe-shared-expert-intermediate-size is set. """ @@ -839,7 +839,7 @@ class TransformerConfig(ModelParallelConfig): no memory: the bias is replaced by the latest global-batch quantile estimate each step.""" moe_router_force_load_balancing: bool = False - """[Experimental] Force load balancing with random logits for MoE router, supports naive topk + """[Experimental] Force load balancing with random logits for MoE router, supports naive topk and group-limited topk. This is an experimental feature and only for benchmark.""" moe_router_force_biased: Optional[float] = None @@ -1202,7 +1202,7 @@ class TransformerConfig(ModelParallelConfig): batch_invariant_mode: bool = False """If true, uses batch-invariant kernels that provide deterministic forward execution regardless of batch size. This ensures bitwise identical results when the same inputs are processed - in different batch configurations. This will significantly affect speed of + in different batch configurations. This will significantly affect speed of training and inference as the kernels are not full optimized. Defaults to False.""" @@ -3435,7 +3435,7 @@ class MLATransformerConfig(TransformerConfig): cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. - This is only for the dynamic inference backend and requires that + This is only for the dynamic inference backend and requires that Flash MLA is installed.""" mla_down_proj_fusion: bool = False diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 90688872a6f..c462c40851f 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -67,11 +67,7 @@ def fake_fp8_model_init(enabled=False, recipe=None, **kwargs): captured["kwargs"] = kwargs return nullcontext() - monkeypatch.setattr( - fp8_utils.transformer_engine.pytorch, - "fp8_model_init", - fake_fp8_model_init, - ) + monkeypatch.setattr(fp8_utils.transformer_engine.pytorch, "fp8_model_init", fake_fp8_model_init) config = TransformerConfig( num_layers=1,