Conversation
jsleep
changed the title
| gather_frequency = n_samples | ||
|
|
||
| gathered = [] | ||
| n_chunks = n_samples // self.gather_frequency + 1 |
There was a problem hiding this comment.
@krishansubudhi For my understanding , why did we have to chunk before? I assumed it was to avoid exceeding GPU memory limit but it looks like we only move tensors to GPU in this loop and never out of it.
There was a problem hiding this comment.
Adding @aminsaied who also initially created the DDP Trainer backend and chunking logic.
There was a problem hiding this comment.
The loop first moves the tensors to GPU, then does all gather op, then moves the gathered tensors back to CPU. I believe this was at the request of @gshruti95 at the time for a specific workload that was being tested (keep me honest Shruti).
There was a problem hiding this comment.
We decided to introduce chunking in case of potential memory or timeout issues when trying to all gather for pretraining workloads.
There was a problem hiding this comment.
I believe this logic will have to change back then, chunking needs to be implemented correctly and not hard coded to 1, but can be for the time being (just will be slow I think).
6 participants
No description provided.