diff --git a/megatron/core/enums.py b/megatron/core/enums.py index c9a715519f9..fcca219badd 100644 --- a/megatron/core/enums.py +++ b/megatron/core/enums.py @@ -20,15 +20,17 @@ def encoder_and_decoder(self): class Fp8Recipe(str, enum.Enum): - """FP8 recipe names: delayed, tensorwise, mxfp8, blockwise.""" + """FP8 recipe names: delayed, tensorwise, mxfp8, blockwise, custom.""" delayed = "delayed" tensorwise = "tensorwise" mxfp8 = "mxfp8" blockwise = "blockwise" + custom = "custom" class Fp4Recipe(str, enum.Enum): - """FP4 recipe names: nvfp4.""" + """FP4 recipe names: nvfp4, custom.""" nvfp4 = "nvfp4" + custom = "custom" diff --git a/megatron/core/fp4_utils.py b/megatron/core/fp4_utils.py index eae4bf91de6..9aebf31c5da 100644 --- a/megatron/core/fp4_utils.py +++ b/megatron/core/fp4_utils.py @@ -7,6 +7,7 @@ import torch from megatron.core.enums import Fp4Recipe +from megatron.core.fp8_utils import _get_custom_recipe from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version @@ -70,9 +71,11 @@ def get_fp4_recipe(config: TransformerConfig): 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: raise ValueError( - "NVFP4BlockScaling is the only supported FP4 recipe. " + "NVFP4BlockScaling and custom are the only supported FP4 recipes. " "Please make sure you are using a compatible TE version >= 2.7.0.dev0." ) else: diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 1c1159f6f7b..9a697981738 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -2,14 +2,15 @@ """Utility functions related to FP8 that are used throughout Megatron core""" +import importlib import weakref from contextlib import nullcontext from functools import wraps -from typing import List, Optional +from typing import List, Optional, Union import torch -from megatron.core.enums import Fp8Recipe +from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.tensor_parallel import ( ColumnParallelLinear, RowParallelLinear, @@ -110,6 +111,53 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor: return fp8_tensor.from_float8() +def _resolve_callable_from_python_import_path(dotted_path: str): + """Resolve a Python import path like 'pkg.mod.func' to a callable. + + Raises ValueError with clear message on failure. + """ + if not isinstance(dotted_path, str) or not dotted_path: + raise ValueError( + "fp8_quantizer_factory must be a non-empty string with format 'pkg.mod.func'." + ) + + parts = dotted_path.rsplit(".", 1) + if len(parts) == 1: + raise ValueError(f"Invalid fp8_quantizer_factory '{dotted_path}'. Expected 'pkg.mod.func'.") + module_path, attr = parts[0], parts[1] + + try: + mod = importlib.import_module(module_path) + except Exception as exc: + raise ValueError( + f"Failed to import module '{module_path}' for fp8_quantizer_factory: {exc}" + ) from exc + + fn = getattr(mod, attr, None) + if fn is None: + raise ValueError( + f"Attribute '{attr}' not found in module '{module_path}' for fp8_quantizer_factory." + ) + if not callable(fn): + raise ValueError( + f"Resolved attribute '{module_path}.{attr}' is not callable for fp8_quantizer_factory." + ) + return fn + + +def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, Fp4Recipe]: + quantizer_factory = _resolve_callable_from_python_import_path(quantizer_factory_python_path) + 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.""" + ) + return custom_recipe + + def get_fp8_align_size(fp8_recipe: Fp8Recipe) -> int: """Get the alignment size required for fp8 GEMM.""" if fp8_recipe == Fp8Recipe.mxfp8: @@ -493,6 +541,8 @@ def get_fp8_recipe(config: TransformerConfig): fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling( fp8_format=fp8_format ) + elif config.fp8_recipe == Fp8Recipe.custom: + fp8_recipe = _get_custom_recipe(config.fp8_quantizer_factory) else: raise ValueError( "Float8CurrentScaling, MXFP8BlockScaling, Float8BlockwiseScaling and " diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 147c5b23b3d..68bdfbcf021 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -7,7 +7,7 @@ import torch import torch.nn.functional as F -from megatron.core.enums import Fp8Recipe +from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.quantization.quant_config import RecipeConfig from megatron.core.transformer.enums import AttnBackend from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout @@ -347,10 +347,10 @@ class TransformerConfig(ModelParallelConfig): activation and weight tensors and e5m2 for all FP8 output activation gradient tensors.""" fp8_recipe: Optional[str] = "delayed" - """If set, enables the use of FP8 precision through Transformer Engine. There are 3 predefined + """If set, enables the use of FP8 precision through Transformer Engine. There are 5 predefined choices (1) 'tensorwise' uses per tensor current scaling recipe, (2) 'delayed' uses delayed scaling recipe, 3) 'mxfp8' for Blackwell architecture only, - 4) 'blockwise' for blockwise scaling recipe.""" + 4) 'blockwise' for blockwise scaling recipe, 5) 'custom' for custom quantization recipe.""" fp8_param: bool = False """If set, keep the parameters in fp8 precision to save memory. This option must be used @@ -359,6 +359,10 @@ class TransformerConfig(ModelParallelConfig): primarily the weights of GEMMs. The specific parameters that will be converted to fp8 are determined by TE.""" + fp8_quantizer_factory: Optional[str] = None + """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. + Required when fp8_recipe is custom.""" + fp8_margin: int = 0 """Margin for the scaling factor computation.""" @@ -420,6 +424,10 @@ class TransformerConfig(ModelParallelConfig): together with fp4 mode (i.e., TransformerConfig.fp4 is not None). Note that not all parameters will be converted to fp4; for example, biases will remain unchanged.""" + fp4_quantizer_factory: Optional[str] = None + """Python import path to a callable quantizer factory, e.g., package.module.quantizer_factory. + Required when fp4_recipe is custom.""" + #################### # MoE related #################### @@ -792,6 +800,14 @@ def __post_init__(self): f"({max_bf16_layers_per_pipeline_stage})." ) + if self.fp8_recipe == Fp8Recipe.custom: + if not self.fp8_quantizer_factory: + raise ValueError( + "fp8_quantizer_factory must be provided when fp8_recipe is 'custom'. " + "Specify a Python import path (e.g., package.module.quantizer_factory) " + "via --fp8-quantizer-factory." + ) + if self.fp8_param and not self.fp8: raise ValueError("fp8_param must be used together with fp8 mode.") @@ -802,6 +818,14 @@ def __post_init__(self): if self.fp4 and self.fp8: raise ValueError("fp4 and fp8 cannot be used simultaneously. Please choose one.") + if self.fp4 and self.fp4_recipe == Fp4Recipe.custom: + if not self.fp4_quantizer_factory: + raise ValueError( + "fp4_quantizer_factory must be provided when fp4_recipe is 'custom'. " + "Specify a Python import path (e.g., package.module.quantizer_factory) " + "via --fp4-quantizer-factory." + ) + if self.apply_query_key_layer_scaling: self.attention_softmax_in_fp32 = True diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0b14140529a..3eb23763625 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1319,9 +1319,13 @@ def _add_transformer_engine_args(parser): dest='fp8') # per tensor current scaling recipe selection group.add_argument('--fp8-recipe', default='delayed', - choices=['tensorwise', 'delayed', 'mxfp8', 'blockwise'], + choices=['tensorwise', 'delayed', 'mxfp8', 'blockwise', 'custom'], help='Which fp8 recipe to use for FP8 tensors in the forward and backward pass', dest='fp8_recipe') + group.add_argument('--fp8-quantizer-factory', default=None, + help='Python import path to a callable quantizer factory, ' + 'e.g., package.module.quantizer_factory.', + dest='fp8_quantizer_factory') # delayed scaling only configs group.add_argument('--fp8-margin', type=int, default=0, help='Scaling margin for fp8', @@ -1358,9 +1362,13 @@ def _add_transformer_engine_args(parser): help='Which nvfp4 format scheme to use for FP4 tensors in the forward and backward pass', dest='fp4') group.add_argument('--fp4-recipe', default='nvfp4', - choices=['nvfp4'], + choices=['nvfp4', 'custom'], help='Which fp4 recipe to use for FP4 tensors in the forward and backward pass', dest='fp4_recipe') + group.add_argument('--fp4-quantizer-factory', default=None, + help='Python import path to a callable quantizer factory, ' + 'e.g., package.module.quantizer_factory.', + dest='fp4_quantizer_factory') group.add_argument('--fp4-param-gather', action='store_true', help='Keep the compute param in fp4 (do not use any other intermediate ' 'dtype) and perform the param all-gather in fp4.',