Fix LoRA save/load to use stable layer paths (#2701)#2774
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates LoRA weight saving and loading to identify layers by their stable path attribute instead of their unstable integer index in a topologically sorted list, preventing issues when models are reconstructed. Backward compatibility is maintained for loading legacy integer-indexed weights, and comprehensive tests have been added. The reviewer recommends using a try-except KeyError block instead of directly checking the internal store.h5_file attribute to make the loading logic more robust.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if f"lora/{layer.path}" in store.h5_file: | ||
| inner_store = store.get(f"lora/{layer.path}") | ||
| else: | ||
| inner_store = store.get(f"lora/{layer_index}") |
There was a problem hiding this comment.
Relying on the internal h5_file attribute of H5IOStore is fragile because it is an implementation detail of Keras's saving library that could change in future Keras releases. Instead of checking membership in store.h5_file directly, it is more robust and idiomatic to use a try-except KeyError block around store.get(). This adheres to the EAFP (Easier to Ask for Forgiveness than Permission) principle and only relies on the public store interface.
| if f"lora/{layer.path}" in store.h5_file: | |
| inner_store = store.get(f"lora/{layer.path}") | |
| else: | |
| inner_store = store.get(f"lora/{layer_index}") | |
| try: | |
| inner_store = store.get(f"lora/{layer.path}") | |
| except KeyError: | |
| inner_store = store.get(f"lora/{layer_index}") |
References
- Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)
`save_lora_weights` / `load_lora_weights` identified each LoRA layer by its integer position in the weight-filtered, flattened layer list. That position is not stable across model reconstructions, so reloading LoRA weights into a freshly built model could place them on the wrong layers and silently change the outputs (e.g. with `bert_base_en`). Identify layers by their unique and stable `layer.path` instead. Loading falls back to the legacy integer keys, so `.lora.h5` files written by earlier versions still load. Added a round-trip test and a path-key regression test to `backbone_test.py`.
ManojPawar2
force-pushed
the
fix-lora-save-load-path-keys
branch
from
8cf4e86 to
786bd37
Compare
laxmareddyp
left a comment
laxmareddyp
left a comment
There was a problem hiding this comment.
Left initial comments and address the gemini comment as it is valid comment.
| # We identify each layer by its `path` rather than by its position in | ||
| # the flattened layer list. Layer names are non-unique, and integer | ||
| # positions are not stable: the topologically sorted list of layers | ||
| # that have weights can be ordered differently when a model is rebuilt, | ||
| # which previously caused LoRA weights to be restored into the wrong | ||
| # layers (see #2701). A layer's `path` is both unique and stable across | ||
| # reconstructions of the same architecture. |
There was a problem hiding this comment.
nit: comment is a bit long ,trimming to 2-3 lines and pointing at #2701 for full context.
| # so the factored weights are always named `kernel` | ||
| layer = all_layers[layer_index] | ||
| inner_store = store.make(f"lora/{layer_index}") | ||
| inner_store = store.make(f"lora/{layer.path}") |
There was a problem hiding this comment.
can you confirm you verified layer.path stability across different model types (BERT, Gemma, etc.)? we want to make sure paths don't include the model name and stay consistent across from_preset vs fresh construction
| legacy_path = os.path.join(self.get_temp_dir(), "legacy.lora.h5") | ||
| store = keras.src.saving.saving_lib.H5IOStore(legacy_path, mode="w") | ||
| lora_store = store.make("lora") | ||
| lora_store["rank"] = backbone._lora_rank | ||
| all_layers = [ | ||
| lyr | ||
| for lyr in backbone._flatten_layers(include_self=False) | ||
| if lyr.weights | ||
| ] | ||
| for layer_index in backbone._lora_enabled_layers: | ||
| layer = all_layers[layer_index] | ||
| inner_store = store.make(f"lora/{layer_index}") | ||
| inner_store["lora_kernel_a"] = layer.lora_kernel_a | ||
| inner_store["lora_kernel_b"] = layer.lora_kernel_b | ||
| store.close() | ||
|
|
||
| reloaded = BertBackbone(**init_kwargs) | ||
| reloaded.enable_lora(4) | ||
| reloaded.load_lora_weights(legacy_path) # must not raise |
There was a problem hiding this comment.
legacy compat test only checks "doesn't crash" worth adding non-zero adapter values and assertAllClose like the other test? with zero-init weights, even assertAllClose wouldn't catch wrong-layer assignment here..
| def test_lora_save_and_reload(self): | ||
| # Regression test for #2701: saving LoRA weights and reloading them | ||
| # into a freshly constructed model must reproduce the same outputs. | ||
| init_kwargs = { |
There was a problem hiding this comment.
nit: init_kwargs is repeated in both tests — can we pull it into setUp?
2 participants
Summary
Fixes #2701.
Backbone.save_lora_weights/load_lora_weightscould restore LoRA adapters into the wrong layers when reloading into a freshly constructed model, which changed the model's outputs. The problem is reproducible onbert_base_enand is not model-specific.Root cause
The save/load logic identified each LoRA-enabled layer by its integer position in the weight-filtered flattened layer list (
[l for l in self._flatten_layers() if l.weights]). Layer names are non-unique, so indices were used as a workaround, but that position is not stable across model reconstructions: the ordering of layers-with-weights can differ between the model that saved the adapters and a freshly built model that loads them. When the ordering shifts,lora_kernel_a/lora_kernel_bare assigned to the wrong layers, and the reloaded model no longer matches the original (thertol=1e-5, atol=1e-5comparison fails).Fix
Identify each layer by its
layer.path, which is unique and stable across reconstructions of the same architecture:save_lora_weightsstores each adapter underlora/{layer.path}.load_lora_weightslooks each adapter up bylayer.path, and falls back to the legacy integer-index key so.lora.h5files written by earlier versions continue to load.Tests
Added to
keras_hub/src/models/backbone_test.py:test_lora_save_and_reload: enables LoRA on a BERT backbone, assigns non-zero adapters, saves the LoRA weights, reloads them into a separately constructed backbone (different auto-generated name), and asserts the outputs match atrtol=1e-5, atol=1e-5, the exact scenario from the issue.test_lora_weights_use_stable_path_keys: asserts the saved file is keyed by layer path rather than integer index, and that a legacy index-keyed file still loads.All existing Gemma/Gemma3/Gemma4 LoRA fine-tuning and save/reload tests continue to pass.