pytorch
c1037d0d - [PT-D][Sharding] Move Partial Tensor to the _shard folder and add logic to remove padding (#76199)

Comment changes are shownComment changes are hidden
Commit
2 years ago
[PT-D][Sharding] Move Partial Tensor to the _shard folder and add logic to remove padding (#76199) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/76199 Since Partial Tensor is somehow isolated to sharded tensor. We now move it to the _shard folder. Also, we added the logic to remove paddings when the size is not divisible by the world size. Modify the unit test to reflect this changes. Finally, we need to consider the placement order for the reshading spec for partial tensor, related logic is added in this change. Futhermore, for sharded linear, we will need to order the placement by rank to get the expected local result. ghstack-source-id: 154853290 Test Plan: CI Reviewed By: pritamdamania87, wanchaol Differential Revision: D35827894 fbshipit-source-id: 58dab77969b8b6557f42afa7e8f5a8a053dd5793 (cherry picked from commit abeb28f16582dcf707c2e165f39df6caf692384d)
Author
Committer
Parents
  • test
    • distributed/_shard
      • sharded_tensor/ops
        • File
          test_linear.py
      • File
        test_partial_tensor.py
    • File
      run_test.py
  • torch/distributed/_shard
    • File
      partial_tensor.py
    • sharded_tensor
      • File
        __init__.py
      • _ops
        • File
          linear.py