Skip to content

Preserve TorchScript builtin identity when wrapping PyTorch APIs - #143

Open
Essoz wants to merge 5 commits into
codex/trace-integrityfrom
codex/torchscript-builtin-identity
Open

Essoz wants to merge 5 commits into
codex/trace-integrityfrom
codex/torchscript-builtin-identity

Conversation

@Essoz

@Essoz Essoz commented Oct 25, 2026 •

Copy link
Copy Markdown
Collaborator

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:

  • Fresh-process CPU reproduction changes both broken wrapper controls from compilation failure to success with aten::tanh in the graph. Outputs, input gradients, and bias gradients exactly match the unwrapped baseline within each platform.
  • All 13 semantic/trace checks and four helper tests pass on macOS/PyTorch 2.2.2 and Linux/PyTorch 2.2.2+cu121. Eager tracing still emits one paired tanh call with arguments and return snapshots.
  • The portable checkout suite also passes against the trace-integrity source plus this integration: four helper tests and one regression covering six fresh subprocesses.

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). See docs/experiments/torchscript-builtin-validation.md for 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.

This branch has not been deployed

No deployments
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