Prefix sizing from per_layer_config ignores KV head count

[DaoyuanLi2816]

System Info

• PEFT main: 051b2c5d
• Reproduction is config-only; no model download or accelerator is required.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

#3471 provisions prefix tuning for the largest per-layer KV footprint, but the selection loop compares only head_dim:

if layer_head_dim > head_dim:
    head_dim = layer_head_dim
    num_key_value_heads = layer_num_kv

A layer can have the same (or smaller) head dimension and more KV heads, making its total KV width larger. In that case the override is ignored. Later, _get_layer_kv_target_shape() returns the wider layer shape and get_prompt() skips that layer because the prefix was provisioned with too few heads.

Minimal reproduction:

from peft import PrefixTuningConfig, TaskType
from peft.utils.other import _prepare_prompt_learning_config

config = PrefixTuningConfig(task_type=TaskType.CAUSAL_LM, num_virtual_tokens=4)
model_config = {
    "num_hidden_layers": 2,
    "hidden_size": 1024,
    "num_attention_heads": 8,
    "num_key_value_heads": 4,
    "head_dim": 128,
    "per_layer_config": {
        1: {"head_dim": 128, "num_key_value_heads": 8},
    },
}

_prepare_prompt_learning_config(config, model_config)
print(config.num_attention_heads, config.token_dim)

Current output:

4 512

The largest per-layer KV footprint is 8 * 128 = 1024, so the prefix should be provisioned as:

8 1024

The comparison should use layer_head_dim * layer_num_kv (including ties in head_dim) and retain the head/count pair that produces the largest footprint. A focused regression can exercise _prepare_prompt_learning_config() directly without depending on a Transformers 5.15 model. I can prepare that fix if this direction matches the intended behavior of #3471.

Expected behavior

Prefix tuning should provision enough KV width for the largest effective entry in per_layer_config, so a valid layer is not skipped solely because another layer has the same head dimension but fewer KV heads.

[DaoyuanLi2816]

@BenjaminBossan when you get a moment — is this one worth fixing at all? I opened #3479 with the fix and a config-level regression test, but I did that before the issue had any maintainer response, which I now realise is the wrong order. Happy to close the PR and leave this in the backlog, or to keep it open if the change looks right to you.

[BenjaminBossan]

I think it only makes sense fixing this type of very niche bug if there are models out there that are actually used and have this configuration. I haven't checked, did you encounter anything like that?

[DaoyuanLi2816]

No, I haven't — and having gone looking, I can't find one. You're right to push back.

I found this by reading the code, not from a model that broke. Two things I checked since your question:

No checkpoint uses per_layer_config yet. That branch is the forward-compat path for transformers >= 5.15. I pulled config.json for the current Gemma 4 releases — google/gemma-4-31B-it, gemma-4-12B-it, gemma-4-26B-A4B-it, gemma-4-E4B-it, plus the unsloth and RedHatAI mirrors — and none of them carry per_layer_config at all. They all still use global_head_dim / num_global_key_value_heads, which is the branch above the one I changed.

On the models that do exist, the existing comparison already picks the right layer. Real Gemma 4 has genuinely different head dims per layer kind, so layer_head_dim > head_dim resolves correctly. Checked against the real google/gemma-4-31B-it text config:

head_dim = 256, num_key_value_heads = 16          (sliding_attention -> footprint 4096)
global_head_dim = 512, num_global_key_value_heads = 4   (full_attention -> footprint 2048)
per_layer_config present = False
_prepare_prompt_learning_config -> num_heads=4, head_dim=512, footprint=2048

The case my patch guards against needs two layer kinds with the same head_dim and different KV head counts, and that combination doesn't appear in anything shipping. So this is a latent inconsistency in a branch nothing reaches, which by your criterion isn't worth a fix. I've closed #3479.

I'd leave this issue open only as a marker for when checkpoints do start shipping per_layer_config, since the branch becomes live then — but if you'd rather not carry that, close it and I won't reopen.

One observation from the same check, which looks intended but I'll mention since it's about a model people actually use. With the prefix provisioned at the global footprint (4 heads x 512), the injection guard in PeftModel skips any layer whose KV is wider, and on gemma-4-31B-it that is 50 of 60 layers — every sliding_attention layer is skipped, and prefix tuning lands on the 10 full_attention layers only. The warning does fire, and the code comment says this matches the default in google-deepmind/gemma#631, so I read it as deliberate rather than broken. Just flagging in case it's worth saying explicitly in the prefix-tuning docs that on heterogeneous-attention models the prefix reaches only the global layers.