Trainable Tokens drops the output-head bias before merge

[DaoyuanLi2816]

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.

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:

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 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.