Skip to content

fix(save_and_load): save base_layer.bias for bias="lora_only"#3307

Open
Anai-Guo wants to merge 1 commit into
huggingface:mainfrom
Anai-Guo:fix/lora-only-base-layer-bias-save
Open

fix(save_and_load): save base_layer.bias for bias="lora_only"#3307
Anai-Guo wants to merge 1 commit into
huggingface:mainfrom
Anai-Guo:fix/lora-only-base-layer-bias-save

Conversation

@Anai-Guo

@Anai-Guo Anai-Guo commented Jun 7, 2026

Copy link
Copy Markdown
Contributor

What

LoraConfig(bias="lora_only") correctly marks a targeted layer's bias as trainable, but get_peft_model_state_dict() / save_pretrained() silently drop it, so a reloaded adapter does not reproduce the trained outputs.

Fixes #3306

Why

For a targeted layer the original module is wrapped, and its bias lives at <module>.base_layer.bias. The bias-key reconstruction was:

bias_name = k.split("lora_")[0] + "bias"   # -> base_model.model.proj.bias

but the real trained key is base_model.model.proj.base_layer.bias, so the lookup never matched and the bias was excluded from the saved state dict. (bias="all" works because it keys off the substring "bias" directly.)

Fix

Look up <prefix>base_layer.bias (current tuner-layer structure) in addition to the legacy <prefix>bias. The same one-line pattern was applied to the identical bias="boft_only" branch for BOFT.

Test

Added TestLoraInitialization::test_lora_only_bias_is_saved_and_reloaded, which perturbs all trainable params (incl. the base_layer bias), asserts the bias key is present in both get_peft_model_state_dict() output and the saved adapter_model.safetensors, and checks that reloading reproduces the pre-save output.

Reproducer from the issue now round-trips with max diff 0.0 (was 1.21).

🤖 Generated with Claude Code

When a tuner targets a layer, the original module is wrapped and its bias
lives at <module>.base_layer.bias. The previous key reconstruction produced
<module>.bias, so get_peft_model_state_dict and save_pretrained silently
dropped the trained bias for bias="lora_only" (and bias="boft_only"),
breaking adapter round-trips. Check the base_layer.bias name as well, keeping
the legacy <module>.bias name for backward compatibility.

Fixes huggingface#3306
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LoraConfig(bias="lora_only") trains base_layer.bias but save_pretrained/get_peft_model_state_dict drops it

1 participant