Fixes KeyError during Gemma3 LoRA fine-tuning initialization on Optimum Habana 1.18.1 #2352
+95
−1
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.
Root cause: PEFT removes LM head
weightentry fromLinear._parametersdict while leaving the attribute intact. When model moves to HPU andtie_weights()runs, PyTorch tries to re-register the parameter and hits the existing attribute.Solution:
_safe_tie_weights()that manually re-ties embeddings without callingregister_parameter_replace_module_parameter()helper that:_parameters[name]entries directlynn.Parameter_parametersinjection ifregister_parameterstill failsChanges:
optimum/habana/transformers/trainer.py:OrderedDict_move_model_to_device()tie_weights call in try/except_safe_tie_weights()method_replace_module_parameter()helperTarget: Optimum Habana 1.18.1 release (standalone backport for users on stable release)
Testing: Verified Gemma3-12B LoRA on ChartQA - training proceeds past initialization with warning log