From 514a1c91edbeb7625a50f18954620e76ba584c4e Mon Sep 17 00:00:00 2001 From: Caio Rocha <164253795+caiomcbr@users.noreply.github.com> Date: Tue, 26 Nov 2024 14:26:57 -0800 Subject: [PATCH] Initialization all the chuncks in AllGather Collective (#26) * implementing allgather collective creating all the chuncks * adding comments --- msccl/language/collectives.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/msccl/language/collectives.py b/msccl/language/collectives.py index c8434da..2e4fa73 100755 --- a/msccl/language/collectives.py +++ b/msccl/language/collectives.py @@ -73,9 +73,12 @@ 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" + # 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 def init_buffers(self): @@ -84,8 +87,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,