[`peft`] Support key_mapping with PEFT models
Author: tomaarsen
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=46766)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46766)
<!-- ci-dashboard-badge:end -->

Hello!

# What does this PR do?

When a `key_mapping` is passed to `from_pretrained` for a model that also loads a PEFT adapter, the mapping was applied to the base model weights but silently **not** to the adapter (LoRA) weights. As a result, an adapter whose keys only line up with the model after the mapping is dropped on load, leaving the LoRA at its fresh/zero init, i.e. the model silently degrades to the untrained base.

This PR makes `load_adapter` reuse the conversion mapping that `from_pretrained` already computed (which includes the user `key_mapping`) instead of recomputing it from scratch.

## Details
`from_pretrained` builds the weight conversions once, including any user `key_mapping`, and stashes them on `load_config.weight_mapping`:

```python
# modeling_utils.py, PreTrainedModel.from_pretrained
weight_conversions = get_model_conversion_mapping(model, key_mapping, hf_quantizer)
load_config = LoadStateDictConfig(..., weight_mapping=weight_conversions, ...)
...
model.load_adapter(_adapter_model_path, ..., load_config=load_config, ...)
```

But `load_adapter` ignored `load_config.weight_mapping` and recomputed the mapping **without** the `key_mapping`:

## The fix

```python
# after
weight_conversions = load_config.weight_mapping or get_model_conversion_mapping(self)
```

`load_config.weight_mapping` already contains the model's built-in conversions plus the user `key_mapping` (and accounts for the quantizer), so the adapter weights now go through the same conversion the base weights did. The `or` fallback preserves behaviour for direct `model.load_adapter(...)` calls, where no `from_pretrained`-supplied mapping exists.

## Code Agent Policy

- [x] I confirm that this is not a pure code agent PR.

The original bug was spotted when I tried to integrate https://huggingface.co/vidore/colpali into Sentence Transformers using agents. This model was trained with `colpali-engine`, whose `ColPali` class wraps PaliGemma under an extra `model.` level. Loading the adapter onto a current PaliGemma backbone therefore needs `key_mapping={"^model\\.": ""}` to strip the wrapper. This works fine for a pure PaliGemma model (e.g. https://huggingface.co/vidore/colpali-v1.2-merged), but not for a PEFT Adapter on a PaliGemma model (e.g. https://huggingface.co/vidore/colpali) as the `key_mapping` is ignored. An Agent also helped write the test.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://huggingface.co/docs/transformers/contributing) and the
      [Pull Request](https://huggingface.co/docs/transformers/pr_checks) checks?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes according to the [guidelines](https://github.com/huggingface/transformers/tree/main/docs)?
- [x] Did you write any new necessary [tests](https://huggingface.co/docs/transformers/testing)?


## Who can review?

@BenjaminBossan @githubnemo

- Tom Aarsen

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46766). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> Thanks for this PR. The fix generally looks good to me. I tried to understand if we can generally replace `get_model_conversion_mapping(self)` by `load_config.weight_mapping`, which would happen with this change; AFAICT, it should be okay, but ideally a maintainer can double-check.

Yeah, we can probably fully drop the `get_model_conversion_mapping(self)` part. cc @vasqu could you perhaps have a look?

> If I have a complaint, it's that the test looks a bit forced with the manual renaming of the keys, I wonder if a more realistic test could not be constructed that follows your report with usage of a nested model.

Oh yeah, it'll be possible to create and upload a tiny-random model that requires `key_mapping`, that's probably preferable. I just didn't want to reuse https://huggingface.co/vidore/colpali as it's a 3B model after loading the base model too. Will have a look at this.

- Tom Aarsen

> Oh yeah, it'll be possible to create and upload a tiny-random model that requires `key_mapping`

Or wouldn't it be enough to adjust the test to create a small custom module:

```python
class WrapperModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = AutoModelForCausalLM.from_pretrained(model_id).to(torch_device)
```

Or is it not enough if the indirection is at the outer-most level?

I uploaded https://huggingface.co/hf-internal-testing/tiny-random-paligemma-lora-key-mapping, which matches roughly my scenario with https://huggingface.co/vidore/colpali, except with a tiny-random base model instead. The test was also simplified a lot, i.e. just loading it with the required `key_mapping`, and then checking if the weights are indeed the expected sentinel values instead of randomly assigned.

I'm also okay to drop the ` or get_model_conversion_mapping(self)` fully, but I'll defer to a maintainer for that.

- Tom Aarsen

Edit: I see there's been changes in #46442 on the weight conversion code, but it still carried the same bug as before. I've pulled that work and applied the same fix on it. The below still applies.

Apologies for the delay, I totally forgot about this. I'd still like to push this through though, as it would allow me to simplify some model implementations a lot (e.g. avoid or shrink this file: https://huggingface.co/tomaarsen/colpali-hard-v1.1-st/blob/main/modeling_st_colpali.py#L21-L23)

I had an agent dig into it and the Trainer actually never reaches this `load_adapter`.

Both Trainer call sites are behind `_is_peft_model(...)`:
https://github.com/huggingface/transformers/blob/3fa92a082eb4b56df6a5e263e64f5b74ba979ce6/src/transformers/trainer.py#L3405-L3420
and
https://github.com/huggingface/transformers/blob/3fa92a082eb4b56df6a5e263e64f5b74ba979ce6/src/transformers/trainer.py#L3473-L3482
and that is `isinstance(model, (PeftModel, PeftMixedModel))`. For a `PeftModel`, `model.load_adapter` resolves to PEFT's own `peft.peft_model.PeftModel.load_adapter`, not `PeftAdapterMixin.load_adapter`:

That said, the fallback is still necessary, just for a different reason. `PeftAdapterMixin.load_adapter` is public API and is commonly called directly (`model.load_adapter(path)`) with no `load_config`, so `weight_mapping` is `None`. Without the fallback, `build_peft_weight_mapping(None, ...)` early-returns `[]`, which drops the mandatory PEFT renamings (`base_model.model.` stripping, the `.default` suffix), and the adapter weights then silently fail to load. Removing it fails `test_peft_from_pretrained_missing_keys_warning` and `test_peft_from_pretrained_unexpected_keys_warning`.

On testing the Trainer case: since the Trainer routes through PEFT's method, a Trainer test would not exercise this code path at all. What is worth pinning is the no-`load_config` path, so I have added `test_peft_load_adapter_without_load_config_recomputes_conversions`. It saves an adapter with sentinel weights, calls `reloaded.load_adapter(tmpdir)` with no `load_config` and no `key_mapping`, and asserts the weights come back. I confirmed it fails if the fallback is removed.
Long story short: it fails if we don't have `or get_model_conversion_mapping(self)`.

So I'll keep that fallback. I think this is ready for a final review (once I fix the merge conflict).

- Tom Aarsen


From my understanding, this change should be good from the PEFT perspective.

> Here it is not the case that we pass already converted mappings. But tbh maybe we should fix all calls to pass key mappings in some way or form

I agree that this might be the cleaner solution. Right now, it makes no difference compared to the suggested change AFAICT, but maybe it's more robust to future changes.

Regarding trainer: IIUC, training a Transformer model with PEFT weights loaded directly (i.e. not a `PeftModel`) would not work correctly when 1) loading from checkpoint, and 2) conversion is needed. This PR doesn't change the picture, so I'd be fine with that.

