From 03d1f475a03a27757bea5b34765d557b31b563b4 Mon Sep 17 00:00:00 2001 From: Zihan Yang Date: Fri, 28 Aug 2026 18:14:55 +0800 Subject: [PATCH 1/2] Add GraphTrainer FP8 validation and metadata --- torchtitan/components/quantization/float8.py | 3 + .../experiments/graph_trainer/common_utils.py | 35 +- .../experiments/graph_trainer/configs.py | 38 + .../experiments/graph_trainer/fp8_passes.py | 377 ++++++++ .../experiments/graph_trainer/passes.py | 49 +- .../graph_trainer/tests/test_fp8.py | 868 ++++++++++++++++++ 6 files changed, 1364 insertions(+), 6 deletions(-) create mode 100644 torchtitan/experiments/graph_trainer/fp8_passes.py create mode 100644 torchtitan/experiments/graph_trainer/tests/test_fp8.py diff --git a/torchtitan/components/quantization/float8.py b/torchtitan/components/quantization/float8.py index 99770ed0c9..4a61a9d73e 100644 --- a/torchtitan/components/quantization/float8.py +++ b/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 diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py index 71a38c10f0..3220c1ed93 100644 --- a/torchtitan/experiments/graph_trainer/common_utils.py +++ b/torchtitan/experiments/graph_trainer/common_utils.py @@ -162,6 +162,8 @@ 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" @@ -169,6 +171,26 @@ def ensure_boxed_graph_module(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: _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: diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 52e94b0b02..40314dbe03 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/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, diff --git a/torchtitan/experiments/graph_trainer/fp8_passes.py b/torchtitan/experiments/graph_trainer/fp8_passes.py new file mode 100644 index 0000000000..34186cf47b --- /dev/null +++ b/torchtitan/experiments/graph_trainer/fp8_passes.py @@ -0,0 +1,377 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""FP8 graph annotation and validation for GraphTrainer compilation.""" + +from __future__ import annotations + +import warnings +from collections import defaultdict + +import torch + +from torchtitan.experiments.graph_trainer.common_utils import ( + _MODULE_FQN, + _QUANTIZATION_EMULATE, + _QUANTIZATION_KIND, +) +from torchtitan.tools.logging import logger + +_FP8_META = "fp8" + + +def _available_scaled_mm_targets() -> frozenset[object]: + targets = [] + for op_name in ("_scaled_mm", "_scaled_grouped_mm"): + op = getattr(torch.ops.aten, op_name, None) + target = getattr(op, "default", None) + if target is not None: + targets.append(target) + return frozenset(targets) + + +_SCALED_MM_TARGETS = _available_scaled_mm_targets() +_FP8_DATA_DTYPES = frozenset( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e5m2", None), + getattr(torch, "float8_e5m2fnuz", None), + ) + if dtype is not None +) + +# Each quantization kind declares the operations that prove its compute was +# lowered to a supported FP8 implementation. Add a new kind here together +# with its lowering targets when extending GraphTrainer FP8 support. +FP8_COMPUTE_TARGETS: dict[str, frozenset[object]] = { + "float8_linear": _SCALED_MM_TARGETS, + "mxfp8_linear": _SCALED_MM_TARGETS, + "float8_grouped_experts": _SCALED_MM_TARGETS, + "mxfp8_grouped_experts": _SCALED_MM_TARGETS, +} + + +def _value_contains_fp8_tensor(value: object) -> bool: + if isinstance(value, torch.fx.Node): + return _value_contains_fp8_tensor(value.meta.get("val")) + if isinstance(value, torch.Tensor): + return value.dtype in _FP8_DATA_DTYPES + if isinstance(value, (tuple, list)): + return any(_value_contains_fp8_tensor(item) for item in value) + if isinstance(value, dict): + return any(_value_contains_fp8_tensor(item) for item in value.values()) + return False + + +def _has_fp8_data_operand(node: torch.fx.Node) -> bool: + # The first two operands of the supported scaled matrix multiplication + # overloads are the input matrices. Scale operands may use FP8 formats too, + # so they cannot establish that the compute itself is FP8. + return any(_value_contains_fp8_tensor(operand) for operand in node.args[:2]) + + +def _classify_fp8_node(node: torch.fx.Node, quantization_kind: str) -> str: + if ( + node.target in FP8_COMPUTE_TARGETS.get(quantization_kind, frozenset()) + and _has_fp8_data_operand(node) + ): + return "compute" + + # GraphPP may compute shared backward quantization in bw_di and pass the + # resulting FP8 tensor to bw_dw. Treat that callable input as an explicit + # boundary, not as a missing cast. The local region starts at its consumer + # (for example, _scaled_mm), while the placeholder remains interpreted. + if node.op == "placeholder" and _value_contains_fp8_tensor( + node.meta.get("val") + ): + return "input_boundary" + + target_name = str(node.target) + if "amax" in target_name or node.target == torch.ops.aten.amax.default: + return "amax" + if "_to_copy" in target_name or "convert_element_type" in target_name: + return "cast" + + value = node.meta.get("val") + if _value_contains_fp8_tensor(value): + return "cast" + return "other" + + +def _inspect_fp8_regions( + gm: torch.fx.GraphModule, + *, + strict: bool, + annotate: bool, +) -> torch.fx.GraphModule: + """Inspect FP8 regions and optionally annotate their nodes.""" + regions: dict[tuple[str, str], dict[str, int]] = {} + emulated_regions: set[tuple[str, str]] = set() + target_inventory: defaultdict[str, set[str]] = defaultdict(set) + + for node in gm.graph.nodes: + custom = node.meta.get("custom", {}) + quantization_kind = custom.get(_QUANTIZATION_KIND) + if quantization_kind is None: + continue + + module_fqn = custom.get(_MODULE_FQN, "") + region_key = (module_fqn, quantization_kind) + region = regions.setdefault( + region_key, + {"forward_compute_ops": 0, "backward_compute_ops": 0}, + ) + if custom.get(_QUANTIZATION_EMULATE, False): + emulated_regions.add(region_key) + target_inventory[quantization_kind].add(str(node.target)) + + role = _classify_fp8_node(node, quantization_kind) + if annotate: + node.meta.setdefault("custom", {})[_FP8_META] = { + "op_role": role, + } + if role == "compute": + phase = "backward" if node.meta.get("autograd_backward", False) else "forward" + region[f"{phase}_compute_ops"] += 1 + + if strict and not regions: + raise RuntimeError( + "GraphTrainer did not find quantized module regions while " + "--compile.fp8.enabled is set." + ) + + missing_regions = [ + region_key + for region_key, region in regions.items() + if region_key not in emulated_regions + if not region["forward_compute_ops"] and not region["backward_compute_ops"] + ] + if strict and missing_regions: + raise RuntimeError( + "GraphTrainer found quantized module regions without a supported " + f"FP8 compute operation: {missing_regions}. Observed targets: " + f"{dict(target_inventory)}." + ) + + summary = { + "regions": {str(key): value for key, value in regions.items()}, + "target_inventory": { + kind: sorted(targets) for kind, targets in target_inventory.items() + }, + } + logger.info("GraphTrainer FP8 analysis: %s", summary) + return gm + + +def _is_regional_fp8_compute_node( + node: torch.fx.Node, + *, + module_fqn: str, + quantization_kind: str, +) -> bool: + if node.op != "call_function": + return False + custom = node.meta.get("custom", {}) + if ( + custom.get(_MODULE_FQN) != module_fqn + or custom.get(_QUANTIZATION_KIND) != quantization_kind + ): + return False + target_name = str(node.target) + if not target_name.startswith("aten."): + return False + value = node.meta.get("val") + return isinstance(value, torch.Tensor) and value.device.type == "cuda" + + +def _identify_fp8_regional_components( + gm: torch.fx.GraphModule, +) -> torch.fx.GraphModule: + """Identify maximal dense FP8 compute components for regional Inductor. + + Each component is seeded by a supported FP8 compute operation and expands + only through CUDA aten nodes with the same module and quantization + provenance. FP8 placeholders are legal callable boundaries, notably when + GraphPP passes shared grad-output quantization from bw_di to bw_dw; they + prove the compute operand dtype but are not compiled as part of the local + region. Communication and host work remain outside these regions. + Grouped-expert FP8 is not supported by regional Inductor compilation. + """ + candidate_nodes: set[torch.fx.Node] = set() + seeds: list[torch.fx.Node] = [] + + for node in gm.graph.nodes: + custom = node.meta.get("custom", {}) + fp8 = custom.get(_FP8_META) + if fp8 is None: + continue + quantization_kind = custom.get(_QUANTIZATION_KIND) + module_fqn = custom.get(_MODULE_FQN, "") + if quantization_kind is None: + continue + if quantization_kind.endswith("grouped_experts"): + raise ValueError( + "FP8 regional compilation does not support grouped experts. " + "Use full Inductor for grouped-expert FP8 graphs." + ) + if _is_regional_fp8_compute_node( + node, + module_fqn=module_fqn, + quantization_kind=quantization_kind, + ): + candidate_nodes.add(node) + if fp8["op_role"] == "compute": + seeds.append(node) + + if not seeds: + return gm + + identified_nodes: set[torch.fx.Node] = set() + num_regions = 0 + for seed in seeds: + if seed in identified_nodes: + continue + seed_custom = seed.meta["custom"] + module_fqn = seed_custom[_MODULE_FQN] + quantization_kind = seed_custom[_QUANTIZATION_KIND] + component = {seed} + pending = [seed] + while pending: + node = pending.pop() + neighbors = (*node.all_input_nodes, *node.users) + for neighbor in neighbors: + if ( + neighbor in component + or neighbor in identified_nodes + or neighbor not in candidate_nodes + ): + continue + if not _is_regional_fp8_compute_node( + neighbor, + module_fqn=module_fqn, + quantization_kind=quantization_kind, + ): + continue + component.add(neighbor) + pending.append(neighbor) + + for node in component: + custom = node.meta.setdefault("custom", {}) + fp8 = custom[_FP8_META] + fp8["regional_region_id"] = num_regions + fp8["regional_region_num_nodes"] = len(component) + identified_nodes.update(component) + num_regions += 1 + + summary = { + "num_regions": num_regions, + "num_region_nodes": len(identified_nodes), + } + logger.info("GraphTrainer FP8 regional annotation: %s", summary) + return gm + + +def annotate_complete_fp8_regions_for_regional_inductor_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Tag only complete FP8 regions identified in the current callable. + + GraphPP re-identifies regions after extracting each callable, so an FP8 + placeholder is a valid local boundary. The node-count check still protects + against graph rewrites between identification and tagging. Other regional + annotations are not modified. + + Each tagged node gets ``compile_with_inductor["inductor_region"]`` set to + its ``regional_region_id``. PyTorch's ``regional_inductor`` uses that key + to keep separately identified FP8 regions from being scooped into one + giant default partition. + """ + del example_inputs + regions: dict[tuple[str, str, int], list[torch.fx.Node]] = defaultdict(list) + expected_num_nodes: dict[tuple[str, str, int], int] = {} + + for node in gm.graph.nodes: + custom = node.meta.get("custom", {}) + fp8 = custom.get(_FP8_META) + if fp8 is None: + continue + region_id = fp8.get("regional_region_id") + region_num_nodes = fp8.get("regional_region_num_nodes") + if not isinstance(region_id, int) or not isinstance(region_num_nodes, int): + continue + region_key = ( + custom.get(_MODULE_FQN, ""), + custom.get(_QUANTIZATION_KIND, ""), + region_id, + ) + regions[region_key].append(node) + expected_num_nodes[region_key] = region_num_nodes + + num_tagged_nodes = 0 + incomplete_regions: list[tuple[str, str, int]] = [] + for region_key, nodes in regions.items(): + if len(nodes) != expected_num_nodes[region_key]: + incomplete_regions.append(region_key) + continue + if not any( + node.meta["custom"][_FP8_META].get("op_role") == "compute" + for node in nodes + ): + continue + _, _, region_id = region_key + for node in nodes: + custom = node.meta.setdefault("custom", {}) + compile_annotation = custom.setdefault("compile_with_inductor", {}) + compile_annotation["inductor_region"] = region_id + num_tagged_nodes += 1 + + if incomplete_regions: + warnings.warn( + "GraphTrainer skipped incomplete FP8 regional Inductor regions: " + f"{incomplete_regions}. The affected nodes will run eagerly.", + stacklevel=2, + ) + if num_tagged_nodes: + gm.meta["fp8_regional_tagged_complete_nodes"] = num_tagged_nodes + logger.info( + "Tagged %d complete FP8 regional Inductor nodes", + num_tagged_nodes, + ) + return gm + + +def validate_fp8_graph_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, + *, + strict: bool, +) -> torch.fx.GraphModule: + """Validate FP8 lowering without changing node-level compilation metadata.""" + del example_inputs + return _inspect_fp8_regions( + gm, + strict=strict, + annotate=False, + ) + + +def annotate_fp8_regions_for_regional_inductor_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, + *, + strict: bool, +) -> torch.fx.GraphModule: + """Validate FP8 lowering and identify dense regions for regional Inductor.""" + del example_inputs + gm = _inspect_fp8_regions( + gm, + strict=strict, + annotate=True, + ) + return _identify_fp8_regional_components(gm) diff --git a/torchtitan/experiments/graph_trainer/passes.py b/torchtitan/experiments/graph_trainer/passes.py index 5a8c2291eb..db795ac333 100644 --- a/torchtitan/experiments/graph_trainer/passes.py +++ b/torchtitan/experiments/graph_trainer/passes.py @@ -29,7 +29,6 @@ import time import warnings from collections.abc import Callable - import torch from torchtitan.experiments.graph_trainer.configs import ( @@ -72,6 +71,11 @@ reassign_collective_pgs_pass, schedule_fsdp_comms_to_dense_regions_pass, ) +from torchtitan.experiments.graph_trainer.fp8_passes import ( + annotate_complete_fp8_regions_for_regional_inductor_pass, + annotate_fp8_regions_for_regional_inductor_pass, + validate_fp8_graph_pass, +) from torchtitan.experiments.graph_trainer.inductor_passes import ( annotate_flex_attention_for_regional_inductor_pass, full_inductor_compilation_pass, @@ -132,6 +136,20 @@ def async_tensor_parallel_pass( return gm +def graph_pp_pre_partition_fp8_passes( + compile_config: GraphTrainerCompileConfig, +) -> list[Callable]: + """Return FP8 validation that requires the complete GraphPP joint graph.""" + if not compile_config.fp8.enabled: + return [] + return [ + functools.partial( + validate_fp8_graph_pass, + strict=compile_config.fp8.strict_validation, + ) + ] + + def _tensor_parallel_degree(config, parallel_dims=None) -> int: """Return TP degree from ``ParallelDims`` when available, else config.""" if parallel_dims is not None and hasattr(parallel_dims, "tp"): @@ -169,6 +187,7 @@ def compile_time_passes( ``include_mandatory_normalization=False`` lets GraphPP run required normalization unconditionally and then append only the optional passes controlled by ``enable_passes``. + """ from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.experiments.graph_trainer.common_utils import ( @@ -356,20 +375,36 @@ def final_inductor_compile_passes( *, use_cudagraph: bool = False, boxed_codegen: bool = False, + fp8_strict_validation: bool | None = None, ) -> list[Callable]: """Return the terminal Inductor passes for a traced graph. GraphTrainer applies these to the full train-step graph. GraphPP applies the same pass list to each extracted stage callable after its PP-specific partitioning has chosen the callable boundary. Terminal Inductor selection - only depends on compile config; model- and parallelism-aware rewrites stay - in ``compile_time_passes``. + only depends on compile config. GraphPP validates the complete stage joint + graph before partitioning, then re-identifies regions in each extracted + callable with strict validation disabled. FP8 pass inclusion is derived + directly from ``compile_config.fp8.enabled``. """ from torchtitan.models.common.attention import FlexAttention passes: list[Callable] = [] inductor_compilation = compile_config.inductor_compilation + fp8_enabled = compile_config.fp8.enabled + strict_validation = ( + compile_config.fp8.strict_validation + if fp8_strict_validation is None + else fp8_strict_validation + ) if inductor_compilation == "full": + if fp8_enabled: + passes.append( + functools.partial( + validate_fp8_graph_pass, + strict=strict_validation, + ) + ) # Compile the entire graph into optimized Triton kernels. Must be # terminal; the FX graph is no longer authoritative after this pass. passes.append( @@ -393,6 +428,14 @@ def final_inductor_compile_passes( ) passes.append(annotate_rmsnorm_for_regional_inductor_pass) + if fp8_enabled: + passes.append( + functools.partial( + annotate_fp8_regions_for_regional_inductor_pass, + strict=strict_validation, + ) + ) + passes.append(annotate_complete_fp8_regions_for_regional_inductor_pass) passes.append( functools.partial( regional_inductor_pass, diff --git a/torchtitan/experiments/graph_trainer/tests/test_fp8.py b/torchtitan/experiments/graph_trainer/tests/test_fp8.py new file mode 100644 index 0000000000..200d5fead7 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_fp8.py @@ -0,0 +1,868 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.nn as nn +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.traceback import preserve_node_meta +from torch.testing._internal.common_utils import TestCase + +from torchtitan.components.quantization import Float8Linear, Float8LinearConverter +from torchtitan.components.quantization.float8 import ( + _get_float8_grouped_experts_cls, +) +from torchtitan.components.loss import CrossEntropyLoss +from torchtitan.experiments.graph_trainer.common_utils import ( + _MODULE_FQN, + _QUANTIZATION_EMULATE, + _QUANTIZATION_KIND, + annotate_module_fqns, +) +from torchtitan.experiments.graph_trainer.cudagraph import ( + CUDAGraphWrapper, + cudagraph_pass, +) +from torchtitan.experiments.graph_trainer.make_fx_tracer import ( + _copy_fwd_metadata_to_bw_nodes, + minimal_fx_tracer, + run_traced, +) +from torchtitan.experiments.graph_trainer.configs import ( + FP8GraphConfig, + GraphTrainerCompileConfig, + validate_fp8_graph_config, +) +from torchtitan.experiments.graph_trainer.fp8_passes import ( + FP8_COMPUTE_TARGETS, + annotate_complete_fp8_regions_for_regional_inductor_pass, + annotate_fp8_regions_for_regional_inductor_pass, + validate_fp8_graph_pass, +) +from torchtitan.experiments.graph_trainer.inductor_passes import ( + regional_inductor_pass, +) +from torchtitan.experiments.graph_trainer.passes import ( + compile_time_passes, + construct_default_graph_passes, + final_inductor_compile_passes, + graph_pp_pre_partition_fp8_passes, +) +from torchtitan.models.common.linear import Linear +from torchtitan.models.common.moe import GroupedExperts +from torchtitan.tools.utils import has_cuda_capability + + +class TestFP8GraphConfig(TestCase): + def test_enabled_config_contract(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="full", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + validate_fp8_graph_config(config) + + def test_fp8_rejects_unknown_inductor_mode(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="unknown", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + with self.assertRaisesRegex(ValueError, "inductor_compilation full or regional"): + validate_fp8_graph_config(config) + + def test_fp8_requires_graph_passes(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + enable_passes=False, + inductor_compilation="full", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + with self.assertRaisesRegex(ValueError, "enable_passes"): + validate_fp8_graph_config(config) + + def test_fp8_allows_cudagraph(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="full", + fp8=FP8GraphConfig(enabled=True), + ) + validate_fp8_graph_config(config) + + def test_fp8_precompile_is_rejected(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="full", + disable_passes=["cudagraph_pass"], + precompile_artifact_dir="/tmp/fp8", + fp8=FP8GraphConfig(enabled=True), + ) + with self.assertRaisesRegex(ValueError, "incompatible with"): + validate_fp8_graph_config(config) + + def test_regional_fp8_precompile_is_rejected(self) -> None: + config = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="regional", + disable_passes=["cudagraph_pass"], + precompile_artifact_dir="/tmp/fp8", + fp8=FP8GraphConfig(enabled=True), + ) + with self.assertRaisesRegex(ValueError, "incompatible with"): + validate_fp8_graph_config(config) + +class TestFP8Provenance(TestCase): + def test_module_annotation_includes_quantization_kind(self) -> None: + model = nn.Sequential(nn.ReLU()) + with patch( + "torchtitan.experiments.graph_trainer.common_utils.get_quantization_kind", + return_value="float8_linear", + ): + annotate_module_fqns(model) + + with preserve_node_meta(): + gm = make_fx(model)(torch.randn(2, 2)) + custom_metadata = [ + node.meta.get("custom", {}) + for node in gm.graph.nodes + if node.meta.get("custom", {}).get(_MODULE_FQN) == "0" + ] + self.assertTrue(custom_metadata) + self.assertTrue( + all( + metadata[_QUANTIZATION_KIND] == "float8_linear" + for metadata in custom_metadata + ) + ) + self.assertTrue( + all( + metadata[_QUANTIZATION_EMULATE] is False + for metadata in custom_metadata + ) + ) + + def test_module_annotation_includes_emulation_mode(self) -> None: + model = nn.Sequential(nn.ReLU()) + model[0]._quantization_emulate = True + with patch( + "torchtitan.experiments.graph_trainer.common_utils.get_quantization_kind", + return_value="float8_linear", + ): + annotate_module_fqns(model) + + with preserve_node_meta(): + gm = make_fx(model)(torch.randn(2, 2)) + + custom_metadata = [ + node.meta.get("custom", {}) + for node in gm.graph.nodes + if node.meta.get("custom", {}).get(_MODULE_FQN) == "0" + ] + self.assertTrue(custom_metadata) + self.assertTrue( + all( + metadata[_QUANTIZATION_EMULATE] is True + for metadata in custom_metadata + ) + ) + + def test_quantization_kind_is_copied_to_backward_nodes(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + forward = graph.call_function(torch.ops.aten.relu.default, args=(x,)) + forward.meta["seq_nr"] = 1 + forward.meta["custom"] = { + _MODULE_FQN: "layers.0.feed_forward.w1", + _QUANTIZATION_KIND: "float8_linear", + } + backward = graph.call_function(torch.ops.aten.relu.default, args=(forward,)) + backward.meta["seq_nr"] = 1 + backward.meta["autograd_backward"] = True + graph.output(backward) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + _copy_fwd_metadata_to_bw_nodes(gm) + + self.assertEqual(backward.meta["custom"], forward.meta["custom"]) + + +class TestFP8ValidationPass(TestCase): + def _graph_with_quantized_node( + self, + target, + *, + backward: bool = False, + quantization_kind: str = "float8_linear", + data_operand_dtype: torch.dtype = torch.float8_e4m3fn, + ): + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(1, dtype=torch.bfloat16) + data_operand = graph.placeholder("data_operand") + data_operand.meta["val"] = torch.empty(1, dtype=data_operand_dtype) + scale_operand = graph.placeholder("scale_operand") + scale_operand.meta["val"] = torch.empty(1, dtype=torch.float8_e4m3fn) + node = graph.call_function(target, args=(x, data_operand, scale_operand)) + node.meta["custom"] = { + _MODULE_FQN: "layers.0.feed_forward.w1", + _QUANTIZATION_KIND: quantization_kind, + } + if backward: + node.meta["autograd_backward"] = True + graph.output(node) + return torch.fx.GraphModule(torch.nn.Module(), graph), node + + def test_records_forward_and_backward_fp8_compute_ops(self) -> None: + targets = FP8_COMPUTE_TARGETS["float8_linear"] + self.assertTrue(targets) + target = next(iter(targets)) + gm, forward = self._graph_with_quantized_node(target) + output = next(node for node in gm.graph.nodes if node.op == "output") + data_operand = next( + node for node in gm.graph.nodes if node.name == "data_operand" + ) + scale_operand = next( + node for node in gm.graph.nodes if node.name == "scale_operand" + ) + with gm.graph.inserting_before(output): + backward = gm.graph.call_function( + target, args=(forward, data_operand, scale_operand) + ) + backward.meta["custom"] = dict(forward.meta["custom"]) + backward.meta["autograd_backward"] = True + output.args = (backward,) + gm.recompile() + + validate_fp8_graph_pass(gm, strict=True) + + self.assertNotIn("fp8_summary", gm.meta) + self.assertNotIn("fp8", forward.meta["custom"]) + self.assertNotIn("fp8", backward.meta["custom"]) + self.assertNotIn("compile_with_inductor", forward.meta["custom"]) + self.assertNotIn("compile_with_inductor", backward.meta["custom"]) + + def test_strict_validation_rejects_missing_fp8_compute_op(self) -> None: + gm, _ = self._graph_with_quantized_node(torch.ops.aten.relu.default) + + with self.assertRaisesRegex(RuntimeError, "without a supported FP8 compute"): + validate_fp8_graph_pass(gm, strict=True) + + def test_compute_targets_are_selected_by_quantization_kind(self) -> None: + targets = FP8_COMPUTE_TARGETS["mxfp8_linear"] + self.assertTrue(targets) + gm, _ = self._graph_with_quantized_node( + next(iter(targets)), + quantization_kind="mxfp8_linear", + ) + + validate_fp8_graph_pass(gm, strict=True) + + self.assertNotIn("fp8_summary", gm.meta) + + def test_grouped_compute_requires_scaled_grouped_mm(self) -> None: + dense_targets = FP8_COMPUTE_TARGETS["float8_linear"] + self.assertTrue(dense_targets) + + for quantization_kind in ( + "float8_grouped_experts", + "mxfp8_grouped_experts", + ): + grouped_targets = FP8_COMPUTE_TARGETS[quantization_kind] + self.assertTrue(grouped_targets) + grouped_gm, _ = self._graph_with_quantized_node( + next(iter(grouped_targets)), + quantization_kind=quantization_kind, + ) + validate_fp8_graph_pass(grouped_gm, strict=True) + + dense_gm, _ = self._graph_with_quantized_node( + next(iter(dense_targets)), + quantization_kind=quantization_kind, + ) + with self.assertRaisesRegex( + RuntimeError, "without a supported FP8 compute" + ): + validate_fp8_graph_pass(dense_gm, strict=True) + + def test_fp8_scale_operand_does_not_prove_fp8_compute(self) -> None: + targets = FP8_COMPUTE_TARGETS["float8_linear"] + self.assertTrue(targets) + gm, _ = self._graph_with_quantized_node( + next(iter(targets)), + data_operand_dtype=torch.bfloat16, + ) + + with self.assertRaisesRegex(RuntimeError, "without a supported FP8 compute"): + validate_fp8_graph_pass(gm, strict=True) + + def test_emulated_float8_region_does_not_require_scaled_compute(self) -> None: + gm, node = self._graph_with_quantized_node(torch.ops.aten.mm.default) + node.meta["custom"][_QUANTIZATION_EMULATE] = True + + validate_fp8_graph_pass(gm, strict=True) + + self.assertNotIn("fp8_summary", gm.meta) + + def test_non_quantized_graph_is_a_noop(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + graph.output(x) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + validate_fp8_graph_pass(gm, strict=False) + + self.assertNotIn("fp8_summary", gm.meta) + + def test_strict_validation_rejects_missing_quantized_region(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + graph.output(x) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with self.assertRaisesRegex(RuntimeError, "did not find quantized module"): + validate_fp8_graph_pass(gm, strict=True) + + +class TestFP8RegionalAnnotation(TestCase): + def _node(self, graph, target, args=()): + node = graph.call_function(target, args=args) + node.meta["val"] = torch.empty(1, device="meta") + node.meta["custom"] = { + _MODULE_FQN: "layers.0.feed_forward.w1", + _QUANTIZATION_KIND: "float8_linear", + "fp8": {"op_role": "cast"}, + } + return node + + def test_identifies_and_tags_connected_fp8_component(self) -> None: + targets = FP8_COMPUTE_TARGETS["float8_linear"] + self.assertTrue(targets) + target = next(iter(targets)) + graph = torch.fx.Graph() + x = graph.placeholder("x") + cast = self._node(graph, torch.ops.aten.clone.default, (x,)) + cast.meta["val"] = torch.empty( + 1, device="meta", dtype=torch.float8_e4m3fn + ) + gemm = self._node(graph, target, (cast,)) + gemm.meta["custom"]["fp8"]["op_role"] = "compute" + graph.output(gemm) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with patch( + "torchtitan.experiments.graph_trainer.fp8_passes._is_regional_fp8_compute_node", + return_value=True, + ): + annotate_fp8_regions_for_regional_inductor_pass(gm, strict=False) + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + self.assertNotIn("fp8_regional_summary", gm.meta) + expected = {"inductor_region": 0} + self.assertEqual(cast.meta["custom"]["compile_with_inductor"], expected) + self.assertEqual(gemm.meta["custom"]["compile_with_inductor"], expected) + + def test_identifies_and_tags_grouped_experts_component(self) -> None: + for quantization_kind in ( + "float8_grouped_experts", + "mxfp8_grouped_experts", + ): + graph = torch.fx.Graph() + x = graph.placeholder("x") + cast = self._node(graph, torch.ops.aten.clone.default, (x,)) + compute = self._node( + graph, + next(iter(FP8_COMPUTE_TARGETS[quantization_kind])), + (cast,), + ) + for node in (cast, compute): + node.meta["custom"][_MODULE_FQN] = ( + "layers.0.moe.routed_experts.inner_experts" + ) + node.meta["custom"][_QUANTIZATION_KIND] = quantization_kind + cast.meta["val"] = torch.empty( + 1, device="meta", dtype=torch.float8_e4m3fn + ) + compute.meta["custom"]["fp8"]["op_role"] = "compute" + graph.output(compute) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with patch( + "torchtitan.experiments.graph_trainer.fp8_passes._is_regional_fp8_compute_node", + return_value=True, + ): + annotate_fp8_regions_for_regional_inductor_pass(gm, strict=False) + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + self.assertNotIn("fp8_regional_summary", gm.meta) + expected = {"inductor_region": 0} + self.assertEqual(cast.meta["custom"]["compile_with_inductor"], expected) + self.assertEqual( + compute.meta["custom"]["compile_with_inductor"], expected + ) + + def test_identifies_fp8_region_without_tagging_inductor(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + cast = self._node(graph, torch.ops.aten.clone.default, (x,)) + cast.meta["val"] = torch.empty( + 1, device="meta", dtype=torch.float8_e4m3fn + ) + compute = self._node( + graph, + next(iter(FP8_COMPUTE_TARGETS["float8_linear"])), + (cast,), + ) + graph.output(compute) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with patch( + "torchtitan.experiments.graph_trainer.fp8_passes._is_regional_fp8_compute_node", + return_value=True, + ): + annotate_fp8_regions_for_regional_inductor_pass( + gm, + strict=False, + ) + + self.assertNotIn("compile_with_inductor", cast.meta["custom"]) + self.assertNotIn("compile_with_inductor", compute.meta["custom"]) + self.assertEqual(cast.meta["custom"]["fp8"]["regional_region_num_nodes"], 2) + + def test_tags_only_complete_identified_fp8_region(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + cast = self._node(graph, torch.ops.aten.clone.default, (x,)) + compute = self._node( + graph, + next(iter(FP8_COMPUTE_TARGETS["float8_linear"])), + (cast,), + ) + for node in (cast, compute): + node.meta["custom"]["fp8"].update( + {"regional_region_id": 3, "regional_region_num_nodes": 2} + ) + compute.meta["custom"]["fp8"]["op_role"] = "compute" + other = graph.call_function(torch.ops.aten.relu.default, args=(compute,)) + other.meta["custom"] = {"compile_with_inductor": {"source": "flex"}} + graph.output(other) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + expected = {"inductor_region": 3} + self.assertEqual(cast.meta["custom"]["compile_with_inductor"], expected) + self.assertEqual(compute.meta["custom"]["compile_with_inductor"], expected) + self.assertEqual( + other.meta["custom"]["compile_with_inductor"], + {"source": "flex"}, + ) + + partial_graph = torch.fx.Graph() + partial_input = partial_graph.placeholder("x") + partial_cast = self._node( + partial_graph, + torch.ops.aten.clone.default, + (partial_input,), + ) + partial_cast.meta["custom"]["fp8"].update( + {"regional_region_id": 3, "regional_region_num_nodes": 2} + ) + partial_graph.output(partial_cast) + partial_gm = torch.fx.GraphModule(torch.nn.Module(), partial_graph) + + with self.assertWarnsRegex(UserWarning, "incomplete FP8 regional"): + annotate_complete_fp8_regions_for_regional_inductor_pass(partial_gm) + + self.assertNotIn("compile_with_inductor", partial_cast.meta["custom"]) + + def test_skips_regions_with_unbound_input_size_symbols(self) -> None: + # EP-padded MoE tensors have sizes like round_up(u2 + u3 + C, 16). + # When those tensors become regional Inductor inputs without a simple + # binding for u2/u3, skip regional compilation (eager fallback). + # With TORCHTITAN_FP8_EP_UNBACKED_PAD=1 the pad dim is a fresh unbacked + # SymInt instead; this skip path remains the default EP fallback. + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + u2 = shape_env.create_unbacked_symint() + u3 = shape_env.create_unbacked_symint() + padded = ((u2 + u3 + 79) // 16) * 16 + padded_val = torch.empty( + (padded, 8), device="meta", dtype=torch.float8_e4m3fn + ) + + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = padded_val + compute = self._node( + graph, + next(iter(FP8_COMPUTE_TARGETS["float8_grouped_experts"])), + (x,), + ) + compute.meta["custom"]["fp8"].update( + { + "op_role": "compute", + "regional_region_id": 0, + "regional_region_num_nodes": 1, + } + ) + graph.output(compute) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with self.assertWarnsRegex(UserWarning, "unbound input size symbols"): + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + self.assertNotIn("compile_with_inductor", compute.meta["custom"]) + + def test_reidentifies_partitioned_fp8_regions(self) -> None: + target = next(iter(FP8_COMPUTE_TARGETS["float8_linear"])) + + def make_partition(region_id: int) -> tuple[torch.fx.GraphModule, list]: + graph = torch.fx.Graph() + x = graph.placeholder("x") + cast = self._node(graph, torch.ops.aten.clone.default, (x,)) + cast.meta["val"] = torch.empty( + 1, device="meta", dtype=torch.float8_e4m3fn + ) + compute = self._node(graph, target, (cast,)) + compute.meta["custom"]["fp8"]["op_role"] = "compute" + graph.output(compute) + for node in (cast, compute): + node.meta["custom"]["fp8"].update( + { + "regional_region_id": region_id, + "regional_region_num_nodes": 4, + } + ) + return torch.fx.GraphModule(torch.nn.Module(), graph), [cast, compute] + + for region_id in (0, 1): + gm, nodes = make_partition(region_id) + with patch( + "torchtitan.experiments.graph_trainer.fp8_passes." + "_is_regional_fp8_compute_node", + return_value=True, + ): + annotate_fp8_regions_for_regional_inductor_pass(gm, strict=False) + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + self.assertNotIn("fp8_regional_summary", gm.meta) + self.assertTrue( + all( + node.meta["custom"]["fp8"]["regional_region_num_nodes"] == 2 + for node in nodes + ) + ) + self.assertTrue( + all("compile_with_inductor" in node.meta["custom"] for node in nodes) + ) + + def test_tags_compute_after_graph_pp_fp8_input_boundary(self) -> None: + target = next(iter(FP8_COMPUTE_TARGETS["float8_linear"])) + graph = torch.fx.Graph() + grad_output_fp8 = graph.placeholder("grad_output_fp8") + weight_fp8 = graph.placeholder("weight_fp8") + for node in (grad_output_fp8, weight_fp8): + node.meta["val"] = torch.empty( + 1, device="meta", dtype=torch.float8_e4m3fn + ) + node.meta["custom"] = { + _MODULE_FQN: "layers.0.feed_forward.w1", + _QUANTIZATION_KIND: "float8_linear", + } + compute = self._node(graph, target, (grad_output_fp8, weight_fp8)) + graph.output(compute) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + with patch( + "torchtitan.experiments.graph_trainer.fp8_passes." + "_is_regional_fp8_compute_node", + side_effect=lambda node, **_kwargs: node.op == "call_function", + ): + annotate_fp8_regions_for_regional_inductor_pass(gm, strict=False) + annotate_complete_fp8_regions_for_regional_inductor_pass(gm) + + for boundary in (grad_output_fp8, weight_fp8): + self.assertEqual( + boundary.meta["custom"]["fp8"]["op_role"], "input_boundary" + ) + self.assertNotIn("compile_with_inductor", boundary.meta["custom"]) + self.assertEqual(compute.meta["custom"]["fp8"]["op_role"], "compute") + self.assertEqual( + compute.meta["custom"]["fp8"]["regional_region_num_nodes"], 1 + ) + self.assertEqual( + compute.meta["custom"]["compile_with_inductor"], + {"inductor_region": 0}, + ) + + +class TestFP8PassOrdering(TestCase): + def test_full_inductor_precedes_cudagraph(self) -> None: + config = SimpleNamespace( + compile=GraphTrainerCompileConfig( + inductor_compilation="full", + fp8=FP8GraphConfig(enabled=True), + ), + loss=CrossEntropyLoss.Config(), + model_spec=SimpleNamespace(model=SimpleNamespace(layers=[0])), + parallelism=SimpleNamespace(enable_async_tensor_parallel=False), + ) + traced_result = SimpleNamespace( + state_fqns=[], + num_static_inputs=1, + tensor_input_indices=[0, 1], + ) + + passes = construct_default_graph_passes(traced_result, config) + pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in passes + ] + + self.assertLess( + pass_names.index("full_inductor_compilation_pass"), + pass_names.index("cudagraph_pass"), + ) + + def test_full_terminal_fp8_pass_inclusion(self) -> None: + config = GraphTrainerCompileConfig( + inductor_compilation="full", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + + fp8_passes = final_inductor_compile_passes(config) + fp8_pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in fp8_passes + ] + self.assertEqual( + fp8_pass_names, + ["validate_fp8_graph_pass", "full_inductor_compilation_pass"], + ) + + config.fp8.enabled = False + skipped_passes = final_inductor_compile_passes(config) + self.assertEqual(len(skipped_passes), 1) + + def test_regional_terminal_fp8_pass_inclusion(self) -> None: + config = GraphTrainerCompileConfig( + inductor_compilation="regional", + numerics_changing_optim=True, + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + + passes = final_inductor_compile_passes(config) + pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in passes + ] + self.assertLess( + pass_names.index("annotate_rmsnorm_for_regional_inductor_pass"), + pass_names.index("annotate_fp8_regions_for_regional_inductor_pass"), + ) + self.assertLess( + pass_names.index("annotate_fp8_regions_for_regional_inductor_pass"), + pass_names.index( + "annotate_complete_fp8_regions_for_regional_inductor_pass" + ), + ) + self.assertLess( + pass_names.index( + "annotate_complete_fp8_regions_for_regional_inductor_pass" + ), + pass_names.index("regional_inductor_pass"), + ) + + config.fp8.enabled = False + skipped_passes = final_inductor_compile_passes(config) + skipped_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in skipped_passes + ] + self.assertNotIn( + "annotate_fp8_regions_for_regional_inductor_pass", skipped_names + ) + self.assertNotIn( + "annotate_complete_fp8_regions_for_regional_inductor_pass", + skipped_names, + ) + + def test_full_validation_follows_graph_rewrites(self) -> None: + config = SimpleNamespace( + compile=GraphTrainerCompileConfig( + enable=True, + inductor_compilation="full", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + enable_async_tensor_parallel=True, + ), + loss=CrossEntropyLoss.Config(), + model_spec=SimpleNamespace(model=SimpleNamespace(layers=[0])), + parallelism=SimpleNamespace(), + ) + traced_result = SimpleNamespace(state_fqns=[]) + + passes = compile_time_passes(traced_result, config) + pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in passes + ] + + self.assertIn("validate_fp8_graph_pass", pass_names) + self.assertNotIn("annotate_fp8_regions_for_regional_inductor_pass", pass_names) + self.assertLess( + pass_names.index("async_tensor_parallel_pass"), + pass_names.index("validate_fp8_graph_pass"), + ) + self.assertLess( + pass_names.index("validate_fp8_graph_pass"), + pass_names.index("full_inductor_compilation_pass"), + ) + + def test_regional_annotation_precedes_regional_inductor(self) -> None: + config = SimpleNamespace( + compile=GraphTrainerCompileConfig( + inductor_compilation="regional", + disable_passes=["cudagraph_pass"], + numerics_changing_optim=True, + fp8=FP8GraphConfig(enabled=True), + ), + loss=CrossEntropyLoss.Config(), + model_spec=SimpleNamespace(model=SimpleNamespace(layers=[0])), + parallelism=SimpleNamespace(enable_async_tensor_parallel=False), + ) + traced_result = SimpleNamespace(state_fqns=[]) + + passes = compile_time_passes(traced_result, config) + pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in passes + ] + + self.assertLess( + pass_names.index("annotate_flex_attention_for_regional_inductor_pass"), + pass_names.index("annotate_rmsnorm_for_regional_inductor_pass"), + ) + self.assertLess( + pass_names.index("annotate_rmsnorm_for_regional_inductor_pass"), + pass_names.index("annotate_fp8_regions_for_regional_inductor_pass"), + ) + self.assertLess( + pass_names.index("annotate_fp8_regions_for_regional_inductor_pass"), + pass_names.index("regional_inductor_pass"), + ) + + def test_graph_pp_validates_before_partitioning(self) -> None: + compile_config = GraphTrainerCompileConfig( + inductor_compilation="regional", + disable_passes=["cudagraph_pass"], + fp8=FP8GraphConfig(enabled=True), + ) + + passes = graph_pp_pre_partition_fp8_passes(compile_config) + pass_names = [ + pass_fn.func.__name__ if hasattr(pass_fn, "func") else pass_fn.__name__ + for pass_fn in passes + ] + + self.assertEqual(pass_names, ["validate_fp8_graph_pass"]) + + +class ScaledGroupedMMMetaTest(TestCase): + def test_soft_meta_falls_back_on_data_dependent_layout_guard(self) -> None: + import os + from unittest import mock + + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ( + GuardOnDataDependentSymNode, + ShapeEnv, + ) + + from torchtitan.components.quantization import scaled_grouped_mm_meta as meta + from torchtitan.components.quantization.scaled_grouped_mm_meta import ( + ENV_FLAG, + install_scaled_grouped_mm_meta, + ) + + previous = os.environ.get(ENV_FLAG) + os.environ[ENV_FLAG] = "1" + # Allow reinstall in this process if a prior test already patched. + meta._installed = False + meta._orig_meta_grouped_mm_common = None + try: + self.assertTrue(install_scaled_grouped_mm_meta()) + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + u = shape_env.create_unbacked_symint() + mat_a = torch.empty((u, 16), device="meta", dtype=torch.float8_e4m3fn) + mat_b = torch.empty( + (4, 32, 16), device="meta", dtype=torch.float8_e4m3fn + ).transpose(-2, -1) + scale_a = torch.empty((u,), device="meta", dtype=torch.float32) + scale_b = torch.empty((4, 32), device="meta", dtype=torch.float32) + offs = torch.empty((4,), device="meta", dtype=torch.int32) + + def _raise_guard(*_args, **_kwargs): + raise GuardOnDataDependentSymNode( + "Could not guard on data-dependent expression u > 1" + ) + + with mock.patch.object( + meta, "_orig_meta_grouped_mm_common", side_effect=_raise_guard + ): + out = torch._scaled_grouped_mm( + mat_a, + mat_b, + scale_a, + scale_b, + offs=offs, + out_dtype=torch.bfloat16, + ) + self.assertEqual(tuple(out.shape), (u, 32)) + finally: + if previous is None: + os.environ.pop(ENV_FLAG, None) + else: + os.environ[ENV_FLAG] = previous + + def test_soft_meta_disabled_by_default(self) -> None: + import os + + from torchtitan.components.quantization.scaled_grouped_mm_meta import ( + ENV_FLAG, + ep_unbacked_pad_enabled, + install_scaled_grouped_mm_meta, + ) + + previous = os.environ.get(ENV_FLAG) + os.environ.pop(ENV_FLAG, None) + try: + self.assertFalse(ep_unbacked_pad_enabled()) + # Already-installed patch from other tests may make this a no-op; + # without the env flag a fresh install must not proceed. + from torchtitan.components.quantization import scaled_grouped_mm_meta as meta + + was_installed = meta._installed + meta._installed = False + try: + self.assertFalse(install_scaled_grouped_mm_meta()) + finally: + meta._installed = was_installed + finally: + if previous is not None: + os.environ[ENV_FLAG] = previous From d61e2b711baa7a21731e1491f22a277938c18d2b Mon Sep 17 00:00:00 2001 From: Zihan Yang Date: Fri, 28 Aug 2026 18:14:55 +0800 Subject: [PATCH 2/2] Add dense FP8 regional compilation --- ...egration_test_8gpu_graph_trainer_h100.yaml | 3 + .../graph_trainer/graph_pp/graph_builder.py | 9 +++ .../graph_trainer/llama3/config_registry.py | 26 +++++++ .../tests/test_graph_pp_passes.py | 54 +++++++++++++++ .../tests/test_graph_pp_runner.py | 69 +++++++++++-------- .../graph_trainer/tests/test_numerics.py | 49 +++++++++++++ .../experiments/graph_trainer/trainer.py | 2 + torchtitan/models/llama3/config_registry.py | 11 +-- 8 files changed, 189 insertions(+), 34 deletions(-) diff --git a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml b/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml index 53a942e9c4..0d62904f8c 100644 --- a/.github/workflows/integration_test_8gpu_graph_trainer_h100.yaml +++ b/.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 diff --git a/torchtitan/experiments/graph_trainer/graph_pp/graph_builder.py b/torchtitan/experiments/graph_trainer/graph_pp/graph_builder.py index 3f4e42c848..d97f459c9b 100644 --- a/torchtitan/experiments/graph_trainer/graph_pp/graph_builder.py +++ b/torchtitan/experiments/graph_trainer/graph_pp/graph_builder.py @@ -82,6 +82,7 @@ deduplicate_fsdp_unshard_chains_pass, eliminate_dead_code_pass, final_inductor_compile_passes, + graph_pp_pre_partition_fp8_passes, ) from torchtitan.protocols.model import BaseModel from torchtitan.tools.logging import logger @@ -675,6 +676,9 @@ def _compile_graph_pp_module( compile_config, use_cudagraph=False, boxed_codegen=True, + # The complete stage joint graph was validated before partitioning. + # Extracted helper callables may legitimately contain no FP8 nodes. + fp8_strict_validation=False, ), compile_config=compile_config, ) @@ -824,6 +828,11 @@ def _apply_graph_pp_pre_partition_passes( include_inductor=False, include_mandatory_normalization=False, ) + passes.extend( + graph_pp_pre_partition_fp8_passes( + compile_config, + ) + ) traced.gm = apply_graph_passes( traced.gm, traced.example_inputs, diff --git a/torchtitan/experiments/graph_trainer/llama3/config_registry.py b/torchtitan/experiments/graph_trainer/llama3/config_registry.py index 041623cc7d..45c6405bc1 100644 --- a/torchtitan/experiments/graph_trainer/llama3/config_registry.py +++ b/torchtitan/experiments/graph_trainer/llama3/config_registry.py @@ -10,6 +10,7 @@ from torchtitan.components.loss import CrossEntropyLoss from torchtitan.components.quantization import MXFP8LinearConverter from torchtitan.experiments.graph_trainer.configs import ( + FP8GraphConfig, GraphTrainerCompileConfig, to_graph_trainer_config, ) @@ -23,6 +24,7 @@ llama3_8b, llama3_debugmodel, llama3_debugmodel_dist_gemm, + llama3_debugmodel_float8, ) from . import model_registry @@ -50,6 +52,30 @@ def graph_trainer_llama3_debugmodel_dist_gemm() -> GraphTrainer.Config: return config +def graph_trainer_llama3_debugmodel_float8() -> GraphTrainer.Config: + config = to_graph_trainer_config( + llama3_debugmodel_float8(model_compile_enabled=True), model_registry + ) + config.compile = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="full", + fp8=FP8GraphConfig(enabled=True), + ) + return config + + +def graph_trainer_llama3_debugmodel_float8_regional() -> GraphTrainer.Config: + config = to_graph_trainer_config( + llama3_debugmodel_float8(model_compile_enabled=True), model_registry + ) + config.compile = GraphTrainerCompileConfig( + enable=True, + inductor_compilation="regional", + fp8=FP8GraphConfig(enabled=True), + ) + return config + + def graph_trainer_llama3_debugmodel_sdpa() -> GraphTrainer.Config: """Debug model on the test-only SDPA backend. diff --git a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py index 2d05525860..39e9912b42 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_passes.py @@ -21,6 +21,8 @@ from torchtitan.config import DebugConfig, ParallelismConfig, TrainingConfig from torchtitan.distributed import ParallelDims from torchtitan.experiments.graph_trainer.common_utils import ( + _MODULE_FQN, + _QUANTIZATION_KIND, maybe_register_blockmask_pytree_node, ) from torchtitan.experiments.graph_trainer.deepseek_v3 import ( @@ -263,6 +265,58 @@ def _assert_tensor_sequence_equal( class GraphPPPartitionTest(unittest.TestCase): + def test_partition_preserves_atomic_fp8_chains_and_metadata(self) -> None: + def stage_step(x: torch.Tensor): + forward_cast = x.sin() + forward_scale = forward_cast.cos() + forward_gemm = forward_scale.tanh() + backward_cast = x.neg() + backward_scale = backward_cast.sin() + backward_gemm = backward_scale.cos() + return [forward_gemm, backward_gemm] + + traced = minimal_fx_tracer(stage_step)(torch.randn(2, 4)) + compute_nodes = [ + node for node in traced.gm.graph.nodes if node.op == "call_function" + ] + self.assertEqual(len(compute_nodes), 6) + expected_names_by_phase = { + "forward": {node.name for node in compute_nodes[:3]}, + "backward": {node.name for node in compute_nodes[3:]}, + } + for phase, nodes in ( + ("forward", compute_nodes[:3]), + ("backward", compute_nodes[3:]), + ): + for node in nodes: + node.meta["custom"] = { + _MODULE_FQN: "layers.0.feed_forward.w1", + _QUANTIZATION_KIND: "float8_linear", + "fp8_chain_phase": phase, + "compile_with_inductor": {}, + } + + fw_module, bw_module, _ = partition_joint_graph( + traced, + num_fwd_outputs=1, + ) + + def fp8_chain_names(gm: fx.GraphModule) -> set[str]: + return { + node.name + for node in gm.graph.nodes + if node.meta.get("custom", {}).get(_QUANTIZATION_KIND) + == "float8_linear" + } + + self.assertEqual(fp8_chain_names(fw_module), expected_names_by_phase["forward"]) + self.assertEqual(fp8_chain_names(bw_module), expected_names_by_phase["backward"]) + for gm in (fw_module, bw_module): + for node in gm.graph.nodes: + custom = node.meta.get("custom", {}) + if custom.get(_QUANTIZATION_KIND) == "float8_linear": + self.assertEqual(custom["compile_with_inductor"], {}) + def test_real_dsv3_moe_block_partition_matches_joint_graph(self) -> None: traced_block = _trace_dsv3_moe_block_stage() diff --git a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py index c2f5d390fc..4cd2cb817c 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py +++ b/torchtitan/experiments/graph_trainer/tests/test_graph_pp_runner.py @@ -31,7 +31,10 @@ ensure_boxed_graph_module, maybe_register_blockmask_pytree_node, ) -from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig +from torchtitan.experiments.graph_trainer.configs import ( + FP8GraphConfig, + GraphTrainerCompileConfig, +) from torchtitan.experiments.graph_trainer.graph_pp import multiplex_fw_bw_graph from torchtitan.experiments.graph_trainer.graph_pp.graph_builder import ( _build_graph_pp_overlap_graphs, @@ -475,39 +478,45 @@ def test_graph_pp_accepts_zero_two_fsdp_reshard_policies(self) -> None: def test_graph_pp_compile_uses_inductor_compilation_with_default_backend( self, ) -> None: - gm = torch.fx.symbolic_trace(lambda x: x + 1) - for node in gm.graph.find_nodes(op="placeholder"): - node.meta["val"] = torch.randn(2) - compile_config = GraphTrainerCompileConfig(enable=True) - def boxed_apply_graph_passes(gm, example_inputs, passes, compile_config): return ensure_boxed_graph_module(gm) - with ( - mock.patch( - "torchtitan.experiments.graph_trainer.graph_pp.graph_builder." - "final_inductor_compile_passes", - return_value=[], - ) as final_inductor_passes, - mock.patch( - "torchtitan.experiments.graph_trainer.graph_pp.graph_builder." - "apply_graph_passes", - side_effect=boxed_apply_graph_passes, - ) as apply_graph_passes, - ): - compiled = _compile_graph_pp_module( - gm, - compile_config=compile_config, - graph_name="test_graph", - ) + for fp8_enabled in (False, True): + with self.subTest(fp8_enabled=fp8_enabled): + gm = torch.fx.symbolic_trace(lambda x: x + 1) + for node in gm.graph.find_nodes(op="placeholder"): + node.meta["val"] = torch.randn(2) + compile_config = GraphTrainerCompileConfig( + enable=True, + fp8=FP8GraphConfig(enabled=fp8_enabled), + ) - self.assertIs(compiled, gm) - final_inductor_passes.assert_called_once_with( - compile_config, - use_cudagraph=False, - boxed_codegen=True, - ) - apply_graph_passes.assert_called_once() + with ( + mock.patch( + "torchtitan.experiments.graph_trainer.graph_pp.graph_builder." + "final_inductor_compile_passes", + return_value=[], + ) as final_inductor_passes, + mock.patch( + "torchtitan.experiments.graph_trainer.graph_pp.graph_builder." + "apply_graph_passes", + side_effect=boxed_apply_graph_passes, + ) as apply_graph_passes, + ): + compiled = _compile_graph_pp_module( + gm, + compile_config=compile_config, + graph_name="test_graph", + ) + + self.assertIs(compiled, gm) + final_inductor_passes.assert_called_once_with( + compile_config, + use_cudagraph=False, + boxed_codegen=True, + fp8_strict_validation=False, + ) + apply_graph_passes.assert_called_once() def test_graph_pp_graph_execution_uses_mutable_boxed_args(self) -> None: gm = torch.fx.symbolic_trace(lambda x, y: x + y) diff --git a/torchtitan/experiments/graph_trainer/tests/test_numerics.py b/torchtitan/experiments/graph_trainer/tests/test_numerics.py index 0755d05678..957b306cda 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_numerics.py +++ b/torchtitan/experiments/graph_trainer/tests/test_numerics.py @@ -22,6 +22,7 @@ from torchtitan.components.loss import cross_entropy_loss from torchtitan.distributed import ParallelDims from torchtitan.experiments.graph_trainer.simple_fsdp import data_parallel +from torchtitan.tools.utils import has_cuda_capability STEPS = 20 @@ -244,6 +245,29 @@ def _run_llama3_loss_compare(test_options_extra: str = "") -> bool: ) +def _run_llama3_fp8_loss_compare( + test_config: str, + test_options_extra: str = "", + *, + rtol: float | None = None, +) -> bool: + """Compare Trainer FP8 with a GraphTrainer FP8 compilation mode.""" + test_options = LLAMA3_PARALLELISM + if test_options_extra: + test_options += f" {test_options_extra}" + compare_fn = run_loss_compare if rtol is None else run_loss_compare_close + compare_kwargs = {} if rtol is None else {"rtol": rtol} + return compare_fn( + baseline_module="llama3", + baseline_config="llama3_debugmodel_float8", + test_module="graph_trainer.llama3", + test_config=test_config, + baseline_options=LLAMA3_PARALLELISM, + test_options=test_options, + **compare_kwargs, + ) + + DSV3_PARALLELISM = ( "--parallelism.data_parallel_shard_degree=4" " --parallelism.tensor_parallel_degree=2" @@ -645,6 +669,31 @@ def test_moe_qwen3_aot_fx_trace_vs_eager(self): ) +@unittest.skipUnless( + torch.cuda.is_available() + and has_cuda_capability(9, 0) + and importlib.util.find_spec("torchao") is not None, + "FP8 numerics tests require TorchAO and an H100-class GPU", +) +class TestGraphTrainerFP8Numerics(unittest.TestCase): + """H100-only dense FP8 loss equivalence against the Trainer path.""" + + def test_dense_llama3_fp8_full_cudagraph_vs_trainer(self): + self.assertTrue( + _run_llama3_fp8_loss_compare( + "graph_trainer_llama3_debugmodel_float8", + rtol=1e-4, + ) + ) + + def test_dense_llama3_fp8_regional_cudagraph_vs_trainer(self): + self.assertTrue( + _run_llama3_fp8_loss_compare( + "graph_trainer_llama3_debugmodel_float8_regional", + ) + ) + + @unittest.skipUnless( importlib.util.find_spec("autoparallel"), "AutoParallel numerics tests require the autoparallel package", diff --git a/torchtitan/experiments/graph_trainer/trainer.py b/torchtitan/experiments/graph_trainer/trainer.py index 2d6105e5a6..385bca1ac1 100644 --- a/torchtitan/experiments/graph_trainer/trainer.py +++ b/torchtitan/experiments/graph_trainer/trainer.py @@ -21,6 +21,7 @@ from torchtitan.experiments.graph_trainer.configs import ( GraphTrainerCompileConfig, trace_input_preparer_keys, + validate_fp8_graph_config, ) from torchtitan.experiments.graph_trainer.cudagraph import cudagraph_teardown from torchtitan.experiments.graph_trainer.make_fx_tracer import ( @@ -116,6 +117,7 @@ class Config(Trainer.Config): def __init__(self, config): super().__init__(config) + validate_fp8_graph_config(self.config.compile) validate_memory_policy_config(self.config.compile) diff --git a/torchtitan/models/llama3/config_registry.py b/torchtitan/models/llama3/config_registry.py index 1dbc1eee9c..4009677006 100644 --- a/torchtitan/models/llama3/config_registry.py +++ b/torchtitan/models/llama3/config_registry.py @@ -103,11 +103,14 @@ def llama3_debugmodel_dist_gemm() -> Trainer.Config: return config -def llama3_debugmodel_float8() -> Trainer.Config: +def llama3_debugmodel_float8( + model_compile_enabled: bool | None = None, +) -> Trainer.Config: config = llama3_debugmodel() - model_compile_enabled = ( - config.compile.enable and "model" in config.compile.components - ) + if model_compile_enabled is None: + model_compile_enabled = ( + config.compile.enable and "model" in config.compile.components + ) config.model_spec = model_registry( "debugmodel", converters=[