Skip to content

Commit

Permalink
#0: Fix test_all_gather_multiple_submeshes (#18319)
Browse files Browse the repository at this point in the history
### Ticket
N/A

### Problem description
`test_all_gather_multiple_submeshes` is broken, as all gather OP assumes
ring connected topology. Physical 2D mesh of devices is as follows (the
first 2x2 submesh of T3K):
```
4 0
5 1
```
... which givers `4 0 5 1` row-major ordering, while all gather expects
`4 0 1 5`.

### What's changed
Reshape submesh to `1x4` to force the the correct ordering.

### Checklist
- [X] Ran the test locally and confirmed it fixes the issue.
  • Loading branch information
omilyutin-tt authored Feb 25, 2025
1 parent 5db78f8 commit cbe0e1a
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tests/ttnn/unit_tests/test_multi_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -681,6 +681,8 @@ def test_all_gather_multiple_submeshes(mesh_device):
pytest.skip()

def model(submesh):
# Reshape to a 1x4 mesh to enforce ring connected topological order.
submesh.reshape(ttnn.MeshShape(1, 4))
full_tensor = torch.ones((1, 1, 32, 32 * submesh.get_num_devices()), dtype=torch.bfloat16)
for i in range(submesh.get_num_devices()):
full_tensor[..., i * 32 : (i + 1) * 32] = i
Expand Down

0 comments on commit cbe0e1a

Please sign in to comment.