make quant receipt flexible

Author: rakkit

torchao/prototype/moe_training/examples/simple_moe_layer.py

@@ -27,10 +27,23 @@
     ) from e
 
 
+from argparse import ArgumentParser
+
+parser = ArgumentParser()
+parser.add_argument(
+    "--scaling_type",
+    type=str,
+    default="fp8_rowwise",
+    choices=["fp8_rowwise", "mxfp8"],
+)
+args = parser.parse_args()
+
+
 # initialize model
 device = torch.device("cuda")
+torch.manual_seed(42)
 model_args = MoEArgs(num_experts=8, top_k=2, use_grouped_mm=True)
-dim = 256
+dim = 1024
 hidden_dim = dim * 4
 model = MoE(model_args, dim, hidden_dim).to(torch.bfloat16).to(device)
 init_std = 0.02
@@ -47,11 +60,15 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
     return False
 
 
-# quantize the model, by default it is rowwise fp8
-config = MoETrainingConfig()
-quantize_(model, config=config, filter_fn=moe_module_filter_fn)
+if args.scaling_type == "fp8_rowwise":
+    config = MoETrainingConfig()
+    alignment_size = 16
+
+elif args.scaling_type == "mxfp8":
+    config = MoETrainingConfig(scaling_type=MoEScalingType.MXFP8)
+    alignment_size = 32
 
-alignment_size = 32 if config.scaling_type == MoEScalingType.MXFP8 else 16
+quantize_(model, config=config, filter_fn=moe_module_filter_fn)
 set_token_group_alignment_size_m(alignment_size)
 
 # training loop