Skip to content

Refactor the JACCL ring and add threads for the multiple rings#3900

Open
angeloskath wants to merge 7 commits into
mainfrom
jaccl-threads
Open

Refactor the JACCL ring and add threads for the multiple rings#3900
angeloskath wants to merge 7 commits into
mainfrom
jaccl-threads

Conversation

@angeloskath

Copy link
Copy Markdown
Member

This is on top of the coordinator changes and I will rebase it when that is merged, for now focus on the changes in the reduction ops and more importantly the ring implementation.

The TL;DR is the following:

  • Refactor the ring to extract the ring pass in a standalone function (with copy op or reduction)
  • Refactor the reduction ops to support out of place reduction that allows us to avoid the initial copy
  • Write a per-wire all reduce and all gather using the ring pass
  • Copy the MLX threadpool and use it to dispatch the other wires to different threads

The performance on a single wire is maintained and there should be 0 overhead since that wire runs in the calling thread.

TODO: Add multi-wire benchmarks to show improved performance compared to before.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant