Repository navigation
Fix subclass tracking of Transformer masks and registered assignments - #154
Merged
Merged
Conversation
Essoz
marked this pull request as draft
October 26, 2026 19:06
This was referenced Oct 26, 2026
Essoz
marked this pull request as ready for review
October 26, 2026 19:45
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.
Subclass tracking crashes when a Transformer forward method assigns an ordinary Tensor attention mask: the old hook treats it as a child model. It also mishandles registered replacements and tied parameters, splitting parameter identity or losing a live alias from observation.
Run PyTorch's native setter once, then inspect the actual registries, including registration-hook replacements. Ordinary tensors and buffers remain ordinary values. Promote exact
nn.Parameterobjects in place so retained references, tied weights, existing optimizer references, gradients, storage and autograd hooks stay consistent.Track live root/path ownership with weak references. Replacement, assignment to
None, retracking and collection preserve surviving aliases and prevent removed parameters from overwriting replacement registry entries. An externally retained finalizer handles cyclic roots. Registry retirement waits until the outermost sample finishes; reentrant locking, stable membership and live stale/type checks preserve record names, order and multiplicity during GC. Malformed internal hooks/bindings raise explicitly.Validation:
035f7f1976c5d84bd01f5d96bb095fb07c9cc059: 246/247 tests pass. The sole failure,test_inject_eval_loop, also fails on unchanged main9ba90acand is unrelated source injection. Pinned hooks and configured mypy across all 78 production files pass; strict installed-Torch checking of eight new/changed contract/test modules passes. Stronger legacy subclass checking retains eight existing errors, with none introduced.Correctness cost: a separate registry-only benchmark with 100 no-op variables and nine alternating repeats measures 8.726 → 16.209 microseconds/sample, adding 7.483 microseconds. Snapshot containers occupy 7,320 bytes excluding existing referents; this is not a training-overhead or incremental-RSS estimate. The final matched training study in #156 includes this repair on both sides.
Limits: custom Parameter subclasses retain legacy out-of-place conversion; their retained external references do not gain the exact-Parameter identity guarantee. Direct
register_parameter/register_moduleanddelattrremain outside setter hooks. Model mutation must be serialized, and distinct tracked roots need distinct name namespaces.All commits are signed. Targets
maindirectly. #155's default change should follow this fix; #156 and #157 are separate concerns. Exact integration evidence, logs and microbenchmark provenance are retained under/Users/matt/gitrepos/traincheck/outputs/instrumenter-final-validation-20261027andoutputs/instrumenter-final-cli-20261027.