Skip to content
This repository has been archived by the owner on Jan 28, 2025. It is now read-only.

Creating All Chuncks in AllGather Collective OOP #31

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Creating All Chuncks in AllGather Collective OOP #31

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
530e835
adding option for create all chuncks in outplace case for allgather c…
caiomcbr Dec 12, 2024
23ac27e
WIP
caiomcbr Dec 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions msccl/language/collectives.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,10 @@ def init_buffers(self):
output_buffer = [None] * (self.num_ranks * self.chunk_factor)
for ch in range(self.chunk_factor):
input_buffer[ch] = Chunk(r, ch, -1, r * self.chunk_factor + ch)
if self.create_all_chunks:
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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The args for Chunk seems incorrect. There are some comments:

  1. original_rank should always be r I think. For example, for GPU 0, the chunks in output buffer should set original rank to 0
  2. For last two args in Chunk constructor, set it to -1, -1?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This make sense, I guess you are right. BTW the same is wrong in the inplace scenario, right? I guess that I did the same mistake there.

buffers = {Buffer.input: input_buffer, Buffer.output: output_buffer}
rank_buffers.append(buffers)
return rank_buffers
Expand Down
Loading