add float8 moe prototype

Author: danielvegamyhre

torchtitan/components/quantization/float8.py

@@ -49,6 +49,7 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
 
         self.enabled = True
         self.filter_fqns = float8_config.filter_fqns
+        self.moe_fqns = float8_config.moe_fqns
 
         if float8_config.recipe_name is not None:
             assert (
@@ -109,6 +110,22 @@ def convert(self, model: nn.Module):
             f"{self.config.enable_fsdp_float8_all_gather}"
         )
 
+        # Mutates the model inplace replacing instances of nn.Parameter with ScaledGroupedMMTensor,
+        # to perform dynamic float8 rowwise quantization + scaled grouped GEMMs for the target MoE FQNs.
+        if self.moe_fqns:
+            from torchao.prototype.scaled_grouped_mm.conversion_utils import (
+                convert_moe_to_float8_training,
+            )
+
+            def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
+                for target_fqn in self.moe_fqns:
+                    if target_fqn in cur_fqn:
+                        return True
+                return False
+
+            convert_moe_to_float8_training(model, module_filter_fn=moe_module_filter_fn)
+            logger.info("Converted MoE to float8")
+
     def post_optimizer_hook(self, model: nn.Module | list[nn.Module]):
         if not self.enabled:
             return

torchtitan/config_manager.py

@@ -450,6 +450,12 @@ class Float8:
     Example: --float8.filter_fqns "attention.wq,attention.wk,attention.wv,output"
     """
 
+    moe_fqns: list[str] | str = field(default_factory=list)
+    """
+    Comma-separated list of fully qualified names of MoE modules to apply float8 rowwise training to.
+    Example: --float8.moe_fqns="experts"
+    """
+
 
 @dataclass
 class MX:

torchtitan/experiments/llama4/train_configs/debug_model.toml

@@ -69,3 +69,4 @@ selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac bas
 enable_fsdp_float8_all_gather = false
 precompute_float8_dynamic_scale_for_fsdp = false
 filter_fqns = ["output", "router.gate"]
+moe_fqns = []