[mxfp8 moe training] Add mxfp8 to FSDP tests (#2849)

Author: liangel-02

test/prototype/moe_training/test_fsdp.py

@@ -27,32 +27,44 @@
         "CUDA not available or compute capability < 8.9", allow_module_level=True
     )
 
+from testing_utils import _validate_model_conversion
+
 from torchao.float8.float8_utils import compute_error
-from torchao.prototype.moe_training.conversion_utils import MoETrainingConfig
+from torchao.prototype.moe_training.conversion_utils import (
+    MoEScalingType,
+    MoETrainingConfig,
+)
 from torchao.quantization.quant_api import quantize_
 
-from .testing_utils import _validate_model_conversion
-
 # this test requires torchtitan
 try:
-    from torchtitan.distributed.expert_parallel import (
-        set_token_group_alignment_size_m,
-    )
+    from torchtitan.distributed.expert_parallel import set_token_group_alignment_size_m
     from torchtitan.models.moe import MoE, MoEArgs
 except ImportError:
     pytest.skip(
         "torchtitan not installed, skipping MoE tests.", allow_module_level=True
     )
 
 
-def test_moe_float8_training_fsdp():
+@pytest.mark.parametrize(
+    "recipe, min_out_sqnr, alignment_size, min_param_grad_sqnr",
+    [
+        (MoEScalingType.FP8_ROWWISE, 29.0, 16, 23.0),
+        (MoEScalingType.MXFP8, 28.0, 32, 21.0),
+    ],
+)
+def test_moe_float8_training_fsdp(
+    recipe: MoEScalingType,
+    min_out_sqnr: float,
+    alignment_size: int,
+    min_param_grad_sqnr: float,
+):
     assert torch.cuda.is_available()
 
     # setup distributed for fsdp
     setup_distributed()
 
-    # token group aligment size must be 16 for fp8
-    set_token_group_alignment_size_m(16)
+    set_token_group_alignment_size_m(alignment_size)
 
     # define model args
     target_fqns = ["experts"]
@@ -83,7 +95,7 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
         return False
 
     # quantize test model
-    config = MoETrainingConfig()
+    config = MoETrainingConfig(recipe)
     quantize_(model, config=config, filter_fn=moe_module_filter_fn)
 
     # validate that only the experts were converted
@@ -109,7 +121,6 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
 
     # validate output
     out_sqnr = compute_error(out, ref_out)
-    min_out_sqnr = 29.0
     assert out_sqnr.item() >= min_out_sqnr, (
         f"SQNR must be >= {min_out_sqnr}, got {out_sqnr.item()}."
     )
@@ -131,7 +142,6 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
     )
 
     # validate param gradients
-    min_param_grad_sqnr = 23.0
     for param1, param2 in zip(model.parameters(), ref_model.parameters()):
         param_grad_sqnr = compute_error(param1.grad, param2.grad)
         assert param_grad_sqnr.item() >= min_param_grad_sqnr, (