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

Commit

1 year ago

Add all-gather and reduce-scatter coalescence support for FSDP. Also allow using reduce-scatter's scale param in FSDP. (revived https://github.com/pytorch/xla/pull/4145)

References

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

Author

jeffhataws

Committer

jeffhataws

Parents

6addbde3

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

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