Repository navigation
Preserve TorchScript builtin identity when wrapping PyTorch APIs - #143
Merged
Merged
Conversation
This was referenced Oct 25, 2026
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.
Wrapping a registered PyTorch builtin changes its function ID, so TorchScript loses the existing operator mapping. A real-file tanh-based bias-GELU function then fails to compile with
TypeError: ... expected, got builtin_function_or_method, although eager execution still works.Before publishing a wrapper, copy the original function's existing TorchScript builtin mapping. Feature-detect the private lookup/registration APIs and leave unregistered functions unchanged. Retain registered wrappers so Python cannot reuse an ID still present in TorchScript's registry.
Validation:
aten::tanhin the graph. Outputs, input gradients, and bias gradients exactly match the unwrapped baseline within each platform.Scripted graphs execute the operator directly and emit no inner Python tanh records. This fixes the focused builtin-identity failure; full Megatron startup and checkpoint-shape compatibility remain unverified. Registered wrappers stay alive for the process lifetime, so repeated creation retains additional objects.
Stacked on #142 (
codex/trace-integrity). Seedocs/experiments/torchscript-builtin-validation.mdfor controls, reproduction, and compact evidence.Engineering hardening
Added a discriminated compatibility-result type, typed registry and wrapper factory, validation of present callbacks/operator metadata, per-wrapper status metadata, and a warning emitted once per process for absent private support (
d4ebb12). Eight added tests cover malformed metadata, callback failures, unsupported capability, and warning behavior. All 13 TorchScript tests (including six fresh subprocess controls) and 15 tracing tests pass. Stronger Mypy 1.9 checks pass for the compatibility module and contract tests. Stronger legacy tracer/dumper diagnostics decrease from 7 to 4; the four remaining diagnostics predate this change and are documented.Separately composed updated production sources pass 140 affected Linux regression tests, whole-source Mypy (77 files), Ruff, and isort. Whole-source Black still reports three unchanged legacy files, down from six in the original base; changed-file hooks pass. Strict checks are scoped and do not establish repository-wide strict typing. Historical performance numbers above retain their original measured revisions; this hardening pass does not re-estimate them.
Integration with current main
Based on #142 (
codex/trace-integrity), with signed merge checkpointf37be1efcb50. The shared base is #150, which reconciles Daniel's parallel-inference tip122e769with mained54c7d. History is preserved; no force push or GitHub PR merge was used. All nine PR heads merge cleanly with their declared bases and current main, and the full stack composes without conflicts.The resolved composition passes 140 production + 28 harness tests locally (Python 3.10.16 / Torch 2.2.2), pinned Mypy over 77 production files, and stricter Mypy over 9 harness/test files. Conflict-file commit hooks pass. Full-stack Ruff/Mypy/isort pass; Black retains the same three pre-existing formatting failures (
config.py,proxy_wrapper/hash.py,proxy_wrapper/subclass.py). Remote CI is ignored because the self-hosted runner is unavailable. These checks validate integration; existing performance measurements retain their original revisions and are not remeasured here.Land the parent in main first, then retarget this PR to main before merging it. Repeat for its children. Use merge commits to preserve stack ancestry, or refresh descendants after a squash/rebase merge. Exact heads and local logs are retained in
outputs/conflict-resolution-20261026/.