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.

Sign up for GitHub

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

Jump to bottom

feat: experimental sharding backend #1544

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

feat: experimental sharding backend #1544

wants to merge 30 commits into from

Commits on Sep 25, 2024

  1. feat: add basic backend layout

    @polvalente
    polvalente committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    25788c0 View commit details
    Browse the repository at this point in the history

  2. proof of concept solution

    @polvalente
    polvalente committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    6585353 View commit details
    Browse the repository at this point in the history

Commits on Sep 26, 2024

  1. feat: add working POC for sharding

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    54e7abf View commit details
    Browse the repository at this point in the history

  2. refactor: use only sharding compiler

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    5fe92ca View commit details
    Browse the repository at this point in the history

  3. refactor: move things into __compile__

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    09b844a View commit details
    Browse the repository at this point in the history

  4. feat: working EXLA example

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    f6bd4ed View commit details
    Browse the repository at this point in the history

  5. wip: initial work on dot (doesn't work)

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    4c8200f View commit details
    Browse the repository at this point in the history

  6. wip

    @polvalente
    polvalente committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    2f9370a View commit details
    Browse the repository at this point in the history

Commits on Sep 28, 2024

  1. wip: refactor input sharding calculation

    @polvalente
    polvalente committed Sep 28, 2024
    Configuration menu
    Copy the full SHA
    918118e View commit details
    Browse the repository at this point in the history

  2. wip: refactor to support broadcasts

    @polvalente
    polvalente committed Sep 28, 2024
    Configuration menu
    Copy the full SHA
    bd06388 View commit details
    Browse the repository at this point in the history

Commits on Sep 29, 2024

  1. wip: rework sharding representation (each slice is a shard)

    @polvalente
    polvalente committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    5d7b106 View commit details
    Browse the repository at this point in the history

  2. feat: deal with broadcasting and re-slicing

    @polvalente
    polvalente committed Sep 29, 2024
    Configuration menu
    Copy the full SHA
    1f7dfb3 View commit details
    Browse the repository at this point in the history

Commits on Oct 10, 2024

  1. refactor: build parents tree into each shard

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    243eee2 View commit details
    Browse the repository at this point in the history

  2. feat: support implicit broadcasting

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    4bb2d97 View commit details
    Browse the repository at this point in the history

  3. chore: remove unused var'

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    87f1c35 View commit details
    Browse the repository at this point in the history

  4. feat: support dot product without contraction sharding

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    bca3943 View commit details
    Browse the repository at this point in the history

  5. feat: support constants

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    73a3553 View commit details
    Browse the repository at this point in the history

  6. add :tensor op to example function

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    52fc860 View commit details
    Browse the repository at this point in the history

  7. feat: support squeeze

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    7706678 View commit details
    Browse the repository at this point in the history

  8. chore: remove empty file

    @polvalente
    polvalente committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    0ddb69a View commit details
    Browse the repository at this point in the history

Commits on Oct 11, 2024

  1. refactor: remove TensorSharding module

    @polvalente
    polvalente committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    489d655 View commit details
    Browse the repository at this point in the history

  2. chore: add stubs for missing callbacks

    @polvalente
    polvalente committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    42b5ca6 View commit details
    Browse the repository at this point in the history

  3. test: add tests

    @polvalente
    polvalente committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    59034c1 View commit details
    Browse the repository at this point in the history

  4. fix: transpose axis

    @polvalente
    polvalente committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    d6da7a8 View commit details
    Browse the repository at this point in the history

  5. chore: remove example .exs files

    @polvalente
    polvalente committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    68565bc View commit details
    Browse the repository at this point in the history

Commits on Oct 14, 2024

  1. Update nx/lib/nx/defn/sharding_compiler.ex 

    Co-authored-by: José Valim <[email protected]>
    @polvalente @josevalim
    polvalente and josevalim authored Oct 14, 2024
    Configuration menu
    Copy the full SHA
    fb86713 View commit details
    Browse the repository at this point in the history

  2. chore: format

    @polvalente
    polvalente committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    0c111df View commit details
    Browse the repository at this point in the history

  3. chore: remove __stream__

    @polvalente
    polvalente committed Oct 14, 2024
    Configuration menu
    Copy the full SHA
    5c5c881 View commit details
    Browse the repository at this point in the history

Commits on Oct 15, 2024

  1. Merge remote-tracking branch 'origin/main' into pv-feat/experimental-… 

    …sharding-backend
    @polvalente
    polvalente committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    5e05bbb View commit details
    Browse the repository at this point in the history

Commits on Oct 17, 2024

  1. feat: add graph splitter for all-gather/all-reduce operations (#1545)

    @polvalente
    polvalente authored Oct 17, 2024
    Configuration menu
    Copy the full SHA
    6eb6fba View commit details
    Browse the repository at this point in the history