ddp_weakref is not set during the backward pass

[lokesh-vr-17773]

Training runs successfully on Torch 2.0 and PyTorch Lightning 2.0 with the same model architecture and dataset. However, using the same configuration and dataset on Torch 2.2 and Lightning 2.2 results in the following error: the training crashes during the 67th epoch. With Torch 2.0, training completes all 150 epochs without issues.

    Trainer.train_from_cfg(cfg)
  File "/home/lokesh/zoho/zspeech_premix/zspeech/zspeech_dev/zspeech/training/trainer/_trainer.py", line 331, in train_from_cfg
    raise ex
  File "/home/lokesh/zoho/zspeech_premix/zspeech/zspeech_dev/zspeech/training/trainer/_trainer.py", line 329, in train_from_cfg
    trainer.train(experiment, data)
  File "/home/lokesh/zoho/zspeech_premix/zspeech/zspeech_dev/zspeech/training/trainer/_trainer.py", line 177, in train
    self.trainer.fit(model=experiment, datamodule=data, ckpt_path=self.ckpt_path)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 543, in fit
    call._call_and_handle_interrupt(
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 43, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
    return function(*args, **kwargs)                                                                                                                                                                                                                                 [690/1733]
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 579, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 986, in _run
    results = self._run_stage()
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1032, in _run_stage
    self.fit_loop.run()
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
    self.advance()
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 138, in run
    self.advance(data_fetcher)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 242, in advance
    batch_output = self.automatic_optimization.run(trainer.optimizers[0], batch_idx, kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 191, in run
    self._optimizer_step(batch_idx, closure)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 269, in _optimizer_step
    call._call_lightning_module_hook(
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 157, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1303, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 152, in step
    step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 270, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 239, in optimizer_step
    return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 80, in optimizer_step
    closure_result = closure()
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 144, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 138, in closure
    self._backward_fn(step_output.closure_loss)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 240, in backward_fn
    call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 213, in backward
    self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision.py", line 72, in backward
    model.backward(tensor, *args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1090, in backward
    loss.backward(*args, **kwargs)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward
    torch.autograd.backward(
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/torch/autograd/__init__.py", line 266, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/torch/autograd/function.py", line 289, in apply
    return user_fn(self, *args)
  File "/home/lokesh/external/miniconda3/envs/zspeech_premix/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 255, in backward
    reducer = ddp_weakref.reducer
AttributeError: 'NoneType' object has no attribute 'reducer'

env info for 2.0 run:

lightning-utilities           0.11.8
numpy                         1.23.5
nvidia-cuda-cupti-cu11        11.7.101
nvidia-cuda-nvrtc-cu11        11.7.99
nvidia-cuda-runtime-cu11      11.7.99
nvidia-cudnn-cu11             8.5.0.96
pytorch-lightning             2.0.0
torch                         2.0.0+cu117
torch_poly_lr_decay           0.0.1
torchaudio                    2.0.0+cu117
torchdata                     0.7.1
torchmetrics                  1.5.1

env info for 2.2 run:

lightning-utilities           0.11.8
numpy                         1.23.5
nvidia-cuda-cupti-cu11        11.7.101
nvidia-cuda-cupti-cu12        12.1.105
nvidia-cuda-nvrtc-cu11        11.7.99
nvidia-cuda-nvrtc-cu12        12.1.105
nvidia-cuda-runtime-cu11      11.7.99
nvidia-cuda-runtime-cu12      12.1.105
nvidia-cudnn-cu11             8.5.0.96
nvidia-cudnn-cu12             8.9.2.26
pytorch-lightning             2.2.0
torch                         2.2.1+cu121
torch_poly_lr_decay           0.0.1
torchaudio                    2.2.1+cu121
torchdata                     0.7.1
torchmetrics                  1.5.1