Lightning-AI / Lightning-AI/litData

Restart training with new data, mid-epoch

Open
#436 1 comment 3 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
614
Forks
106
Avg merge
15h 8m
Merged PRs (30d)
22

Description

## 🚀 Feature

Right now, you can change the datasets used in a `CombinedStreamingDataset` and resume training on epoch boundaries. It would be great if you could resume training with new datasets mid-epoch.

### Motivation

When we're doing curriculum learning, we don't know the right number of steps or epochs to train. If we reach a sufficient validation loss, we kill the training and resume training with a new group of datasets (i.e. adjust the curriculum). Accordingly, we often have to kill training mid-epoch and restart with new datasets.

### Pitch

If you are training with N datasets and kill training K steps into epoch N, change the underlying datasets, and resume training from a checkpoint that was saved mid-epoch, the trainer should jump to epoch N + 1 with the new datasets, the old optimizer state, the correct global batch index.

### Alternatives

Right now, we have two workarounds:
- **Delete the `loops` part of the last saved mid-epoch checkpoint, before we resume training with different datasets.** This isn't a great solution - because you resume training at epoch 0, which messes with any learning rate schedulers we have.
- **Copy the `loops` part of the last saved epoch checkpoint into the more recent epoch that was saved mid-epoch, before we resume training with different datasets**. This approximately mitigates the learning rate scheduler issue -- but isn't the cleanest solution and is a pain to do manually every time.

### Additional context

Contributor guide

Open the contributing guide

Research direction

Start with the CombinedStreamingDataset behavior and the trainer's checkpoint-resume path described in the issue. Trace how mid-epoch checkpoints store loops, optimizer state, epoch, and global batch index; done means replacing datasets resumes at the next epoch with the old optimizer state and correct global batch index.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.