microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-04-19 22:39:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

NCCL API Executor Integration (#331)

Browse Source 
Co-authored-by: Changho Hwang <changhohwang@microsoft.com>
This commit is contained in:
caiomcbr
2024-07-25 15:05:02 -07:00
committed by GitHub
parent f131fae3ec
commit 67eb9b04cc
6 changed files with 255 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/test/test_mscclpp.py
View File

@@ -613,7 +613,10 @@ def test_executor(mpi_group: MpiGroup, filename: str):
cp.random.seed(42)
buffer = cp.random.random(nelems).astype(cp.float16)
sub_arrays = cp.split(buffer, mpi_group.comm.size)
sendbuf = sub_arrays[mpi_group.comm.rank]
nelems_per_rank = int(nelems / mpi_group.comm.size)
sendbuf = cp.empty(nelems_per_rank).astype(cp.float16)
for i in range(nelems_per_rank):
sendbuf[i] = sub_arrays[mpi_group.comm.rank][i]
expected = cp.zeros_like(sendbuf)
for i in range(mpi_group.comm.size):
expected += sub_arrays[i]