float8 moe training conversion API prototype

[danielvegamyhre]

Stacked PRs:

• add float8 moe training readme and example #2348
• ->float8 moe training conversion API prototype #2275

---

float8 moe training conversion API prototype

• quantize_ will recursively swap nn.Parameter data tensors to a tensor subclass, which has an override for grouped_mm => dynamic quant + scaled grouped mm prototype. Context: see implementation of GroupedExperts here.
• Prototype name changed from "scaled_grouped_mm" to "moe_training" since the scaled grouped mm is now just a building block in the larger MoE training prototype.

Testing

• Added unit tests using the torchtitan Llama4 MoE layer to perform a forward + backward pass, then validating reasonable SQNR for model outputs, input gradients and parameter gradients when comparing the quantized vs baseline bf16 model.
• Tested E2E training with torchtitan (see PR with MoE conversion API) and confirmed single GPU training works as expected with Llama4.

Limitations

• Only supports single GPU training. I tried with FSDP=2 and hit this issue which seems to be related to a known issue that is being addressed.
• Only performs grouped_mm override for routed experts (see condition here). For shared experts, I'll need to update the torchao prototype to support 3d A tensor.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2275

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit c80dfa6 with merge base 83663b8 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_quant_semi_sparse_compile_False

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[danielvegamyhre]

@drisspg @vkuzo for review when you have a chance

[drisspg]

@bdhirsh I suggested this Torch function approach to stay above autograd in order call into our autograd func, AFAIK compile's fucntion subclass support has gotten much better but just wanted to double check w/ you

[vkuzo]

+1. How's compile looking on this now?

[danielvegamyhre]

• llama4 (what I tested with in torchtitan) cannot compile yet due to backward hooks implementing auxiliary loss not playing nicely with compile. Tianyu is aware of this and working on a new approach I believe.
• In this minimal GroupedExperts example, it cannot compile with fullgraph=True due to _grouped_mm not being traceable yet: torch._dynamo.exc.Unsupported: .... Explanation: Dynamo does not know how to trace the builtin `torch._VariableFunctionsClass._grouped_mm. (I believe Brian is aware of this and working on it).
• The minimal example can compile with graph breaks.

this was tested with the latest nightly pytorch.

[bdhirsh]

(1) yep, torch_function subclass compile support has improved in the last few months (thanks to Ryan / Lazos)

(2) on the grouped_mm issue, pytorch/pytorch#153384 should help, im going to land early next week but its a small change so feel free to rebase on it if you want to test sooner

[vkuzo]

you can also use the quantize_ API to match the rest of torchao. We want to eventually migrate the float8 training API to that.

[danielvegamyhre]

Added todo for this. Hoping to get the functionality working w/ key parallelisms first and then will look into migrating to quantize_.

[danielvegamyhre]

Updated to use quantize_ and added tests. The prototype works for both forward, however, the reference bf16 moe hits an issue during backward, so I can't use it to validate the fp8 results until it is resolved.

[drisspg]

Is the fp8 version training ?

[danielvegamyhre]

Is the fp8 version training ?

In the e2e llama4 training with this PR, yes.

However, in the unit test in this PR test_moe_training_conversion.py where I import and use the torchtitan MoE layer for testing, I hit this issue during backward in the bf16 model (it also happens in the fp8 version, but I mention bf16 since it rules out fp8 conversion specific issues).

I hit the same issue when I just use GroupedExperts or the whole MoE layer.

I might try just importing the transformer itself and seeing if that works.

[danielvegamyhre]

I managed to get it working by using a MSE loss instead of the out.sum().backward(), but the root cause is still a mystery

[danielvegamyhre]

@vkuzo @drisspg this is ready for another look, I (1) migrated to the quantize_ API, (2) added unit tests using the torchtitan Llama4 MoE implementation, and (3) validated E2E training with torchtitan again as well.

[drisspg]

Some small comments, looks good, re-req me when ready

[danielvegamyhre]

Some small comments, looks good, re-req me when ready

@drisspg Thanks, I addressed the comments

[danielvegamyhre]

Confirmed failing tests are unrelated to this change (they are autoquant, sparsity, etc), see below:

FAILED test/integration/test_integration.py::TestSubclass::test_int4_weight_only_quant_subclass_grouped_5_cuda - torch._inductor.exc.InductorError: AssertionError: -201885665299331/250000000000000

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
FAILED test/integration/test_integration.py::TestAutoQuant::test_autoquant_compile_09_cuda - RuntimeError: self.size(0) needs to be greater than 16, but got 1
FAILED test/integration/test_integration.py::TestAutoQuant::test_autoquant_compile_15_cuda - RuntimeError: self.size(0) needs to be greater than 16, but got 1
FAILED test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_quant_semi_sparse_compile_False - AssertionError: Tensor-likes are not close!