Lightning-AI / litgpt

Pretrain, finetune, deploy 20+ LLMs on your own data. Uses state-of-the-art techniques: flash attention, FSDP, 4-bit, LoRA, and more.
https://lightning.ai
Apache License 2.0
6.85k stars 726 forks source link

Will CycleIterator forward to dataset on resume for pretrain? #1386

Open calvintwr opened 1 week ago

calvintwr commented 1 week ago

When resuming finetuning, I see that the CycleIterator is forwarded to the dataset where the iteration is to continue from:

https://github.com/Lightning-AI/litgpt/blob/f3343784bbd192490e2a70aa5ef75c52608b1d35/litgpt/finetune/full.py#L208-L219

However, for pretrain, this does not exist and the training seems to resume from the begining:

https://github.com/Lightning-AI/litgpt/blob/f3343784bbd192490e2a70aa5ef75c52608b1d35/litgpt/pretrain.py#L217-L271

Can I check in this case, it looks like when resuming, the pretraining will start from the first dataset, and not forwarded?

awaelchli commented 1 week ago

The pretraining code uses a stateful dataloader from LitData: https://github.com/Lightning-AI/litgpt/blob/f3343784bbd192490e2a70aa5ef75c52608b1d35/litgpt/pretrain.py#L192-L198