[GraphTrainer] Enhance FP8 region compilation

.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml

@@ -98,6 +98,9 @@ jobs:
         # Run the MoE numerics tests
         NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerNumerics -v -k "moe"
 
+        # Run dense FP8 full/regional Inductor and CUDA Graph numerics tests.
+        NCCL_NVLS_ENABLE=0 pytest torchtitan/experiments/graph_trainer/tests/test_numerics.py::TestGraphTrainerFP8Numerics -v
+
         # Run precompile integration tests (DSv3 with EP; Llama3 runs in default workflow)
         NCCL_NVLS_ENABLE=0 python -m torchtitan.experiments.graph_trainer.tests.run_precompile_tests $RUNNER_TEMP/artifacts-to-be-uploaded/precompile --ngpu 8 --test_name aot_fx_trace_deepseek_v3_precompile_fsdp_tp_ep
 

tests/unit_tests/test_quantization.py

@@ -10,7 +10,9 @@
 from torchtitan.components.quantization import Float8Linear
 from torchtitan.components.quantization.float8 import _get_float8_grouped_experts_cls
 from torchtitan.components.quantization.mx import _get_mxfp8_grouped_experts_cls
-from torchtitan.components.quantization.utils import has_quantization
+from torchtitan.components.quantization.utils import (
+    has_quantization,
+)
 from torchtitan.config import ConfigManager
 from torchtitan.models.common.decoder_sharding import colwise_config, rowwise_config
 from torchtitan.models.common.linear import Linear
@@ -41,11 +43,15 @@ def test_float8_applied_by_model_registry():
     assert has_quantization(model_config)
     # Some Linear.Config instances should be swapped to Float8Linear
     converted = [
-        fqn
-        for fqn, lc, _parent, _attr in model_config.traverse(Linear.Config)
-        if isinstance(lc, Float8Linear.Config)
+        linear_config
+        for _fqn, linear_config, _parent, _attr in model_config.traverse(
+            Linear.Config
+        )
+        if isinstance(linear_config, Float8Linear.Config)
     ]
     assert len(converted) > 0
+    assert all(config._quantization_recipe_name == "rowwise" for config in converted)
+    assert all(config._quantization_emulate for config in converted)
 
 
 @pytest.mark.parametrize(

torchtitan/components/quantization/float8.py

@@ -36,6 +36,8 @@ class Config(Linear.Config):
             """Drop-in replacement for Linear.Config that builds Float8Linear."""
 
             _torchao_config: object = None
+            _quantization_recipe_name: str = ""
+            _quantization_emulate: bool = False
 
         def __init__(self, config: Config):
             TorchAOFloat8Linear.__init__(
@@ -45,6 +47,10 @@ def __init__(self, config: Config):
                 bias=config.bias,
                 config=config._torchao_config,
             )
+            self._torchtitan_quantization_recipe_name = (
+                config._quantization_recipe_name
+            )
+            self._torchtitan_quantization_emulate = config._quantization_emulate
 
 except ImportError:
     Float8Linear = None
@@ -161,6 +167,8 @@ def convert(self, model_config):
                     bias=linear_config.bias,
                     param_init=linear_config.param_init,
                     _torchao_config=self.torchao_config,
+                    _quantization_recipe_name=self.config.recipe_name,
+                    _quantization_emulate=self.config.emulate,
                 )
                 if parent is None:
                     model_config = new_config

torchtitan/components/quantization/mx.py

@@ -33,7 +33,7 @@ class MXFP8Linear(TorchAOMXFP8Linear, Module):
         class Config(Linear.Config):
             """Drop-in replacement for Linear.Config that builds MXFP8Linear."""
 
-            pass
+            _quantization_recipe_name: str = "mxfp8_rceil"
 
         def __init__(self, config: Config):
             TorchAOMXFP8Linear.__init__(
@@ -42,6 +42,9 @@ def __init__(self, config: Config):
                 config.out_features,
                 bias=config.bias,
             )
+            self._torchtitan_quantization_recipe_name = (
+                config._quantization_recipe_name
+            )
 
 except ImportError:
     MXFP8Linear = None
@@ -86,6 +89,7 @@ def convert(self, model_config):
                     out_features=config.out_features,
                     bias=config.bias,
                     param_init=config.param_init,
+                    _quantization_recipe_name="mxfp8_rceil",
                 )
                 if parent is None:
                     model_config = new_config

torchtitan/components/quantization/utils.py

@@ -4,6 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from dataclasses import dataclass
+
+import torch.nn as nn
+
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.moe import GroupedExperts
 from torchtitan.models.common.token_dispatcher import (
@@ -92,3 +96,70 @@ def has_quantization(model_config) -> bool:
         for _fqn, config, _parent, _attr in model_config.traverse(GroupedExperts.Config)
     )
     return has_quant_linear or has_quant_moe
