Enable TE custom quantization recipe

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"

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,17 @@ 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:
+                if not config.fp4_quantizer_factory:
+                    raise ValueError(
+                        "Python import path, e.g. package.module.quantizer_factory, "
+                        "must be provided via --fp4-quantizer-factory when "
+                        "--fp4-recipe custom is selected."
+                    )
+                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:

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.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import get_te_version, is_te_min_version
 
@@ -104,6 +105,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:
@@ -466,6 +514,14 @@ def get_fp8_recipe(config: TransformerConfig):
                 fp8_recipe = transformer_engine.common.recipe.MXFP8BlockScaling(
                     fp8_format=fp8_format
                 )
+            elif config.fp8_recipe == Fp8Recipe.custom:
+                if not config.fp8_quantizer_factory:
+                    raise ValueError(
+                        "Python import path, e.g. package.module.quantizer_factory, "
+                        "must be provided via --fp8-quantizer-factory when "
+                        "--fp8-recipe custom is selected."
+                    )
+                fp8_recipe = _get_custom_recipe(config.fp8_quantizer_factory)
             else:
                 raise ValueError(
                     "Float8CurrentScaling, MXFP8BlockScaling, Float8BlockwiseScaling and "

megatron/core/transformer/transformer_config.py

@@ -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
     ####################

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.',