Repository: pytorch
Issue: [Design][RFC] RemoteModule API Design
Labels: oncall: jit,triaged,module: rpc

# [Design][RFC] RemoteModule API Design

FB internal link, [[Design] RemoteModule](https://fb.quip.com/Rb5bAKee8TQ4)

## Goal

* Provide a wrapper `nn.Module` API that serves these purposes
   * Users should feel like using remote module is like using a local module. RPCs are done transparently for users.
   * Users can look up if a the parameters and modules that sit on the remote side rather than doing it in an ad-hoc way. (Internal diff that does it in an ad-hoc way, D19954520)
   * RemoteModule should be a RRef + helper functions container. It should be able to be copied and sent over RPC. Here is an example use case. Master creates a list of EmbeddingBag Modules on a set of Embedding Servers and refers to them using RemoteModules on Master. Master would want to share this list of RemoteModules to all the Trainers.
* This API as a modeling primitive should be TorchScript compatible.
   * The reason is that RPC-supported distributed training (Model-parallel) largely leverages TorchScript to remove Python dependency to avoid Python GIL contention overhead.
* This API need to be able to integrated with autograd profiler
   * In distributed model parallel training, how to place remote modules and parameters is better a result of profiling results. Since uses of this API are marks of local/remote pairs in distributed model parallel training, it is naturally the good hooks for profiler.
   * A RemoteModule should hold a global unique name for identifying an instance among all remote modules.

## RemoteModule API Example

```
class MyModule(nn.Module):
    pass


# Local side is aware of this, Remote side runs this.
def module_creator(args, kwargs) -> ModuleInterfaceXYZ: (An important type var)
    my_module = MyModule(*args, **kwargs)
    scripted_my_moduel = torch.jit.script(my_module)
    return scripted_my_moduel


to = "worker1" 
my_module_remote: RemoteModuleInterfaceXYZ = RemoteModule(
    to, module_creator, args=(), kwargs={}
)

parameter_rrefs: List[RRef[Tensor]] = my_module_remote.remote_prameters()

ret_fut = my_module_remote.forward_async()
ret = ret_fut.wait()

ret= my_module_remote.forward()
```
## nn.Module New API

* Add API to nn.Module
   * nn.Module.**remote_modules**()
      * It’s aware of how to get RemoteModules.
   * nn.Module.**remote_parameters**()
      * It gets all Parameter RRefs from all sub-RemoteModules.
* Note, for the first phase of this project, we don’t want to expose it in PyTorch code API nn.Module. Instead we provide helper functions that serve the same purpose as experimental API under RPC module.

## Challenges imposed on scripting

* RemoteModule initialization should allowed be done in Python. But when it’s done in TorchScript or the Python RemoteModule object got scripted, need to do the following type inference and generation.
* A requirement: Type of the contained ScriptModule must be known on both sides. Expect it’s provided on the “module_creater” return type by user. Check it’s provided by the user. See “module_creator” return type “ModuleInterfaceXYZ” in the above example.
* Generate a “forward” method and necessary schema/types on both local side and remote side, see “LocalForwardTemplate” and “RemoteForwardTemplate” below. Notice JIT Compilation Unit lifetime management and thread safety, because when local side initiates type registration on remote side, remote side might be compiling other global scope ScriptFunctions initiated from importing.
* LocalForwardTemplate doesn’t do computation, but it initiates an RPC. It needs to know how to refer to the RemoteForwardTemplate instance.
* RemoteForwardTemplate is a dynamically-generated RRef method call helper. It’s signature depends on the “forward” method schema of the custom ScriptModule, “ModuleInterfaceXYZ”, which is constructed in the initialization phase above. 

## Template

Need to have 2 TorchScript templates 

* LocalForwardTemplate

```
def local_forward_async(self: RemoteModuleInterfaceXYZ, ...args_0) -> Future[...args_1]:
    return rpc.rpc_async(
        to,
        remote_forward,
        args,
        kwargs,
        )

def local_foward(self: RemoteModuleInterfaceXYZ, ...args_0) -> ...args_1:
    fut = local_forward_async(...)
    return fut.wait()
```
* RemoteForwardTemplate

```
def remote_forward(self_rref: RRef[ModuleInterfaceXYZ], ...args_0) -> ...args_1:
    self: ModuleInterfaceXYZ = self_rref.local_value()
    return self.forward(...args_0)
```
## Working Items

* ~~Get ModuleInterface type from module_creator. Done~~
* ~~Get type strs from ModuleInterface. Done~~
* ~~Instantiate Python code template with type strs inserted. Done~~
* ~~Instantiate template as generated output files. Done~~
* ~~Test scripting the Remote Module. Done~~
* ~~Test 3 ways of constructing RemoteModule. Done~~
* ~~Test forward_async. Done~~
* ~~Test forward_sync. Done~~
* ~~Test foward with kwargs. Done~~
* ~~Setup OSS tests. Done~~
* ~~Resolve ModuleImportError when import CodeTemplate class. Done~~
* ~~Resolve FileNotFoundError seeing after python setup.py install. Done~~
* ~~Cleaning up generated module files to avoid torch package size exploding. Done~~
* ~~Added docstrings and examples to explain how to use the API. Done~~

cc @suo @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar