fix: make convert_from_litgpt output loadable with AutoModel.from_pretrained#2265
Open
discobot wants to merge 1 commit into
Open
fix: make convert_from_litgpt output loadable with AutoModel.from_pretrained#2265discobot wants to merge 1 commit into
discobot wants to merge 1 commit into
Conversation
…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.
discobot
requested review from
andyland,
k223kim,
lianakoleva and
t-vi
as code owners
June 13, 2026 13:19
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1871.
litgpt convert_from_litgptproduced a directory transformers cannot load: onlymodel.pth, nopytorch_model.bin, no config/tokenizer files. This PR saves theconverted weights directly as
pytorch_model.binand copies the config and tokenizerfiles from the source checkpoint, so
AutoModel.from_pretrained(converted_dir)justworks — the same thing
litgpt evaluatealready did internally before handing thedirectory to HFLM.
One extra finding: the torch.load/torch.save re-save hack in
litgpt/eval/evaluate.py(added in #1357 when the
incremental_savepickle wasn'tweights_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.loadhas effectivelybeen 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.pthtopytorch_model.bin; both tutorialsthat referenced
model.pth(and the manualcp config.json/cp tokenizer*steps,now unnecessary) are updated.
litgpt evaluateoutput directories are unaffected —they already contained
pytorch_model.bin.Tested with an updated
tests/convert/test_lit_checkpoint.pyplus a new regressiontest that converts a tiny pythia-14m-config checkpoint and checks
AutoModelForCausalLM.from_pretrainedloads it with matching logits. Also verified theissue's full pretrain-convert-load repro on CPU.