Add a _RemoteDevice structure for ShardedTensor/ShardingSpec. (#62927)
authorPritam Damania <pritam.damania@fb.com>
Wed, 11 Aug 2021 18:22:48 +0000 (11:22 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Wed, 11 Aug 2021 18:27:32 +0000 (11:27 -0700)
commitb8e6144e0a328fe49067ae7d18d5582c8f1a61b0
treefd07e0ee5cbe0fa02972c55b8bda00140da4e678
parentb746fed16482a4f11f0fdaa514fda807e138ff7f
Add a _RemoteDevice structure for ShardedTensor/ShardingSpec. (#62927)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62927

As part of the ShardedTensor work, we realized we do need some sort of
_RemoteDevice structure that deals with our format of "workername/device" so
that users don't have to worry about parsing this string directly.

Right now this structure is just the bare minimum and is mostly a container for
describing a remote device. It is currently only used in ShardedTensor,
ShardingSpec and RemoteModule.

Once we actually have a consolidated remote device proposal, this class can be
extended appropriately if needed.
ghstack-source-id: 135534086

Test Plan:
1) unit tests
2) waitforbuildbot

Reviewed By: SciPioneer

Differential Revision: D30170689

fbshipit-source-id: 1ac2e81c7a597dc40bf3fbf2c1168c382c66649f
test/distributed/_sharded_tensor/test_sharded_tensor.py
test/distributed/_sharding_spec/test_sharding_spec.py
torch/distributed/__init__.py
torch/distributed/_sharded_tensor/api.py
torch/distributed/_sharding_spec/_internals.py
torch/distributed/_sharding_spec/api.py
torch/distributed/nn/api/remote_module.py
torch/distributed/remote_device.py [new file with mode: 0644]
torch/distributed/utils.py [deleted file]