Repository: pytorch
Issue: Internal Assert During Distributed Autograd Backprop
Labels: oncall: distributed

### 🐛 Describe the bug

#### Problem
An INTERNAL ASSERT error occurs during the backward call of distributed autograd when doing remote calculations via torch.distributed.rpc. Notes:
1. Specifically, the remote calculation is a single call to `torchode` to solve a trivial differential equation.
1. The precise time to fail varies from 5-60 seconds. The failure iteration number varies widely, even with a fixed random seed. 
1. It fails 100% of the time if multiple calculation nodes are used (i.e. world_size > 2), and never fails if only 1 calculation node is used (world_size=2).
1. The failure occurs regardless of whether the workers are all run on the same machine or different machines
1. It works fine with multiple clients if I replace the `torchode` call with a dummy calculation.
1. It is not clear to me whether this is really a distributed autograd issue or a downstream consequence of a problem with torchode, but the message requested that I report it :)

#### Source code
The following standalone program demonstrates the failure with a contrived case:
```python
import os, torch, torchode
import torch.distributed
import torch.distributed.nn
import torch.distributed.optim
import torch.distributed.rpc
import torch.multiprocessing


def SolveRPC(model, times):
    return model.forward(times)


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.r = torch.nn.Parameter(torch.randn(3))

    def forward(self, times):
        def RHS(t, Y):
            return -(self.r**2) * Y
        Y0 = torch.ones(5, 3).double()
        solns = torchode.solve_ivp(RHS, Y0, t_eval=times)
        return solns.ys


def RunClient(worker_id, num_clients):
    torch.distributed.rpc.init_rpc(f"worker{worker_id}", rank=worker_id,
                                   world_size=num_clients+1)
    torch.distributed.rpc.shutdown()


def RunServer(num_clients):
    world_size = num_clients + 1
    torch.distributed.rpc.init_rpc('worker0', rank=0, world_size=world_size)
    workers = [f'worker{i}' for i in range(1, world_size)]

    # setup remote model and optimizer
    remote_model = torch.distributed.nn.RemoteModule('worker0/cpu', Model, ())
    par_refs = remote_model.remote_parameters()
    opt = torch.distributed.optim.DistributedOptimizer(
       torch.optim.Adam, par_refs, lr=1e-4)

    for i in range(1000000):
        with torch.distributed.autograd.context() as grad_context:
            # make RPC solve calls
            times = torch.arange(10).double()
            task_futures= [
                torch.distributed.rpc.rpc_async(worker, SolveRPC,
                                                args=(remote_model, times))
                for worker in workers
            ]

            # collect results
            solutions = torch.futures.wait_all(task_futures)
            Y_hat = torch.stack(solutions)

            # score and update pars
            loss = (Y_hat**2).mean()
            torch.distributed.autograd.backward(grad_context, [loss])
            opt.step(grad_context)
            print(i, loss.item())
    torch.distributed.rpc.shutdown()


def Launch(node_id, num_clients):
    torch.manual_seed(3)
    if node_id == 0:
        RunServer(num_clients)
    else:
        RunClient(node_id, num_clients)


def main():
    #num_clients = 1 # works
    num_clients = 2  # fails

    os.environ['MASTER_PORT'] = '65001'
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    torch.multiprocessing.spawn(Launch, nprocs=num_clients+1,
                                args=(num_clients,))


if __name__ == '__main__':
    main()
```

#### Traceback
```
RuntimeError: Error on Node 0: res[i].defined() INTERNAL ASSERT FAILED at "../torch/csrc/autograd/functions/tensor.cpp":111, please report a bug to PyTorch. 
Exception raised from apply at ../torch/csrc/autograd/functions/tensor.cpp:111 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7f58ad1a2457 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x68 (0x7f58ad16c4b5 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #2: torch::autograd::CopySlices::apply(std::vector<at::Tensor, std::allocator<at::Tensor> >&&) + 0x86f (0x7f58de04a64f in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #3: <unknown function> + 0x48689cb (0x7f58de0409cb in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #4: torch::autograd::Engine::evaluate_function(std::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptr<torch::autograd::ReadyQueue> const&) + 0x1638 (0x7f58de03a2b8 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #5: torch::distributed::autograd::DistEngine::execute_graph_task_until_ready_queue_empty(torch::autograd::NodeTask&&, bool) + 0x434 (0x7f58dea774a4 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #6: <unknown function> + 0x52a02e0 (0x7f58dea782e0 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #7: <unknown function> + 0x196e271 (0x7f58db146271 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so)
frame #8: c10::ThreadPool::main_loop(unsigned long) + 0x285 (0x7f58ad193d85 in /home/foo/venv/lib/python3.10/site-packages/torch/lib/libc10.so)
frame #9: <unknown function> + 0xd6de4 (0x7f59075d9de4 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #10: <unknown function> + 0x8609 (0x7f5925c3e609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #11: clone + 0x43 (0x7f5925d78133 in /lib/x86_64-linux-gnu/libc.so.6)
```

### Versions

Versions of relevant libraries:
[pip3] functorch==1.13.1
[pip3] numpy==1.23.4
[pip3] torch==1.13.1
[pip3] torchode==0.1.6
[pip3] torchtyping==0.1.4


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu