Skip to content

ppocrv6 + quit=early crashes: OverflowError in Muon cosine schedule #792

Description

@johnlockejrr

Training with --arch ppocrv6 and early stopping (-q early) fails before the first step:

OverflowError: cannot convert float infinity to integer

at kraken/train/ppocr.py in configure_optimizers, on:

total_steps = max(int(self.trainer.estimated_stepping_batches), 1)

With quit=early, ketos train sets max_epochs=-1 (kraken/ketos/recognition.py). Lightning then reports estimated_stepping_batches as infinity, and casting that to int blows up.

Fixed schedule (-q fixed -N 100) is fine. Same setup otherwise: AdamW+Muon + cosine (the ppocrv6 defaults).

Repro: any small dataset is enough —

ketos train --arch ppocrv6 -q early --lag 10 -N 100 \
  -f binary -t train.lst -e val.lst --base-dir R

(-N is ignored for max_epochs when quit=early, which is part of the problem.)

Expected: early stopping should work with the default Muon cosine schedule. Either use --epochs as a ceiling when it’s set, or don’t rely on estimated_stepping_batches when max_epochs < 0.

kraken 7.0.2.post32+g5494db67 (mittagessen/kraken @ 5494db67)
Python 3.13.11
torch 2.13.0+cu126
lightning 2.6.1
OS: Pop!_OS 24.04 (Linux x86_64)
GPU: RTX 3090 (repro is not GPU-specific)

Happy to send a PR if useful.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions