remove excessive logging

Author: danielvegamyhre

torchao/prototype/moe_training/scaled_grouped_mm.py

@@ -60,7 +60,7 @@ def forward(
         offs: Optional[torch.Tensor] = None,
         out_dtype: Optional[torch.dtype] = torch.bfloat16,
     ) -> torch.Tensor:
-        # torchao _scaled_grouped_mm only supports A=2D|3D + B=3D.
+        # torchao _scaled_grouped_mm only supports A=2D|3D and B=3D.
         assert A.ndim == 2 or A.ndim == 3, "A must be 2D or 3D"
         assert B_t.ndim == 3, "B must be 3D"
 
@@ -150,17 +150,6 @@ def forward(
         assert _is_column_major(B_t_fp8_col_major), (
             "B must be column-major for output = A @ B"
         )
-
-        # TODO: remove excessive logging once prototype is more mature.
-        logger.debug(
-            (
-                f"forward scaled_grouped_mm: A_fp8_row_major.shape={A_fp8_row_major.shape}, "
-                f"A_scale.shape={A_scales.squeeze(-1).shape}, "
-                f"B_t_fp8_col_major.shape={B_t_fp8_col_major.shape}, "
-                f"B_t_scale.shape={B_t_scales.squeeze(1).shape}, "
-                f"offs={offs if offs is not None else None}"
-            )
-        )
         return torch._scaled_grouped_mm(
             A_fp8_row_major,
             B_t_fp8_col_major,
@@ -205,14 +194,6 @@ def backward(ctx, grad_output: torch.Tensor):
         assert _is_column_major(B_fp8_col_major), (
             "B must be column-major for grad_A = grad_output @ B"
         )
-        logger.debug(
-            (
-                f"backward grad_A: grad_output_fp8_row_major.shape={grad_output_fp8_row_major.shape}, "
-                f"grad_output_scale.shape={grad_output_scales.shape}, "
-                f"B_fp8_col_major.shape={B_fp8_col_major.shape}, "
-                f"B_scale.shape={B_scales.shape}, "
-            )
-        )
         grad_A = torch._scaled_grouped_mm(
             grad_output_fp8_row_major,
             B_fp8_col_major,
@@ -258,15 +239,6 @@ def backward(ctx, grad_output: torch.Tensor):
         assert _is_column_major(A_fp8_col_major), (
             "A must be column-major for grad_B = grad_output_t @ A"
         )
-
-        logger.debug(
-            (
-                f"backward grad_B: grad_output_t_fp8_row_major.shape={grad_output_t_fp8_row_major.shape}, "
-                f"grad_output_t_scale.shape={grad_output_t_scales.shape}, "
-                f"A_fp8_col_major.shape={A_fp8_col_major.shape}, "
-                f"A_scale.shape={A_scales.shape}, "
-            )
-        )
         grad_B = torch._scaled_grouped_mm(
             grad_output_t_fp8_row_major,
             A_fp8_col_major,

torchao/prototype/moe_training/tensor.py

@@ -13,8 +13,6 @@
 
 from torchao.prototype.moe_training import _scaled_grouped_mm
 
-logger: logging.Logger = logging.getLogger(__name__)
-
 _ops_to_preserve_subclass = {
     torch.ops.aten.empty_like.default,
     torch.ops.aten.new_zeros.default,
@@ -77,9 +75,6 @@ def __torch_function__(cls, func, types, args, kwargs={}):
             A, B = args[0], args[1]
             A_is_2d_or_3d = A.dim() in (2, 3)
             B_is_3d = B.dim() == 3
-            has_offs = kwargs.get(cls.offs_arg_name) is not None
-            logger.debug(f"A.shape={A.shape}, B.shape={B.shape}, has_offs={has_offs}")
-
             if A_is_2d_or_3d and B_is_3d:
                 return _scaled_grouped_mm(
                     *args,