zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

mpi send use input as output (#1750)

Browse Source 
* mpi send use input as output

* move earlier
This commit is contained in:
Awni Hannun
2025-01-06 06:08:43 -08:00
committed by GitHub
parent eab93985b8
commit 058d6ce683
5 changed files with 19 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/tests/mpi_test_distributed.py
View File

@@ -111,6 +111,14 @@ class TestDistributed(mlx_tests.MLXTestCase):
self.assertTrue(mx.all(x == (1024 if pairs.rank() == 0 else 512)))
# Check recv and computation in same eval:
y = mx.ones((5, 5)) + mx.array(2.0)
if send:
x = mx.distributed.send(2 * x, neighbor, group=pairs)
else:
x = mx.distributed.recv_like(x, neighbor, group=pairs)
mx.eval(y, x)
def test_average_gradients(self):
original_all_sum = mx.distributed.all_sum
n_calls = 0