Remove device maps from TensorPipe for v1.7 release (pytorch#45353)

Summary:
Pull Request resolved: pytorch#45353

Temporarily removing this feature, will add this back after branch cut.

Test Plan: Imported from OSS

Reviewed By: rohan-varma

Differential Revision: D23939865

Pulled By: mrshenli

fbshipit-source-id: 7dceaffea6b9a16512b5ba6036da73e7f8f83a8e

Author: mrshenli

torch/csrc/distributed/rpc/init.cpp

@@ -483,28 +483,20 @@ PyObject* rpc_init(PyObject* /* unused */) {
               optional<std::vector<std::string>>,
               optional<std::vector<std::string>>,
               float,
-              std::string,
-              std::unordered_map<std::string, tensorpipe::DeviceMap>>(),
+              std::string>(),
           py::arg("num_worker_threads") = kDefaultNumWorkerThreads,
           py::arg("_transports") = optional<std::vector<std::string>>(),
           py::arg("_channels") = optional<std::vector<std::string>>(),
           py::arg("rpc_timeout") = kDefaultRpcTimeoutSeconds,
-          py::arg("init_method") = kDefaultInitMethod,
-          py::arg("device_maps") =
-              std::unordered_map<std::string, tensorpipe::DeviceMap>())
+          py::arg("init_method") = kDefaultInitMethod)
       .def_readwrite(
           "num_worker_threads",
           &TensorPipeRpcBackendOptions::numWorkerThreads,
           R"(
               The number of threads in the thread-pool used by
               :class:`~torch.distributed.rpc.TensorPipeAgent` to execute
               requests.
-          )")
-      .def_readwrite(
-          "device_maps",
-          &TensorPipeRpcBackendOptions::deviceMaps,
-          R"(The device map locations.)")
-      .def("set_device_map", &TensorPipeRpcBackendOptions::setDeviceMap);
+          )");
 
   module.attr("_DEFAULT_NUM_WORKER_THREADS") =
       py::cast(kDefaultNumWorkerThreads);

torch/csrc/distributed/rpc/tensorpipe_agent.cpp

