Skip to content

fix(peft): patch the LayerNorm copies, not the ModulesToSaveWrapper itself - #1425

Open
yurekami wants to merge 1 commit into
linkedin:mainfrom
yurekami:fix-peft-layernorm-patching
Open

yurekami wants to merge 1 commit into
linkedin:mainfrom
yurekami:fix-peft-layernorm-patching

Conversation

@yurekami

Copy link
Copy Markdown

Root cause

The PEFT branch of _patch_layer_norm_module (src/liger_kernel/transformers/monkey_patch.py:103) diverges from its _patch_rms_norm_module sibling in three ways, all defects:

  1. It binds LigerLayerNorm.forward/extra_repr onto the ModulesToSaveWrapper itself. The instance-level forward shadows PEFT's _forward_wrapped dispatch (which routes to the trainable modules_to_save.default copy when adapters are enabled, and to original_module under disable_adapters), silently breaking PEFT semantics. The RMSNorm branch from the same PR correctly leaves the wrapper untouched.
  2. hidden_size is set on the wrapper and (twice) on original_module, but never on modules_to_save.default — yet extra_repr (which reads self.hidden_size) is bound to it, so repr()/print() of a patched model raises AttributeError: 'LayerNorm' object has no attribute 'hidden_size'.
  3. The duplicated original_module.hidden_size assignment is dead weight.

The branch is reached whenever apply_liger_kernel_to_*(model=...) patches a LayerNorm wrapped by a PEFT config's modules_to_save (e.g. vision-tower layernorms in mllama/InternVL/SmolVLM-style setups). It was introduced by #632 (which fixed #631 for the RMSNorm path) and has zero test coverage — grep -c peft test/transformers/test_monkey_patch.py on main is 0.

Fix

Mirror _patch_rms_norm_module: patch only modules_to_save.default and original_module, set hidden_size on both, and leave the wrapper's forward alone so PEFT's adapter dispatch stays intact. Net −3 lines in src. Adds a regression test covering the PEFT wrapper path, including the invariant that "forward" not in wrapper.__dict__.

Verification

Executed on CPU (WSL, torch 2.13, transformers 4.56.2, peft 0.20.0):

  • Red: on main, the new test fails with AttributeError: 'LayerNorm' object has no attribute 'hidden_size' on print(wrapper), and the wrapper's __dict__ carries the shadowing forward.
  • Green: with the fix, full test/transformers/test_monkey_patch.py passes — 44 passed, 19 skipped, 0 failed. ruff check and ruff format clean.

Disclosure: the GPU kernel suite (make test) and convergence tests were not run — no GPU on this box. The change touches only patch-time attribute/binding structure, no kernel code.

…tself

_patch_layer_norm_module bound LigerLayerNorm.forward/extra_repr onto the
PEFT ModulesToSaveWrapper, bypassing PEFT's adapter dispatch, and set
hidden_size on the wrapper and (twice) on original_module but never on
modules_to_save.default - so repr() of a patched model raised
AttributeError. Mirror _patch_rms_norm_module: patch only
modules_to_save.default and original_module, and set hidden_size on both.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Error when training using FSDP and LoRA with modules_to_save parameter

1 participant