Skip to content

fix: make convert_from_litgpt output loadable with AutoModel.from_pretrained#2265

Open
discobot wants to merge 1 commit into
Lightning-AI:mainfrom
discobot:fix/1871-convert-hf-loadable
Open

fix: make convert_from_litgpt output loadable with AutoModel.from_pretrained#2265
discobot wants to merge 1 commit into
Lightning-AI:mainfrom
discobot:fix/1871-convert-hf-loadable

Conversation

@discobot

Copy link
Copy Markdown

Fixes #1871.

litgpt convert_from_litgpt produced a directory transformers cannot load: only
model.pth, no pytorch_model.bin, no config/tokenizer files. This PR saves the
converted weights directly as pytorch_model.bin and copies the config and tokenizer
files from the source checkpoint, so AutoModel.from_pretrained(converted_dir) just
works — the same thing litgpt evaluate already did internally before handing the
directory to HFLM.

One extra finding: the torch.load/torch.save re-save hack in litgpt/eval/evaluate.py
(added in #1357 when the incremental_save pickle wasn't weights_only=True-compatible)
is no longer needed. On every supported torch version (>=2.7) the incremental_save file
loads cleanly with weights_only=True (verified on 2.7.0 and 2.12, incl. mmap=True,
which transformers uses) — evaluate.py's own un-flagged torch.load has effectively
been doing a weights-only load of it in CI since torch 2.6 flipped the default. So the
weights are written once, with no full-state-dict re-save, keeping conversion
memory-friendly for large models, and evaluate.py now just calls the conversion.

Note the output filename changes from model.pth to pytorch_model.bin; both tutorials
that referenced model.pth (and the manual cp config.json / cp tokenizer* steps,
now unnecessary) are updated. litgpt evaluate output directories are unaffected —
they already contained pytorch_model.bin.

Tested with an updated tests/convert/test_lit_checkpoint.py plus a new regression
test that converts a tiny pythia-14m-config checkpoint and checks
AutoModelForCausalLM.from_pretrained loads it with matching logits. Also verified the
issue's full pretrain-convert-load repro on CPU.

…trained

convert_lit_checkpoint wrote only a model.pth file and copied no config or
tokenizer files, so transformers could not load the converted directory
(no file named pytorch_model.bin, no config.json). Save the converted
weights as pytorch_model.bin and copy the config and tokenizer files from
the source checkpoint, mirroring what litgpt/eval/evaluate.py already did.
The torch.load/torch.save re-save hack in evaluate.py is dropped because
the incremental_save file loads fine with weights_only=True on all
supported torch versions (>=2.7). Fixes Lightning-AI#1871.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

failure converting pretrained litgpt checkpoints to HF format: a reproducible example

1 participant