Skip to content

Fix subclass tracking of Transformer masks and registered assignments - #154

Merged
Essoz merged 4 commits into
mainfrom
codex/subclass-tracker-assignment
Oct 27, 2026
Merged

Essoz merged 4 commits into
mainfrom
codex/subclass-tracker-assignment

Conversation

@Essoz

@Essoz Essoz commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

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.Parameter objects 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:

  • 29 assignment/identity/ownership regressions and six lifecycle tests pass. Baseline regressions reproduce the mask failure, alias loss, stale retained-source state, cyclic-root leakage and two bounded deadlocks during sampling.
  • Independent review reproduced both GC failures and confirmed their fixes, including collection of either shared owner and GC during both sampling modes.
  • Final composition of Fix subclass tracking of Transformer masks and registered assignments #154–Snapshot constructor receivers generally without invoking application getters #157 at 035f7f1976c5d84bd01f5d96bb095fb07c9cc059: 246/247 tests pass. The sole failure, test_inject_eval_loop, also fails on unchanged main 9ba90ac and 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.
  • On labmach, six final training cases pass exact initial/final state, every loss and Python/NumPy/PyTorch RNG checks: Transformer plain/API/VAR/combined and MNIST plain/combined. Every tracked parameter has changing data hashes.
  • Final real two-step Transformer CLI runs also pass ordinary, selective and full-flag collection, preserving losses, parameters, gradients, mask and Torch CPU RNG. All 12 parameter paths appear, with complete API pairs, no duplicate raw records, and two active passing invariant decisions in each mode.

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_module and delattr remain outside setter hooks. Model mutation must be serialized, and distinct tracked roots need distinct name namespaces.

All commits are signed. Targets main directly. #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-20261027 and outputs/instrumenter-final-cli-20261027.

@Essoz
Essoz marked this pull request as ready for review October 26, 2026 19:45
@Essoz
Essoz merged commit b66b39a into main Oct 27, 2026
0 of 2 checks passed
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