feat(native): add param-free pre-norm (LNPre/RMSPre) and fold_ln support - #1580
Conversation
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Hey @MdSadiqMd, thanks for taking this on!
The param-free norm half of this PR is solid. The math is good, and the code is well structured. The fold_ln half is not quite there yet, and I have a handful of requests below. Let me know if you have any questions!
| self.supports_center_writing_weights = False | ||
| # Support fold_ln (LN weight folds into downstream layers) and | ||
| # center_writing_weights (residual-stream-writing weights are centered). | ||
| self.supports_fold_ln = True |
There was a problem hiding this comment.
Setting supports_fold_ln true turns on a path that mis-folds: with non-identity LN weights the folded model's logits diverge relative to the unfolded model, and enable_compatibility_mode() on an LNPre model crashes with an EinopsError inside fold_layer_norm. The root cause is that weight_processing_conversions stays {}, so convert_tensor_to_tl_format passes native (out, in) Linear weights through unchanged while the fold formulas assume TL [head, d_model, d_head] format. Issue #1570 called for populating the conversions alongside the flag flip. Could you wire the conversions? The base adapter ships a shared helper _qkvo_weight_conversions that other split-QKV adapters use.
| if force_rms or _uses_rms(cfg): | ||
| norm_type = (getattr(cfg, "normalization_type", None) or "LN").upper() | ||
|
|
||
| if _uses_param_free_norm(cfg): |
There was a problem hiding this comment.
Returning a plain GeneralizedComponent for LNPre/RMSPre drops hook_scale and hook_normalized. HookedTransformer's LayerNormPre exposes those attributes and ActivationCache.apply_ln_to_stack consumes them. run_with_cache on an LNPre native model records zero scale/normalized hooks and ActivationCache.apply_ln_to_stack raises KeyError on ln_final.hook_scale. Could the pre-norms get a norm-shaped bridge component that keeps these hooks, the way NormalizationBridge does?
| super().__init__() | ||
| self.eps = eps | ||
|
|
||
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
NativeRMSNormPre normalizes in fp32 and casts back (lines 76–80), matching HT LayerNormPre's reduced-precision handling, but NativeLayerNormPre computes its mean and variance in the input dtype. Should the two new classes share the same dtype policy so half-precision native models normalize consistently?
| if force_rms or _uses_rms_norm(cfg): | ||
| norm_type = _normalization_type(cfg) | ||
|
|
||
| if force_rms or norm_type in ("RMS", "RMSPRE"): |
There was a problem hiding this comment.
With normalization_type="LNPre" and final_rms=True, this branch builds a parameterized NativeRMSNorm for ln_out, but HookedTransformer's contract for that config is the param-free RMSNormPre. The bridge side already treats it as param-free (supported_architectures/native.py:48 fires before its force_rms check), so the two sides now disagree about whether ln_out has weights. Could this branch return the param-free RMS norm whenever the config is a pre-norm type, so the module matches both the bridge wrapper and HT?
| optimizer.zero_grad() | ||
|
|
||
|
|
||
| def test_boot_native_fold_ln_output_invariant(): |
There was a problem hiding this comment.
This test cannot fail. boot_native initializes norms to identity, making the scale-fold a no-op, and the 0.01 loss tolerance hides the logit error the current fold introduces. Can we randomize the norm parameters and compare at the logit level, similar to how tests/unit/test_weight_processing.py:838-877 works? It would also be nice to add a LNPre-config case, since that workflow currently crashes.
|
Hey @jlarson4 implemented the changes. Let me know if anything is required from my side |
|
Looks great! Merging now, thank you @MdSadiqMd! |
2 participants
Description
Add param-free pre-norm support (LNPre/RMSPre) and enable fold_ln on the native bridge adapter.
Problem: Silent aliasing where "RMSPre" was treated as "RMS" (creating parameterized norms) and "LNPre" fell through to
nn.LayerNorm. Both silently produced parameterized norms instead of param-free ones.Solution:
NativeRMSNormPreandNativeLayerNormPreparam-free norm classes_make_normto distinguish RMS vs RMSPRE, LN vs LNPREsupports_fold_lnandsupports_center_writing_weightson native adapterMotivation: Deprecating HookedTransformer requires its toy-model demos (
demos/Othello_GPT.ipynb, in the CI notebook matrix) to run on TransformerBridge, which needsnormalization_type="LNPre"+ folding.Fixes #1570
Type of change
Screenshots
N/A
Checklist: