Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Support group specification #616

Merged
merged 4 commits into from
Jan 11, 2024
Merged

[Feature] Support group specification #616

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fa85ab7
[feat] support group specification
Jan 11, 2024
7c53d62
add tests
vmoens Jan 11, 2024
ccc06ca
fix tests
vmoens Jan 11, 2024
5d3787f
amend
vmoens Jan 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 23 additions & 5 deletions tensordict/_lazy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -1418,22 +1425,31 @@ 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
_futures = []
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(
Expand All @@ -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:
Expand All @@ -1455,6 +1472,7 @@ def _irecv(
_tag=_tag,
_future_list=_future_list,
pseudo_rand=pseudo_rand,
group=group,
)

if not root:
Expand Down
Loading
Loading