fix(analysis): reject unfolded final norm in DLA - #1806
Merged
jlarson4 merged 1 commit intoSep 24, 2026
Merged
Conversation
Collaborator
|
Solid fix to this bug @emerardd, this is good to go and will be merged now! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
direct_logit_attribution()currently checkscompatibility_modebut acceptsenable_compatibility_mode(fold_ln=False)andno_processing=True. In those states a learned final-normalization scale is still applied during the forward pass, whileActivationCache.logit_attrs()projects againstW_Uwithout that scale. A tiny native Bridge with nonuniform final-LN weights produced an attribution sum of-2.2847for a token difference whose actual logit difference was-3.0057; the folded control reconstructed the difference.This change checks the final norm's current affine parameters before DLA: its effective weight must be identity and its bias zero. It fails clearly when either learned parameter remains active, while allowing parameter-free final norms and a folded norm. Checking live weights also catches changes made after compatibility mode was enabled. The API docstring and error messages describe the supported state and how to obtain it. No separate Issue was filed.
Type of change
Validation
fold_ln=False,no_processing=True, an active final-LN bias, post-fold edits, folded numerical reconstruction, and parameter-free LNPre/RMSPre.pycln --check --all, pinnedisort --check-only, Black--check,git diff --check, andmypy .passed.tests/unit/toolsrun on this Windows environment: 962 passed, 1 skipped, 2 failed intest_sparse_probing.py(test_default_tolerance_accepts_large_scale_activationsandtest_stop_reason_distinguishes_converged_from_capped_fits). Both nodes also showed failure in an isolated, unmodifiedupstream/devcheckout using the same locked environment; their optimizer reachedtolerance_gradrather than the test's expected refinement/max_evalroute. The DLA patch does not touch that solver or its tests. I am not claiming the complete unit suite passed locally.Checklist