1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
|
# mypy: allow-untyped-defs
# Copyright (c) Meta Platforms, Inc. and affiliates
from contextlib import contextmanager
from torch.distributed.tensor._api import DTensor
from torch.distributed.tensor.experimental._attention import context_parallel
from torch.distributed.tensor.experimental._func_map import local_map
from torch.distributed.tensor.experimental._register_sharding import register_sharding
__all__ = ["context_parallel", "implicit_replication", "local_map", "register_sharding"]
@contextmanager
def implicit_replication():
"""
This context manager allows :class:`DTensor` to implicitly treat all non-DTensors (``torch.Tensor``)
in the program be replicate :class:`DTensor` s during the operator computation.
.. warning:: This might possible lead to incorrect results if ``torch.Tensor`` s are not replicated
in practice, please use it at your discretion.
"""
try:
DTensor._op_dispatcher._allow_implicit_replication = True
yield
finally:
DTensor._op_dispatcher._allow_implicit_replication = False
# Set namespace for exposed private names
context_parallel.__module__ = "torch.distributed.tensor.experimental"
implicit_replication.__module__ = "torch.distributed.tensor.experimental"
local_map.__module__ = "torch.distributed.tensor.experimental"
register_sharding.__module__ = "torch.distributed.tensor.experimental"
|
