microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-05-03 21:21:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing RegisterMemory Allocation for ProxyChannels (#353)

Browse Source 
Co-authored-by: Binyang Li <binyli@microsoft.com>
Co-authored-by: Changho Hwang <changhohwang@microsoft.com>
This commit is contained in:
Caio Rocha
2024-09-24 23:01:41 -07:00
committed by GitHub
parent 8a330f9135
commit 08a0cec2eb
3 changed files with 71 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/test/executor_test.py
View File

@@ -78,12 +78,13 @@ def dtype_to_mscclpp_dtype(dtype):
def main(
execution_paln_name: str,
execution_plan_name: str,
execution_plan_path: str,
size: int,
in_place: bool = True,
dtype: cp.dtype = cp.float16,
packet_type: PacketType = PacketType.LL16,
seed: int = 42 + MPI.COMM_WORLD.rank,
seed: int = 42,
):
mscclpp_group = mscclpp_comm.CommGroup(MPI.COMM_WORLD)
cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()
@@ -91,21 +92,33 @@ def main(
npkit_dump_dir = os.getenv("NPKIT_DUMP_DIR")
if npkit_dump_dir is not None:
npkit.init(mscclpp_group.my_rank)
execution_plan = ExecutionPlan(execution_paln_name, execution_plan_path)
execution_plan = ExecutionPlan(execution_plan_name, execution_plan_path)
cp.random.seed(seed)
nelems = size // cp.dtype(dtype).itemsize
sendbuf = cp.random.random(nelems).astype(dtype)
expected = cp.asnumpy(sendbuf)
expected = MPI.COMM_WORLD.allreduce(expected, op=MPI.SUM)
buffer = cp.random.random(nelems * mscclpp_group.nranks, dtype=cp.float32).astype(dtype)
sub_arrays = cp.split(buffer, MPI.COMM_WORLD.size)
sendbuf = cp.zeros(nelems, dtype=dtype)
for i in range(nelems):
sendbuf[i] = sub_arrays[MPI.COMM_WORLD.rank][i]
if "allgather" in execution_plan_name:
recvbuf = cp.zeros(nelems * mscclpp_group.nranks, dtype=dtype)
expected = buffer
else:
cp.random.seed(seed)
recvbuf = cp.zeros(nelems, dtype=dtype)
expected = cp.zeros_like(sendbuf, dtype=dtype)
for i in range(mscclpp_group.nranks):
expected += sub_arrays[i]
mscclpp_group.barrier()
executor_func = lambda stream: executor.execute(
MPI.COMM_WORLD.rank,
sendbuf.data.ptr,
sendbuf.data.ptr,
sendbuf.nbytes,
sendbuf.data.ptr if in_place else recvbuf.data.ptr,
sendbuf.nbytes,
sendbuf.nbytes if in_place else recvbuf.nbytes,
dtype_to_mscclpp_dtype(dtype),
execution_plan,
stream.ptr,
@@ -115,7 +128,8 @@ def main(
stream = cp.cuda.Stream(non_blocking=True)
executor_func(stream)
stream.synchronize()
assert cp.allclose(sendbuf, expected, atol=1e-2 * mscclpp_group.nranks)
assert cp.allclose(sendbuf if in_place else recvbuf, expected, atol=1e-2 * mscclpp_group.nranks)
mscclpp_group.barrier()
execution_time = bench_time(100, 10, executor_func)
@@ -136,6 +150,7 @@ if __name__ == "__main__":
parser.add_argument("-n", "--execution_plan_name", type=str, required=True)
parser.add_argument("-path", "--execution_plan_path", type=str, required=True)
parser.add_argument("--size", type=str, required=True)
parser.add_argument("--in_place", action="store_true", help="flag to define an in-place operation")
parser.add_argument("--dtype", type=str, default="float16", help="Choose from float16, float32, int32")
parser.add_argument("--packet_type", type=str, default="LL16", help="Choose from LL8, LL16")
parser.add_argument("--seed", type=int, default=42)
@@ -151,6 +166,7 @@ if __name__ == "__main__":
args.execution_plan_name,
args.execution_plan_path,
buffer_size,
args.in_place,
dtype,
packet_type,
args.seed,