Running out of memory when resuming the training from a checkpoint

Nobody has claimed this yet.

### Bug description

When I try to resume the training from a checkpoint, program runs out of GPU memory. This is an unexpected behavior because when I set trainer's `ckpt_path` parameter to `None`, training works perfectly fine.

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

v2.0

### How to reproduce the bug

```python
trainer = pl.Trainer(
accelerator=args.accelerator,
devices=args.devices,
num_nodes=args.num_nodes,
max_steps=args.max_steps,
max_epochs=args.max_epochs,
log_every_n_steps=args.log_every_n_steps,
gradient_clip_val=args.gradient_clip_val,
gradient_clip_algorithm=args.gradient_clip_algorithm,
val_check_interval=args.val_check_interval,
precision=args.precision,
default_root_dir=args.output_dir,
strategy=strategy,
callbacks=callbacks,
logger=loggers,
)
ckpt_path = None
if args.ckpt_path:
ckpt_path = args.ckpt_path

trainer.fit(model_module, datamodule=data_module, ckpt_path=ckpt_path)
```

### Error messages and logs

Resuming from checkpoint: (OOM error)
```
Global seed set to 42
Using 16bit Automatic Mixed Precision (AMP)
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
[rank: 0] Global seed set to 42
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/4
[rank: 2] Global seed set to 42
[rank: 1] Global seed set to 42
[rank: 3] Global seed set to 42
[rank: 3] Global seed set to 42
[rank: 1] Global seed set to 42
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/4
Initializing distributed: GLOBAL_RANK: 3, MEMBER: 4/4
[rank: 2] Global seed set to 42
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/4
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 4 processes
----------------------------------------------------------------------------------------------------

You are using a CUDA device ('NVIDIA A100-SXM4-40GB') 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
/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:612: UserWarning: Checkpoint directory /...path.../output exists and is not empty.
rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")
Restoring states from the checkpoint path at /...path.../output/latest-hourly-epoch=0-step=35392.ckpt
LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]
LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3]

| Name | Type | Params
-------------------------------------------
0 | model | Model | 135 M
1 | loss | CrossEntropyLoss | 0
-------------------------------------------
135 M Trainable params
0 Non-trainable params
135 M Total params
543.127 Total estimated model params size (MB)
Restored all states from the checkpoint at /...path.../output/latest-hourly-epoch=0-step=35392.ckpt
Training: 0it [00:00, ?it/s]
Training: 0%| | 0/81497 [00:00
main(args)
File "/...path.../model_dir/pretrain_model.py", line 185, in main
trainer.fit(model_module, datamodule=data_module, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
call._call_and_handle_interrupt(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 92, in launch
return function(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
results = self._run_stage()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 978, in _run_stage
self.fit_loop.run()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
self.advance()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
self.advance(data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in advance
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 185, in run
self._optimizer_step(kwargs.get("batch_idx", 0), closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 261, in _optimizer_step
call._call_lightning_module_hook(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 142, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1265, in optimizer_step
optimizer.step(closure=optimizer_closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 158, in step
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 259, in optimizer_step
optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 224, in optimizer_step
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 70, in optimizer_step
closure_result = closure()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in __call__
self._result = self.closure(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 135, in closure
self._backward_fn(step_output.closure_loss)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 233, in backward_fn
call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 199, in backward
self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 67, in backward
model.backward(tensor, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1054, in backward
loss.backward(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
torch.autograd.backward(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.25 GiB (GPU 1; 39.41 GiB total capacity; 24.62 GiB already allocated; 5.81 GiB free; 31.72 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Traceback (most recent call last):
File "/...path.../model_dir/pretrain_model.py", line 354, in
main(args)
File "/...path.../model_dir/pretrain_model.py", line 185, in main
trainer.fit(model_module, datamodule=data_module, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
call._call_and_handle_interrupt(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 92, in launch
return function(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
results = self._run_stage()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 978, in _run_stage
self.fit_loop.run()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
self.advance()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
self.advance(data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in advance
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 185, in run
self._optimizer_step(kwargs.get("batch_idx", 0), closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 261, in _optimizer_step
call._call_lightning_module_hook(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 142, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1265, in optimizer_step
optimizer.step(closure=optimizer_closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 158, in step
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 259, in optimizer_step
optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 224, in optimizer_step
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 70, in optimizer_step
closure_result = closure()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in __call__
self._result = self.closure(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 135, in closure
self._backward_fn(step_output.closure_loss)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 233, in backward_fn
call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 199, in backward
self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 67, in backward
model.backward(tensor, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1054, in backward
loss.backward(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
torch.autograd.backward(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.25 GiB (GPU 0; 39.41 GiB total capacity; 24.62 GiB already allocated; 5.83 GiB free; 31.72 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Traceback (most recent call last):
File "/...path.../model_dir/pretrain_model.py", line 354, in
main(args)
File "/...path.../model_dir/pretrain_model.py", line 185, in main
trainer.fit(model_module, datamodule=data_module, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
call._call_and_handle_interrupt(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 92, in launch
return function(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
Traceback (most recent call last):
File "/...path.../model_dir/pretrain_model.py", line 354, in
results = self._run_stage()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 978, in _run_stage
self.fit_loop.run()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
main(args)
File "/...path.../model_dir/pretrain_model.py", line 185, in main
self.advance()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
trainer.fit(model_module, datamodule=data_module, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 520, in fit
self.advance(data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in advance
call._call_and_handle_interrupt(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 185, in run
return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 92, in launch
return function(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 559, in _fit_impl
self._optimizer_step(kwargs.get("batch_idx", 0), closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 261, in _optimizer_step
self._run(model, ckpt_path=ckpt_path)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 935, in _run
call._call_lightning_module_hook(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 142, in _call_lightning_module_hook
results = self._run_stage()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 978, in _run_stage
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1265, in optimizer_step
optimizer.step(closure=optimizer_closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 158, in step
self.fit_loop.run()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 201, in run
self.advance()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 133, in run
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 259, in optimizer_step
self.advance(data_fetcher)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 218, in advance
optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 224, in optimizer_step
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 185, in run
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 70, in optimizer_step
self._optimizer_step(kwargs.get("batch_idx", 0), closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 261, in _optimizer_step
closure_result = closure()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in __call__
self._result = self.closure(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 135, in closure
self._backward_fn(step_output.closure_loss)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 233, in backward_fn
call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
call._call_lightning_module_hook(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 142, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 199, in backward
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1265, in optimizer_step
self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 67, in backward
optimizer.step(closure=optimizer_closure)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/optimizer.py", line 158, in step
model.backward(tensor, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1054, in backward
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 259, in optimizer_step
loss.backward(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
optimizer_output = super().optimizer_step(optimizer, closure, model, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 224, in optimizer_step
torch.autograd.backward(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/amp.py", line 70, in optimizer_step
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.25 GiB (GPU 2; 39.41 GiB total capacity; 24.62 GiB already allocated; 5.81 GiB free; 31.72 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
closure_result = closure()
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 140, in __call__
self._result = self.closure(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 135, in closure
self._backward_fn(step_output.closure_loss)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/loops/optimization/automatic.py", line 233, in backward_fn
call._call_strategy_hook(self.trainer, "backward", loss, optimizer)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 288, in _call_strategy_hook
output = fn(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py", line 199, in backward
self.precision_plugin.backward(closure_loss, self.lightning_module, optimizer, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 67, in backward
model.backward(tensor, *args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1054, in backward
loss.backward(*args, **kwargs)
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/_tensor.py", line 487, in backward
torch.autograd.backward(
File "/...path.../.conda/envs/environment/lib/python3.10/site-packages/torch/autograd/__init__.py", line 200, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 6.25 GiB (GPU 3; 39.41 GiB total capacity; 24.62 GiB already allocated; 5.83 GiB free; 31.72 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
Epoch 0: 0%| | 0/81497 [00:09
Current environment

```
- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer
- PyTorch Lightning Version (e.g., 1.5.0): 2.0.2
- PyTorch Version (e.g., 2.0): 2.0.1
- Python version (e.g., 3.9): 3.10.11
- OS (e.g., Linux): Linux
- CUDA/cuDNN version: 11.6
- GPU models and configuration: 4 x NVIDIA A100-SXM4-40GB
- How you installed Lightning(`conda`, `pip`, source): conda
```

### More info

It is worth noting that GPUs are definitely "empty" (nothing else is being run on them).

cc @awaelchli @borda

Start with the reproducer in pretrain_model.py around trainer.fit(..., ckpt_path=ckpt_path), then trace the checkpoint-resume path through the training and optimization stack shown in the traceback. Compare resuming from the checkpoint with ckpt_path=None under the reported DDP and 16-bit AMP setup; done means resuming does not trigger the CUDA out-of-memory failure.

Written by the indexing model from the issue text.