xla
Implement mark_sharding as a custom op to support dynamo spmd activation sharding
#5712

Merged

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

wonjoo-wj merged 14 commits into master from wonjoo/dynamo-custom-op

wonjoo-wj force pushed from e1b09a7d to 47d34635 2 years ago

wonjoo-wj force pushed from 47d34635 to 31baf053 2 years ago

wonjoo-wj force pushed from 63f26731 to d6b58523 2 years ago

yeounoh commented on 2023-10-25

yeounoh added distributed

yeounoh added dynamo

yeounoh commented on 2023-10-25

wonjoo-wj force pushed from d6b58523 to 2b5f14a8 2 years ago

wonjoo-wj force pushed from 81bb5722 to ad4cc4f8 2 years ago

wonjoo-wj force pushed from ad4cc4f8 to e88c7d39 2 years ago

Implement mark_sharding as a custom op to support dynamo spmd activat…

53e7ec0b

Update to include OpSharding as an input

f0e8a94b

Rebase with master and run linter

7891b42c

Update unit tests

6aeeecff

Refine custom marking sharding op

dc19b9bc

Re-run linter due to wrong version

a20d710f

Add new API for custom mark sharding op and update tests

bd169c2e

Add torch pin

ae05c9a4

wonjoo-wj force pushed from 1bceb4a6 to 3c51d5ed 2 years ago

Clean up some code

08d6296c

wonjoo-wj force pushed from 3c51d5ed to 08d6296c 2 years ago

Update code to transfer pylist to xla::OpSharding

9aaa533b

wonjoo-wj force pushed from 4b453c73 to 9aaa533b 1 year ago

wonjoo-wj changed the title ~~[WIP] Implement mark_sharding as a custom op to support dynamo spmd activation sharding~~ Implement mark_sharding as a custom op to support dynamo spmd activation sharding 1 year ago

wonjoo-wj marked this pull request as ready for review 1 year ago

wonjoo-wj force pushed from c8a76f93 to 281d49a9 1 year ago

wonjoo-wj force pushed from 281d49a9 to fbe5c797 1 year ago

Update unit test and run linter

a98bfb22

wonjoo-wj force pushed from fbe5c797 to a98bfb22 1 year ago

wonjoo-wj requested a review from

yeounoh 1 year ago

wonjoo-wj requested a review from

JackCaoG 1 year ago

wonjoo-wj commented on 2023-11-07

JackCaoG commented on 2023-11-07

wonjoo-wj force pushed from 1808663f to 05763711 1 year ago

wonjoo-wj force pushed from 05763711 to f468ed62 1 year ago

wonjoo-wj force pushed from f468ed62 to 3604c159 1 year ago

Address comments -- fix typos and variable names

a4318a6c

wonjoo-wj force pushed from 3604c159 to a4318a6c 1 year ago

wonjoo-wj requested a review from

JackCaoG 1 year ago

Run linter

e35ca641

JackCaoG commented on 2023-11-08

Add metric assertions to unit tests

ae721ae2

wonjoo-wj requested a review from

JackCaoG 1 year ago

JackCaoG approved these changes on 2023-11-09

yeounoh commented on 2023-11-11

yeounoh approved these changes on 2023-11-11

wonjoo-wj merged 367f47f5 into master 1 year ago

Reviewers

JackCaoG

yeounoh

Assignees

No one assigned

Labels

distributed dynamo

Milestone

No milestone

xla Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712 Merged

Implement mark_sharding as a custom op to support dynamo spmd activation sharding #5712

xla
Implement mark_sharding as a custom op to support dynamo spmd activation sharding
#5712

Merged