From 6ec5661315cd26832a342df6a63d00834f5bd8a7 Mon Sep 17 00:00:00 2001 From: Caio Date: Tue, 19 Nov 2024 09:11:10 +0000 Subject: [PATCH 1/2] implementing allgather collective creating all the chuncks --- msccl/language/collectives.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/msccl/language/collectives.py b/msccl/language/collectives.py index 6dc20cc..f291bad 100755 --- a/msccl/language/collectives.py +++ b/msccl/language/collectives.py @@ -62,9 +62,10 @@ def check(self, prog): class AllGather(Collective): - def __init__(self, num_ranks, chunk_factor, inplace): + def __init__(self, num_ranks, chunk_factor, inplace, create_all_chunks=False): Collective.__init__(self, num_ranks, chunk_factor, inplace) self.name = "allgather" + self.create_all_chunks = create_all_chunks # Initializes input buffer for an allgather def init_buffers(self): @@ -73,8 +74,13 @@ def init_buffers(self): # Inplace AllGather only uses the output buffer for r in range(self.num_ranks): output_buffer = [None] * (self.num_ranks * self.chunk_factor) - for ch in range(self.chunk_factor): - output_buffer[r * self.chunk_factor + ch] = Chunk(r, ch, -1, r * self.chunk_factor + ch) + if not self.create_all_chunks: + for ch in range(self.chunk_factor): + output_buffer[r * self.chunk_factor + ch] = Chunk(r, ch, -1, r * self.chunk_factor + ch) + else: + for rank in range(self.num_ranks): + for ch in range(self.chunk_factor): + output_buffer[rank * self.chunk_factor + ch] = Chunk(rank, ch, -1, rank * self.chunk_factor + ch) buffers = { Buffer.input: output_buffer[r * self.chunk_factor : (r + 1) * self.chunk_factor], Buffer.output: output_buffer, From 93c7629bf5cdc2037ede2a9c903845bf21252328 Mon Sep 17 00:00:00 2001 From: Caio Date: Mon, 25 Nov 2024 22:38:41 +0000 Subject: [PATCH 2/2] adding comments --- msccl/language/collectives.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/msccl/language/collectives.py b/msccl/language/collectives.py index f291bad..41bae5e 100755 --- a/msccl/language/collectives.py +++ b/msccl/language/collectives.py @@ -65,6 +65,8 @@ class AllGather(Collective): def __init__(self, num_ranks, chunk_factor, inplace, create_all_chunks=False): Collective.__init__(self, num_ranks, chunk_factor, inplace) self.name = "allgather" + # This flag is a temporary solution, which initialize all the chuncks only for inputbuffer + # In this future we need to remove this flag and always initialize all the chunks self.create_all_chunks = create_all_chunks # Initializes input buffer for an allgather