fix(shadow): sync lm_head trainability on set_adapter for CAUSAL_LM

[Sravanjangam]

Hi @BenjaminBossan — thanks for ShadowPEFT and for the quick nod on #3626! This is the small trainability-sync fix we discussed, Only following your green light here.

Fixes #3626

Problem

ShadowModel._sync_shadow_module_trainability (src/peft/tuners/shadow/model.py:731-749) only walks shadow_backbone / shadow_projection / shadow_head to flip requires_grad. For task_type="CAUSAL_LM" ShadowPEFT intentionally does not create a shadow_head — it reuses the frozen base lm_head via _resolve_shadow_head (model.py:787-793) and lets users train it via modules_to_save=["lm_head"] (docstring 491-497, test 295-308).

lm_head lives on self.model, not in any shadow_* container, so it is never visited by the sync. After the first set_adapter("other") or set_adapter("default", inference_mode=True), the shadow modules are flipped but lm_head stays in its previous requires_grad state — frozen when it should be trainable (or vice versa). Loss still decreases (adapter trains), so the bug is invisible until eval — checkpoint saves an untouched head.

Who can trigger: Every ShadowConfig(CAUSAL_LM, modules_to_save=["lm_head"]) workflow that calls set_adapter (multi-adapter training, inference_mode toggling). Single-adapter without set_adapter is fine, which is why test_save_includes_trainable_lm_head passes.

Concrete failure:

• Before set_adapter: lm_head.weight.requires_grad=True, grad.norm()≈4.27
• After set_adapter("other") (bug): requires_grad=False, grad=None, norm 0 — head frozen the whole fine-tune

Solution

In _sync_shadow_module_trainability, also sync the base lm_head when it is the resolved shadow head for CAUSAL_LM. The approach is minimal and matches the existing shadow-container logic:

• should_train_head = not inference_mode and any active CAUSAL_LM adapter requests lm_head via modules_to_save
• Only touch the head if some adapter ever requested lm_head for CAUSAL_LM (to avoid touching unrelated SEQ_CLS or non-lm models)
• head = self.model.get_output_embeddings(); for p in head.parameters(): p.requires_grad = should_train_head

No API change, no new config, only fixes the missing walk. Single-adapter and SEQ_CLS paths unchanged.

Changes

