Fix silent hook_normalized/hook_scale no-op on NormalizationBridge native-autograd path

[Hetens]

Description

Fixes #1526.

On the native-autograd path (use_native_layernorm_autograd=True, hardwired in 13 adapters incl. Gemma 3, NeoX/Pythia, OLMo, Phi, GPT-OSS), NormalizationBridge fired hook_scale / hook_normalized on detached observation values and discarded their return values: forward edits silently no-oped and backward hooks never fired, while the default python-norm path honored both.

Per the dispatch design agreed in #1526:

• Backward hooks attached: falls back to the python-norm path so hook tensors stay in the autograd graph.
• Forward hook edits: detected by object identity (a hook returning non-None replaces the tensor, matching HookPoint semantics); the output is reconstructed from the hooked values and stays grad-connected to the input. When hook_scale is edited, hook_normalized is recomputed from the edited scale first, mirroring the python path's ordering.
• Observation-only hooks (run_with_cache caching hooks return None) and hook-free forwards: unchanged, bit-identical native HF output and exact autograd. Guarded by a dedicated regression test that fails on any warning.
• Both fallbacks emit a UserWarning (the "solution 1" stopgap), since fallback numerics differ from HF's at float-rounding scale.

Known limitation (documented in a comment): a hook that mutates the tensor in place and returns None is not detected as an edit. Return the tensor from the hook to edit.

Based on main per maintainer request in #1526 (patch release). normalization.py is currently identical on main and dev, so this merges cleanly upward. The _last_input_before_norm dead-code cleanup from #1526 will follow as a separate PR to dev.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Notes on the checklist:

• No new warnings: no warnings on any previously-working configuration; the two UserWarnings introduced fire only in the situations that previously silently produced wrong results, which is the point of the fix.
• Docs: no docs page covers these hook points' per-path semantics today; happy to add one if wanted.

Validation

• New test file red on unfixed main (verified by swapping in origin/main's normalization.py): 5 failed, 4 passed. The 5 failures are exactly the bug-demonstrating tests; the 4 passes are the behavior-preservation guards.
• On this branch: all 9 new tests pass; full tests/unit -m "not slow" suite green: 3035 passed, 0 failed, 32 skipped, 9 xfailed (20 min, CPU).
• make format / make check-format clean, uv run mypy . clean (no # type: ignore added).

Found, diagnosed, and fixed with the assistance of Claude (Anthropic's AI assistant), working under my direction; I reviewed the changes and the validation runs are from my machine.

[jlarson4]

@Hetens great work on this! Just this one small note from me.

The Gemma family of models applies an offset of +1 to their weight. Without a check for that here, we will produce incorrect results. Should be an easy fix, just add w = (1.0 + self.weight) if getattr(self.config, "rmsnorm_uses_offset", False) else self.weight. We will want to also add a test that covers this specific case.

[Hetens]

Good catch, done!