Repository: pytorch/tensordict

Issue:
[Feature Request] Support specifying group in `isend` and `irecv`
## Motivation

Currently, `isend` and `irecv` do not have a `group` argument, making it impossible to specify the desired communication group when there are multiple groups.

An example is when using a NCCL-based group for GPU communication, while also using a GLOO-based group for CPU communication.

## Solution

Add a `group` argument to `isend` and `irecv`, and pass the parameter down to the underlying `torch.distributed` methods.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)



Discussion:
Good point will take care of vat today unless you want to do it :)

Resolving patch:
diff --git a/tensordict/_lazy.py b/tensordict/_lazy.py
--- a/tensordict/_lazy.py
+++ b/tensordict/_lazy.py
@@ -19,6 +19,7 @@
 
 import numpy as np
 import torch
+import torch.distributed as dist
 from functorch import dim as ftdim
 from tensordict._td import _SubTensorDict, _TensorDictKeysView, TensorDict
 from tensordict._tensordict import _unravel_key_to_tuple, unravel_key_list
@@ -1407,9 +1408,15 @@ def any(self, dim: int = None) -> bool | TensorDictBase:
             )
         return any(value.any() for value in self.tensordicts)
 
-    def _send(self, dst: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
+    def _send(
+        self,
+        dst: int,
+        _tag: int = -1,
+        pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
+    ) -> int:
         for td in self.tensordicts:
-            _tag = td._send(dst, _tag=_tag, pseudo_rand=pseudo_rand)
+            _tag = td._send(dst, _tag=_tag, pseudo_rand=pseudo_rand, group=group)
         return _tag
 
     def _isend(
@@ -1418,6 +1425,7 @@ def _isend(
         _tag: int = -1,
         _futures: list[torch.Future] | None = None,
         pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
     ) -> int:
         if _futures is None:
             is_root = True
@@ -1425,15 +1433,23 @@ def _isend(
         else:
             is_root = False
         for td in self.tensordicts:
-            _tag = td._isend(dst, _tag=_tag, pseudo_rand=pseudo_rand, _futures=_futures)
+            _tag = td._isend(
+                dst, _tag=_tag, pseudo_rand=pseudo_rand, _futures=_futures, group=group
+            )
         if is_root:
             for future in _futures:
                 future.wait()
         return _tag
 
-    def _recv(self, src: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
+    def _recv(
+        self,
+        src: int,
+        _tag: int = -1,
+        pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
+    ) -> int:
         for td in self.tensordicts:
-            _tag = td._recv(src, _tag=_tag, pseudo_rand=pseudo_rand)
+            _tag = td._recv(src, _tag=_tag, pseudo_rand=pseudo_rand, group=group)
         return _tag
 
     def _irecv(
@@ -1443,6 +1459,7 @@ def _irecv(
         _tag: int = -1,
         _future_list: list[torch.Future] = None,
         pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
     ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None:
         root = False
         if _future_list is None:
@@ -1455,6 +1472,7 @@ def _irecv(
                 _tag=_tag,
                 _future_list=_future_list,
                 pseudo_rand=pseudo_rand,
+                group=group,
             )
 
         if not root:
diff --git a/tensordict/base.py b/tensordict/base.py
--- a/tensordict/base.py
+++ b/tensordict/base.py
@@ -2862,11 +2862,17 @@ def del_(self, key: NestedKey) -> T:
         ...
 
     # Distributed functionality
-    def gather_and_stack(self, dst: int) -> T | None:
+    def gather_and_stack(
+        self, dst: int, group: "dist.ProcessGroup" | None = None
+    ) -> T | None:
         """Gathers tensordicts from various workers and stacks them onto self in the destination worker.
 
         Args:
             dst (int): the rank of the destination worker where :func:`gather_and_stack` will be called.
+            group (torch.distributed.ProcessGroup, optional): if set, the specified process group
+                will be used for communication. Otherwise, the default process group
+                will be used.
+                Defaults to ``None``.
 
         Example:
             >>> from torch import multiprocessing as mp
@@ -2917,24 +2923,37 @@ def gather_and_stack(self, dst: int) -> T | None:
             ...     secondary_worker.join()
         """
         output = (
-            [None for _ in range(dist.get_world_size())]
-            if dst == dist.get_rank()
+            [None for _ in range(dist.get_world_size(group=group))]
+            if dst == dist.get_rank(group=group)
             else None
         )
-        dist.gather_object(self, output, dst=dst)
-        if dst == dist.get_rank():
+        dist.gather_object(self, output, dst=dst, group=group)
+        if dst == dist.get_rank(group=group):
             # remove self from output
             output = [item for i, item in enumerate(output) if i != dst]
             self.update(torch.stack(output, 0), inplace=True)
             return self
         return None
 
-    def send(self, dst: int, init_tag: int = 0, pseudo_rand: bool = False) -> None:
+    def send(
+        self,
+        dst: int,
+        *,
+        group: "dist.ProcessGroup" | None = None,
+        init_tag: int = 0,
+        pseudo_rand: bool = False,
+    ) -> None:  # noqa: D417
         """Sends the content of a tensordict to a distant worker.
 
         Args:
             dst (int): the rank of the destination worker where the content
                 should be sent.
+
+        Keyword Args:
+            group (torch.distributed.ProcessGroup, optional): if set, the specified process group
+                will be used for communication. Otherwise, the default process group
+                will be used.
+                Defaults to ``None``.
             init_tag (int): the initial tag to be used to mark the tensors.
                 Note that this will be incremented by as much as the number of
                 tensors contained in the TensorDict.
@@ -3003,15 +3022,21 @@ def send(self, dst: int, init_tag: int = 0, pseudo_rand: bool = False) -> None:
             ...     secondary_worker.join()
 
         """
-        self._send(dst, _tag=init_tag - 1, pseudo_rand=pseudo_rand)
+        self._send(dst, _tag=init_tag - 1, pseudo_rand=pseudo_rand, group=group)
 
-    def _send(self, dst: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
+    def _send(
+        self,
+        dst: int,
+        _tag: int = -1,
+        pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
+    ) -> int:
         for key in self.sorted_keys:
             value = self._get_str(key, NO_DEFAULT)
             if isinstance(value, Tensor):
                 pass
             elif _is_tensor_collection(value.__class__):
-                _tag = value._send(dst, _tag=_tag, pseudo_rand=pseudo_rand)
+                _tag = value._send(dst, _tag=_tag, pseudo_rand=pseudo_rand, group=group)
                 continue
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
@@ -3019,17 +3044,30 @@ def _send(self, dst: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
                 _tag += 1
             else:
                 _tag = int_generator(_tag + 1)
-            dist.send(value, dst=dst, tag=_tag)
+            dist.send(value, dst=dst, tag=_tag, group=group)
 
         return _tag
 
-    def recv(self, src: int, init_tag: int = 0, pseudo_rand: bool = False) -> int:
+    def recv(
+        self,
+        src: int,
+        *,
+        group: "dist.ProcessGroup" | None = None,
+        init_tag: int = 0,
+        pseudo_rand: bool = False,
+    ) -> int:  # noqa: D417
         """Receives the content of a tensordict and updates content with it.
 
         Check the example in the `send` method for context.
 
         Args:
             src (int): the rank of the source worker.
+
+        Keyword Args:
+            group (torch.distributed.ProcessGroup, optional): if set, the specified process group
+                will be used for communication. Otherwise, the default process group
+                will be used.
+                Defaults to ``None``.
             init_tag (int): the ``init_tag`` used by the source worker.
             pseudo_rand (bool): if True, the sequence of tags will be pseudo-
                 random, allowing to send multiple data from different nodes
@@ -3038,17 +3076,22 @@ def recv(self, src: int, init_tag: int = 0, pseudo_rand: bool = False) -> int:
                 slow down the runtime of your algorithm.
                 This value must match the one passed to :func:`send`.
                 Defaults to ``False``.
-
         """
-        return self._recv(src, _tag=init_tag - 1, pseudo_rand=pseudo_rand)
+        return self._recv(src, _tag=init_tag - 1, pseudo_rand=pseudo_rand, group=group)
 
-    def _recv(self, src: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
+    def _recv(
+        self,
+        src: int,
+        _tag: int = -1,
+        pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
+    ) -> int:
         for key in self.sorted_keys:
             value = self._get_str(key, NO_DEFAULT)
             if isinstance(value, Tensor):
                 pass
             elif _is_tensor_collection(value.__class__):
-                _tag = value._recv(src, _tag=_tag, pseudo_rand=pseudo_rand)
+                _tag = value._recv(src, _tag=_tag, pseudo_rand=pseudo_rand, group=group)
                 continue
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
@@ -3056,17 +3099,30 @@ def _recv(self, src: int, _tag: int = -1, pseudo_rand: bool = False) -> int:
                 _tag += 1
             else:
                 _tag = int_generator(_tag + 1)
-            dist.recv(value, src=src, tag=_tag)
+            dist.recv(value, src=src, tag=_tag, group=group)
             self._set_str(key, value, inplace=True, validated=True)
 
         return _tag
 
-    def isend(self, dst: int, init_tag: int = 0, pseudo_rand: bool = False) -> int:
+    def isend(
+        self,
+        dst: int,
+        *,
+        group: "dist.ProcessGroup" | None = None,
+        init_tag: int = 0,
+        pseudo_rand: bool = False,
+    ) -> int:  # noqa: D417
         """Sends the content of the tensordict asynchronously.
 
         Args:
             dst (int): the rank of the destination worker where the content
                 should be sent.
+
+        Keyword Args:
+            group (torch.distributed.ProcessGroup, optional): if set, the specified process group
+                will be used for communication. Otherwise, the default process group
+                will be used.
+                Defaults to ``None``.
             init_tag (int): the initial tag to be used to mark the tensors.
                 Note that this will be incremented by as much as the number of
                 tensors contained in the TensorDict.
@@ -3139,7 +3195,7 @@ def isend(self, dst: int, init_tag: int = 0, pseudo_rand: bool = False) -> int:
             ...     secondary_worker.join()
 
         """
-        return self._isend(dst, init_tag - 1, pseudo_rand=pseudo_rand)
+        return self._isend(dst, _tag=init_tag - 1, pseudo_rand=pseudo_rand, group=group)
 
     def _isend(
         self,
@@ -3147,6 +3203,7 @@ def _isend(
         _tag: int = -1,
         _futures: list[torch.Future] | None = None,
         pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
     ) -> int:
         root = False
         if _futures is None:
@@ -3156,7 +3213,11 @@ def _isend(
             value = self._get_str(key, NO_DEFAULT)
             if _is_tensor_collection(value.__class__):
                 _tag = value._isend(
-                    dst, _tag=_tag, pseudo_rand=pseudo_rand, _futures=_futures
+                    dst,
+                    _tag=_tag,
+                    pseudo_rand=pseudo_rand,
+                    _futures=_futures,
+                    group=group,
                 )
                 continue
             elif isinstance(value, Tensor):
@@ -3167,7 +3228,7 @@ def _isend(
                 _tag += 1
             else:
                 _tag = int_generator(_tag + 1)
-            _future = dist.isend(value, dst=dst, tag=_tag)
+            _future = dist.isend(value, dst=dst, tag=_tag, group=group)
             _futures.append(_future)
         if root:
             for _future in _futures:
@@ -3177,6 +3238,8 @@ def _isend(
     def irecv(
         self,
         src: int,
+        *,
+        group: "dist.ProcessGroup" | None = None,
         return_premature: bool = False,
         init_tag: int = 0,
         pseudo_rand: bool = False,
@@ -3187,6 +3250,12 @@ def irecv(
 
         Args:
             src (int): the rank of the source worker.
+
+        Keyword Args:
+            group (torch.distributed.ProcessGroup, optional): if set, the specified process group
+                will be used for communication. Otherwise, the default process group
+                will be used.
+                Defaults to ``None``.
             return_premature (bool): if ``True``, returns a list of futures to wait
                 upon until the tensordict is updated. Defaults to ``False``,
                 i.e. waits until update is completed withing the call.
@@ -3204,7 +3273,11 @@ def irecv(
                 upon until the tensordict is updated.
         """
         return self._irecv(
-            src, return_premature, _tag=init_tag - 1, pseudo_rand=pseudo_rand
+            src,
+            return_premature=return_premature,
+            _tag=init_tag - 1,
+            pseudo_rand=pseudo_rand,
+            group=group,
         )
 
     def _irecv(
@@ -3214,6 +3287,7 @@ def _irecv(
         _tag: int = -1,
         _future_list: list[torch.Future] = None,
         pseudo_rand: bool = False,
+        group: "dist.ProcessGroup" | None = None,
     ) -> tuple[int, list[torch.Future]] | list[torch.Future] | None:
         root = False
         if _future_list is None:
@@ -3228,6 +3302,7 @@ def _irecv(
                     _tag=_tag,
                     _future_list=_future_list,
                     pseudo_rand=pseudo_rand,
+                    group=group,
                 )
                 continue
             elif isinstance(value, Tensor):
@@ -3238,7 +3313,7 @@ def _irecv(
                 _tag += 1
             else:
                 _tag = int_generator(_tag + 1)
-            _future_list.append(dist.irecv(value, src=src, tag=_tag))
+            _future_list.append(dist.irecv(value, src=src, tag=_tag, group=group))
         if not root:
             return _tag, _future_list
         elif return_premature:
@@ -3248,13 +3323,20 @@ def _irecv(
                 future.wait()
             return
 
-    def reduce(self, dst, op=dist.ReduceOp.SUM, async_op=False, return_premature=False):
+    def reduce(
+        self,
+        dst,
+        op=dist.ReduceOp.SUM,
+        async_op=False,
+        return_premature=False,
+        group=None,
+    ):
         """Reduces the tensordict across all machines.
 
         Only the process with ``rank`` dst is going to receive the final result.
 
         """
-        return self._reduce(dst, op, async_op, return_premature)
+        return self._reduce(dst, op, async_op, return_premature, group=group)
 
     def _reduce(
         self,
@@ -3263,6 +3345,7 @@ def _reduce(
         async_op=False,
         return_premature=False,
         _future_list=None,
+        group=None,
     ):
         root = False
         if _future_list is None:
@@ -3282,7 +3365,9 @@ def _reduce(
                 pass
             else:
                 raise NotImplementedError(f"Type {type(value)} is not supported.")
-            _future_list.append(dist.reduce(value, dst=dst, op=op, async_op=async_op))
+            _future_list.append(
+                dist.reduce(value, dst=dst, op=op, async_op=async_op, group=group)
+            )
         if not root:
             return _future_list
         elif async_op and return_premature:
@@ -3332,6 +3417,8 @@ def apply(
                 unnamed inputs as the number of tensordicts, including self.
                 If other tensordicts have missing entries, a default value
                 can be passed through the ``default`` keyword argument.
+
+        Keyword Args:
             batch_size (sequence of int, optional): if provided,
                 the resulting TensorDict will have the desired batch_size.
                 The :obj:`batch_size` argument should match the batch_size after


Test patch:
diff --git a/test/test_distributed.py b/test/test_distributed.py
--- a/test/test_distributed.py
+++ b/test/test_distributed.py
@@ -16,7 +16,7 @@
 
 from packaging.version import parse
 
-from tensordict import MemmapTensor, TensorDict
+from tensordict import MemoryMappedTensor, TensorDict
 from torch import distributed as dist, multiprocessing as mp, nn
 from torch.distributed._tensor import (
     DeviceMesh,
@@ -215,7 +215,7 @@ def client(memmap_filename):
             {
                 ("a", "b"): torch.randn(2),
                 "c": torch.randn(2),
-                ("d", "e", "f"): MemmapTensor.from_tensor(
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(
                     torch.randn(2, 2), filename=memmap_filename
                 ),
             },
@@ -237,7 +237,7 @@ def server(queue):
                 {
                     ("a", "b"): torch.zeros(2),
                     "c": torch.zeros(2),
-                    ("d", "e", "f"): MemmapTensor.from_tensor(torch.zeros(2, 2)),
+                    ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.zeros(2, 2)),
                 },
                 [2],
             )
@@ -287,7 +287,7 @@ def client(memmap_filename, rank, op, async_op, return_premature):
             {
                 ("a", "b"): torch.ones(2),
                 "c": torch.ones(2),
-                ("d", "e", "f"): MemmapTensor.from_tensor(
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(
                     torch.ones(2, 2), filename=memmap_filename
                 ),
             },
@@ -310,7 +310,7 @@ def server(queue, op, async_op, return_premature):
                 {
                     ("a", "b"): torch.ones(2),
                     "c": torch.ones(2),
-                    ("d", "e", "f"): MemmapTensor.from_tensor(torch.ones(2, 2)),
+                    ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)),
                 },
                 [2],
             )
@@ -378,34 +378,40 @@ def make_td(ones):
         raise NotImplementedError
 
     @classmethod
-    def client(cls, pseudo_rand):
+    def client(cls, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=1,
             world_size=2,
             init_method="tcp://localhost:10017",
         )
+        if group is not None:
+            group = dist.new_group(group)
         td = cls.make_td(ones=True)
-        td.send(0, pseudo_rand=pseudo_rand)
+        td.send(0, pseudo_rand=pseudo_rand, group=group)
 
     @classmethod
-    def server(cls, queue, pseudo_rand):
+    def server(cls, queue, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=0,
             world_size=2,
             init_method="tcp://localhost:10017",
         )
+        if group is not None:
+            group = dist.new_group(group)
         td = cls.make_td(ones=False)
-        td.recv(1, pseudo_rand=pseudo_rand)
+        td.recv(1, pseudo_rand=pseudo_rand, group=group)
         assert (td == 1).all()
         queue.put("yuppie")
 
     @pytest.mark.flaky(reruns=5, reruns_delay=5)
-    def test_send(self, pseudo_rand, set_context):
+    @pytest.mark.parametrize("group", [[0, 1], None])
+    @pytest.mark.parametrize("pseudo_rand", [True, False])
+    def test_send(self, pseudo_rand, group, set_context):
         queue = mp.Queue(1)
-        main_worker = mp.Process(target=type(self).server, args=(queue, pseudo_rand))
-        secondary_worker = mp.Process(target=type(self).client, args=(pseudo_rand,))
+        main_worker = mp.Process(target=self.server, args=(queue, pseudo_rand, group))
+        secondary_worker = mp.Process(target=self.client, args=(pseudo_rand, group))
 
         main_worker.start()
         secondary_worker.start()
@@ -417,7 +423,6 @@ def test_send(self, pseudo_rand, set_context):
             secondary_worker.join()
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestSend(SendBase):
     """Test send for tensordict as root."""
 
@@ -431,7 +436,7 @@ def make_td(ones):
             {
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -439,7 +444,6 @@ def make_td(ones):
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestSendLazyStackRoot(SendBase):
     """Test send for lazy-stack as root."""
 
@@ -453,7 +457,7 @@ def make_td(ones):
             {
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -466,7 +470,6 @@ def make_td(ones):
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestSendLazyStackNest(SendBase):
     """Test send for tensordict as root with lazy stacked field."""
 
@@ -480,7 +483,7 @@ def make_td(ones):
             {
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -505,7 +508,7 @@ def make_td(ones):
         raise NotImplementedError
 
     @classmethod
-    def client(cls, pseudo_rand):
+    def client(cls, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=1,
@@ -513,10 +516,12 @@ def client(cls, pseudo_rand):
             init_method="tcp://localhost:10017",
         )
         td = cls.make_td(ones=True)
-        td.send(0, pseudo_rand=pseudo_rand)
+        if group is not None:
+            group = dist.new_group(group)
+        td.send(0, pseudo_rand=pseudo_rand, group=group)
 
     @classmethod
-    def server(cls, queue, return_premature, pseudo_rand):
+    def server(cls, queue, return_premature, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=0,
@@ -524,20 +529,28 @@ def server(cls, queue, return_premature, pseudo_rand):
             init_method="tcp://localhost:10017",
         )
         td = cls.make_td(ones=False)
-        out = td.irecv(1, return_premature=return_premature, pseudo_rand=pseudo_rand)
+        if group is not None:
+            group = dist.new_group(group)
+        out = td.irecv(
+            1, return_premature=return_premature, pseudo_rand=pseudo_rand, group=group
+        )
         if return_premature:
             for fut in out:
                 fut.wait()
         assert (td == 1).all()
         queue.put("yuppie")
 
+    @pytest.mark.parametrize("group", [None, [0, 1]])
+    @pytest.mark.parametrize("pseudo_rand", [True, False])
     @pytest.mark.parametrize("return_premature", [True, False])
-    def test_irecv(self, pseudo_rand, return_premature, set_context):
+    def test_irecv(self, pseudo_rand, return_premature, set_context, group):
         queue = mp.Queue(1)
         main_worker = mp.Process(
-            target=type(self).server, args=(queue, return_premature, pseudo_rand)
+            target=type(self).server, args=(queue, return_premature, pseudo_rand, group)
+        )
+        secondary_worker = mp.Process(
+            target=type(self).client, args=(pseudo_rand, group)
         )
-        secondary_worker = mp.Process(target=type(self).client, args=(pseudo_rand,))
 
         main_worker.start()
         secondary_worker.start()
@@ -549,7 +562,6 @@ def test_irecv(self, pseudo_rand, return_premature, set_context):
             secondary_worker.join()
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiRecv(iRecvBase):
     @staticmethod
     def make_td(ones):
@@ -562,14 +574,13 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiRecvLazyStackRoot(iRecvBase):
     @staticmethod
     def make_td(ones):
@@ -582,7 +593,7 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -593,7 +604,6 @@ def make_td(ones):
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiRecvLazyStackNest(iRecvBase):
     @staticmethod
     def make_td(ones):
@@ -606,7 +616,7 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -630,7 +640,7 @@ def make_td(ones):
         raise NotImplementedError
 
     @classmethod
-    def client(cls, pseudo_rand):
+    def client(cls, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=1,
@@ -639,10 +649,12 @@ def client(cls, pseudo_rand):
         )
 
         td = cls.make_td(True)
-        td.isend(0, pseudo_rand=pseudo_rand)
+        if group is not None:
+            group = dist.new_group(group)
+        td.isend(0, pseudo_rand=pseudo_rand, group=group)
 
     @classmethod
-    def server(cls, queue, pseudo_rand):
+    def server(cls, queue, pseudo_rand, group):
         torch.distributed.init_process_group(
             "gloo",
             rank=0,
@@ -650,15 +662,23 @@ def server(cls, queue, pseudo_rand):
             init_method="tcp://localhost:10017",
         )
         td = cls.make_td(False)
-        td.recv(1, pseudo_rand=pseudo_rand)
+        if group is not None:
+            group = dist.new_group(group)
+        td.recv(1, pseudo_rand=pseudo_rand, group=group)
         assert (td == 1).all()
         queue.put("yuppie")
 
+    @pytest.mark.parametrize("group", [[0, 1], None])
+    @pytest.mark.parametrize("pseudo_rand", [True, False])
     @pytest.mark.flaky(reruns=5, reruns_delay=5)
-    def test_isend(self, pseudo_rand, set_context):
+    def test_isend(self, pseudo_rand, set_context, group):
         queue = mp.Queue(1)
-        main_worker = mp.Process(target=type(self).server, args=(queue, pseudo_rand))
-        secondary_worker = mp.Process(target=type(self).client, args=(pseudo_rand,))
+        main_worker = mp.Process(
+            target=type(self).server, args=(queue, pseudo_rand, group)
+        )
+        secondary_worker = mp.Process(
+            target=type(self).client, args=(pseudo_rand, group)
+        )
 
         main_worker.start()
         secondary_worker.start()
@@ -673,7 +693,6 @@ def test_isend(self, pseudo_rand, set_context):
             secondary_worker.join()
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiSend(iSendBase):
     @staticmethod
     def make_td(ones):
@@ -686,14 +705,13 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiSendLazyStackRoot(iSendBase):
     @staticmethod
     def make_td(ones):
@@ -706,7 +724,7 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )
@@ -717,7 +735,6 @@ def make_td(ones):
         return td
 
 
-@pytest.mark.parametrize("pseudo_rand", [True, False])
 class TestiSendLazyStackNest(iSendBase):
     @staticmethod
     def make_td(ones):
@@ -730,7 +747,7 @@ def make_td(ones):
                 ("a", "b"): fun(2),
                 "c": fun(2, 3),
                 "_": fun(2, 1, 5),
-                ("d", "e", "f"): MemmapTensor.from_tensor(fun(2, 2)),
+                ("d", "e", "f"): MemoryMappedTensor.from_tensor(fun(2, 2)),
             },
             [2],
         )