[wip] sharding strategy for dim1 kernel

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: ddb1f80
ghstack-comment-id: 3001601112
Pull Request resolved: #2436

Author: vkuzo

test/float8/test_dtensor.py

@@ -183,7 +183,7 @@ def _test_dtensor_fp8_autograd(mesh: DeviceMesh, size=16):
     loss.backward()
 
 
-def _test_fp8_mlp_tensor_parallelism_eager(mesh: DeviceMesh, size=16):
+def _test_fp8_mlp_tensor_parallelism_eager(mesh: DeviceMesh, size=32):
     tensorwise_config = Float8LinearConfig(emulate=True)
     _test_lowp_mlp_tensor_parallelism_base(
         mesh, tensorwise_config, size, compile=False, allgather_in_lowp=True
@@ -198,7 +198,7 @@ def _test_fp8_mlp_tensor_parallelism_eager(mesh: DeviceMesh, size=16):
     )
 
 
-def _test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=16):
+def _test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=32):
     tensorwise_config = Float8LinearConfig(emulate=True)
     _test_lowp_mlp_tensor_parallelism_base(
         mesh, tensorwise_config, size, compile=True, allgather_in_lowp=True

test/float8/test_fsdp2_tp.py

@@ -61,7 +61,7 @@ def _test_fp8_mlp_tensor_parallelism_base(
         enable_fsdp_float8_all_gather=True,
     )
 
-    toy_model = ToyModel().to(device)
+    toy_model = ToyModel(size).to(device)
 
     tp_model = copy.deepcopy(toy_model)
     tp_model = convert_to_float8_training(tp_model, config=config)
@@ -94,11 +94,11 @@ def _test_fp8_mlp_tensor_parallelism_base(
     # TODO(future PR): test numerics, and add more cases
 
 
-def _test_fp8_mlp_tensor_parallelism_eager(mesh: DeviceMesh, size=16):
+def _test_fp8_mlp_tensor_parallelism_eager(mesh: DeviceMesh, size=32):
     _test_fp8_mlp_tensor_parallelism_base(mesh, size, compile=False)
 
 
-def _test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=16):
+def _test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=32):
     _test_fp8_mlp_tensor_parallelism_base(mesh, size, compile=True)
 
 

test/prototype/mx_formats/test_mx_dtensor.py

@@ -68,21 +68,22 @@ def _test_dtensor_cast_to_mxfp8(mesh: DeviceMesh, size=4):
     )
 
 
-def _test_mxfp8_mlp_tensor_parallelism(mesh: DeviceMesh, size=16):
+def _test_mxfp8_mlp_tensor_parallelism(mesh: DeviceMesh, size=128):
     config = MXLinearConfig.from_recipe_name("mxfp8_emulated")
-    config.block_size = 16
+    config.block_size = 32
+    config.use_fp8_dim1_cast_triton_kernel = True
     _test_lowp_mlp_tensor_parallelism_base(
         mesh, config, size, compile=False, allgather_in_lowp=False
     )
-    _test_lowp_mlp_tensor_parallelism_base(
-        mesh, config, size, compile=True, allgather_in_lowp=False
-    )
+    # _test_lowp_mlp_tensor_parallelism_base(
+    #     mesh, config, size, compile=True, allgather_in_lowp=False
+    # )
 
 
 if __name__ == "__main__":
     device_mesh = setup_distributed()
     tests = [
-        _test_dtensor_cast_to_mxfp8,
+        # _test_dtensor_cast_to_mxfp8,
         _test_mxfp8_mlp_tensor_parallelism,
     ]
 

torchao/prototype/mx_formats/kernels.py

@@ -1315,7 +1315,8 @@ def triton_to_mxfp8_dim1(
         * `col_scale`: the `e8m0` values of `x_scale` used to cast `x` to mxfp8 across dim1
         """
         assert x.is_contiguous(), "`x` must be contiguous"
-        assert x.dtype == torch.bfloat16
+        # TODO(before land): maybe gate by FakeTensor below?
+        # assert x.dtype == torch.bfloat16
         assert inner_block_size <= 32
 
         # Get tensor shape
@@ -1362,6 +1363,10 @@ def triton_to_mxfp8_dim1(
             output_col_major.t(),
             col_scale.view(torch.float8_e8m0fnu),
         )
+        
+    print('ASDFASDFASDF')
+    from torchao import triton_to_mxfp8_dim1
+    print(triton_to_mxfp8_dim1)
 
     def triton_to_mxfp8_dim1_reference(
         x_hp: torch.Tensor, block_size

torchao/testing/training/dtensor_utils.py

@@ -32,11 +32,11 @@
 class FeedForward(nn.Module):
     """MLP based model"""
 
-    def __init__(self):
+    def __init__(self, size):
         super(FeedForward, self).__init__()
-        self.w1 = nn.Linear(16, 32, bias=False)
-        self.w2 = nn.Linear(16, 32, bias=False)
-        self.out_proj = nn.Linear(32, 16, bias=False)
+        self.w1 = nn.Linear(size, size * 2, bias=False)
+        self.w2 = nn.Linear(size, size * 2, bias=False)
+        self.out_proj = nn.Linear(size * 2, size, bias=False)
 
     def forward(self, x):
         x = F.silu(self.w1(x)) * self.w2(x)
@@ -45,9 +45,9 @@ def forward(self, x):
 
 
 class ToyModel(nn.Module):
-    def __init__(self):
+    def __init__(self, size):
         super(ToyModel, self).__init__()
-        self.ffn = FeedForward()
+        self.ffn = FeedForward(size)
 
     def forward(self, x):
         return self.ffn(x)
@@ -56,7 +56,7 @@ def forward(self, x):
 def _test_lowp_mlp_tensor_parallelism_base(
     mesh: DeviceMesh,
     config: Union[Float8LinearConfig, MXLinearConfig],
-    size=16,
+    size=32,
     compile: bool = False,
     allgather_in_lowp: bool = False,
 ):
@@ -67,7 +67,7 @@ def _test_lowp_mlp_tensor_parallelism_base(
     if isinstance(config, MXLinearConfig):
         convert_model_func = quantize_
 
-    toy_model = ToyModel().to(device)
+    toy_model = ToyModel(size).to(device)
     toy_model_fp8 = copy.deepcopy(toy_model)
     convert_model_func(toy_model_fp8, config=config)