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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
#!/usr/bin/python3
def get_remote_module_template(enable_moving_cpu_tensors_to_cuda: bool):
return _TEMPLATE_PREFIX + (
_REMOTE_FORWARD_TEMPLATE_ENABLE_MOVING_CPU_TENSORS_TO_CUDA
if enable_moving_cpu_tensors_to_cuda
else _REMOTE_FORWARD_TEMPLATE
)
_TEMPLATE_PREFIX = """from typing import *
import torch
import torch.distributed.rpc as rpc
from torch import Tensor
from torch._jit_internal import Future
from torch.distributed.rpc import RRef
from typing import Tuple # pyre-ignore: unused import
{assign_module_interface_cls}
def forward_async(self, {arg_types}){arrow_and_future_return_type}:
args = (self.module_rref, self.device, self.is_device_map_set, {args})
kwargs = {{{kwargs}}}
return rpc.rpc_async(
self.module_rref.owner(),
_remote_forward,
args,
kwargs,
)
def forward(self, {arg_types}){arrow_and_return_type}:
args = (self.module_rref, self.device, self.is_device_map_set, {args})
kwargs = {{{kwargs}}}
ret_fut = rpc.rpc_async(
self.module_rref.owner(),
_remote_forward,
args,
kwargs,
)
return ret_fut.wait()
_generated_methods = [
forward_async,
forward,
]
{jit_script_decorator}
"""
# This template may cause typing error (the mismatch between ``Tuple[()]`` and ``Tuple[Any]``)
# even if the code is only used for instaniation but not execution.
# Therefore, only include handling moving CPU tensors to a cuda device if necessary.
# TODO: Merge these two templates together in the future once TorchScript syntax is improved.
_REMOTE_FORWARD_TEMPLATE_ENABLE_MOVING_CPU_TENSORS_TO_CUDA = """
def _remote_forward(
module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, {arg_types}){arrow_and_return_type}:
module = module_rref.local_value()
device = torch.device(device)
if device.type != "cuda":
return module.forward({args}, {kwargs})
# If the module is on a cuda device,
# move any CPU tensor in args or kwargs to the same cuda device.
# Since torch script does not support generator expression,
# have to use concatenation instead of
# ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``.
args = ({args},)
out_args: Tuple[()] = ()
for arg in args:
arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,)
out_args = out_args + arg
kwargs = {{{kwargs}}}
for k, v in kwargs.items():
if isinstance(v, Tensor):
kwargs[k] = kwargs[k].to(device)
if is_device_map_set:
return module.forward(*out_args, {kwargs})
# If the device map is empty, then only CPU tensors are allowed to send over wire,
# so have to move any GPU tensor to CPU in the output.
# Since torch script does not support generator expression,
# have to use concatenation instead of
# ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, {kwargs}))``.
ret: Tuple[()] = ()
for i in module.forward(*out_args, {kwargs}):
i = (i.cpu(),) if isinstance(i, Tensor) else (i,)
ret = ret + i
return ret
"""
_REMOTE_FORWARD_TEMPLATE = """
def _remote_forward(
module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, {arg_types}){arrow_and_return_type}:
module = module_rref.local_value()
return module.forward({args}, {kwargs})
"""
|
