huggingface / huggingface/peft

Trainable Tokens drops the output-head bias before merge

Open
#3,649 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
21.7k
Forks
2.5k
Avg merge
4d 16h
Merged PRs (30d)
60

Description

### System Info

PEFT `main` at `3d881e97426db449902478fa545ed6fa3086991f` (`0.20.1.dev0`), Python 3.12.13, PyTorch 2.13.0+cpu, Transformers 5.14.1, Linux/WSL2 CPU.

### Reproduction

Trainable Tokens omits an existing output-head bias in its unmerged Linear forward path. On a pretrained BERT model, adding a fresh adapter changes the logits without any training, while disabling or merging it restores the original output. Both standalone Trainable Tokens and LoRA's `trainable_token_indices` are affected.

```python
import copy

import torch
from transformers import BertForMaskedLM, BertTokenizer
from peft import LoraConfig, TrainableTokensConfig, get_peft_model

model_id = "prajjwal1/bert-tiny"
revision = "6f75de8b60a9f8a2fdf7b69cbd86d9e64bcb3837"
base = BertForMaskedLM.from_pretrained(model_id, revision=revision).eval()
tokenizer = BertTokenizer.from_pretrained(model_id, revision=revision)
inputs = tokenizer("The capital of France is [MASK].", return_tensors="pt")
with torch.no_grad():
expected = base(**inputs).logits
bias = base.get_output_embeddings().bias.detach().clone()

configs = [
TrainableTokensConfig(token_indices=[tokenizer.mask_token_id]),
LoraConfig(r=2, target_modules=["query"],
trainable_token_indices=[tokenizer.mask_token_id]),
]
for config in configs:
model = get_peft_model(copy.deepcopy(base), config).eval()
with torch.no_grad():
active = model(**inputs).logits
with model.disable_adapter():
disabled = model(**inputs).logits
model.merge_adapter()
merged = model(**inputs).logits
print(type(config).__name__)
print("active error:", (active - expected).abs().max().item())
print("active + bias error:", (active + bias - expected).abs().max().item())
print("disabled error:", (disabled - expected).abs().max().item())
print("merged error:", (merged - expected).abs().max().item())
```

Both configurations produce:

```text
active error: 9.51423454284668
active + bias error: 0.0
disabled error: 0.0
merged error: 0.0
```

No training or replacement of pretrained parameters is involved. The explicit BERT classes avoid model-type inference for this older checkpoint. Loading reports unrelated NSP/pooler keys as unexpected, not missing decoder weights.

The [Linear branch in TrainableTokensLayer](https://github.com/huggingface/peft/blob/3d881e97426db449902478fa545ed6fa3086991f/src/peft/tuners/trainable_tokens/layer.py#L263-L268) calls `F.linear(input=x, weight=W)` without the base layer's bias. The merged and disabled paths call the original layer, which includes it.

### Expected behavior

With default initialization, a fresh adapter should preserve the original model output. Active and merged inference should agree, including the existing output-head bias, without changing which bias parameters are trainable.

### Proposed scope and ownership

I would like to take ownership of this fix. Would it be okay for me to submit a focused PR preserving the Linear bias and adding regression coverage for both standalone Trainable Tokens and the LoRA auxiliary path? The tests can use a tiny locally initialized BERT with nonzero-bias and zero-bias controls, so CI would not need to download this checkpoint.

I did not find an overlapping issue or open PR; #2863 fixed Linear initialization rather than this forward-path behavior. I will wait for maintainer confirmation before implementing the patch.

Contributor guide

Open the contributing guide

Research direction

Start with the Linear branch in src/peft/tuners/trainable_tokens/layer.py at the referenced forward path, then trace how the standalone Trainable Tokens and LoRA trainable_token_indices paths use it. Add regression coverage with a tiny locally initialized BERT using nonzero- and zero-bias controls, and verify active, disabled, and merged inference preserve the original output.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
75/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.