[1/2] [GraphTrainer] Support dense FP8 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
 

torchtitan/components/quantization/float8.py

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

torchtitan/experiments/graph_trainer/common_utils.py

@@ -162,13 +162,35 @@ def ensure_boxed_graph_module(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
 # Tuple of unique parameter FQNs naming one gradient value. Consumers must
 # treat it as a set: tied parameters can associate multiple FQNs with one node.
 PARAMETER_GRADIENT_FQNS_META = "parameter_gradient_fqns"
+_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"
 _EP_TOKEN_EXCHANGE_WAIT = "EP_token_exchange_wait"
 _NOT_IN_LAYERS = -1
 
 
+def get_quantization_kind(module: nn.Module) -> str | None:
+    """Return the quantization category used by GraphTrainer annotations."""
+    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
+
+
 def compute_parameter_gradients(
     loss: torch.Tensor,
     named_parameters: Iterable[tuple[str, torch.Tensor]],
@@ -257,14 +279,21 @@ 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.
     """
     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, "_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
@@ -178,6 +189,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:
+        raise ValueError(
+            "GraphTrainer FP8 is incompatible with --compile.precompile_artifact_dir. "
+            "CooR precompile tracing fails on torchao FP8 layout guards and "
+            "FlexAttention BlockMask serialization. Use online "
+            "--compile.mode aot_fx_trace without precompile, or a non-FP8 "
+            "config for CooR precompile."
+        )
 
 def validate_autoparallel_config(
     compile_config: GraphTrainerCompileConfig,