Formalize progress tracking inside of the trainer internals
🚀 Feature
We should better enforce progress tracking across these dimensions:
Stage: training, evaluation (validation and test), and prediction loops
Granularity: batches, steps, epochs

batches vs steps: steps = optimizer steps (parameter updates) and applies to training loop only. this will differ from batches when gradient accumulation is used, for instance

Motivation

Provide consistency across trainer stages for users
Better debugging
Create mid-epoch resumable state for training


Pitch
See the full example here

Create a dataclass that each Loop in Lightning is responsible for maintaining. All of this state is local to a particular trainer rank. These dataclasses live as attributes on the corresponding loop.

@dataclass
class LoopProgress:
    # this also serves as the index of the current epoch
    total_epochs_processed: int = 0
    # this monotonically increases and is summed across epochs
    total_batches_processed: int = 0
    # this resets at the end of the epoch back to 0
    batches_processed_this_epoch: int = 0
    
    # convenience utils for accessing state
    def bump_batch(self, increment: int = 1):
        self.total_batches_processed += increment
        self.batches_processed_this_epoch += increment
    
    def bump_epoch(self, increment: int = 1):
        self.total_epochs_processed += increment
    
    def reset_batch_in_epoch(self):
        self.batches_processed_this_epoch = 0

The train loop extends this to track optimizer steps
@dataclass
class TrainLoopProgress(LoopProgress):
    total_optimizer_steps_processed: int = 0
    optimizer_steps_processed_this_epoch: int = 0

    def bump_step(self, increment: int = 1):
        self.total_optimizer_steps_processed += increment
        self.optimizer_steps_processed_this_epoch += increment
    
    def reset_step_in_epoch(self):
        self.optimizer_steps_processed_this_epoch = 0


The Trainer maintains its own tracker that has references to the individual loop progress trackers

@dataclass
class Progress:
    train_progress: TrainLoopProgress
    val_progress: LoopProgress
    test_progress: LoopProgress
    predict_progress: LoopProgress


For convenience, we could offer synchronization utilities to sum the progress state across ranks to get totals
We can also offer convenience utilities to get the totals across different stages
We update the loops to populate this state, and make backwards compatible changes to reference this state
We save the progress state into checkpoints as part of the trainer state
We handle loading this state when resuming from checkpoints