PR #5950 Add all-gather coalescing for FSDP/ZeRO1

Add all-gather coalescing for FSDP/ZeRO1 #5950

JackCaoG merged 13 commits into master from cc_coalesce_revival

Add all-gather and reduce-scatter coalescence support for FSDP.

57e6b366

clang-format-7 and python lint fixes

63431828

Fix "SyntaxError: 'return' outside function" error

cb5b6fbe

Code/test fixes to get run_tests.sh to run on CPU

3e97aa84

Fix allgather to be compatible with openxla allgather tuple change wi…

ec27f90d

Fix reduce-scatter-coalesce to be compatible with openxla reduce-scat…

f9f3a712

Separate out the reduce-scatter-coalesce changes into a separate PR

5843b434

Some cleanups

8e6715e7

Add separate BuildAllGatherCoalesced builder and AllGatherCoalesced c…

3cd39ea5

jeffhataws requested a review from

JackCaoG 1 year ago

Use token_handler.GetInput to capture token

49f4ba4e

Clean up

6fd0d0a7

Clean up

fae4d071

JackCaoG commented on 2023-11-30

Switch to GetOperandListWithToken naming for func GetOperandList

8dc7af94

JackCaoG added backport_2.2

JackCaoG approved these changes on 2023-12-02

JackCaoG merged 12719646 into master 1 year ago

Reviewers

JackCaoG

Assignees

No one assigned

Labels

backport_2.2

Milestone

No milestone