Repository navigation
Restore the LoRA adapter when resuming from a checkpoint - #166
Merged
Merged
Conversation
HF Trainer's _load_from_checkpoint loads model.safetensors and, when that file exists, never reaches its PEFT branch. Our model.safetensors holds only the projector (LoRA is saved separately as adapter_model.safetensors), so a resumed run kept the trained projector but restarted LoRA from its initialization under the trained optimizer state. Seen resuming granite-qwen-verbatim-aug-2 at step 52,000: loss 0.28 -> 1.7, grad norm 0.4 -> 1.9; the run was stopped before its next Hub push. ASRTrainer._load_from_checkpoint now also calls the new ASRModel.load_lora_checkpoint, which loads adapter_model.safetensors into the PEFT-wrapped decoder and raises if a saved tensor matches no LoRA parameter. A test reproduces the stock Trainer's behaviour so it fails if transformers changes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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.
Summary
Resuming training silently threw away the LoRA adapter.
_load_from_checkpointloadsmodel.safetensorsand, when that file exists, never reaches its PEFT branch. Ourmodel.safetensorsholds only the projector; LoRA is saved separately asadapter_model.safetensors. A resumed run therefore kept the trained projector but restarted LoRA from its initialization, under optimizer moments that belonged to the trained adapter.granite-qwen-verbatim-aug-2at step 52,000 after a pod reboot, loss jumped from 0.28 to 1.7 and grad norm from 0.4 to 1.9. I stopped the run before its next Hub push, so the good step-52,000 model on the Hub was not overwritten.ASRModel.load_lora_checkpointloads the adapter file into the PEFT-wrapped decoder. It raises if a saved tensor matches no LoRA parameter.ASRTrainer._load_from_checkpointcalls it after the stock load.Test plan
tests/test_resume_checkpoint.py: save a "trained" LoRA model, scramble its weights, resume throughASRTrainer, and check that both the projector and the adapter are restored exactly. A second test reproduces the stock Trainer skipping the adapter, so it fails if transformers changes this. Also covers the edge cases (no adapter file; tensor that matches no parameter).ta dev check,ta dev test🤖 Generated with Claude Code