Checkpoints not saving: global_step and current_epoch not updating despite correct progress bar

[homjay]

Bug description

During training, I noticed that no checkpoints are being saved. Upon investigation, I found that although the progress bar shows the correct step and epoch values after each step, the model's internal global_step and current_epoch remain unchanged.

I am using the model.fit() method and suspect the issue may be related to using multiple optimizers and manual optimization operations. There are no errors in the logs, and I haven't found any documentation addressing this behavior.

Has anyone encountered this or know what might be causing it? Any help is appreciated.

---

I discovered that the global_step wasn't updating, causing _should_skip_saving_checkpoint to always return true:

def _should_skip_saving_checkpoint(self, trainer: "pl.Trainer") -> bool:
    from lightning.pytorch.trainer.states import TrainerFn

    return (
        bool(trainer.fast_dev_run)  # disable checkpointing with fast_dev_run
        or trainer.state.fn != TrainerFn.FITTING  # don't save anything during non-fit
        or trainer.sanity_checking  # don't save anything during sanity check
        or self._last_global_step_saved == trainer.global_step  # already saved at the last step
    )

After examining the code, I believe this is an optimizer-related issue:

@property
def global_step(self) -> int:
    lightning_module = self.trainer.lightning_module
    if lightning_module is None or lightning_module.automatic_optimization:
        return self.automatic_optimization.optim_progress.optimizer_steps
    return self.manual_optimization.optim_step_progress.total.completed

A temporary solution is to save the checkpoint manually:

self.trainer.save_checkpoint("./model.ckpg")

I'm still uncertain why the model's global_step isn't updating, especially since the optimization has been rewritten and is quite complex.

What version are you seeing the problem on?

master, v2.5

Reproduced in studio

No response

How to reproduce the bug

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0): 2.5.2
#- PyTorch Version (e.g., 2.5): 2.19
#- Python version (e.g., 3.12): 3.12.9
#- OS (e.g., Linux): Ubuntu 22.04
#- CUDA/cuDNN version: 12.9
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source): pip

More info

No response

cc @lantiga

[SkafteNicki]

@homjay thanks for raising this issue. Could you please provide some more details on what your optimization setup is looking like and your training_step in particular? If you can provide a fully reproducible example that would also be great.

[homjay]

Thank you for your reply. Below is the optimized configure_optimizers method:

def configure_optimizers(self):
    params_vid = (
        list(self.ctx_dec_8.parameters())
        + list(self.ctx_enc_8.parameters())
        + list(self.frm_gen.parameters())
    )
    params_flo = list(self.flo_enc_8.parameters()) + list(
        self.flo_dec_8.parameters()
    )
    params_aux = list(self.aux_8.parameters())

    if self.cfg.use_muon_aux_adam:
        vid_w = [p for p in params_vid if p.ndim >= 2]
        vid_b = [p for p in params_vid if p.ndim < 2]
        flo_w = [p for p in params_flo if p.ndim >= 2]
        flo_b = [p for p in params_flo if p.ndim < 2]
        aux_w = [p for p in params_aux if p.ndim >= 2]
        aux_b = [p for p in params_aux if p.ndim < 2]

        vid_pgs = [
            dict(params=vid_w, use_muon=True, lr=1e-4, weight_decay=0.01),
            dict(params=vid_b, use_muon=False, lr=1e-4, weight_decay=0.01),
        ]
        flo_pgs = [
            dict(params=flo_w, use_muon=True, lr=1e-4, weight_decay=0.01),
            dict(params=flo_b, use_muon=False, lr=1e-4, weight_decay=0.01),
        ]
        aux_pgs = [
            dict(params=aux_w, use_muon=True, lr=1e-4, weight_decay=0.01),
            dict(params=aux_b, use_muon=False, lr=1e-4, weight_decay=0.01),
        ]

        opt_vid = MuonWithAuxAdam(vid_pgs)
        opt_flo = MuonWithAuxAdam(flo_pgs)
        opt_aux = MuonWithAuxAdam(aux_pgs)
    else:
        opt_vid = optim.Adam(params_vid, lr=1e-4, weight_decay=0.01)
        opt_flo = optim.Adam(params_flo, lr=1e-4, weight_decay=0.01)
        opt_aux = optim.Adam(params_aux, lr=1e-4, weight_decay=0.01)

    return opt_vid, opt_flo, opt_aux

For context, here's a simplified example of the training step implementation:

