handle out != None

Author: danielvegamyhre

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -40,16 +40,8 @@ def _scaled_grouped_mm(
         offs (int32 torch.Tensor): The offsets to use to mark the starting index of each group along dim0 of the A tensor.
         out_dtype (Optional[torch.dtype]): The dtype of the output tensor. Currently only torch.bfloat16 is supported.
     """
-<<<<<<< HEAD
-<<<<<<< HEAD
-    # logger.info("Using scaled_grouped_mm")
-=======
-    #logger.info("Using scaled_grouped_mm")
->>>>>>> 6ca070de (handle out != None)
-=======
     # TODO: Remove once prototype is more mature. This is currently very useful for development and debugging.
     logger.info("Using scaled_grouped_mm")
->>>>>>> 2f3bb137 (add tp support for fp8 moe training)
     return _Float8GroupedMM.apply(
         A,
         B_t,

torchao/prototype/moe_training/tensor.py

@@ -183,6 +183,7 @@ def fsdp_post_all_gather(
                     f"{out_data.dtype} {param_dtype}"
                 )
                 out_data.copy_(data)
+
             return
 
         # For training step 0, out=None, so we need to return a new ScaledGroupedMMTensor.