PR #7216 Enable bucketized all-reduce for gradients

Enable bucketized all-reduce for gradients #7216

JackCaoG merged 11 commits into master from bucketing_gradients

JackCaoG requested a review from

jeffhataws 1 year ago

JackCaoG requested a review from

alanwaketan 1 year ago

jeffhataws requested changes on 2024-06-07

JackCaoG approved these changes on 2024-06-11

jeffhataws approved these changes on 2024-06-11

Gradient bucketing using a pre-defined bucket size cap

c491f05c

Fix linter issues

3f7cc775

Added ALLREDUCE_BUCKET_SIZE_MB to turn on bucketing for allreduce

4c1233b0

Fix import

45ba8857

Fixing API for allreduce bucketized gradients

d7ecbbaa

fix linter checks

35991a69

Fixing test case

7fba20aa

Add bucketized all-reduce test to run_tests.sh; move test into torch_…

a7193b4a

Add init_method='xla://' to bucketized allreduce test

c63f00cc

Lint fix

f332a8fe

jeffhataws force pushed from b1bf93c7 to f332a8fe 1 year ago

Add bucket_cap_mb arguments; fix bucketized allreduce test

85c6ec17

jeffhataws changed the title ~~Bucketing gradients~~ Enable bucketized all-reduce for gradients 1 year ago

JackCaoG merged 28f9887b into master 1 year ago

jeffhataws deleted the bucketing_gradients branch 1 year ago

Reviewers

jeffhataws

JackCaoG

alanwaketan

Assignees

No one assigned

Labels

None yet

Milestone

No milestone