[1/2] [GraphTrainer] Support dense FP8 compilation

[WhoisZihan]

This is a refactor and enhancement of #3989

RFC available in #3990

• This is the first of two stacked PRs, with the second PR available in my forked repo [2/2] Add grouped-expert FP8 GraphTrainer support WhoisZihan/torchtitan#2 (will create in torchtitan repo once current PR is merged.)

The goal is to adds the foundation for compiling and running quantized TorchAO FP8 models with GraphTrainer, which currently does not support regional inductor for fp8 nodes.

Implementation

The main problem is that TorchAO FP8 quantization is implemented through a dynamic sequence of casts, scaling operations, and scaled matrix multiplications. Detecting FP8 regions solely from individual operators is fragile because the traced decomposition can vary across PyTorch and TorchAO versions. This change preserves quantization provenance from the converted modules and uses it during GraphTrainer validation and regional annotation.

Another challenge is PP which partitions graph and might cut a complete FP8 region into two separated parts. To handle such situation, we should

• perform strict check of fp8 validness before graph partition, see graph_pp_pre_partition_fp8_passes, which means we ensure the fp8 region contains complete chains of nodes, instead of silent fallback to bf16.
• after partition, we perform non-strict scan of fp8 nodes and include as many FP8-compute nodes as possible into the regional inductor candidates. Non-compute nodes, such as placeholders/cast, are kept eager after partition.

Detailed changes are

• Adds FP8 quantization metadata in relevant nodes.
• Annotates complete dense FP8 compute regions for regional Inductor compilation.
• Adds Llama3 debug FP8 GraphTrainer configurations for full and regional compilation.
• Adds unit, integration, and numerics coverage for the new paths.

Code Scope

Almost all of the changes are GraphTrainer only, with only one exception that we need to keep _quantization_emulate field in Float8Linear module so that fp8 passes can identify emulated nodes.

Known Limits

This PR only supports the fp8 linear nodes, then MoE grouped experts part can be found in [2/2] Add grouped-expert FP8 GraphTrainer support

[WhoisZihan]

@tianyu-l Hi, I've refactored and enhanced the old PR #3989, hopefully the new one is ready for review. Thanks.