Understandable. Should we move forward with this fix in the meantime, however? It should simplify some things on my end.

- Tom Aarsen

From my perspective, it would be fine.

Also fine with this, just need to merge with main and recheck that CI doesnt complain :saluting_face: 

This new test might be failing as the file system is read-only: https://github.com/huggingface/transformers/actions/runs/29009986379/job/86091350053?pr=46766
Is that expected? I can look into a quick fix before we queue this for merge.

- Tom Aarsen

Ah crap it's about the read-only cache. I think it was fixed on main otherwise we need to ping in slack 

This branch was still outdated, I merged from `main`, we'll see how it does now.

- Tom Aarsen

Even more unrelated failures :cry: Edit: rerunning CI, let's see if they are flaky (or at least some of them)

The error was now
```
Can't find 'adapter_config.json' at 'hf-internal-testing/tiny-random-paligemma-lora-key-mapping'
```
But it seems that this error is always shown when `hf_hub_download` fails for whatever reason: https://github.com/huggingface/peft/blob/cea8213158c8b682acc0839405c2062d57fdf867/src/peft/config.py#L254-L259

So it's hiding the real error with the catch-all try-except, and surprise-surprise: it's our friend again:
```
>           mkdir(name, mode)
E           OSError: [Errno 30] Read-only file system: '/mnt/cache/hub/models--hf-internal-testing--tiny-random-paligemma-lora-key-mapping'
```

So I'll try setting `cache_dir` to avoid the `/mnt/cache/...`, I'll try and see if it helps.

- Tom Aarsen

run-slow: peft

Trying run slow iirc it does cache afterwards but not sure if it works on peft

P.s. when https://github.com/huggingface/transformers/pull/47338 is merged, I'll remove the cherry-picked commits, pull from main, and queue the merge for this PR.

- Tom Aarsen

run-slow: peft

Hey, at least the `tests_peft_integration` passed now! 🙃

- Tom Aarsen

Should be ready now, sorry release broke main a bit :grimacing: 

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46766)
**Latest run:** [29475698778:2](https://github.com/huggingface/transformers/actions/runs/29501723518)
**Result:** `success` | **Jobs:** 15 | **Tests:** 171,597 | **Failures:** 163 | **Duration:** 1h 41m

<!-- ci-dashboard-recap:end -->