Skip to content

Fix batchless accumulated residual normalization - #1678

Open
emerardd wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/batchless-accumulated-resid-ln
Open

Fix batchless accumulated residual normalization#1678
emerardd wants to merge 1 commit into
TransformerLensOrg:dev-4.xfrom
emerardd:fix/batchless-accumulated-resid-ln

Conversation

@emerardd

Copy link
Copy Markdown
Contributor

Description

Fix ActivationCache.accumulated_resid(..., apply_ln=True) for caches whose batch dimension has been removed before selecting a scalar position.

The final-normalization recomputation path previously inferred batch and position semantics only from tensor rank. A batchless scalar residual stack has shape [components, d_model], so each component reached ln_final as a one-dimensional tensor and then failed when the code attempted squeeze(1).

accumulated_resid() now forwards the cache's explicit has_batch_dim state. The recomputation path temporarily restores the structural batch and position dimensions expected by ln_final, then reshapes each result back to its original component shape.

The new download-free regression matrix uses tiny native LN and RMS models. It compares batchless outputs against the corresponding row of the batched result across scalar, sliced, and unsliced positions, with and without normalization, for both cached-scale and final-normalization recomputation paths.

Fixes #1677

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Validation

  • 66 passed across the new unit regression matrix and the existing tests/acceptance/test_activation_cache.py surface
  • mypy .: success across 434 source files
  • pycln, isort, Black, and git diff --check: clean

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

No documentation text changes are required because this restores the existing batchless ActivationCache shape contract without changing the public API.

@emerardd
emerardd marked this pull request as ready for review August 15, 2026 10:57
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.

1 participant