File size: 2,355 Bytes
9b96ca6 |
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 |
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
module_interface_cls = None
def forward_async(self, *args, **kwargs):
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, *args, **kwargs):
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,
]
def _remote_forward(
module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs):
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
|