• src/peft/tuners/shadow/model.py:742-761 → 17 lines added in _sync_shadow_module_trainability to sync base lm_head for CAUSAL_LM + modules_to_save=["lm_head"]
• tests/test_shadow.py:310-342 → 32 lines added: test_shadow_lm_head_trainable_after_set_adapter (repro from [Bug] ShadowPEFT LM-head trainability freeze after set_adapter — modules_to_save=["lm_head"] silently stops receiving grads #3626) and test_shadow_lm_head_frozen_in_inference_mode

Verification

Exact gates run fresh from the committed branch fix/shadow-lm-head-trainability@d531dbee after final commit, before PR:

• Repro before fix (5 runs, peft@9c16ee66): 0/5 OK — after_req=False, grad_after=False, norm 0.00 (BUG)

  Run 1: init_req=True grad_before=True norm_before=4.27 | after_req=False grad_after=False norm_after=0.00 -> BUG
  (x5 same)
  

• Repro after fix (5 runs, this branch): 5/5 OK — after_req=True, grad_after=True, norm ~4.27

  Run 1: init_req=True grad_before=True norm_before=4.28 | after_req=True grad_after=True norm_after=4.27 -> OK
  Run 2: init_req=True grad_before=True norm_before=4.28 | after_req=True grad_after=True norm_after=4.28 -> OK
  Run 3: init_req=True grad_before=True norm_before=4.28 | after_req=True grad_after=True norm_after=4.27 -> OK
  Run 4: init_req=True grad_before=True norm_before=4.28 | after_req=True grad_after=True norm_after=4.27 -> OK
  Run 5: init_req=True grad_before=True norm_before=4.27 | after_req=True grad_after=True norm_after=4.29 -> OK
  

• Additional metrics (after fix):
  • inference_mode=True → lm_head.requires_grad=False (expected)
  • inference_mode=False → True (expected)
  • Without modules_to_save → False before and after set_adapter (expected, no touch)
• New tests (HF_HUB_CACHE=/tmp/hf_cache):

  pytest tests/test_shadow.py::TestShadowCausalLM::test_shadow_lm_head_trainable_after_set_adapter -v --no-cov
  → 1 passed in 12.07s
  pytest tests/test_shadow.py::TestShadowCausalLM::test_shadow_lm_head_trainable_after_set_adapter + test_shadow_lm_head_frozen_in_inference_mode -q --no-cov
  → 2 passed in 8.93s
  

• Style: ruff check src/peft/tuners/shadow/model.py tests/test_shadow.py → All checks passed; ruff format --check → 2 files already formatted
• Existing tests: test_shadow.py baseline parity — 3 passed, 8 skipped, 29 failed pre-existing (all PermissionError cache on peft-internal-testing/tiny-random-LlamaForCausalLM when run without HF_HUB_CACHE=/tmp/hf_cache; same failures on main, not introduced by this diff). With HF_HUB_CACHE=/tmp/hf_cache, new tests pass.

Environment

• OS: macOS 15.x arm64 (Darwin CBG5APLTJ2FHG9JVXV.local)
• Python: 3.12.11, torch 2.13.0, transformers 5.15.1, safetensors, peft 0.20.1.dev0
• Branch: fix/shadow-lm-head-trainability@d531dbee (fix bba8d99f, test d531dbee), base main@9c16ee66, 0 behind upstream/main at creation
• Commands: HF_HUB_CACHE=/tmp/hf_cache pytest tests/test_shadow.py -k "lm_head" and ruff check/format as above

[BenjaminBossan]

@SeanLee97 Do you have time to check this?

[Sravanjangam]

Hi @BenjaminBossan and @SeanLee97 thanks for taking a look! Happy to hear any suggestions and ready to improve on the change let me know what you'd like to see.

[BenjaminBossan]

Thanks for the fix, it's almost good to go, just two small comments from my side.

[BenjaminBossan]

Wouldn't this work?

Suggested change

	for p in head.parameters():
	p.requires_grad = should_train_head
	head.requires_grad_(should_train_head)

[BenjaminBossan]

This pretty much duplicates the logic just 3 lines above. Let's rewrite this block to determine this once.

[Sravanjangam]

Hi @BenjaminBossan — both addressed in 8f2015c2, thanks for the suggestions.

1. head.requires_grad_(should_train_head) (model.py:758) — done, replaces the per-parameter loop. Verified nn.Module.requires_grad_ covers all head parameters; the 3 lm_head tests still pass.

2. Single-pass block (model.py:747-755) — done. One loop over peft_config.items() now sets both manages_head (any adapter routes lm_head via modules_to_save for CAUSAL_LM) and should_train_head (such an adapter is active and not in inference mode). Behavior is identical: head untouched when no adapter manages it, frozen in inference mode, trainable otherwise.

Verification on fix/shadow-lm-head-trainability@8f2015c2:

• pytest tests/test_shadow.py -k "lm_head" → 3 passed
• ruff check + ruff format --check → clean; doc-builder style --check_only → exit 0

[HuggingFaceDocBuilderDev]

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.

[BenjaminBossan]

Thanks for the updates. I did another review and checked the tests you added. It turns out they're already passing on the current main branch. So I'm not sure if the fix is unnecessary or if the tests are too weak. Remember that for bugfix PRs, the tests should first be written so that they fail on main and then ensure that they pass with the fix.

[Sravanjangam]

Hi @BenjaminBossan — full update on your 09-07 review, with a correction to my earlier note below.

Your review was right about the set_adapter paths: I verified end to end (10 scenarios on main@9c16ee66 vs this branch, tiny Llama on CPU) — the issue's exact repro, both-adapters param maps, inference-construct back-to-train, double toggle, disable-toggle, asymmetric adapters, roundtrip grad flow. All identical and correct on main: ModulesToSaveWrapper.set_adapter already handles per-adapter granularity there. The old lm_head block in this PR never fired (always-False str() == str() comparison), and both old tests passed on main — all removed.

But one path over, the invariant genuinely breaks on main — delete-fallback (FAILS on main, PASSES with fix): default manages lm_head, other does not; set_adapter("other") then delete_adapter("other") falls back to default, yet the head stays frozen while its managing adapter is active:

main:  wrapper_active=['default'] copy_requires_grad=False   # wrong
fixed: wrapper_active=['default'] copy_requires_grad=True    # + grad flows, norm > 0

Root cause (ModulesToSaveWrapper.delete_adapter, src/peft/utils/other.py): the early return when the deleted adapter has no copy drops the fallback sync the normal path performs, leaving a stale active. The fix (8 lines, pushed 62d982da, no force-push) mirrors the normal path, so deleted-with-copy and deleted-without-copy behave identically — benefits every tuner, not just shadow.

Verification: new bug-fix test asserts requires_grad + grad-norm (true FAIL-on-main → PASS); full test_shadow.py 40 passed / 1 skipped; adjacent delete_adapter/modules_to_save selections 926 passed (7 FrodConfig failures reproduce identically on clean main — pre-existing); ruff + doc-builder clean. Net diff vs base: other.py +8, test_shadow.py +18, shadow/model.py untouched. The issue's original set_adapter-only expectation remains wrong (that matrix passes on main) — the delete-fallback is the case that doesn't. Leaving the close call on #3626 to you.