multi-node training runs crash because ddp_weakref is None during backward

[mishooax]

(also reported here: #20706)

Multi-node / multi-GPU training fails mid-way through because ddp_weakref is not being set correctly during the backward pass. This appears to be similar to the issue reported in #20390. I was unable to reproduce this with a small model. Also the exact moment it fails (epoch/step) can vary between training runs. Any ideas? 🙏

[rank13]: Traceback (most recent call last):
[rank13]:              ^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/hydra/main.py", line 94, in decorated_main
[rank13]:     _run_hydra(
[rank13]:   File "/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
[rank13]:     _run_app(
[rank13]:   File "/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457, in _run_app
[rank13]:     run_and_report(
[rank13]:   File "/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
[rank13]:     raise ex
[rank13]:   File "/lib/python3.11/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
[rank13]:   File "/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
[rank13]:     lambda: hydra.run(
[rank13]:     _ = ret.return_value
[rank13]:     DOPTrainer(config).train()
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 539, in fit
[rank13]:     call._call_and_handle_interrupt(
[rank13]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank13]:     return function(*args, **kwargs)
[rank13]:     self._run(model, ckpt_path=ckpt_path)
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 982, in _run
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/trainer/trainer.py", line 1026, in _run_stage
[rank13]:     self.fit_loop.run()
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/loops/fit_loop.py", line 216, in run
[rank13]:     self.advance(data_fetcher)
[rank13]:     self._optimizer_step(batch_idx, closure)
[rank13]:     output = fn(*args, **kwargs)
[rank13]:              ^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/core/optimizer.py", line 154, in step
[rank13]:     step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
[rank13]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/strategies/ddp.py", line 270, in optimizer_step
[rank13]:     optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
[rank13]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank13]:     return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
[rank13]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 146, in __call__
[rank13]:     self._result = self.closure(*args, **kwargs)
[rank13]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank13]:     return func(*args, **kwargs)
[rank13]:            ^^^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in closure
[rank13]:     self._backward_fn(step_output.closure_loss)
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 241, in backward_fn
[rank13]:     call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/trainer/call.py", line 323, in _call_strategy_hook
[rank13]:     output = fn(*args, **kwargs)
[rank13]:              ^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/strategies/strategy.py", line 213, in backward
[rank13]:     self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/plugins/precision/precision.py", line 73, in backward
[rank13]:     model.backward(tensor, *args, **kwargs)
[rank13]:   File "/lib/python3.11/site-packages/pytorch_lightning/core/module.py", line 1097, in backward
[rank13]:     loss.backward(*args, **kwargs)
[rank13]:   File "/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
[rank13]:     torch.autograd.backward(
[rank13]:   File "/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank13]:     _engine_run_backward(
[rank13]:   File "/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
[rank13]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank13]:   File "/lib/python3.11/site-packages/torch/autograd/function.py", line 307, in apply
[rank13]:     return user_fn(self, *args)
[rank13]:            ^^^^^^^^^^^^^^^^^^^^
[rank13]:   File "/lib/python3.11/site-packages/torch/nn/parallel/distributed.py", line 260, in backward
[rank13]:     reducer = ddp_weakref.reducer
[rank13]:               ^^^^^^^^^^^^^^^^^^^
[rank13]: AttributeError: 'NoneType' object has no attribute 'reducer'
[rank12]:[E410 06:07:50.769590035 ProcessGroupNCCL.cpp:629] [Rank 4] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=253173, OpType=ALLGATHER, NumelIn=91356, NumelOut=730848, Timeout(ms)=600000) ran for 600009 milliseconds before timing out.