Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: experimental sharding backend #1544

Open
wants to merge 30 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 25 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
25788c0
feat: add basic backend layout
polvalente Sep 25, 2024
6585353
proof of concept solution
polvalente Sep 25, 2024
54e7abf
feat: add working POC for sharding
polvalente Sep 26, 2024
5fe92ca
refactor: use only sharding compiler
polvalente Sep 26, 2024
09b844a
refactor: move things into __compile__
polvalente Sep 26, 2024
f6bd4ed
feat: working EXLA example
polvalente Sep 26, 2024
4c8200f
wip: initial work on dot (doesn't work)
polvalente Sep 26, 2024
2f9370a
wip
polvalente Sep 26, 2024
918118e
wip: refactor input sharding calculation
polvalente Sep 28, 2024
bd06388
wip: refactor to support broadcasts
polvalente Sep 28, 2024
5d7b106
wip: rework sharding representation (each slice is a shard)
polvalente Sep 29, 2024
1f7dfb3
feat: deal with broadcasting and re-slicing
polvalente Sep 29, 2024
243eee2
refactor: build parents tree into each shard
polvalente Oct 10, 2024
4bb2d97
feat: support implicit broadcasting
polvalente Oct 10, 2024
87f1c35
chore: remove unused var'
polvalente Oct 10, 2024
bca3943
feat: support dot product without contraction sharding
polvalente Oct 10, 2024
73a3553
feat: support constants
polvalente Oct 10, 2024
52fc860
add :tensor op to example function
polvalente Oct 10, 2024
7706678
feat: support squeeze
polvalente Oct 10, 2024
0ddb69a
chore: remove empty file
polvalente Oct 10, 2024
489d655
refactor: remove TensorSharding module
polvalente Oct 11, 2024
42b5ca6
chore: add stubs for missing callbacks
polvalente Oct 11, 2024
59034c1
test: add tests
polvalente Oct 11, 2024
d6da7a8
fix: transpose axis
polvalente Oct 11, 2024
68565bc
chore: remove example .exs files
polvalente Oct 11, 2024
fb86713
Update nx/lib/nx/defn/sharding_compiler.ex
polvalente Oct 14, 2024
0c111df
chore: format
polvalente Oct 14, 2024
5c5c881
chore: remove __stream__
polvalente Oct 14, 2024
5e05bbb
Merge remote-tracking branch 'origin/main' into pv-feat/experimental-…
polvalente Oct 15, 2024
6eb6fba
feat: add graph splitter for all-gather/all-reduce operations (#1545)
polvalente Oct 17, 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
Loading
Loading