@Vibhatha_Abeykoon @mrshenli
I am reviewing this topic today and I have a few suggestions.
I will use “process/device” (Eg: “worker:0/cuda:0”) as a location descriptor of a tensor, and tensor is the only holder of all data.
The first thing is:



 Vibhatha_Abeykoon:

How to partition the model


I have a primitive design for this purpose: an assigner and a simple wrapper, assigner currently won’t partition a model auto matically, instead it will just assign user specified partitions based on a series of heuristics (GPU mem, GPU power, CPU mem, CPU power, Model complexity, bandwidth of connection between models), but it could be also reworked to fulfill your purpose, and you just need to wrap all of your submodules in the wrapper, the wrapper just stores input/output location descritors, nothing more.
Simply speaking, partitioning just require users to specify the input and output (process/device) requirements for a module/model, and then a smart assigner to assign partitions to process/device.
Theoretically a dynamic profiler is much better than a static heuristic assigner, since it actively detects hotspots and try to even the load on all of your nodes, but this introduces additional issues: Is evening the load across nodes increasing the performance? How much cost have the additional transmission introduced? Is the decreasing the load not pushing your GPUs to their full capacity (kernel launching cost should be considered)? So there are possibilities that this solution does not meet with the theoretical standard.
I beilieve @mrshenli has studied about this issue, from his profile page. I think many more experiments are needed to determine the best scheme.
The second thing is:



 Vibhatha_Abeykoon:

I have a few suggestions, if we can make use of to(device) call into to(machine:device) kind of an API endpoint, it will be much easier to work, but I am not quite sure how the changes should reflect internally.


I think it won’t be too defficult if rpc.pair in #41546 2 is implemented,
tensor.to("worker:1/cuda:0")

is equivalent to
# suppose there is a tensor on process "worker:0" and device "cuda:0" of this process
# move to process "worker:1" and device "cuda:1" of that process
# take care when torch.cuda.set_device is used

def pair_and_move_to(tensor, device, uuid):
    # uuid should be a unique identifier to identify this tensor
    # could be process_name:tensor_ptr
    tensor = tensor.to(device)
    rpc.pair(uuid, tensor)
    return RRef(tensor)

# on worker:0 when .to is invoked:
rpc.rpc_sync("worker:1", rpc.pair, args=(tensor, "cuda:1", tensor.some_uuid))

And for implementing Distributed Model Parallel Using Distributed RPC, there are many model parallel methods, it depends on your model, algorithm framework and application. For DDP compatible RPC solutions specifically, I have implementations of a gradient reduction server in my framework, which should be able to do the exact same thing as DDP does, however, this server implementation is also based on new API RFC #41546 2, which haven’t been implemented in pytorch. I have made a simple wrapper upon current primitive RPC APIs for this RFC, it is not efficient since two primitive RPC requests have to be made per wrapped high level RPC API, but it is tested and workable, if you would like to take a look.
From my personal point of view, if torch could provide a way to “expose” a resource or service upon the RPC module, even RemoteModule could be easily implemented, since it is basically create a module on a remote process, then expose its “__call__()” method as a service on the global scope,  #41546 2 could solve this problem.
Summary:
You can achieve all functions, using current torch APIs, if you don’t mind 20% ~ 30% efficiency loss, and spend a little?(much) time to construct your wheel, if you don’t, you can also use mine . It would be definetly better if torch could just provide these functions, with more optimizations.
And @mrshenli @Kiuk_Chung, please chime in and offer some feedback and precious suggestions on  #41546 2, there are some torchelastic issues I am not very familiar with and need some help