xla

a805505d - Backport ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter (#6025) (#6806)

Backport ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter (#6025) (#6806) Co-authored-by: jeffhataws <jthuynh@amazon.com> Co-authored-by: Rahul Solanki <rhsoln@amazon.com> Co-authored-by: guangtai <guangtai@amazon.com> Co-authored-by: Amithrajith Mamidala <amithrm@amazon.com>