Deadlock when manually logging from on_train_epoch_end

Nobody has claimed this yet.

### Bug description

When using DDP and manually logging a TorchMetric by calling `metric.update` inside of `training_step` and then `metric.compute` inside of `on_training_epoch_end`, training appears to end up in a deadlock and eventually crashes with NCCL timeout errors.

This sort of logging setup is desirable when a TorchMetric returns a non-scalar data structure, i.e., a confusion matrix or dictionary, and some further processing is required at the end of the epoch before logging with `self.log`.

A minimal example is attached. Deadlock only occurs when using DDP, multiple GPUs, and also logging the training loss with `on_epoch`.

### What version are you seeing the problem on?

v2.1, v2.2

### How to reproduce the bug

```python
import os, sys

import torch
from torchmetrics import MeanSquaredError
from pytorch_lightning import LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset

class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)

def __getitem__(self, index):
return self.data[index]

def __len__(self):
return self.len

class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(32, 32)

# Use MSE for simple demo but this might be a complex
# metric that doesn't return a scalar.
self.metric = MeanSquaredError()

def forward(self, x):
return self.layer(x)

def training_step(self, batch, batch_idx):
targs = batch
preds = self(batch)
loss = torch.mean((preds - targs)**2)

# Log the loss at each step
self.log('train_loss/step', loss, on_step=True, on_epoch=False)

# Log the loss averaged over each epoch.
# IF THIS IS COMMENTED OUT, EVERYTHING WORKS!!
self.log('train_loss', loss, on_step=False, on_epoch=True, sync_dist=True)

# Update our custom metric. We don't want to just log it directly with
# `on_epoch=True` because it returns a complex data structure, e.g.,
# a confusion matrix or a dictionary of values that need processing at the
# end of the epoch. Instead, we want to update it and later handle the
# logging in `on_train_epoch_end`. That way, we can inject cutom processing
# at the end of the epoch.
self.metric.update(preds, targs)

return {'loss': loss, 'preds': preds, 'targs': targs}

def on_train_epoch_end(self):
print(f'logging metrics for rank `{os.environ.get("LOCAL_RANK", 0)}`.')
sys.stdout.flush()

# At the end of the epoch, we compute the metric and then reset it.
metric_value = self.metric.compute()
self.metric.reset()

# We can do sophisticated steps here to convert `metric_value` into something
# that can be logged.
self.log('metric', metric_value, on_step=False, on_epoch=True, sync_dist=True)

def validation_step(self, batch, batch_idx):
targs = batch
preds = self(batch)
loss = torch.mean((preds - targs)**2)

self.log('valid_loss', loss, on_step=False, on_epoch=True, sync_dist=True)

def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.01)

def run():
train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
model = BoringModel()

trainer = Trainer(
default_root_dir=os.getcwd(),
limit_train_batches=5,
limit_val_batches=5,
log_every_n_steps=2,
num_sanity_val_steps=0,
max_epochs=5,
enable_model_summary=False,
strategy='ddp', # Everything works fine when using a single GPU
accelerator='gpu',
devices=[0, 1],
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

if __name__ == '__main__':
run()
```

### Error messages and logs

This is the output, it hangs for a very long time before actually segfaulting.
```
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
/home/eforney/venv/lightning_bug/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
You are using a CUDA device ('NVIDIA RTX A6000') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/2
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 2 processes
----------------------------------------------------------------------------------------------------

LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1]
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
/home/eforney/venv/lightning_bug/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:441: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=9` in the `DataLoader` to improve performance.
/home/eforney/venv/lightning_bug/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:441: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=9` in the `DataLoader` to improve performance.
Epoch 0: 100%|█████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 28.00it/s, v_num=23logging metrics for rank `1`.%|██████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 440.52it/s]
Epoch 0: 100%|█████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 20.92it/s, v_num=23][rank0]:[E ProcessGroupNCCL.cpp:523] [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=33, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800562 milliseconds before timing out.
[rank0]:[E ProcessGroupNCCL.cpp:537] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank0]:[E ProcessGroupNCCL.cpp:543] To avoid data inconsistency, we are taking the entire process down.
[rank0]:[E ProcessGroupNCCL.cpp:1182] [Rank 0] NCCL watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=33, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800562 milliseconds before timing out.
Exception raised from checkTimeout at ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:525 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fe4691ffd87 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional > >) + 0x1e6 (0x7fe46a3a76e6 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::workCleanupLoop() + 0x19d (0x7fe46a3aac3d in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x119 (0x7fe46a3ab839 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #4: + 0xd6df4 (0x7fe4b40bbdf4 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #5: + 0x8609 (0x7fe4b5264609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #6: clone + 0x43 (0x7fe4b539e353 in /lib/x86_64-linux-gnu/libc.so.6)

terminate called after throwing an instance of 'c10::DistBackendError'
what(): [Rank 0] NCCL watchdog thread terminated with exception: [Rank 0] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=33, OpType=ALLREDUCE, NumelIn=1, NumelOut=1, Timeout(ms)=1800000) ran for 1800562 milliseconds before timing out.
Exception raised from checkTimeout at ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:525 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fe4691ffd87 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional > >) + 0x1e6 (0x7fe46a3a76e6 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::workCleanupLoop() + 0x19d (0x7fe46a3aac3d in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x119 (0x7fe46a3ab839 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #4: + 0xd6df4 (0x7fe4b40bbdf4 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #5: + 0x8609 (0x7fe4b5264609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #6: clone + 0x43 (0x7fe4b539e353 in /lib/x86_64-linux-gnu/libc.so.6)

Exception raised from ncclCommWatchdog at ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1186 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x57 (0x7fe4691ffd87 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libc10.so)
frame #1: + 0xdf6b11 (0x7fe46a101b11 in /home/eforney/venv/lightning_bug/lib/python3.8/site-packages/torch/lib/libtorch_cuda.so)
frame #2: + 0xd6df4 (0x7fe4b40bbdf4 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #3: + 0x8609 (0x7fe4b5264609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #4: clone + 0x43 (0x7fe4b539e353 in /lib/x86_64-linux-gnu/libc.so.6)

Aborted (core dumped)
```

### Environment

Current environment

```
- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer / LightningModule
- PyTorch Lightning Version (e.g., 1.5.0): 2.2.1
- PyTorch Version (e.g., 2.0): 2.2.1
- Python version (e.g., 3.9): 3.8.10
- OS (e.g., Linux): Linux, Ubuntu 20.04.6 LTS
- CUDA/cuDNN version: 12.2
- GPU models and configuration: 2x RTX A6000
- How you installed Lightning(`conda`, `pip`, source): pip
- Running environment of LightningApp (e.g. local, cloud): local
```

### More info

_No response_

Start with the attached minimal BoringModel example, focusing on training_step, on_train_epoch_end, TorchMetric.update/compute, and self.log with on_epoch and sync_dist under DDP. Reproduce with two GPUs and the stated Lightning/PyTorch versions, then trace the epoch-end collectives. Done means the configuration completes without an NCCL deadlock or timeout while retaining the described custom metric and epoch loss logging.

Written by the indexing model from the issue text.