[CCL] Add negative dim support (#15305)

### Ticket #15288 ### Problem description Negative dim input throws error ### What's changed Adds negative dim support across ccl ops ### Checklist - [ ] Post commit CI passes - [ ] Blackhole Post commit (if applicable) - [ ] Model regression CI testing passes (if applicable) - [ ] Device performance regression CI testing passes (if applicable) - [ ] New/Existing tests provide coverage for changes
tenstorrent · Nov 25, 2024 · aeca52f · aeca52f
1 parent f73dc7f
commit aeca52f
Show file tree

Hide file tree

Showing 24 changed files with 125 additions and 100 deletions.
diff --git a/models/demos/llama3/tt/llama_attention.py b/models/demos/llama3/tt/llama_attention.py
@@ -357,7 +357,7 @@ def forward_decode(
         if self.is_multichip and not self.use_fused_all_gather_matmul:
             dense_out_reduced = ttnn.reduce_scatter(
                 dense_out,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.L1_MEMORY_CONFIG,
@@ -530,7 +530,7 @@ def forward_prefill(self, x_11SH, rot_mats, transformation_mats, user_id: int =
         if self.is_multichip and not self.use_fused_all_gather_matmul:
             dense_out_reduced = ttnn.reduce_scatter(
                 output_11SH,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,

diff --git a/models/demos/llama3/tt/llama_mlp.py b/models/demos/llama3/tt/llama_mlp.py
@@ -137,7 +137,7 @@ def forward(self, x: ttnn.Tensor, mode) -> ttnn.Tensor:
         if self.args.is_multichip:
             w2_out_reduced = ttnn.reduce_scatter(
                 w2_out,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG if mode == "prefill" else ttnn.L1_MEMORY_CONFIG,

diff --git a/models/demos/llama3/tt/multimodal/llama_cross_attention.py b/models/demos/llama3/tt/multimodal/llama_cross_attention.py
@@ -271,7 +271,7 @@ def forward_decode(self, x_11SH, xattn_mask, full_text_row_masked_out_mask_1NSH,
         if self.is_multichip:
             dense_out_reduced = ttnn.reduce_scatter(
                 output,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.L1_MEMORY_CONFIG,
@@ -358,7 +358,7 @@ def forward_prefill(
         if self.is_multichip:  # TODO use_fused_all_gather_matmul
             dense_out_reduced = ttnn.reduce_scatter(
                 output,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,

diff --git a/models/demos/qwen/tt/qwen_attention.py b/models/demos/qwen/tt/qwen_attention.py
@@ -414,7 +414,7 @@ def forward_decode(
         if self.is_multichip and not self.use_fused_all_gather_matmul:
             dense_out_reduced = ttnn.reduce_scatter(
                 dense_out,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.L1_MEMORY_CONFIG,
@@ -598,7 +598,7 @@ def forward_prefill(self, x_11SH, rot_mats, transformation_mats, user_id: int =
         if self.is_multichip and not self.use_fused_all_gather_matmul:
             dense_out_reduced = ttnn.reduce_scatter(
                 output_11SH,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG,

diff --git a/models/demos/qwen/tt/qwen_mlp.py b/models/demos/qwen/tt/qwen_mlp.py
@@ -142,7 +142,7 @@ def forward(self, x: ttnn.Tensor, mode) -> ttnn.Tensor:
         if self.args.is_multichip:
             w2_out_reduced = ttnn.reduce_scatter(
                 w2_out,
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,
                 memory_config=ttnn.DRAM_MEMORY_CONFIG if mode == "prefill" else ttnn.L1_MEMORY_CONFIG,

diff --git a/models/demos/t3000/falcon40b/tt/falcon_mlp.py b/models/demos/t3000/falcon40b/tt/falcon_mlp.py
@@ -124,7 +124,7 @@ def fwd_decode(self, x: List[ttnn.Tensor]) -> List[ttnn.Tensor]:
         hidden_states = ttnn.get_device_tensors(
             ttnn.reduce_scatter(
                 ttnn.aggregate_as_tensor(hidden_states),
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,  # only unidirectional supported for now
                 memory_config=self.model_config["DEFAULT_MEMCFG"],
@@ -200,7 +200,7 @@ def fwd_prefill(self, x: List[ttnn.Tensor]) -> List[ttnn.Tensor]:
         hidden_states = ttnn.get_device_tensors(
             ttnn.reduce_scatter(
                 ttnn.aggregate_as_tensor(hidden_states),
-                scatter_dim=3,
+                dim=3,
                 math_op=ttnn.ReduceType.Sum,
                 num_links=1,  # only one link supported for now
                 memory_config=self.model_config["DEFAULT_MEMCFG"],

diff --git a/models/demos/t3000/llama2_70b/tt/llama_mlp_optimized.py b/models/demos/t3000/llama2_70b/tt/llama_mlp_optimized.py
@@ -219,7 +219,7 @@ def prefill_forward(self, x: List[ttnn.Tensor]) -> List[ttnn.Tensor]:
 
         hidden_states_reduced = ttnn.reduce_scatter(
             hidden_states_mm,
-            scatter_dim=3,
+            dim=3,
             math_op=ttnn.ReduceType.Sum,
             num_links=1,
             memory_config=ttnn.DRAM_MEMORY_CONFIG,
@@ -268,7 +268,7 @@ def decode_forward(self, x: List[ttnn.Tensor]) -> List[ttnn.Tensor]:
 
         hidden_states_reduced = ttnn.reduce_scatter(
             hidden_states,
-            scatter_dim=3,
+            dim=3,
             math_op=ttnn.ReduceType.Sum,
             num_links=1,
             memory_config=self.model_config["RESIDUAL_16_CORES_OUTPUT_MEMCFG"],

diff --git a/models/demos/tg/llama3_70b/tt/llama_common.py b/models/demos/tg/llama3_70b/tt/llama_common.py
@@ -93,7 +93,7 @@ def tt_composite_sharded_all_reduce(
     input_mem_cfg = input_tensor.memory_config()
     reduce_scattered_tensor = ttnn.reduce_scatter(
         input_tensor,
-        scatter_dim=dim,
+        dim=dim,
         math_op=ttnn.ReduceType.Sum,
         num_links=num_links,
         cluster_axis=cluster_axis,

diff --git a/tests/ttnn/unit_tests/operations/ccl/perf/test_ccl_perf.py b/tests/ttnn/unit_tests/operations/ccl/perf/test_ccl_perf.py
@@ -141,7 +141,7 @@ def test_all_gather_on_t3000(
     ],
 )
 @pytest.mark.parametrize(
-    "per_chip_output_shape, scatter_dim, layout",
+    "per_chip_output_shape, dim, layout",
     [
         ([1, 8, 1024, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 4, 1024, 1024], 3, ttnn.TILE_LAYOUT),
@@ -171,7 +171,7 @@ def test_reduce_scatter_on_t3000(
     t3k_mesh_device,
     num_devices,
     per_chip_output_shape,
-    scatter_dim,
+    dim,
     num_links,
     math_op,
     input_dtype,
@@ -187,7 +187,7 @@ def test_reduce_scatter_on_t3000(
         t3k_mesh_device,
         num_devices,
         per_chip_output_shape,
-        scatter_dim,
+        dim,
         num_links,
         math_op,
         input_dtype,
@@ -210,7 +210,7 @@ def test_reduce_scatter_on_t3000(
     ],
 )
 @pytest.mark.parametrize(
-    "per_chip_output_shape, scatter_dim, layout",
+    "per_chip_output_shape, dim, layout",
     [
         ([1, 1, 32, 4096], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 2048], 3, ttnn.TILE_LAYOUT),
@@ -239,7 +239,7 @@ def test_reduce_scatter_on_n300(
     n300_mesh_device,
     num_devices,
     per_chip_output_shape,
-    scatter_dim,
+    dim,
     num_links,
     math_op,
     input_dtype,
@@ -254,7 +254,7 @@ def test_reduce_scatter_on_n300(
         n300_mesh_device,
         num_devices,
         per_chip_output_shape,
-        scatter_dim,
+        dim,
         num_links,
         math_op,
         input_dtype,

diff --git a/tests/ttnn/unit_tests/operations/ccl/test_all_gather_nightly.py b/tests/ttnn/unit_tests/operations/ccl/test_all_gather_nightly.py
@@ -22,15 +22,15 @@
     [
         (4, 1, [4, 1, 33, 256], 0, ttnn.ROW_MAJOR_LAYOUT),
         (8, 1, [8, 1, 33, 256], 0, ttnn.ROW_MAJOR_LAYOUT),
-        (8, 1, [8, 1, 256, 32], 0, ttnn.TILE_LAYOUT),
+        (8, 1, [8, 1, 256, 32], -4, ttnn.TILE_LAYOUT),
         (8, 1, [8, 8, 256, 384], 1, ttnn.ROW_MAJOR_LAYOUT),
         # (4, 2, [8, 8, 256, 384], 1, ttnn.TILE_LAYOUT),
         (8, 1, [8, 8, 256, 384], 1, ttnn.TILE_LAYOUT),
-        (4, 1, [8, 5, 13, 384], 3, ttnn.ROW_MAJOR_LAYOUT),
-        (8, 1, [8, 5, 13, 512], 3, ttnn.ROW_MAJOR_LAYOUT),
+        (4, 1, [8, 5, 13, 384], -1, ttnn.ROW_MAJOR_LAYOUT),
+        (8, 1, [8, 5, 13, 512], -1, ttnn.ROW_MAJOR_LAYOUT),
         (4, 1, [8, 5, 32, 384], 3, ttnn.TILE_LAYOUT),
         (8, 1, [8, 5, 32, 512], 3, ttnn.TILE_LAYOUT),
-        (4, 1, [1, 1, 32, 16384], 3, ttnn.TILE_LAYOUT),
+        (4, 1, [1, 1, 32, 16384], -1, ttnn.TILE_LAYOUT),
     ],
 )
 @pytest.mark.parametrize(

diff --git a/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_N300_post_commit.py b/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_N300_post_commit.py
@@ -20,7 +20,7 @@
     ],
 )
 @pytest.mark.parametrize(
-    "per_chip_output_shape, scatter_dim, layout",
+    "per_chip_output_shape, dim, layout",
     [
         ([1, 1, 32, 4096], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 2048], 3, ttnn.TILE_LAYOUT),
@@ -50,7 +50,7 @@ def test_ring_reduce_scatter_n300_post_commit(
     n300_mesh_device,
     num_devices,
     per_chip_output_shape,
-    scatter_dim,
+    dim,
     num_links,
     math_op,
     input_dtype,
@@ -65,7 +65,7 @@ def test_ring_reduce_scatter_n300_post_commit(
         n300_mesh_device,
         num_devices,
         per_chip_output_shape,
-        scatter_dim,
+        dim,
         num_links,
         math_op,
         input_dtype,

diff --git a/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_TG_nightly.py b/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_TG_nightly.py
@@ -145,7 +145,7 @@ def run_line_reduce_scatter_on_TG_with_mesh_tensor_along_rows(
     # ttnn.visualize_mesh_device(mesh_device, tensor=ttnn_tensor)
     ttnn_tensor_out = ttnn.reduce_scatter(
         ttnn_tensor,
-        scatter_dim=dim,
+        dim=dim,
         cluster_axis=cluster_axis,
         mesh_device=mesh_device,
         math_op=math_op,
@@ -158,7 +158,7 @@ def run_line_reduce_scatter_on_TG_with_mesh_tensor_along_rows(
     for _ in range(num_iters):
         ttnn_tensor_out = ttnn.reduce_scatter(
             ttnn_tensor,
-            scatter_dim=dim,
+            dim=dim,
             cluster_axis=cluster_axis,
             mesh_device=mesh_device,
             math_op=math_op,

diff --git a/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_nightly.py b/tests/ttnn/unit_tests/operations/ccl/test_reduce_scatter_nightly.py
@@ -19,23 +19,23 @@
     ],
 )
 @pytest.mark.parametrize(
-    "per_chip_output_shape, scatter_dim, layout",
+    "per_chip_output_shape, dim, layout",
     [
         ([1, 8, 1024, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 4, 1024, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 4, 2048, 1024], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 32, 32], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 32, 64], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 32], -1, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 64], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 64, 64], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 32, 128], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 128], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 256], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 512], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 2048], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 128, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 128, 1024], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 128, 8192], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 2048, 1024], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 2048, 8192], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 2048, 8192], -1, ttnn.TILE_LAYOUT),
     ],
 )
 @pytest.mark.parametrize(
@@ -58,7 +58,7 @@ def test_reduce_scatter_t3k_8chip_nightly(
     t3k_mesh_device,
     num_devices,
     per_chip_output_shape,
-    scatter_dim,
+    dim,
     num_links,
     math_op,
     input_dtype,
@@ -73,7 +73,7 @@ def test_reduce_scatter_t3k_8chip_nightly(
         t3k_mesh_device,
         num_devices,
         per_chip_output_shape,
-        scatter_dim,
+        dim,
         num_links,
         math_op,
         input_dtype,
@@ -95,16 +95,16 @@ def test_reduce_scatter_t3k_8chip_nightly(
     ],
 )
 @pytest.mark.parametrize(
-    "per_chip_output_shape, scatter_dim, layout",
+    "per_chip_output_shape, dim, layout",
     [
         ([1, 8, 1024, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 4, 1024, 1024], 3, ttnn.TILE_LAYOUT),
-        ([1, 4, 2048, 1024], 3, ttnn.TILE_LAYOUT),
+        ([1, 4, 2048, 1024], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 512], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 32, 1024], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 32, 2048], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 32, 2048], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 128, 1024], 3, ttnn.TILE_LAYOUT),
-        ([1, 1, 128, 8192], 3, ttnn.TILE_LAYOUT),
+        ([1, 1, 128, 8192], -1, ttnn.TILE_LAYOUT),
         ([1, 1, 2048, 1024], 3, ttnn.TILE_LAYOUT),
         ([1, 1, 2048, 8192], 3, ttnn.TILE_LAYOUT),
         # These shapes result in some workers with no work, which is currently
@@ -136,7 +136,7 @@ def test_reduce_scatter_t3k_4chip_nightly(
     pcie_mesh_device,
     num_devices,
     per_chip_output_shape,
-    scatter_dim,
+    dim,
     num_links,
     math_op,
     input_dtype,
@@ -151,7 +151,7 @@ def test_reduce_scatter_t3k_4chip_nightly(
         pcie_mesh_device,
         num_devices,
         per_chip_output_shape,
-        scatter_dim,
+        dim,
         num_links,
         math_op,
         input_dtype,