float8 moe training conversion API prototype

[danielvegamyhre]

Stacked PRs:

• ->float8 moe training conversion API prototype #2275

---

float8 moe training conversion API prototype

• convert_moe_to_float8_training 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.

Testing

• Tested with torchtitan (see PR with MoE conversion API) and confirmed single GPU training works as expected.

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.

✅ No Failures

As of commit 2b4361d with merge base d963a88 ():
💚 Looks good so far! There are no failures yet. 💚

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