use quantize_ for moe

Author: danielvegamyhre

torchtitan/components/quantization/float8.py

@@ -118,17 +118,19 @@ def convert(self, model: nn.Module):
         # 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_
             from torchao.prototype.scaled_grouped_mm.conversion_utils import (
-                convert_moe_to_float8_training,
+                MoETrainingConfig,
             )
 
             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)
+            
+            config = MoETrainingConfig(module_filter_fn=moe_module_filter_fn)
+            quantize_(model, config=config)
             logger.info("Converted MoE to float8")
 
     def post_optimizer_hook(self, model: nn.Module | list[nn.Module]):

torchtitan/train.py

@@ -228,6 +228,7 @@ def __init__(self, job_config: JobConfig):
             model.to_empty(device=init_device)
             with torch.no_grad():
                 model.init_weights(buffer_device=buffer_device)
+            model = model.to(torch.bfloat16)
             model.train()
 
             self.model_parts = [model]