@@ -32,12 +32,9 @@ const std::string kServerActiveAsyncCalls = "agent.server_active_async_calls";
 inline void checkCPUTensor(const torch::Tensor& tensor) {
   TORCH_CHECK(
       tensor.device() == at::kCPU,
-      "TensorPipeAgent only supports CPU tensors by default. Sending "
-      "GPU tensors using RPC requires explicitly configurations using "
-      "`set_device_map` on `TensorPipeRpcBackendOptions`. Got a tensor "
-      "with device ",
-      tensor.device(),
-      ", but no device map is specified.");
+      "TensorPipe RPC backend only supports CPU tensors, please move your ",
+      "tensors to CPU before sending them over RPC. Found tensor on device: ",
+      tensor.device());
 }
 
 std::vector<c10::DeviceIndex> getDevicesForTensors(
@@ -480,41 +477,16 @@ void TensorPipeAgent::sendCompletedResponseMessage(
   Message&& responseMessage = std::move(*futureResponseMessage).moveValue();
   responseMessage.setId(messageId);
   if (!error) {
-    const auto& iter = reverseDeviceMaps_.find(pipe->getRemoteName());
-    if (iter == opts_.deviceMaps.end()) {
-      for (const auto& t : responseMessage.tensors()) {
-        if (!t.device().is_cpu()) {
-          responseMessage = createExceptionResponse(
-              c10::str(
-                  "TensorPipe RPC backend only supports CPU tensors by default,"
-                  " please move your tensors to CPU before sending them over "
-                  "RPC, or call `set_device_map` on "
-                  "`TensorPipeRpcBackendOptions` to explicitly configure "
-                  "device mapping. Response device mapping is not available for "
-                  "destination ",
-                  pipe->getRemoteName(),
-                  ", but found tensor on device: ",
-                  t.device()),
-              responseMessage.id());
-          break;
-        }
-      }
-    } else {
-      const auto& deviceMap = iter->second;
-      for (const auto& t : responseMessage.tensors()) {
-        if (!t.device().is_cpu() &&
-            deviceMap.find(t.device().index()) == deviceMap.end()) {
-          responseMessage = createExceptionResponse(
-              c10::str(
-                  "TensorPipe RPC backend only supports CPU tensors by default."
-                  " Response device mapping is not available for destination ",
-                  pipe->getRemoteName(),
-                  " for device ",
-                  t.device(),
-                  " but received a tensor on that device."),
-              responseMessage.id());
-          break;
-        }
+    for (const auto& tensor : responseMessage.tensors()) {
+      if (!tensor.device().is_cpu()) {
+        responseMessage = createExceptionResponse(
+            c10::str(
+                "TensorPipe RPC backend only supports CPU tensors, please ",
+                "move your tensors to CPU before sending them over RPC. Found ",
+                "tensor on device: ",
+                tensor.device()),
+            responseMessage.id());
+        break;
       }
     }
 

torch/distributed/rpc/__init__.py

@@ -19,8 +19,6 @@ def is_available():
     raise RuntimeError("Failed to initialize torch.distributed.rpc")
 
 
-
-
 if is_available():
     from . import api, backend_registry, functions, _set_profiler_node_id
     from . import (

torch/distributed/rpc/backend_registry.py

@@ -3,10 +3,8 @@
 from datetime import timedelta
 import enum
 
-import torch
 import torch.distributed as dist
 
-from . import api
 from . import constants as rpc_constants
 
 
@@ -185,57 +183,6 @@ def _tensorpipe_construct_rpc_backend_options_handler(
     )
 
 
-# detect if any worker has invalid device_map configurations, and return
-# names of failed workers
-def _tensorpipe_check_device_maps(agent, device_maps):
-    if device_maps is None:
-        device_maps = {}
-
-    def check_one_worker(name, device_maps, all_device_counts):
-        device_count = all_device_counts[name]
-        wrong_worker_names = set(device_maps) - set(all_device_counts)
-        if wrong_worker_names:
-            raise ValueError(f"Wrong worker names: {wrong_worker_names}")
-        for worker_name in all_device_counts:
-            remote_device_count = all_device_counts[worker_name]
-            if worker_name in device_maps:
-                device_map = device_maps[worker_name]
-                key_set = set(device_map.keys())
-                val_set = set(device_map.values())
-                if not all([
-                    len(device_map) == len(key_set),
-                    len(device_map) == len(val_set),  # check 1-to-1 mapping
-                    min(key_set) >= 0,
-                    max(key_set) < device_count,  # check local range
-                    min(val_set) >= 0,
-                    max(val_set) < remote_device_count  # check remote range
-                ]):
-                    raise ValueError(
-                        f"Invalid device_map configuration on {name}:\n"
-                        f"device_maps = {device_maps}"
-                    )
-
-    gathered = api._all_gather([torch.cuda.device_count(), device_maps])
-    all_device_counts = {name: gathered[name][0] for name in gathered}
-    all_device_maps = {name: gathered[name][1] for name in gathered}
-    for worker_name in all_device_maps:
-        worker_device_maps = all_device_maps[worker_name]
-        check_one_worker(worker_name, worker_device_maps, all_device_counts)
-
-    # passed all checked, construct reverse mapping for return values
-    reverse_device_maps = {}
-    local_name = api.get_worker_info().name
-    for worker_name in all_device_maps:
-        remote_device_maps = all_device_maps[worker_name]
-        if local_name in remote_device_maps:
-            remote_device_map = remote_device_maps[local_name]
-            reverse_device_maps[worker_name] = {
-                remote_device_map[k]: k for k in remote_device_map
-            }
-
-    agent._set_reverse_device_maps(reverse_device_maps)
-
-
 def _tensorpipe_init_backend_handler(store, name, rank, world_size, rpc_backend_options):
     from . import TensorPipeRpcBackendOptions
     from . import TensorPipeAgent
@@ -259,21 +206,10 @@ def _tensorpipe_init_backend_handler(store, name, rank, world_size, rpc_backend_
     group = _init_process_group(store, rank, world_size)
 
     # TODO: add try-except and destroy _agent in all processes if any fails.
-    agent = TensorPipeAgent(
+    return TensorPipeAgent(
         store, name, rank, world_size, group, rpc_backend_options
     )
 
-    api._init_rpc_states(agent)
-
-    try:
-        _tensorpipe_check_device_maps(agent, rpc_backend_options.device_maps)
-        agent.join()
-    except Exception:
-        api.shutdown()
-        raise
-
-    return agent
-
 
 register_backend(
     "TENSORPIPE",

torch/distributed/rpc/options.py

@@ -1,9 +1,7 @@
 from . import _TensorPipeRpcBackendOptionsBase
 from . import constants as rpc_contants
 
-import torch
-
-from typing import Dict, List
+from typing import List
 
 
 class TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):
@@ -27,19 +25,13 @@ class TensorPipeRpcBackendOptions(_TensorPipeRpcBackendOptionsBase):
             store used for rendezvous. It takes any value accepted for the
             same argument of :meth:`~torch.distributed.init_process_group`
             (default: ``env://``).
-        device_maps (Dict[str, Dict]): Device placement mappings from this
-            worker to the callee. Key is the callee worker name and value the
-            dictionary (``Dict`` of ``int``, ``str``, or ``torch.device``) that
-            maps this worker's devices to the callee worker's devices.
-            (default: ``None``)
     """
     def __init__(
         self,
         *,
         num_worker_threads: int = rpc_contants.DEFAULT_NUM_WORKER_THREADS,
         rpc_timeout: float = rpc_contants.DEFAULT_RPC_TIMEOUT_SEC,
         init_method: str = rpc_contants.DEFAULT_INIT_METHOD,
-        device_maps: Dict = None,
         _transports: List = None,
         _channels: List = None,
     ):
@@ -48,71 +40,5 @@ def __init__(
             _transports,
             _channels,
             rpc_timeout,
-            init_method,
-            device_maps if device_maps else {}
+            init_method
         )
-
-    def set_device_map(self, to: str, device_map: Dict):
-        r"""
-        Set device mapping between each RPC caller and callee pair. This
-        function can be called multiple times to incrementally add
-        device placement configurations.
-
-        Arguments:
-            worker_name (str): Callee name.
-            device_map (Dict of int, str, or torch.device): Device placement
-                mappings from this worker to the callee. This map must be
-                invertible.
-
-        Example::
-            >>> # both workers
-            >>> def add(x, y):
-            >>>     print(x)  # tensor([1., 1.], device='cuda:1')
-            >>>     return x + y, (x + y).to(2)
-            >>>
-            >>> # on worker 0
-            >>> options = TensorPipeRpcBackendOptions(
-            >>>     num_worker_threads=8,
-            >>>     device_maps={"worker1": {0, 1}}
-            >>>     # maps worker0's cuda:0 to worker1's cuda:1
-            >>> )
-            >>> options.set_device_map("worker1", {1, 2})
-            >>> # maps worker0's cuda:1 to worker1's cuda:2
-            >>>
-            >>> rpc.init_rpc(
-            >>>     "worker0",
-            >>>     rank=0,
-            >>>     world_size=2
-            >>>     backend=rpc.BackendType.TENSORPIPE,
-            >>>     rpc_backend_options=options
-            >>> )
-            >>>
-            >>> x = torch.ones(2)
-            >>> rets = rpc.rpc_sync("worker1", add, args=(x.to(0), 1))
-            >>> # The first argument will be moved to cuda:1 on worker1. When
-            >>> # sending the return value back, it will follow the invert of
-            >>> # the device map, and hence will be moved back to cuda:0 and
-            >>> # cuda:1 on worker0
-            >>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
-            >>> print(rets[0])  # tensor([2., 2.], device='cuda:1')
-        """
-        device_index_map = {}
-        curr_device_maps = super().device_maps
-        for k in device_map:
-            v = device_map[k]
-            k, v = torch.device(k), torch.device(v)
-            if k.type != 'cuda' or v.type != 'cuda':
-                raise ValueError(
-                    "`set_device_map` only supports CUDA devices, "
-                    f"but got device pair {k}: {v}"
-
-                )
-            if to in curr_device_maps and k.index in curr_device_maps[to]:
-                curr_v = super().device_maps[to][k.index]
-                if curr_v != v.index:
-                    raise ValueError(
-                        "`set_device_map` only supports 1-to-1 mapping, "
-                        f"trying to map {k} to {v} and {curr_v}"
-                    )
-            device_index_map[k.index] = v.index
-        super().set_device_map(to, device_index_map)

torch/testing/_internal/distributed/rpc/rpc_test.py

@@ -2876,6 +2876,38 @@ def _return_gpu_tensor_list():
     def _gpu_tensor_list_arg(tensor_list):
         return torch.rand(3, 3)
 
+    @skip_if_lt_x_gpu(2)
+    @dist_init
+    def test_cuda(self):
+        dst = worker_name((self.rank + 1) % self.world_size)
+        t1 = torch.rand(3, 3).cuda(0)
+        t2 = torch.rand(3, 3).cuda(1)
+        t3 = torch.rand(3, 3)
+
+        # cuda tensors as args fail.
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(dst, torch.add, args=(t1, t2))
+
+        # mix of cpu and cuda tensors as args fail.
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(dst, torch.add, args=(t1, t3))
+
+        # gpu tensor list as args fails.
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(dst, RpcTest._gpu_tensor_list_arg, args=([t1, t2]))
+
+        # cuda tensors as return values fail.
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(dst, RpcTest._return_gpu_tensor, args=())
+
+        # cuda tensors as a list of return value fails
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(dst, RpcTest._return_gpu_tensor_list, args=())
+
+        # Sending to self should fail too.
+        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
+            rpc.rpc_sync(worker_name(self.rank), torch.add, args=(t1, t2))
+
     def _create_rref(self):
         owner_rank = (self.rank + 2) % self.world_size
         return rpc.remote(
@@ -3609,39 +3641,6 @@ def test_logs_deprecation_warning(self):
             "\n".join(cm.output),
         )
 
-    @skip_if_lt_x_gpu(2)
-    @dist_init
-    def test_cuda(self):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        t1 = torch.rand(3, 3).cuda(0)
-        t2 = torch.rand(3, 3).cuda(1)
-        t3 = torch.rand(3, 3)
-
-        # cuda tensors as args fail.
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(dst, torch.add, args=(t1, t2))
-
-        # mix of cpu and cuda tensors as args fail.
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(dst, torch.add, args=(t1, t3))
-
-        # gpu tensor list as args fails.
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(dst, RpcTest._gpu_tensor_list_arg, args=([t1, t2]))
-
-        # cuda tensors as return values fail.
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(dst, RpcTest._return_gpu_tensor, args=())
-
-        # cuda tensors as a list of return value fails
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(dst, RpcTest._return_gpu_tensor_list, args=())
-
-        # Sending to self should fail too.
-        with self.assertRaisesRegex(RuntimeError, "RPC backend only supports CPU tensors.*Found tensor on device: cuda:0"):
-            rpc.rpc_sync(worker_name(self.rank), torch.add, args=(t1, t2))
-
-
     def test_single_threaded_rref_owner(self):
         # We need a process group in order to perform a barrier at the end.
         dist.init_process_group(
@@ -4288,319 +4287,3 @@ def test_tensorpipe_options_throw_on_timedelta_timeout(self):
                 num_worker_threads=self.rpc_backend_options.num_worker_threads,
                 rpc_timeout=timeout,
             )
-
-    def _test_device_maps(self, options, errMsg="Invalid device_map"):
-        with self.assertRaisesRegex(ValueError, errMsg):
-            rpc.init_rpc(
-                name=worker_name(self.rank),
-                backend=self.rpc_backend,
-                rank=self.rank,
-                world_size=self.world_size,
-                rpc_backend_options=options,
-            )
-
-        self.assertFalse(rpc.api._is_current_rpc_agent_set())
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_wrong_worker_name(self):
-        options = self.rpc_backend_options
-        options.set_device_map("none_exist", {0: 1})
-        self._test_device_maps(options, "Wrong worker names")
-
-    @skip_if_lt_x_gpu(1)
-    def test_device_maps_invalid_max_local_device(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options.set_device_map(dst, {torch.cuda.device_count(): 0})
-
-        self._test_device_maps(options)
-
-    @skip_if_lt_x_gpu(1)
-    def test_device_maps_invalid_max_remote_device(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options.set_device_map(dst, {0: torch.cuda.device_count()})
-
-        self._test_device_maps(options)
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_many_to_one(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options.set_device_map(dst, {1: 0})
-        options.set_device_map(dst, {0: 0})
-
-        self._test_device_maps(options)
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_one_to_many(self):
-        if self.rank == 0:
-            options = self.rpc_backend_options
-            dst = worker_name((self.rank + 1) % self.world_size)
-            options.set_device_map(dst, {0: 1})
-            with self.assertRaisesRegex(
-                ValueError, "`set_device_map` only supports 1-to-1 mapping"
-            ):
-                options.set_device_map(dst, {0: 0})
-
-    @skip_if_lt_x_gpu(1)
-    def test_device_maps_invalid_min_device(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        with self.assertRaisesRegex(
-            RuntimeError, "Device index must not be negative"
-        ):
-            options.set_device_map(dst, {-1: 0})
-
-        with self.assertRaisesRegex(
-            RuntimeError, "Device index must not be negative"
-        ):
-            options.set_device_map(dst, {0: -1})
-
-    @staticmethod
-    def _gpu_add(x, y):
-        if all([x.is_cuda, x.device.index == 1, y.is_cuda, y.device.index == 1]):
-            return (x + y).to(0)
-        else:
-            raise ValueError("Wrong device affinity")
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_gpu(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options.set_device_map(dst, {0: 1, 1: 0})
-
-        rpc.init_rpc(
-            name=worker_name(self.rank),
-            backend=self.rpc_backend,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_backend_options=options,
-        )
-
-        ret = rpc.rpc_sync(
-            dst,
-            TensorPipeAgentRpcTest._gpu_add,
-            args=(torch.zeros(2).to(0), torch.ones(2).to(0))
-        )
-        self.assertEqual(ret.device, torch.device(1))
-        self.assertEqual(ret, (torch.zeros(2) + torch.ones(2)).to(1))
-        rpc.shutdown()
-
-    @staticmethod
-    def _gpu_add_multi_gpu(x, y):
-        if all([x.is_cuda, x.device.index == 0, y.is_cuda, y.device.index == 1]):
-            return x + y.to(0), x.to(1) - y
-        else:
-            raise ValueError("Wrong device affinity")
-
-    def _test_device_maps_multi_gpu(self, dst):
-        options = self.rpc_backend_options
-        options.set_device_map(dst, {1: 0})
-        options.set_device_map(dst, {0: 1})
-
-        rpc.init_rpc(
-            name=worker_name(self.rank),
-            backend=self.rpc_backend,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_backend_options=options,
-        )
-
-        rets = rpc.rpc_sync(
-            dst,
-            TensorPipeAgentRpcTest._gpu_add_multi_gpu,
-            args=(torch.zeros(2).to(1), torch.ones(2).to(0))
-        )
-        self.assertEqual(rets[0].device, torch.device(1))
-        self.assertEqual(rets[1].device, torch.device(0))
-        self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1))
-        self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0))
-        rpc.shutdown()
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_multi_gpu(self):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        self._test_device_maps_multi_gpu(dst)
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_multi_gpu_self(self):
-        dst = worker_name(self.rank)
-        self._test_device_maps_multi_gpu(dst)
-
-    @staticmethod
-    def _gpu_add_return_to_gpu(x, y):
-        if x.device.type == 'cpu' and y.device.type == 'cpu':
-            return (x + y).to(0), (x - y).to(1), (x * y).to(2), (x / y).to(3)
-        else:
-            raise ValueError("Wrong device affinity")
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_in_options(self):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options = self.rpc_backend_options
-
-        rpc.init_rpc(
-            name=worker_name(self.rank),
-            backend=self.rpc_backend,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
-                init_method=options.init_method,
-                num_worker_threads=options.num_worker_threads,
-                device_maps={dst: {0: 1, 1: 0}}
-            )
-        )
-
-        rets = rpc.rpc_sync(
-            dst,
-            TensorPipeAgentRpcTest._gpu_add_multi_gpu,
-            args=(torch.zeros(2).to(1), torch.ones(2).to(0))
-        )
-        self.assertEqual(rets[0].device, torch.device(1))
-        self.assertEqual(rets[1].device, torch.device(0))
-        self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1))
-        self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0))
-        rpc.shutdown()
-
-    def _test_device_maps_return_to_gpu(self, dst):
-        options = self.rpc_backend_options
-
-        options.set_device_map(dst, {0: 1})
-        options.set_device_map(dst, {1: 2})
-        options.set_device_map(dst, {2: 3})
-        options.set_device_map(dst, {3: 0})
-
-        rpc.init_rpc(
-            name=worker_name(self.rank),
-            backend=self.rpc_backend,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_backend_options=options,
-        )
-
-        rets = rpc.rpc_sync(
-            dst,
-            TensorPipeAgentRpcTest._gpu_add_return_to_gpu,
-            args=(torch.zeros(2), torch.ones(2))
-        )
-        for i in range(len(rets)):
-            self.assertEqual(rets[i].device, torch.device((3 + i) % 4))
-        self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(3))
-        self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0))
-        self.assertEqual(rets[2], (torch.zeros(2) * torch.ones(2)).to(1))
-        self.assertEqual(rets[3], (torch.zeros(2) / torch.ones(2)).to(2))
-        rpc.shutdown()
-
-    @skip_if_lt_x_gpu(4)
-    def test_device_maps_return_to_gpu(self):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        self._test_device_maps_return_to_gpu(dst)
-
-    @skip_if_lt_x_gpu(4)
-    def test_device_maps_return_to_gpu_self(self):
-        dst = worker_name(self.rank)
-        self._test_device_maps_return_to_gpu(dst)
-
-    @staticmethod
-    def _add_to_gpu(x, y):
-        return (x + y).to(0)
-
-    def _test_device_maps_missing_config(self, mode):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        errMsg = (
-            "TensorPipeAgent only supports CPU tensors by default.*"
-            "`set_device_map` on `TensorPipeRpcBackendOptions`"
-        )
-
-        with self.assertRaisesRegex(RuntimeError, errMsg):
-            if mode == RPCExecMode.SYNC:
-                rpc.rpc_sync(dst, torch.add, args=(torch.zeros(2).to(0), 1))
-            elif mode == RPCExecMode.REMOTE:
-                rpc.remote(dst, torch.add, args=(torch.zeros(2).to(0), 1)).to_here()
-            else:
-                raise ValueError(f"unexpected mode {mode}")
-
-        # make sure RPC is still functioning
-        ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1))
-        self.assertEqual(ret, torch.ones(2) + 1)
-
-    def _test_device_maps_missing_config_response(self, mode):
-        dst = worker_name((self.rank + 1) % self.world_size)
-        errMsg = "Response device mapping is not available"
-
-        with self.assertRaisesRegex(RuntimeError, errMsg):
-            if mode == RPCExecMode.SYNC:
-                rpc.rpc_sync(
-                    dst,
-                    TensorPipeAgentRpcTest._add_to_gpu,
-                    args=(torch.zeros(2), 1)
-                )
-            elif mode == RPCExecMode.REMOTE:
-                rpc.remote(
-                    dst,
-                    TensorPipeAgentRpcTest._add_to_gpu,
-                    args=(torch.zeros(2), 1)
-                ).to_here()
-            else:
-                raise ValueError(f"unexpected mode {mode}")
-
-        # make sure RPC is still functioning
-        ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1))
-        self.assertEqual(ret, torch.ones(2) + 1)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config(self):
-        self._test_device_maps_missing_config(RPCExecMode.SYNC)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config_loop(self):
-        for _ in range(self.rpc_backend_options.num_worker_threads + 5):
-            self._test_device_maps_missing_config(RPCExecMode.SYNC)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config_response(self):
-        self._test_device_maps_missing_config_response(RPCExecMode.SYNC)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config_response_loop(self):
-        for _ in range(self.rpc_backend_options.num_worker_threads + 5):
-            self._test_device_maps_missing_config_response(RPCExecMode.SYNC)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config_remote(self):
-        self._test_device_maps_missing_config(RPCExecMode.REMOTE)
-
-    @skip_if_lt_x_gpu(1)
-    @dist_init
-    def test_device_maps_missing_config_remote_response(self):
-        self._test_device_maps_missing_config_response(RPCExecMode.REMOTE)
-
-    @skip_if_lt_x_gpu(2)
-    def test_device_maps_remote(self):
-        options = self.rpc_backend_options
-        dst = worker_name((self.rank + 1) % self.world_size)
-        options.set_device_map(dst, {1: 0})
-
-        rpc.init_rpc(
-            name=worker_name(self.rank),
-            backend=self.rpc_backend,
-            rank=self.rank,
-            world_size=self.world_size,
-            rpc_backend_options=options,
-        )
-
-        rref = rpc.remote(
-            dst,
-            TensorPipeAgentRpcTest._add_to_gpu,
-            args=(torch.zeros(2), 1)
-        )
-
-        self.assertEqual(rref.to_here(), torch.ones(2).to(1))
-
-        rpc.shutdown()