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

Comment changes are shownComment changes are hidden
Commit
1 year ago
ZeRO1: Add bucketting logic to control the size of tensors for all-gather/reduce-scatter (#6025) Co-authored-by: Rahul Solanki <rhsoln@amazon.com> Co-authored-by: guangtai <guangtai@amazon.com> Co-authored-by: Amithrajith Mamidala <amithrm@amazon.com>
Author
Parents
  • test
    • File
      test_mp_all_gather.py
    • File
      test_mp_reduce_scatter.py
  • torch_xla
    • core
      • File
        xla_model.py
    • distributed
      • File
        zero_redundancy_optimizer.py