huge memory leak and execution stuck after on_train_epoch_start but before training_step

[malfonsoarquimea]

This is a follow-up issue of #12522
I am facing an issue where, when using more than one GPU and ddp strategy, the code gets stuck and starts consuming huge amounts of RAM after the execution of on_train_epoch_start but before training_step.
I managed to track the memory leak to https://github.com/Lightning-AI/lightning/blob/dbb5ca8d436a917fc9c2bdd6e9c00e4fd6187735/src/pytorch_lightning/loops/epoch/training_epoch_loop.py#L147
and then to
https://github.com/Lightning-AI/lightning/blob/dbb5ca8d436a917fc9c2bdd6e9c00e4fd6187735/src/pytorch_lightning/utilities/fetching.py#L179

I am using PyTorch lightning version 1.7.7
Regarding RAM consumption, when using one GPU this training consumes (due to data being preloaded in the Dataset):

• batch size 3000: 12GB
• batch size 6000: 20GB

With 2 GPUs and DDP strategy running on a docker container:

• batch size 3000 per GPU (6000 total) : 296 GB
• batch size 6000 per GPU (12000 total) : 587 GB

EDIT:
Afer using a distributed sampler as indicated here #15164 the memory consumption is by 40% or so:

• batch size 6000 per GPU (12000 total) : 380 GB

I would expect the second case to be twice the ram consumption of the first one, as preloading is happening once per GPU and there are 2 GPUs instead of one.
Also, I noticed that this problem does not occur on the sanity check steps.

Sadly, I can not provide a code example to reproduce but I can give any information about my system or what my code does

Any help on how to solve this would be greatly appreciated. Thanks in advance and have a nice day!

[malfonsoarquimea]

@rohitgr7 can you provide some insight?