[2/2] Add grouped-expert FP8 GraphTrainer support

[WhoisZihan]

This is the second of two stacked PRs. It builds on PR #4384 and adds grouped-expert FP8 support for GraphTrainer.

The goal is to extend GraphTrainer FP8 support from dense linear layers to MoE grouped experts running with expert parallelism (EP), while keeping FP8 limited to local expert computation and preserving the existing communication path.

Implementation

Grouped-expert FP8 is more difficult than dense FP8 because the number of tokens assigned to each expert is dynamic. Under EP, tokens are exchanged and then reordered into expert-major layout before grouped GEMM. FP8/MXFP8 grouped GEMM also requires each expert's token group to satisfy a kernel-specific alignment constraint. These dynamic padded shapes can produce compound symbolic expressions that FakeTensor and regional Inductor cannot codegen directly.

This PR addresses these issues in three parts:

• Convert supported DeepSeek-V3 grouped experts to TorchAO FP8/MXFP8 grouped GEMM modules.
• Pad each local expert token group to the required alignment while preserving the original token indices and unpadding the result after computation.
• Add a focused metadata fallback for the traced aten._scaled_grouped_mm path when a data-dependent symbolic layout guard prevents FakeTensor from constructing output metadata.

The GraphTrainer FP8 pass is also extended to distinguish dense scaled GEMM from grouped scaled GEMM. The grouped-expert pad SymInt path is traceable. With TORCHTITAN_FP8_EP_UNBACKED_PAD=1, the custom path creates an unbacked padded length and installs a shape-only grouped-GEMM metadata fallback, allowing the supported path to reach regional Inductor. Regions with unbound compound EP padding symbols that remain after this workaround are skipped with a warning rather than generating an invalid regional kernel. Full Inductor remains available as the fallback for those cases.

The workaround is opt-in through TORCHTITAN_FP8_EP_UNBACKED_PAD=1 and is installed only when the grouped-expert padding path needs it. The eager execution path continues to use TorchAO's regular permute_and_pad implementation.

Detailed changes

• Add FP8 and MXFP8 grouped-expert conversion for the DeepSeek-V3 debug model.
• Add GraphTrainer FP8 configuration and validation for grouped-expert graphs.
• Add expert token padding for EP and EP=1 debugging/numerics runs.
• Recognize aten._scaled_grouped_mm as the compute target for grouped FP8.
• Re-check local FP8 regions after GraphPP callable extraction.
• Skip regional regions with unbound symbolic input-size expressions.
• Add FakeTensor metadata fallback for the supported scaled grouped GEMM path.
• Add grouped-expert unit, integration, tracing, and numerical tests.

Code Scope

Most of the graph analysis and tracing changes are inside torchtitan/experiments/graph_trainer. A few shared changes are required to connect GraphTrainer with the existing TorchAO grouped-expert implementation:

• quantization converters expose grouped-expert modules and their alignment;
• the shared token dispatcher applies EP token padding;
• DeepSeek-V3 model registries provide the FP8 grouped-expert presets.

No FP8 communication or change to the general high-precision EP communication semantics is introduced by this PR.

Tests

The PR adds or updates coverage for:

• grouped FP8 provenance and compute-target validation;
• grouped-expert token padding and EP=1 behavior;
• GraphPP callable boundaries and local FP8 region analysis;
• FakeTensor tracing with symbolic padded shapes;
• DeepSeek-V3 FP8 grouped-expert integration runs;
• grouped-expert numerical comparisons.

The GPU integration cases use the existing H100 test suite and 10 training steps for smoke coverage.

Known Limits

• Grouped-expert FP8 regional Inductor is supported, including the unbacked pad SymInt path enabled by TORCHTITAN_FP8_EP_UNBACKED_PAD=1. Regions with unbound symbols that remain after this workaround are skipped and run eagerly; full Inductor remains the fallback.
• FP8 EP communication and FP8 FSDP all-gather are not implemented; EP communication remains in the existing higher-precision format.
• CUDA Graph capture remains disabled for GraphPP.
• FP8 precompile artifacts remain unsupported by the current CooR precompile path.
• The metadata fallback targets the current TorchAO/PyTorch aten._scaled_grouped_mm implementation and may need adjustment when that implementation changes.