[float8] add float8 rowwise MoE prototype (#1245)

# Summary
- Adds `--float8.moe_fqns_prototype="..."` option to float8 training API
- API accepts a comma-separated list of FQNs to apply MoE float8
training conversion to.
- `quanttize_` with the `MoETrainingConfig` will recursively swap
nn.Parameter data tensors to a tensor subclass, which has an override
for grouped_mm => [dynamic quant + scaled grouped
mm](https://github.com/pytorch/ao/blob/d963a8840e3c228e303fe14aff5d9be7017c92b6/torchao/prototype/scaled_grouped_mm/scaled_grouped_mm.py#L20)
prototype. Context: see implementation of GroupedExperts
[here](https://github.com/pytorch/torchtitan/blob/ca10545e41582fed4ebb00db4c13db71194a0dfa/torchtitan/experiments/llama4/model/moe.py#L85-L87).

# Testing
- Tested via manual testing with torchao
`convert_moe_to_float8_training` prototype
([PR](pytorch/ao#2275)) and confirmed single GPU
training works as expected.

# Limitations
- Only supports single GPU training so far.
- Only performs grouped_mm override for routed experts (see condition
[here](https://github.com/pytorch/ao/pull/2275/files#diff-c529b94621368096076db6bec8a6fc058d7f7595c39cd59965c657ed5dea861cR29-R33)).
For shared experts, I'll need to update the torchao prototype to support
3d A tensor (see torchtitan
[here](https://github.com/pytorch/torchtitan/blob/ca10545e41582fed4ebb00db4c13db71194a0dfa/torchtitan/experiments/llama4/model/moe.py#L316)).

Author: danielvegamyhre

torchtitan/components/quantization/float8.py

@@ -53,6 +53,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_prototype
 
         if float8_config.recipe_name is not None:
             assert (
@@ -114,6 +115,30 @@ 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.quantization.quant_api import quantize_
+
+            try:
+                from torchao.prototype.moe_training.conversion_utils import (
+                    MoETrainingConfig,
+                )
+            except ImportError as e:
+                raise ImportError(
+                    "torchao installation does not have MoE training support. Please install torchao nightly build."
+                ) from e
+
+            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
+
+            config = MoETrainingConfig()
+            quantize_(model, config=config, 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

@@ -513,6 +513,13 @@ class Float8:
     Not compatible with torch.compile.
     """
 
+    moe_fqns_prototype: list[str] | str = field(default_factory=list)
+    """
+    Comma-separated list of fully qualified names of MoE modules to apply float8 rowwise training to.
+    This is a prototype feature that requires the torchao nightly build.
+    Example: --float8.moe_fqns_prototype="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 = ["experts"]