frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use Devices instead of DeviceIndexes in TensorPipe agent (#57294)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/57294

With the advent of CPUs in the device maps, and to be more generic (e.g., to support AMD GPUs), and to avoid conversions when passing to Future and RRef and such, it's easier to use Devices instead of DeviceIndices. This started by just migrating the TensorPipe agent but the RPC layer is quite intertwined so I had to migrate a lot of stuff.
ghstack-source-id: 127916562

Test Plan: CI

Reviewed By: mrshenli

Differential Revision: D28092733

fbshipit-source-id: 024dcb3648c5898ab13e770413c43958f04f1a8a
This commit is contained in:
Luca Wehrstedt
2021-05-01 16:09:43 -07:00
committed by Facebook GitHub Bot
parent 0c3e79b5b9
commit 0422e67336
26 changed files with 176 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/csrc/distributed/rpc/rpc_agent.cpp
View File

@ -287,7 +287,7 @@ bool RpcAgent::isGILProfilingEnabled() {
return profilingEnabled_.load();
}
std::unordered_map<c10::DeviceIndex, c10::DeviceIndex> RpcAgent::getDeviceMap(
std::unordered_map<c10::Device, c10::Device> RpcAgent::getDeviceMap(
const WorkerInfo& /* unused */) const {
// Default implementation has no device map.
return {};