use filter_fn in quantize_

Author: danielvegamyhre

torchtitan/components/quantization/float8.py

@@ -129,8 +129,8 @@ def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
                         return True
                 return False
 
-            config = MoETrainingConfig(module_filter_fn=moe_module_filter_fn)
-            quantize_(model, config=config)
+            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]):

torchtitan/experiments/llama4/model/moe.py

@@ -82,6 +82,7 @@ def forward(
         assert (
             x.dtype == self.w1.dtype == self.w2.dtype == self.w3.dtype == torch.bfloat16
         ), "torch._grouped_mm only supports bf16 dtypes"
+
         h = F.silu(torch._grouped_mm(x, self.w1, offs=offsets))
         h = h * torch._grouped_mm(x, self.w3, offs=offsets)
         out = torch._grouped_mm(h, self.w2, offs=offsets)