Fix LoRA adapter loading and finetune model serialization - #93
Merged
Merged
Conversation
- load_lora_adapter() now mirrors the InterpreterModule replacement done before training, so PEFT can find target modules by the names recorded in a saved adapter's config. - FinetuneModelConfig now encodes/decodes base_model as a string so the dynamic server CLI (which only passes plain strings) can round-trip it through `command`.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new CLI round-trip path needs stricter validation/quoting and a corrected internal reference to avoid fragile failures and confusing documentation.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
This PR improves LoRA finetuned-model portability by (1) ensuring adapter loading reconstructs the same module tree used during training so PEFT can resolve target_modules, and (2) enabling FinetuneModelConfig.base_model to round-trip through the dynamic server CLI by encoding it as a string in command and decoding it in the constructor.
Changes:
- Decode
base_modelfrom an encoded string inFinetuneModelConfig.__init__and include--base-modelinFinetuneModelConfig.command. - Call
_replace_interpreter_modules(model)duringload_lora_adapter()to mirror the pre-training replacement behavior.
File summaries
| File | Description |
|---|---|
cellmap_flow/models/models_config.py |
Adds base model encode/decode support so finetune model configs can be serialized through CLI command strings. |
cellmap_flow/finetune/lora_wrapper.py |
Ensures adapter loading mirrors pre-training module replacement so PEFT can match saved target_modules. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+784
to
788
| if isinstance(base_model, str): | ||
| from cellmap_flow.utils.web_utils import decode_to_json | ||
|
|
||
| base_model = decode_to_json(base_model) | ||
| self.base_model_dict = base_model |
Comment on lines
+809
to
+813
| encoded_base_model = encode_to_str(self.base_model_dict) | ||
| return ( | ||
| f"finetune --lora-adapter-path {self.lora_adapter_path} " | ||
| f"--base-model {encoded_base_model}" | ||
| ) |
Comment on lines
+375
to
+380
| # Replace any non-standard leaf modules (e.g. InterpreterModule from | ||
| # torch.export unflatten) with real nn.Conv*/Linear so PEFT's dispatch | ||
| # can find the target modules named in the saved adapter config. Must | ||
| # mirror create_lora_model()'s call to this before training, since the | ||
| # adapter's target_modules names were recorded against the post-replacement | ||
| # module tree. |
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.
@stuarteberg
Summary
load_lora_adapter()now mirrors theInterpreterModulereplacement done before training, so PEFT can find target modules by the names recorded in a saved adapter's config.FinetuneModelConfignow encodes/decodesbase_modelas a string so the dynamic server CLI (which only passes plain strings) can round-trip it throughcommand.