xla icon indicating copy to clipboard operation
xla copied to clipboard

Published 20 hours ago • verified publisher icon pytorch

ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter

Open jeffhataws opened this issue 1 year ago • 0 comments

This PR updates XLA ZeRO1 implementation to use allgather coalesed and reduce-scatter coalesced.

jeffhataws avatar Dec 05 '23 05:12 jeffhataws