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

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

jeffhataws committed 2 years ago

clang-format-7 and python lint fixes

jeffhataws committed 2 years ago

Fix "SyntaxError: 'return' outside function" error

jeffhataws committed 2 years ago

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

jeffhataws committed 2 years ago

Fix allgather to be compatible with openxla allgather tuple change without token

jeffhataws committed 2 years ago

Fix reduce-scatter-coalesce to be compatible with openxla reduce-scatter tuple change without token

jeffhataws committed 2 years ago

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

jeffhataws committed 2 years ago

Some cleanups

jeffhataws committed 2 years ago

Add separate BuildAllGatherCoalesced builder and AllGatherCoalesced class

jeffhataws committed 2 years ago

Use token_handler.GetInput to capture token

jeffhataws committed 2 years ago

Clean up

jeffhataws committed 2 years ago

Clean up

jeffhataws committed 2 years ago

Switch to GetOperandListWithToken naming for func GetOperandList

jeffhataws committed 2 years ago

xla Add all-gather coalescing for FSDP/ZeRO1 #5950 Merged