def training_step(self, batch, batch_idx):
    frame, flow = batch
    x_hat = self.forward(frame)
    flow_hat = self.flo_enc_8(frame)
    pred_psnr = self.aux_8(x_hat.clamp(0, 1))

    x_hat_clamped = x_hat.clamp(0, 1)
    video_psnr = self.psnr(x_hat_clamped.float(), frame.float()).detach().cpu().numpy()
    video_ssim = self.ssim(x_hat_clamped.float(), frame.float()).detach().cpu().numpy()
    video_lpips = self.lpips(x_hat_clamped.float(), frame.float()).detach().cpu().numpy()

    loss_flo = self.l1charbonnierloss(flow_hat, flow.detach())
    opt_flo.zero_grad()
    self.manual_backward(self.lambda_flo * loss_flo)
    opt_flo.step()

    loss_vid = self.l1charbonnierloss(x_hat, frame)
    opt_vid.zero_grad()
    self.manual_backward(self.lambda_recon * loss_vid)
    opt_vid.step()

    loss_aux = self.l1charbonnierloss(pred_psnr, video_psnr / 50.0)
    opt_aux.zero_grad()
    self.manual_backward(self.lambda_qem * loss_aux)
    opt_aux.step()

Hope this could be helpful.

[SkafteNicki]

@homjay below are my effort to create a reproduction script based on what information you have provided. I CANNOT reproduce the issue e.g. current epoch and global step increases as expected. Are you setting self.automatic_optimization=False in your LightningModule’s __init__?.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from lightning.pytorch import LightningModule, Trainer, Callback
from lightning.pytorch.callbacks import ModelCheckpoint


class ExactReproductionModel(LightningModule):
    """Model that exactly matches the user's setup"""
    
    def __init__(self):
        super().__init__()
        self.automatic_optimization = False  # Manual optimization
        
        # Simple model
        self.layer = nn.Linear(10, 1)
        self.loss_fn = nn.MSELoss()
        
    def forward(self, x):
        return self.layer(x)
        
    def configure_optimizers(self):
        # Multiple optimizers as mentioned in the issue
        opt1 = torch.optim.Adam([self.layer.weight], lr=1e-3)
        opt2 = torch.optim.Adam([self.layer.bias], lr=1e-3)
        return opt1, opt2
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        opt1, opt2 = self.optimizers()
        
        # Similar to user's approach - one step per training_step
        y_hat = self.forward(x)
        loss = self.loss_fn(y_hat, y)
        
        # Only step one optimizer per training step (like user described)
        if batch_idx % 2 == 0:
            opt1.zero_grad()
            self.manual_backward(loss)
            opt1.step()
        else:
            opt2.zero_grad()
            self.manual_backward(loss)
            opt2.step()
        
        self.log('train_loss', loss)
        return loss
        
    def train_dataloader(self):
        x = torch.randn(20, 10)  # Small dataset
        y = torch.randn(20, 1)
        dataset = TensorDataset(x, y)
        return DataLoader(dataset, batch_size=4)


class ExactMonitor(Callback):
    """Monitor exactly what the user is seeing"""
    
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        # Check the exact values mentioned in the issue
        progress_bar_global_step = trainer.progress_bar_metrics.get('step', 'N/A')
        actual_global_step = trainer.global_step
        
        print(f"Batch {batch_idx}:")
        print(f"  Progress bar step: {progress_bar_global_step}")
        print(f"  trainer.global_step: {actual_global_step}")
        print(f"  trainer.current_epoch: {trainer.current_epoch}")
        
        # Check manual optimization counter directly
        manual_opt = trainer.fit_loop.epoch_loop.manual_optimization
        manual_steps = manual_opt.optim_step_progress.total.completed
        print(f"  Manual optimization steps completed: {manual_steps}")
        
        # Check if checkpoint would be saved
        for cb in trainer.callbacks:
            if isinstance(cb, ModelCheckpoint):
                # Check the internal logic
                print(f"  Checkpoint callback every_n_train_steps: {cb._every_n_train_steps}")
                should_save = cb._should_skip_saving_checkpoint(trainer)
                print(f"  Should skip checkpoint: {should_save}")


def main():
    print("Exact reproduction of issue #21096")
    print("=" * 60)
    
    model = ExactReproductionModel()
    
    # Checkpoint callback with every_n_train_steps like in the issue
    checkpoint_callback = ModelCheckpoint(
        monitor="train_loss",
        dirpath="./checkpoints_exact",
        filename="model-{epoch:02d}-{step:02d}",
        save_top_k=-1,  # Save all
        every_n_train_steps=1,  # Try to save every step
        save_on_train_epoch_end=False,
    )
    
    trainer = Trainer(
        max_epochs=2,
        limit_train_batches=8,  # 8 batches to see the pattern
        callbacks=[ExactMonitor(), checkpoint_callback],
        enable_progress_bar=True,
        log_every_n_steps=1,
    )
    
    print(f"Before training - global_step: {trainer.global_step}")
    trainer.fit(model)
    print(f"After training - global_step: {trainer.global_step}")
    
    # Check saved checkpoints
    import os
    checkpoint_dir = "./checkpoints_exact"
    if os.path.exists(checkpoint_dir):
        checkpoints = os.listdir(checkpoint_dir)
        print(f"\nCheckpoints saved: {len(checkpoints)}")
        for ckpt in sorted(checkpoints):
            print(f"  - {ckpt}")
    else:
        print("\nNo checkpoints were saved!")


if __name__ == "__main__":
    main()

[SkafteNicki]

Going to close issue for now. Feel free to ping and reopen if necessary.