`automatic_optimization=True` will terminate the training process after few iterations.

Nobody has claimed this yet.

### Bug description

When I use Lightning 2.0 to do the experiments, it is strange that I can only use `manual_backward` to train the model.
If I modified the code and adapt `automatic_optimization=True`, the training steps will be terminated in several steps and automatically skip to the next epoch. However, after skipping to the next epoch, the problem still exists, and again.

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

v2_0

### How to reproduce the bug

```python
"""Run R3D with lightning. (Enable colossal-ai, deepspeed and etc.)
"""
import os
import pathlib
import shutil
from collections import defaultdict
from typing import Any, List

import gin
import numpy as np
import torch
from colossalai.nn.optimizer import HybridAdam
from deepspeed.ops.adam import DeepSpeedCPUAdam
from lightning import LightningModule, Trainer, seed_everything
from lightning.pytorch import loggers, logging
from lightning.pytorch.callbacks import (LearningRateMonitor, ModelCheckpoint,
TQDMProgressBar)

from r3d.system.interface import R3D
from r3d.tools import get_callbacks, get_inst_from_str
from r3d.utils import auto_fix_name
from r3d.utils.optim import Adan
from r3d.utils.store_image import make_grid, reshape_rays_to_images

@gin.configurable()
class R3DLightning(LightningModule, R3D):
"""
R3D manages model~(parameterized by theta), differential renderer ops~(in order to speed up) and various guidances.

We will build two optimizers for model and guidance respectively, since model and
guidance contains different networks.
Differential renderer may contains a lot operations that speed up the sampling procedure, such as cuda based or taichi.
"""
# the warmup steps
warmup_step: int

def __init__(
self,
model: str = 'TorchFreqNeRF',
renderer: str = 'TorchNeRFRenderer',
guiders: List[str] = ['PixelMSEGuider'],
lr_init: float = 5e-4,
lr_final: float = 5e-6,
lr_delay_steps: int = 1000,
lr_delay_mult: float = 1.0,
logdir='logs',
warmup_step=0,
optim: str = 'adam',
accumulate_steps: int = 10,
):
super().__init__()
# don't instantiate layers here
# move the creation of layers to `configure_sharded_model`
self.model_name = auto_fix_name(model, 'model')
self.renderer_name = auto_fix_name(renderer, 'renderer')
self.guiders_name = [auto_fix_name(g, 'guider') for g in guiders]
self.lr_init = lr_init
self.lr_final = lr_final
self.lr_delay_steps = lr_delay_steps
self.lr_delay_mult = lr_delay_mult
self.logdir = logdir
self.warmup_step = warmup_step
self.optim = optim
self.accumulate_steps = accumulate_steps

# self.automatic_optimization = False

def configure_sharded_model(self):
"""
This function will be called by Trainer in suitable place.
"""
self.model = get_inst_from_str(self.model_name)
self.renderer = get_inst_from_str(self.renderer_name)
self.guiders = get_inst_from_str('r3d.guide.Guiders',
guiders=self.guiders_name)

def configure_optimizers(self):
"""
We will optimize the Model and Renderer in one optimizer because they always work together.
We put the Guider in another optimizer to optimize it separately, as the Guider may be the discriminator in the adversarial generative network and needs to be optimized separately.
"""
network: Any = self.trainer.model
param_groups = []
adv_param_groups = []

# add encoder, model and renderer's parameters to param_groups_one
param_groups.append({'params': network.model.parameters()})
param_groups.append({'params': network.renderer.parameters()})

# add guider to param_groups_one or param_groups_two depending on detach_training flag.
for guide in network.guiders.guiders:
if guide.detach_training:
adv_param_groups.append({'params': guide.parameters()})
else:
param_groups.append({'params': guide.parameters()})

strategy_name = self.trainer.strategy.strategy_name # type: ignore
if 'colossalai' in strategy_name:
optimizer = HybridAdam
elif 'deepspeed' in strategy_name:
optimizer = DeepSpeedCPUAdam
elif self.optim == 'adan':
optimizer = Adan
elif self.optim == 'adam':
optimizer = torch.optim.Adam
else:
raise NotImplementedError

optim_g = optimizer(param_groups, lr=self.lr_init)
lr_scheduler_g = torch.optim.lr_scheduler.CosineAnnealingLR(
optim_g, self.trainer.max_steps, self.lr_final)

if len(adv_param_groups):
optim_d = optimizer(adv_param_groups, lr=self.lr_init)
lr_scheduler_d = torch.optim.lr_scheduler.CosineAnnealingLR(
optim_d, self.trainer.max_steps, self.lr_final)
return [optim_g, optim_d], [{
'scheduler': lr_scheduler_g,
'interval': 'step'
}, {
'scheduler': lr_scheduler_d,
'interval': 'step'
}]
else:
return {
"optimizer": optim_g,
"lr_scheduler": lr_scheduler_g,
"interval": "step"
}

def training_step(self, batch, batch_idx):
"""
detach_training_step == True indicates that we are training the adv part guider.
detach_training_step == False indicates that we are training a generator.
"""

result = self.renderer(self.model, batch)

# adversarial loss is binary cross-entropy
g_loss, g_loss_details = self.guiders(
result,
batch,
detach_training_step=False,
)

self.log(
'g_loss',
g_loss,
prog_bar=True,
on_step=True,
rank_zero_only=True,
)
for k, v in g_loss_details.items():
self.log(f'g_loss_details/{k}',
v,
on_step=True,
rank_zero_only=True)

return g_loss

def render_rays(self, batch, batch_idx):
ret = {}
rendered = self.renderer(self.model, batch)

rendered = reshape_rays_to_images(rendered, batch)

# determine image size
if 'image_size' in batch:
h, w = batch['image_size']
else:
h, w = batch['H'], batch['W']

rgb_fine = rendered[-1]['image']
depth_fine = rendered[-1]['depth']
if 'normal_image' in rendered[-1] and rendered[-1][
'normal_image'] is not None:
ret['normal_image'] = rendered[-1]['normal_image']

if 'target' in batch:
target = batch['target']
ret['target'] = target.reshape(-1, h, w, 3)
ret['image'] = rgb_fine
ret['depth'] = depth_fine

# convert to cpu
return {k: v.cpu() for k, v in ret.items()}

def on_validation_epoch_start(self) -> None:
self.results_buffers = []

def validation_step(self, batch, batch_idx):
"""Validation step will log images and depths to tensorboard.
"""
result = self.render_rays(batch, batch_idx)
self.results_buffers.append(result)
return result

def on_validation_epoch_end(self) -> None:
"""Log image to tensorboard
"""
results = defaultdict(list)
for result in self.results_buffers:
for k, v in result.items():
results[k].append(v)
results = {k: torch.cat(v) for k, v in results.items()}

writer = self.logger.experiment # type: ignore

image = make_grid(results['image'])
writer.add_image('image', image, self.trainer.current_epoch)

depth = make_grid(results['depth'])
writer.add_image('depth', depth, self.trainer.current_epoch)
if 'normal_image' in results:
normal_image = make_grid(results['normal_image'])
writer.add_image('normal_image', normal_image,
self.trainer.current_epoch)

if 'target' in results:
target = make_grid(results['target'])
writer.add_image('target', target, self.trainer.current_epoch)

# clear buffer
self.results_buffers.clear()

def test_step(self, batch, batch_idx):
"""Test step will save the image and depth as video in local disk,
"""
return self.render_rays(batch, batch_idx)

def predict_step(self, batch, batch_idx):
"""Prediction step will save image and depths as video in local disk
and wil also export mesh from model. (Exporting mesh need more computation.)
"""
return self.render_rays(batch, batch_idx)

def export_mesh(self):
pass

def export_video(self):
pass

@gin.configurable()
def run(
ckpt_path='last.ckpt',
logbase: str = 'logs',
exp_name: str = 'trail',
# Optimization
strategy: str = 'auto', # colossalai or deepspeed tricks.
max_steps: int = 100000,
max_epochs: int = 100,
# Logging
log_every_n_steps: int = 50,
progressbar_refresh_rate: int = 1,
# Run Mode
run_train: bool = True,
run_eval: bool = True,
run_render: bool = False,
num_devices: int = -1,
precision='16-mixed',
num_sanity_val_steps: int = 0,
seed: int = 777,
save_last: bool = True,
grad_max_norm=1.,
grad_clip_algorithm='norm',
debug: bool = True):

logging.getLogger('lightning').setLevel(
logging.ERROR if not debug else logging.DEBUG)

num_devices = num_devices if num_devices >= 0 else torch.cuda.device_count(
)

# model_name = gin.query_parameter('R3D.model_name')

# if model_name in ['plenoxel']:
# num_devices = 1

pathlib.Path(logbase).mkdir(exist_ok=True)
logdir = pathlib.Path(logbase).joinpath(exp_name)
logdir.mkdir(exist_ok=True)
(logdir / exp_name).mkdir(exist_ok=True)

logger = loggers.TensorBoardLogger(save_dir=logdir, name=exp_name)

seed_everything(seed, workers=True)

callbacks = []
lr_monitor = LearningRateMonitor(logging_interval='step')
tqdm_progress = TQDMProgressBar(refresh_rate=progressbar_refresh_rate)
callbacks.append(tqdm_progress)
callbacks.append(lr_monitor)

# callbacks += [model_checkpoint, tqdm_progress]
# append callbacks of renderer
callbacks += get_callbacks()

trainer = Trainer(
logger=logger if run_train else None,
log_every_n_steps=log_every_n_steps,
devices=num_devices,
max_epochs=max_epochs,
max_steps=max_steps,
accelerator='gpu',
strategy=strategy, # colossalai
check_val_every_n_epoch=1,
precision=precision,
num_sanity_val_steps=num_sanity_val_steps,
callbacks=callbacks,
gradient_clip_algorithm=grad_clip_algorithm,
gradient_clip_val=grad_max_norm,
sync_batchnorm=True,
limit_val_batches=64,
)

ckpt_path = logdir / ckpt_path
logging.info('Loading dataset...')
data_module = get_inst_from_str('r3d.dataset.LitData')

r3d = R3DLightning()

if run_train:
best_ckpt = logdir / 'best.ckpt'
if best_ckpt.exists():
os.remove(best_ckpt)
version0 = logdir / exp_name / 'version_0'
if version0.exists():
shutil.rmtree(version0, ignore_errors=True)

if not ckpt_path.exists():
logging.warning('No checkpoint find! Train from scratch.')
ckpt_path = None
else:
ckpt_path = str(ckpt_path)

trainer.fit(r3d, data_module, ckpt_path=ckpt_path)

# if run_eval:
# # render image and depth only
# ckpt_path = logdir / 'best.ckpt'
# trainer.test(r3d, data_module, ckpt_path=str(ckpt_path))

# if run_render:
# # render image and depth, export mesh.
# ckpt_path = logdir / 'best.ckpt'
# trainer.predict(r3d, data_module, ckpt_path=str(ckpt_path))
```

