allow other axiswise dims so we can pass in 3D B tensor tranposed

Author: danielvegamyhre

torchao/float8/float8_ops.py

@@ -151,6 +151,9 @@ def float8_transpose(aten_op, args, kwargs=None):
     else:
         new_scale = args[0]._scale
 
+    if aten_op == aten.transpose.int:
+        _assert_tensorwise_scale(aten_op, args[0]._scale)
+
     old_axiswise_dim = args[0]._axiswise_dim
     new_axiswise_dim = old_axiswise_dim
     if old_axiswise_dim is not None:

torchao/float8/float8_tensor.py

@@ -313,7 +313,7 @@ def __new__(
             linear_mm_config if linear_mm_config is not None else LinearMMConfig()
         )
         self._gemm_input_role = gemm_input_role
-        assert axiswise_dim in (None, 0, -1), f"unsupported axiswise_dim {axiswise_dim}"
+        assert axiswise_dim in (None, 0, 1, -1), f"unsupported axiswise_dim {axiswise_dim}"
         self._axiswise_dim = axiswise_dim
 
         return self

torchao/prototype/grouped_mm/__init__.py

@@ -47,7 +47,7 @@ class _Float8GroupedMM(torch.autograd.Function):
     def forward(
         ctx,
         A: torch.Tensor,
-        B: torch.Tensor,
+        B_t: torch.Tensor,
         float8_recipe_name: Float8LinearRecipeName,
         offs: Optional[torch.Tensor] = None,
         out_dtype: Optional[torch.dtype] = None,
@@ -60,7 +60,7 @@ def forward(
 
         # perform dynamic float8 quantization using the given recipe, if specified
         assert 2 <= A.ndim <= 3, "A must be 2D or 3D"
-        assert B.ndim == 3, "B must be 3D"
+        assert B_t.ndim == 3, "B must be 3D"
         if A.ndim == 2:
             assert offs is not None, "offs must be specified for 2D A tensor"
         else:
@@ -83,34 +83,33 @@ def forward(
         )
 
         # Convert high precision weight tensor to float8.
-        B_fp8 = hp_tensor_to_float8_dynamic(
-            B,
+        B_t_fp8 = hp_tensor_to_float8_dynamic(
+            B_t,
             float8_config.cast_config_input.target_dtype,
             linear_mm_config=LinearMMConfig(),
             gemm_input_role=GemmInputRole.WEIGHT,
             scaling_granularity=float8_config.cast_config_weight.scaling_granularity,
             axiswise_dim=get_maybe_axiswise_dim(
-                -1, float8_config.cast_config_input.scaling_granularity
+                1, float8_config.cast_config_input.scaling_granularity
             ),
             round_scales_to_power_of_2=float8_config.round_scales_to_power_of_2,
         )
-        B_fp8_t = B_fp8.transpose(-2, -1)
 
         # Store what we need for backward.
-        ctx.save_for_backward(A, B)
+        ctx.save_for_backward(A, B_t)
         ctx.float_config = float8_config
         ctx.offs = offs
 
         # For rowwise scaling, torch._scaled_grouped_mm requires scales without any empty dims.
         A_fp8._scale = A_fp8._scale.squeeze()
-        B_fp8_t._scale = B_fp8_t._scale.squeeze()
+        B_t_fp8._scale = B_t_fp8._scale.squeeze()
 
         # Perform scaled grouped GEMM and return result.
         return torch._scaled_grouped_mm(
             A_fp8._data,
-            B_fp8_t._data,
+            B_t_fp8._data,
             A_fp8._scale,
-            B_fp8_t._scale,
+            B_t_fp8._scale,
             offs,
             out_dtype=out_dtype,
             use_fast_accum=use_fast_accum,

torchao/prototype/grouped_mm/test_grouped_mm.py

@@ -30,7 +30,7 @@ def test_grouped_gemm_2d_3d(use_fast_accum, strided):
     offs = torch.arange(m, n_groups * m + 1, m, device="cuda", dtype=torch.int32)
     result = _grouped_scaled_mm(
         a,
-        b,
+        b.transpose(-2, -1),
         offs=offs,
         float8_recipe=float8_recipe_name,
         out_dtype=out_dtype,
@@ -41,7 +41,7 @@ def test_grouped_gemm_2d_3d(use_fast_accum, strided):
     validate_grouped_mm(
         result,
         a,
-        b,
+        b.transpose(-2, -1),
         n_groups,
         out_dtype,
         use_fast_accum,
@@ -67,7 +67,7 @@ def test_grouped_gemm_3d_3d(use_fast_accum, strided):
     ]
     result = _grouped_scaled_mm(
         a,
-        b,
+        b.transpose(-2, -1),
         float8_recipe=float8_recipe_name,
         out_dtype=out_dtype,
         use_fast_accum=use_fast_accum,
@@ -77,7 +77,7 @@ def test_grouped_gemm_3d_3d(use_fast_accum, strided):
     validate_grouped_mm(
         result,
         a,
-        b,
+        b.transpose(-2, -1),
         n_groups,
         out_dtype,
         use_fast_accum,
@@ -95,7 +95,7 @@ def test_tensorwise_scaling_not_supported():
     with pytest.raises(AssertionError):
         _grouped_scaled_mm(
             a,
-            b,
+            b.transpose(-2, -1),
             offs=offs,
             float8_recipe=Float8LinearRecipeName.TENSORWISE,
             out_dtype=torch.bfloat16,
@@ -131,8 +131,8 @@ def validate_grouped_mm(
         round_scales_to_power_of_2=float8_config.round_scales_to_power_of_2,
     )
 
-    B_fp8 = hp_tensor_to_float8_dynamic(
-        B,
+    B_t_fp8 = hp_tensor_to_float8_dynamic(
+        B.transpose(-2, -1),
         float8_config.cast_config_input.target_dtype,
         linear_mm_config=LinearMMConfig(),
         gemm_input_role=GemmInputRole.WEIGHT,
@@ -142,11 +142,10 @@ def validate_grouped_mm(
         ),
         round_scales_to_power_of_2=float8_config.round_scales_to_power_of_2,
     )
-    B_fp8_t = B_fp8.transpose(-2, -1)
 
     # grouped_scaled_mm doesn't support empty dims
     scale_A = A_fp8._scale.squeeze()
-    scale_B = B_fp8_t._scale.squeeze()
+    scale_B = B_t_fp8._scale.squeeze()
 
     A_list, B_list, A_scale_list, B_scale_list, result_list = [], [], [], [], []
     start = 0
@@ -160,7 +159,7 @@ def validate_grouped_mm(
             start = offs_cpu[i]
     else:
         A_list = A_fp8._data
-        B_list = B_fp8_t._data
+        B_list = B_t_fp8._data
 
     A_scale_list = scale_A
     B_scale_list = scale_B