[mxfp8 moe training] Add mxfp8 to FSDP tests

[liangel-02]

addressing #2833

updating test to include mxfp8: torchrun --nproc_per_node=${NUM_GPUS} -m pytest test_fsdp.py , all tests pass

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 7b0add0 with merge base df7bf37 ():
💚 Looks good so far! There are no failures yet. 💚

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

[danielvegamyhre]

@liangel-02 we need to use this torchtitan API:

set_token_group_alignment_size_m(16) # fp8

or

set_token_group_alignment_size_m(32) # mxfp8

This is because TMA (some background here) requires the slowest moving dim (stride of 1) to be 128 bit (16 byte) aligned.

In the backward pass, when grad_weight = grad_output_t @ input, the “M” dimension (flattened token groups) become this “stride 1” dim. Therefore, each token group must be 16 byte aligned for this grouped gemm.

• For bf16, this means groups divisible by 8 elements (16 bytes / 2 bytes per elem).
• For fp8 groups divisible by 16 elements (16 bytes / 1 byte per elem).
• For mxfp8, I enforce they must be divisible by block_size (32) not because of TMA, but to ensure when we quantize along that dim, we have (1) no scaling groups crossing logically distinct tensor boundaries, and (2) no “partial” scaling group at the the end (e..g, 1x17 chunk because the tensor ends).

This might all be a bit confusing without background knowledge on GPU architecture and MoE models, but don't worry! Plenty of time to learn.

For now, just implement this setting in some clean/interpretable way in the test code and it should (hopefully) work.

[danielvegamyhre]

lgtm! 1 minor comment before landing, thanks

[danielvegamyhre]

remove commented code before landing