Skip to content

fix(monkey_patch): let the norm flags cover the whole vision tower - #1443

Open
yupengtang wants to merge 2 commits into
linkedin:mainfrom
yupengtang:fix-vision-tower-layer-norm-flag
Open

yupengtang wants to merge 2 commits into
linkedin:mainfrom
yupengtang:fix-vision-tower-layer-norm-flag

Conversation

@yupengtang

@yupengtang yupengtang commented Sep 7, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Five patchers gate the per-layer norms of a vision tower on the corresponding flag, but patch the tower's own norms unconditionally. Passing the flag as False together with a model= instance therefore still swaps Liger norms in:

patcher flag patched regardless
llama4, mllama layer_norm vision_model.layernorm_pre, layernorm_post
gemma3, paligemma layer_norm vision_model.post_layernorm
glm4v_moe rms_norm visual.post_conv_layernorm, post_layernorm

The class-level branch of the very same functions gets it right (if layer_norm and model is None: in mllama/gemma3/paligemma, if rms_norm: in glm4v_moe), so one function answers the same flag differently depending on whether you hand it a model. And smolvlm, internvl, muse_glimmer and glm4v already gate the tower norm exactly like their encoder layers, which is why these five read as oversights rather than deliberate exceptions.

Both flags default to True on all five, so nothing changes unless a caller explicitly opts out.

Testing Done

Reproduced on transformers 5.9.0 / torch 2.11.0 before the change, patching an instance with the flag off:

paligemma (layer_norm=False)
  vision_model.post_layernorm             PATCHED  <-- ignored the flag
  encoder.layers[0].layer_norm1           untouched
mllama (layer_norm=False)
  vision_model.layernorm_pre              PATCHED  <-- ignored the flag
  vision_model.layernorm_post             PATCHED  <-- ignored the flag
  transformer.layers[0].input_layernorm   untouched

gemma3, llama4 and glm4v_moe behave the same way. After the change every one of those reads untouched.

  • pytest test/transformers/test_monkey_patch.py: 63 passed / 1 skipped on main, 68 passed / 1 skipped here. The five added tests each fail on main and pass with the fix; the skip is muse_glimmer missing from the installed transformers.
  • ruff check . and ruff format --check . with the v0.14.11 that .pre-commit-config.yaml pins: clean.
  • Convergence is untouched by construction. Both test_mini_models_multimodal.py files pin "rms_norm": True and kwargs["layer_norm"] = True for these models, glm4v_moe is not in the convergence set at all, and this change only alters the False branch.

Also ran on GPU. This branch against main, both on the same A100 40GB in one job:

this branch main
test_mini_models_multimodal.py bf16 (-k "mllama or paligemma or gemma3 or llama4 or glm4v") 5 passed 5 passed
the same, fp32 4 passed, 1 xfailed 4 passed, 1 xfailed
test_monkey_patch.py 68 passed, 1 skipped 63 passed, 1 skipped

Full make test scope, separately on an A100 80GB: 4496 passed, 18 failed, 931 skipped, 15 xfailed. Re-running just test/chunked_loss/test_grpo_loss.py on main gives the same 18 failing test ids (luspo/cispo cases that miss tolerance on sm80), and that file does not import monkey_patch.

Hardware Type: A100 40GB and A100 80GB, plus CPU for the unit suite. The patchers are pure Python and test_monkey_patch.py has no CUDA gating, so that suite is the one that actually exercises them.

yupengtang and others added 2 commits September 7, 2026 08:00
When a model instance is passed in, five patchers gate the per-layer norms
on the corresponding flag but patch the vision tower's own norms
unconditionally:

- llama4, mllama          layernorm_pre / layernorm_post   (layer_norm)
- gemma3, paligemma       post_layernorm                   (layer_norm)
- glm4v_moe               post_conv_layernorm / post_layernorm (rms_norm)

So layer_norm=False (or rms_norm=False for glm4v_moe) still leaves Liger
norms on those models.

The class-level branch of the same functions already gets this right
(`if layer_norm and model is None:`, `if rms_norm:`), and smolvlm,
internvl, muse_glimmer and glm4v gate the tower norm the same way as their
encoder layers, so the instance path is the odd one out rather than a
deliberate exception.

Defaults are unchanged - both flags default to True on all five - and the
convergence tests pass layer_norm=True explicitly, so nothing that runs
today changes behaviour.

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.

2 participants