Repository navigation
Filter adapter-name slots by tuner-prefix position - #3623
Sravanjangam wants to merge 10 commits into
Conversation
The negative filter in get_peft_model_state_dict matches by plain string containment, so an adapter whose NAME equals a module path segment of the saved adapter's tensors (e.g. an adapter called 'mlp' while a submodule is named 'mlp') silently removed those tensors from the checkpoint — up to and including writing a completely empty adapter_model.safetensors with no warning at all. Detect this collision before filtering and emit a RuntimeWarning naming the affected tensors, so silent truncation becomes visible. A structural (positional) filter is the cleaner long-term fix and is proposed in the issue; this makes the failure mode observable without changing behavior. Fixes huggingface#3584
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for tackling this issue.
I found a couple of problems with the current PR and flagged them. Could you please share which LLM you are using, it feels rather weak to me.
| self.mlp = nn.Linear(32, 32) | ||
| self.act = nn.ReLU() | ||
|
|
||
| def forward(self, x): |
There was a problem hiding this comment.
forward is not used and can thus be removed.
|
|
||
| return Net() | ||
|
|
||
| def _save_and_load(self, model, save_dir, selected_adapters): |
There was a problem hiding this comment.
Why do we need this? A simple save_pretrained -> from_pretrained should be enough. The returned config is not being used in the tests. To get the PEFT-specific weights from the model, use get_peft_model_state_dict.
| tensors = load_file(str(save_dir / "adapter_model.safetensors")) | ||
| return config, tensors | ||
|
|
||
| def test_partial_collision_saves_all_default_tensors(self, mlp_net, tmp_path, recwarn): |
There was a problem hiding this comment.
The recwarn fixture here and below is not used.
| torch.manual_seed(0) | ||
| model = get_peft_model( | ||
| mlp_net, | ||
| LoraConfig(r=4, lora_alpha=8, target_modules=["lin0", "mlp"], lora_dropout=0.0), |
There was a problem hiding this comment.
For a proper test, this should set init_lora_weights=False as B is otherwise always 0 and thus the values are always equal.
| } | ||
| assert set(tensors.keys()) == expected_keys | ||
|
|
||
| def test_save_load_round_trip_with_colliding_adapter_name(self, mlp_net, tmp_path, recwarn): |
There was a problem hiding this comment.
This whole test doesn't contribute much. The other tests already establish that the key is missing from the checkpoint, so the fact that the weight is wrong trivially follows from that.
| "base_model.model.mlp.lora_A.weight", | ||
| "base_model.model.mlp.lora_B.weight", | ||
| } | ||
| assert set(tensors.keys()) == expected_keys |
There was a problem hiding this comment.
We should also check the right keys when selecting the adapter name "mlp".
| using modules_to_save. | ||
| """Filter the state dict to remove keys that correspond to the unwanted adapter. | ||
|
|
||
| A naive `f".{name}." in k` check is too permissive: if an unwanted adapter's NAME happens to equal a |
There was a problem hiding this comment.
This refers to the previous implementation, but when reading the code, the previous implementation is not visible. The explanation must make sense on its own and not require users to dig into the code history.
| state_dict: dict[str, torch.Tensor], unwanted_adapter_names: list[str] | ||
| state_dict: dict[str, torch.Tensor], | ||
| unwanted_adapter_names: list[str], | ||
| selected_adapter_name: str | None = None, |
There was a problem hiding this comment.
You changed the signature of the function but didn't update all the call sites. This function is also called in tuners/shira/model.py.
| # Keep the historical behaviour to avoid behaviour changes in that niche path. | ||
| return f".{name}." in key or key.endswith(f".{name}") | ||
|
|
||
| return {k: v for k, v in state_dict.items() if not any(_is_unwanted(k, n) for n in unwanted_adapter_names)} |
There was a problem hiding this comment.
This all seems too complicated to me. I tried a different version which takes advantage of the already existing _get_tuner_state_dict_key_prefixes and the tests pass for me:
def _filter_state_dict_for_adapter_name(
state_dict: dict[str, torch.Tensor],
unwanted_adapter_names: list[str],
model,
) -> dict[str, torch.Tensor]:
# avoid circular import
from peft.tuners.tuners_utils import _get_tuner_state_dict_key_prefixes
prefixes = _get_tuner_state_dict_key_prefixes(model)
unwanted = set(unwanted_adapter_names)
filtered_state_dict = {}
for key, value in state_dict.items():
keep = True
for prefix in prefixes:
prefix = prefix + "."
if not key.startswith(prefix):
continue
suffix = key[len(prefix) :]
adapter_name = suffix.partition(".")[0]
if adapter_name in unwanted:
keep = False
break
if keep:
filtered_state_dict[key] = value
return filtered_state_dict…efixes An adapter whose name equals a base-model module segment (e.g. adapter 'mlp') silently truncated the checkpoint: the containment filter dropped the selected adapter's own tensors. Match positionally using _get_tuner_state_dict_key_prefixes instead of mining prefixes from keys. Per BenjaminBossan on huggingface#3623: adopt the simpler model-derived implementation (both call sites incl. shira), standalone comments, tests via get_peft_model_state_dict + plain save/load round trip, mlp-first coverage, init_lora_weights=False where B=0 would trivialize. Verification: - 5 collision tests pass (2x); 4 fail on upstream 5d602fd, foo-guard passes - full test_initialization 349 passed (3 FrodConfig fail identically on base) - test_shira 11 passed; mixed modules_to_save same-as-base (7 Frod pre-existing) - ruff check + format clean; doc-builder style --check_only exit 0 PR huggingface#3623.
|
Hi @BenjaminBossan — all ten points addressed in Your implementation (save_and_load.py:113): adopted as proposed, using the existing Missed call site (save_and_load.py:79): Standalone comment (save_and_load.py:83): rewritten with no history references. Tests (custom loader deleted;
|
There was a problem hiding this comment.
🤗 Serge says:
Verdict
The fix is correct: the old substring filter (f".{adapter_name}." in k) conflated adapter-name slots with base-model module path segments; the new positional check — drop only when the segment immediately after a tuner prefix (lora_A, lora_embedding_A, shira_indices, …) is an unwanted adapter name — resolves #3584 without weakening multi-adapter filtering.
Correctness
- Verified
_get_tuner_state_dict_key_prefixes(tuners_utils.py:1673): prefixes are derived structurally from tuner layers/BaseTunerattributes, excludebase_layer/modeland auxiliary wrappers, and handle the FSDP prefix — so the positional match is robust and consistent withBaseTuner._get_adapter_state_dict's own key selection. - Keys outside any tuner prefix are now always kept. Safe:
modules_to_save/AuxiliaryTrainingWrapperentries are re-selected per adapter downstream viaadapter_state_dict(adapter_name, ...)(other.py:789,save_and_load.py:216-221), which only readsmodules_to_save.{adapter_name}.*, so unwanted adapters' copies cannot leak into the checkpoint. - The
len(...) > 0guard atsave_and_load.py:163still behaves correctly for adaption_prompt: that tuner produces no structural prefixes, so the filter becomes a no-op and the guard logic is unchanged; its_get_adapter_state_dictselectsadaption_*keys independently. - The SHIRA call site correctly passes
model:shira_indicesis a tuner-owned buffer dict keyed by adapter name, which the positional filter now handles properly.
Tests
TestAdapterNameCollisionFiltering covers partial and full collision, reversed adapter order with value attribution, an end-to-end save/load round trip, and a no-regression guard for a non-colliding second adapter; init_lora_weights=False making the adapters' weights distinguishable is the right call. Non-blocking gap: all tests exercise LoRA's lora_A/lora_B prefixes; a case with a different tuner prefix shape (e.g. IA³'s ia3_l, or an adapter named like a tuner attribute such as lora_A) would harden the positional matching further.
Style
- The new
modelparameter lacks a type annotation while its siblings have one.
The description's claims (supersedes #3591, pre-existing FrodConfig failures, test counts) are consistent with the diff. Approving.
serge v0.1.0 · model: zai-org/GLM-5.3-Flash · 48 LLM turns · 53 tool calls · 379.0s · 1472048 in / 22456 out tokens
| state_dict: dict[str, torch.Tensor], unwanted_adapter_names: list[str] | ||
| state_dict: dict[str, torch.Tensor], | ||
| unwanted_adapter_names: list[str], | ||
| model, |
There was a problem hiding this comment.
Nit: the other parameters are annotated; keep this consistent.
| model, | |
| model: nn.Module, |
(nn.Module is already imported in this module.)
BenjaminBossan
left a comment
There was a problem hiding this comment.
@Sravanjangam Thanks for the update. I did my own review but also requested an LLM review from our Serge bot. I generally agree with that review (which is why I submitted it): The fix looks good, let's just extend the tests even more. This is a core functionality that affects all of PEFT, so I'd rather go the extra mile with unit tests to ensure it works as expected. Please add the test cases that the bot suggested.
|
Hi @BenjaminBossan — both bot-suggested items are addressed (plain pushes, no force-push), all gates re-run on the pushed tip: 1. Annotation ( 2. New hardening tests (
Verification (all on pushed tip
Nothing else in the diff changed — the fix itself is untouched since your last review. P.S. I did one final end-to-end verification on the latest tip after adding the additional test coverage. Verification gates:
I also re-ran the broader verification: Looks good to me from my side. just let me know any changes required |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update. I found one test that I'm not sure if it really contributes to the coverage, please check.
| model.add_adapter("foo", LoraConfig(r=4, lora_alpha=8, target_modules=["lin0"], lora_dropout=0.0)) | ||
|
|
||
| tensors = get_peft_model_state_dict(model, adapter_name="default") | ||
| assert all("foo" not in k for k in tensors.keys()) |
There was a problem hiding this comment.
Hmm, isn't this assert trivially correct? "foo" is the adapter name, which is not part of the output of get_peft_model_state_dict. This makes the whole test somewhat moot.
Per BenjaminBossan on huggingface#3623 at tests/test_initialization.py:6469: 'foo' is not part of get_peft_model_state_dict output, so the all('foo' not in k) assert alone is tautological. Fix: prove filtering actually removed something — assert the raw state_dict does contain 'foo', then assert the filtered view's keys are exactly the expected default set, plus value attribution that the returned tensor equals default's live weight (not foo's). Verification: - TestAdapterNameCollisionFiltering 7/7 pass twice - 6/7 fail on base 5d602fd (1 pass is this guard, now non-trivial) - full test_initialization 351 passed 3 failed (same Frod pre-existing) - ruff check + format clean; doc-builder style exit 0 PR huggingface#3623.
|
Hi @BenjaminBossan — addressed your Trivial assert (
Verification: |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update. I still have a few issues around the tests, but with those out of the way, the PR should be ready.
| "base_model.model.mlp.lora_B.weight", | ||
| } | ||
| # Value attribution: the returned tensor must be default's weight, not foo's. | ||
| assert torch.equal( |
There was a problem hiding this comment.
Okay, I see where this test is going now, but I think it's not necessary, please remove the test.
| ) | ||
| model.add_adapter("mlp", LoraConfig(r=4, lora_alpha=8, target_modules=["lin0"], lora_dropout=0.0)) | ||
|
|
||
| tensors = get_peft_model_state_dict(model, adapter_name="default") |
There was a problem hiding this comment.
Here and in the tests below, let's not call the variable tensors, as that sounds like a list of tensors. state_dict is a better fit.
| "base_model.model.mlp.lora_A.weight", | ||
| "base_model.model.mlp.lora_B.weight", | ||
| } | ||
| live = model.base_model.model.mlp.lora_A[name].weight |
There was a problem hiding this comment.
live is a bit of an odd name, let's just go with weight. Also, how about comparing all the weights, not just this one example?
Same argument applies to the tests below.
Per BenjaminBossan on huggingface#3623 (09-17): remove unnecessary IA3 test, rename tensors -> state_dict, live -> weight, compare all weights instead of one example (same for the tests below). - TestAdapterNameCollisionFiltering 6/6 pass (IA3 test removed per review)
|
Hi @BenjaminBossan — all four points addressed in
Verification: |
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update. Just a small suggestion left from my side, otherwise the PR is ready.
| module_name, param_name = key.split(".")[2], key.split(".")[3] | ||
| weight = getattr(getattr(model.base_model.model, module_name), param_name)[name].weight |
There was a problem hiding this comment.
IMO this is simpler:
module_name = key.rpartition(".")[0]
weight = model.get_submodule(module_name)[name].weightSame applies below.
There was a problem hiding this comment.
Done in ce8fd040 (plain push) — applied verbatim at all three sites, keeping each test's own index key ([name] / ["default"] / ["lora_A"]). Gates re-run on the pushed tip: TestAdapterNameCollisionFiltering 6/6 pass, ruff check + format clean.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@Sravanjangam Please run |
|
There are failing SHiRA tests, could you please investigate? |
The positional match alone missed keys that carry the adapter name as a suffix rather than after a tuner prefix. SHiRA keeps its indices in a plain dict on the layer, so they are neither parameters nor buffers and thus absent from the structural prefixes; the other adapters' entries leaked into the checkpoint (saving selected_adapters=["default"] wrote 6 keys instead of 2). Five SHiRA cases of test_save_pretrained_adapter_name_substring failed on CI for this reason and pass again with the suffix check restored, which was the previous behaviour. Also reflow the docstring per make style.
|
Hi @BenjaminBossan — both points handled.
Failing SHiRA tests — root-caused and fixed in Cause: the positional match depends on The fix restores the suffix check next to the positional one ( Verification at
|
Fixes #3584
Problem
save_pretrainedsilently writes an empty/truncated adapter checkpoint when an adapter's name equals a base-model module path segment (e.g. an adapter calledmlpwhile a submodule is literally namedmlp): the negative name filter matched by raw string containment, so savingdefaultdropped every key containing.mlp.— includingdefault's own tensors.Fix
_filter_state_dict_for_adapter_namenow matches positionally using the existing_get_tuner_state_dict_key_prefixes(model): a key is dropped only when the segment immediately after a tuner prefix (lora_A,ia3_l, …) is an unwanted adapter name. Keys outside any tuner prefix are kept;modules_to_savecopies are resolved per adapter downstream, so multi-adapter behavior there is unchanged.Supersedes #3591 — that one only added a warning; the thread on #3584 pointed at fixing the filter itself, which is what's here.
Tests
TestAdapterNameCollisionFiltering(5 tests): partial + full collision key sets viaget_peft_model_state_dict; reversed order (mlpfirst) with per-adapter value attribution; plainsave_pretrained→from_pretrainedround trip withinit_lora_weights=False; non-colliding second adapter still filtered.test_initialization.py349 passed (3 FrodConfig failures reproduce identically on the base commit — pre-existing);test_shira.py11 passed (call site updated); mixedmodules_to_saveselections match base (7 FrodConfig pre-existing failures, identical both sides).Test environment: Python 3.12, torch 2.13.0+cpu, transformers 5.15.1, peft @ b147931, ruff 0.16.4 (check + format clean).