Lightning-AI / Lightning-AI/litgpt

On converting checkpoints to weights

Open
#830 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

conversion enhancement
Dominant language
Python
Stars
13.7k
Forks
1.5k
Avg merge
15h 37m
Merged PRs (30d)
1

Description

## Motivation

Whenever a script wants to load model weights, there are different variations of it that could be loaded depending on which script we are referring to:

1. A lit model weights file `lit_model.pth`. This is the output of `scripts/convert_hf_checkpoint.py`
2. A Fabric weights-only checkpoint. This is the output of `finetune/*.py`. It will include the lit model checkpoint under the `model` key. Example: https://github.com/Lightning-AI/lit-gpt/blob/main/finetune/lora.py#L310-L312
3. A Fabric training checkpoint. This is the output of `pretrained/*.py`. It will include the lit model checkpoint under the `model` key plus extra training state (e.g. optimizer state).
4. A Trainer checkpoint. This is the output of `pretrained/openwebtext_trainer.py`, the only script using the Trainer.

Most of our scripts support loading (1) and (2): https://github.com/search?q=repo%3ALightning-AI%2Flit-gpt%20.get(%22model&type=code
https://github.com/Lightning-AI/lit-gpt/pull/803 added support for loading (3) after a conversion step.

Currently (4) cannot be loaded anywhere other than the pertaining script itself.

## Pitch

This issue suggests unifying these cases by having a single interface to "process" checkpoints and `checkpont_dirs`. There are two ways to do it

### With a previous conversion step:

Roughly:
```bash
python convert_checkpoint.py out/foobar/ converted_checkpoint/
python generate/base.py --checkpoint_dir converted_checkpoint/
```

Cons:
- You have to remember to call this step
- It will create a duplicate version of the weights. This can be very annoying for large checkpoints in environments with limited disk size such as cloud instances.

### With a in-memory conversion function:

Inside `generate/base.py`, we call

```python
from lit_gpt.utils import get_weights_from

state_dict = get_weights_from(ckpt_dir)
```

Cons:
- Some users might prefer to have a clean checkpoint_dir to read from

## What about configs?

If we implement #483, all weights should have a config file beside it so that it can be carried over?

## What about the tokenizer vocabulary?

This will need to be manually copied over. Unless we choose to carry it over as with the configs.

Tutorials such as https://github.com/Lightning-AI/lit-gpt/blob/main/tutorials/finetune_lora.md#merging-lora-weights already indicate the need for this `cp` step

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the checkpoint-loading paths in generate/base.py and the Trainer, finetuning, and pretraining scripts, then compare them with the proposed convert_checkpoint.py and lit_gpt.utils.get_weights_from interface. The issue does not choose between conversion approaches or settle config and tokenizer handling; done would require an agreed design that handles all four checkpoint types consistently.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.