+
+
+def get_quantization_kind(module: nn.Module) -> str | None:
+    """Return the stable quantization category for a runtime module."""
+    from torchtitan.components.quantization.float8 import (
+        _float8_experts_cache,
+        Float8Linear,
+    )
+    from torchtitan.components.quantization.mx import _mxfp8_experts_cache, MXFP8Linear
+
+    if Float8Linear is not None and isinstance(module, Float8Linear):
+        return "float8_linear"
+    if MXFP8Linear is not None and isinstance(module, MXFP8Linear):
+        return "mxfp8_linear"
+
+    if any(isinstance(module, cls) for cls in _float8_experts_cache.values()):
+        return "float8_grouped_experts"
+    if any(isinstance(module, cls) for cls in _mxfp8_experts_cache.values()):
+        return "mxfp8_grouped_experts"
+    return None
+
+
+@dataclass(frozen=True)
+class QuantizationSignature:
+    """Stable lowering-relevant identity for one quantized runtime module."""
+
+    module_fqn: str
+    kind: str
+    recipe_name: str
+    emulate: bool
+
+
+def get_quantization_signature(
+    model: nn.Module,
+) -> tuple[QuantizationSignature, ...]:
+    """Return sorted stable signatures for quantized runtime modules.
+
+    Grouped-expert modules are rejected because their regional compilation
+    path does not yet have a complete precompile artifact signature.
+    """
+    signatures = []
+    for module_fqn, module in model.named_modules():
+        kind = get_quantization_kind(module)
+        if kind is None:
+            continue
+        if kind.endswith("grouped_experts"):
+            raise ValueError(
+                "FP8 precompile does not support grouped experts. "
+                "Use dense FP8 modules for precompiled graphs."
+            )
+        recipe_name = getattr(module, "_torchtitan_quantization_recipe_name", "")
+        if not recipe_name:
+            raise ValueError(
+                "Quantized module is missing stable recipe metadata: "
+                f"{module_fqn} ({kind})."
+            )
+        signatures.append(
+            QuantizationSignature(
+                module_fqn=module_fqn,
+                kind=kind,
+                recipe_name=recipe_name,
+                emulate=bool(
+                    getattr(module, "_torchtitan_quantization_emulate", False)
+                ),
+            )
+        )
+    return tuple(sorted(signatures, key=lambda signature: signature.module_fqn))

torchtitan/experiments/graph_trainer/common_utils.py

@@ -135,6 +135,8 @@ def ensure_boxed_graph_module(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
 
 
 _MODULE_FQN = "module_fqn"
+_QUANTIZATION_KIND = "quantization_kind"
+_QUANTIZATION_EMULATE = "quantization_emulate"
 _EP_TOKEN_COUNT_EXCHANGE = "EP_token_count_exchange"
 _EP_TOKEN_COUNT_SYNC = "EP_token_count_sync"
 _EP_TOKEN_EXCHANGE = "EP_token_exchange"
@@ -205,14 +207,23 @@ def annotate_module_fqns(model: nn.Module) -> None:
     """Annotate all modules' forward with their fully-qualified names.
 
     Every named submodule (excluding the root) gets its forward method wrapped
-    with ``annotate_fn`` so that FX nodes carry ``module_fqn`` in
-    ``node.meta["custom"]``.
+    with ``annotate_fn`` so that FX nodes carry ``module_fqn`` and, for
+    quantized modules, ``quantization_kind`` in ``node.meta["custom"]``.
 
     Call once after model construction, before tracing/compilation.
     """
+    from torchtitan.components.quantization.utils import get_quantization_kind
+
     for fqn, submodule in model.named_modules():
         if fqn:  # skip root module
-            submodule.forward = annotate_fn({_MODULE_FQN: fqn})(submodule.forward)
+            metadata = {_MODULE_FQN: fqn}
+            quantization_kind = get_quantization_kind(submodule)
+            if quantization_kind is not None:
+                metadata[_QUANTIZATION_KIND] = quantization_kind
+                metadata[_QUANTIZATION_EMULATE] = bool(
+                    getattr(submodule, "_torchtitan_quantization_emulate", False)
+                )
+            submodule.forward = annotate_fn(metadata)(submodule.forward)
 
 
 def matches_module_fqn_pattern(pattern: str, fqn: str) -> bool:

torchtitan/experiments/graph_trainer/configs.py

@@ -25,6 +25,17 @@
 SUPPORTED_EP_OVERLAP_MODULE_FQNS = frozenset({TRANSFORMER_BLOCK_FQN, MOE_BLOCK_FQN})
 
 
+@dataclass(kw_only=True, slots=True)
+class FP8GraphConfig:
+    """Controls GraphTrainer support for already-quantized FP8 modules."""
+
+    enabled: bool = False
+    """Enable FP8 graph analysis and validation."""
+
+    strict_validation: bool = True
+    """Require a supported FP8 compute operation for each quantized module region."""
+
+
 @dataclass(kw_only=True, slots=True)
 class EpOverlapConfig:
     enabled: bool = False
@@ -169,6 +180,33 @@ class GraphTrainerCompileConfig(CompileConfig):
     """Use AutoParallelGraph (ILP solver-based SPMD sharding) instead of
     manual TP/FSDP/EP. Forces the AOT compilation path internally."""
 
+    fp8: FP8GraphConfig = field(default_factory=FP8GraphConfig)
+    """FP8 graph analysis and validation configuration."""
+
+
+def validate_fp8_graph_config(compile_config: GraphTrainerCompileConfig) -> None:
+    """Validate the supported FP8 GraphTrainer execution contracts."""
+    if not compile_config.fp8.enabled:
+        return
+    if not compile_config.enable:
+        raise ValueError("--compile.fp8.enabled requires --compile.enable")
+    if not compile_config.enable_passes:
+        raise ValueError("--compile.fp8.enabled requires --compile.enable_passes")
+    if compile_config.mode != "aot_fx_trace":
+        raise ValueError("--compile.fp8.enabled requires --compile.mode aot_fx_trace")
+    if compile_config.inductor_compilation not in {"full", "regional"}:
+        raise ValueError(
+            "--compile.fp8.enabled requires --compile.inductor_compilation full "
+            "or regional"
+        )
+    if (
+        compile_config.precompile_artifact_dir
+        and compile_config.inductor_compilation != "regional"
+    ):
+        raise ValueError(
+            "FP8 precompile requires --compile.inductor_compilation regional"
+        )
+
 
 def validate_autoparallel_config(
     compile_config: GraphTrainerCompileConfig,