### Error messages and logs

```
Global seed set to 777
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
WARNING:root:No checkpoint find! Train from scratch.
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
/docker/software/anaconda3/envs/r3d/lib/python3.8/site-packages/transformers/models/clip/feature_extraction_clip.py:28: FutureWarning: The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please use CLIPImageProcessor instead.
warnings.warn(
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
/docker/software/anaconda3/envs/r3d/lib/python3.8/site-packages/lightning/pytorch/core/optimizer.py:360: RuntimeWarning: Found unsupported keys in the optimizer configuration: {'interval'}
rank_zero_warn(

| Name | Type | Params
-------------------------------------------------------------------------------------------------------------------------------------------------
0 | model | TorchFreqNeRF | 8.4 K
1 | renderer | TorchNeRFRenderer | 0
2 | guiders | ScoreDistillationSamplingGuider,OpacityGeometryGuider,EntropyGeometryGuider,NormalSmoothGuider,OrientGeometryGuider | 1.3 B
-------------------------------------------------------------------------------------------------------------------------------------------------
8.4 K Trainable params
1.3 B Non-trainable params
1.3 B Total params
5,159.843 Total estimated model params size (MB)
Epoch 1: 5%|█▊ | 5/100 [00:01<00:37, 2.55it/s, v_num=0, g_loss=1.18e-7]
```

### Environment

Current environment

```
#- Trainer, LightningModule
#- PyTorch Lightning Version (e.g., 2.0):
#- PyTorch Version (e.g., 1.13.1):
#- Python version (e.g., 3.8):
#- OS (e.g., Linux):
#- CUDA/cuDNN version: 11.7
#- How you installed Lightning(`pip`):
#- Running environment of LightningApp (e.g. local):
```

### More info

_No response_

Start with the R3DLightning.training_step and configure_optimizers entry points shown in the report, then inspect the Trainer configuration using automatic optimization, mixed precision, and the colossalai or deepspeed strategy. Reproduce the run and compare it with manual_backward; done means automatic optimization completes training steps without prematurely skipping to the next epoch.

Written by the indexing model from the issue text.