IIUC, what we’re looking for here is pipeline parallelism. PyTorch currently doesn’t have native support for pipeline parallelism. There are a few projects that have built something similar on top of PyTorch: https://github.com/kakaobrain/torchgpipe 2 and https://github.com/msr-fiddle/pipedream 5. You could also use the Distributed RPC Framework 1 to build something like this.