From 9ac61d1140cf8679c24442a174ec12f4de069980 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 15:16:20 +0800 Subject: [PATCH 1/3] Preserve TorchScript builtin identity through tracing wrappers --- .../torchscript-builtin-validation.json | 208 ++++++++++++ .../torchscript-builtin-validation.md | 38 +++ tests/test_torchscript_compat.py | 305 ++++++++++++++++++ traincheck/instrumentor/torchscript_compat.py | 27 ++ traincheck/instrumentor/tracer.py | 2 + 5 files changed, 580 insertions(+) create mode 100644 docs/experiments/torchscript-builtin-validation.json create mode 100644 docs/experiments/torchscript-builtin-validation.md create mode 100644 tests/test_torchscript_compat.py create mode 100644 traincheck/instrumentor/torchscript_compat.py diff --git a/docs/experiments/torchscript-builtin-validation.json b/docs/experiments/torchscript-builtin-validation.json new file mode 100644 index 00000000..13020e1c --- /dev/null +++ b/docs/experiments/torchscript-builtin-validation.json @@ -0,0 +1,208 @@ +{ + "historical_cpu_study": { + "linux": { + "controls": [ + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 0, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "baseline", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": "TypeError", + "graph_has_aten_tanh": false, + "mode": "buggy-unprimed", + "registered_operator": null, + "script_compilation": "failed", + "scripted_output_and_gradients_equal_baseline": null, + "scripted_records": null + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": "TypeError", + "graph_has_aten_tanh": false, + "mode": "buggy-warm", + "registered_operator": null, + "script_compilation": "failed", + "scripted_output_and_gradients_equal_baseline": null, + "scripted_records": null + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "fixed-unprimed", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "fixed-warm", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "integrated-fixed-warm", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + } + ], + "device": "cpu", + "helper_test_log_sha256": "feabb4add9e2ca3374e538560634d1ec97dc4f4ce89aab12f3d726e36ea9b2f6", + "helper_unit_tests_passed": 4, + "native_threads": 1, + "python": "3.10.16", + "registry_already_initialized_after_import": true, + "semantic_trace_checks_passed": 13, + "summary_sha256": "68c1888b7df4a9a0248e6d218ea5e82d2daeb90c3efcf50b852faaf7e3bfa7c2", + "torch": "2.2.2+cu121" + }, + "macos": { + "controls": [ + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 0, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "baseline", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": "TypeError", + "graph_has_aten_tanh": false, + "mode": "buggy-unprimed", + "registered_operator": null, + "script_compilation": "failed", + "scripted_output_and_gradients_equal_baseline": null, + "scripted_records": null + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": "TypeError", + "graph_has_aten_tanh": false, + "mode": "buggy-warm", + "registered_operator": null, + "script_compilation": "failed", + "scripted_output_and_gradients_equal_baseline": null, + "scripted_records": null + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "fixed-unprimed", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "fixed-warm", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + }, + { + "eager_output_and_gradients_equal_baseline": true, + "eager_records": 2, + "exception_type": null, + "graph_has_aten_tanh": true, + "mode": "integrated-fixed-warm", + "registered_operator": "aten::tanh", + "script_compilation": "passed", + "scripted_output_and_gradients_equal_baseline": true, + "scripted_records": 0 + } + ], + "device": "cpu", + "helper_test_log_sha256": "d094b1bd2c8318c96097f80c50da282add5b2972970a77046de8db6c8935f6ca", + "helper_unit_tests_passed": 4, + "native_threads": 1, + "python": "3.10.16", + "registry_already_initialized_after_import": true, + "semantic_trace_checks_passed": 13, + "summary_sha256": "4e8775e72f31a344f6e82f59f6b3009fddebc9f71c9463f6f5a009b5885d8c43", + "torch": "2.2.2" + } + }, + "historical_source_sha256": { + "bias_gelu_fixture.py": "c6d7503996e5fa2872114aa7ba8b3b37151f03e8f7cc2302a76e6c451fb70059", + "builtin_compat_candidate.py": "744506a0c035396aea443f39394a700ef7908bc7d36b8117ba76fae97c92e845", + "candidate-src/traincheck/instrumentor/torchscript_compat.py": "744506a0c035396aea443f39394a700ef7908bc7d36b8117ba76fae97c92e845", + "candidate-src/traincheck/instrumentor/tracer.py": "a8e39a9dea2c4940543174c2bca97146665f0f08068bdf350b3f16802d4e79bd", + "probe.py": "58b0ce8f8a1fb65b96cd5a756bb82ef5ff8d315b3673081ea58d49ebffb0b29c", + "run.py": "f4c549ea1dd21c47f4e98f35df4126b23895391eacb2f89c182aaf7f2ccf8d12", + "test_candidate.py": "6f5fbbab5053e2303a53209ec25d71c1c240e9c06b9038af844adf762f636eff", + "torchscript-builtin-identity.patch": "8ae809f16bf6464d5a152d9612e37bd81edf04ef6680fee7e3042b42c8f026a5" + }, + "limits": [ + "Registry was already initialized after imports; no genuinely cold-registry result is claimed.", + "Scripted operators bypass Python recorder hooks; successful graphs emit no inner tanh API records.", + "Registered wrappers remain alive for the process lifetime; repeated creation retains additional objects.", + "Full Megatron startup, checkpoint-shape compatibility, and other PyTorch versions/operators are not established.", + "The converted portable checkout test ran locally on macOS; Linux evidence comes from the prior equivalent focused study." + ], + "linux_saved_evidence_archive_sha256": "81f5ede08024c3e099ad61b7489ffa420e4da490c743cded9d9236fd6453a339", + "portable_checkout_validation": { + "fresh_probe_processes": 6, + "fresh_process_regression_methods": 1, + "helper_test_methods": 4, + "log_sha256": "4e23c888ea4641c17ac45c1a413ba0a9b74afbbda23c87e364458b283a4b2438", + "platform": "macOS", + "python": "3.10.16", + "semantic_trace_checks": 13, + "source_includes_trace_integrity_fix": true, + "test_methods_passed": 5, + "tested_base_commit": "7d6c683db92d83ffd6db9fc9c11fbc0ebdd2c5d4", + "tested_base_pr": 142, + "timeout_seconds_per_probe": 90, + "torch": "2.2.2" + }, + "recorded_date": "2026-10-25", + "schema_version": 1, + "source": { + "base_dumper_sha256": "dd813a342b4531a6543995a23591f558db00d04142611a2d9656b1258833fac3", + "base_tracer_sha256": "bce6546cfa13a263150dff891f2f8d7d4af8c5ccb28dc2faedcedc691c5857da", + "candidate_dumper_sha256": "dd813a342b4531a6543995a23591f558db00d04142611a2d9656b1258833fac3", + "candidate_tracer_sha256": "ad57320f0a3764c02c73cf4bfa623f1ef73b9835dfa6cb1365dc2988551de428", + "helper_implementation_ast_unchanged_except_docstring": true, + "helper_sha256": "0d0c348fa02d5f93bb74e874829e5fb4c57365b82ce973ae8ff1f21a42c2d64e", + "integration_patch_sha256": "b374501d22408ebace2d1e2e69c03d2a98abcf3ebfa2e4e056b62adf27874397", + "portable_test_sha256": "653be1032409579030ebd3ec5694cc3e43f34df8e4f5750a76a9ffc64dd0a2e2", + "validated_helper_before_docstring_sha256": "744506a0c035396aea443f39394a700ef7908bc7d36b8117ba76fae97c92e845" + } +} diff --git a/docs/experiments/torchscript-builtin-validation.md b/docs/experiments/torchscript-builtin-validation.md new file mode 100644 index 00000000..71c9f4b3 --- /dev/null +++ b/docs/experiments/torchscript-builtin-validation.md @@ -0,0 +1,38 @@ +# TorchScript builtin identity validation + +TorchScript records builtin operators by function ID. Replacing `torch.tanh` with a TrainCheck wrapper loses the `aten::tanh` lookup. TorchScript attempts to inspect Python source, follows the wrapper metadata to the C builtin, and raises `TypeError: ... expected, got builtin_function_or_method`. + +The compatibility helper looks up the original mapping before the wrapper is published and registers the wrapper with that same operator. It guesses no operator names. Missing private APIs and unregistered Python functions are unchanged. Registered wrappers are retained because the registry stores only IDs; otherwise a later unrelated callable could reuse a stale ID. + +## Controls and results + +The 2026-10-25 CPU experiment uses a real Python source file containing a representative tanh-based bias-GELU expression. Each control runs in a fresh interpreter, with a 90-second timeout and one native/inter-op thread. + +| Control | Compilation | Eager tanh records | Scripted operator | +|---|---|---:|---| +| Unwrapped baseline | Pass | 0 | `aten::tanh` | +| Wrapper without compatibility, no explicit priming | TypeError | 2, one pair | Unregistered wrapper | +| Wrapper without compatibility, explicitly primed | TypeError | 2, one pair | Unregistered wrapper | +| Wrapper plus helper, no explicit priming | Pass | 2, one pair | `aten::tanh` | +| Wrapper plus helper, explicitly primed | Pass | 2, one pair | `aten::tanh` | +| Helper integrated into the actual wrapper | Pass | 2, one pair | `aten::tanh` | + +The original study passed all 13 semantic/trace checks and four helper unit tests on macOS with PyTorch 2.2.2 and Linux with PyTorch 2.2.2+cu121; both used Python 3.10.16 and CPU tensors. Candidate eager outputs and both gradients exactly matched the unwrapped baseline within each platform. Successful scripted outputs and both gradients also exactly matched that platform's scripted baseline. Every eager pair retained arguments, return snapshots, and a shared call ID. + +The portable checkout test was additionally run on macOS against the trace-integrity source plus this integration: all five test methods passed, comprising the four helper contracts and six fresh-process controls with the same 13 semantic/trace checks. The helper implementation is unchanged from the Linux-tested version apart from its module docstring and formatting. Linux's earlier study is separate evidence; the converted checkout test was not rerun there. + +Run the portable tests from the repository root: + +```bash +python -m unittest discover -s tests -p test_torchscript_compat.py -v +``` + +Compact source hashes, control outcomes, and validation details are recorded in [torchscript-builtin-validation.json](torchscript-builtin-validation.json). + +## Limits + +TrainCheck imports had already initialized the registry in both tested environments. “No explicit priming” therefore does not prove a genuinely uninitialized registry path. + +Scripted graphs call `aten::tanh` directly and produce no inner Python tanh records. Eager tracing remains intact; complete compiled-graph tracing is outside this change. Full instrument-before-import Megatron startup and checkpoint-shape compatibility were not validated by this focused fixture. + +Registered wrappers remain alive for the process lifetime; repeatedly creating new wrappers retains additional objects. The helper depends on feature-detected private APIs and does not establish compatibility with every future PyTorch version or wrapped operator. diff --git a/tests/test_torchscript_compat.py b/tests/test_torchscript_compat.py new file mode 100644 index 00000000..05f2b8bb --- /dev/null +++ b/tests/test_torchscript_compat.py @@ -0,0 +1,305 @@ +"""Builtin-registration contracts and fresh-process bias-GELU regressions.""" + +import argparse +import copy +import gc +import json +import os +import subprocess +import sys +import tempfile +import unittest +import weakref +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from traincheck.instrumentor import torchscript_compat +from traincheck.instrumentor.torchscript_compat import ( + _REGISTERED_WRAPPERS, + preserve_torchscript_builtin, +) + + +def bias_gelu(bias: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """Real-source-file fixture for the tanh-based fused bias-GELU expression.""" + x = bias + value + return ( + x * 0.5 * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) + ) + + +def evaluate(function): + bias = torch.linspace(-0.7, 0.9, 8, dtype=torch.float64, requires_grad=True) + value = ( + torch.linspace(-2.3, 2.1, 24, dtype=torch.float64) + .reshape(3, 8) + .requires_grad_() + ) + output = function(bias, value) + grad_bias, grad_value = torch.autograd.grad(output.square().sum(), (bias, value)) + return { + "output": output.detach().tolist(), + "grad_bias": grad_bias.tolist(), + "grad_value": grad_value.tolist(), + } + + +def probe(mode): + """Run one control in its own interpreter so JIT caches cannot leak.""" + from traincheck.instrumentor import tracer + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + records = [] + registry = torch.jit._builtins + original = torch.tanh + result = { + "mode": mode, + "torch": torch.__version__, + "initial_registry_initialized": registry._builtin_table is not None, + } + if mode.endswith("warm"): + result["original_operator_before_wrapping"] = registry._find_builtin(original) + + with ( + patch.object(tracer.config, "DISABLE_WRAPPER", False), + patch.object( + tracer, + "dump_trace_API", + side_effect=lambda record: records.append(copy.deepcopy(record)), + ), + ): + if mode != "baseline": + if mode == "integrated-fixed-warm": + wrapped = tracer.wrapper( + original, + False, + False, + False, + dump_args=True, + dump_ret=True, + handle_proxy=False, + ) + else: + # Disable only registration to reproduce the former wrapper; + # the helper controls then apply the production helper directly. + with patch.object( + tracer, "preserve_torchscript_builtin", return_value=None + ): + wrapped = tracer.wrapper( + original, + False, + False, + False, + dump_args=True, + dump_ret=True, + handle_proxy=False, + ) + if mode.startswith("fixed"): + result["candidate"] = preserve_torchscript_builtin( + original, wrapped, torch + ) + torch.tanh = wrapped + try: + result["eager"] = evaluate(bias_gelu) + result["eager_records"] = copy.deepcopy(records) + records.clear() + scripted = torch.jit.script(bias_gelu) + result["graph_has_aten_tanh"] = "aten::tanh" in str(scripted.graph) + result["scripted"] = evaluate(scripted) + result["scripted_records"] = copy.deepcopy(records) + result["script_compilation"] = "passed" + except Exception as error: + result["script_compilation"] = "failed" + result["exception_type"] = type(error).__name__ + result["exception"] = str(error) + finally: + result["current_operator_after_compilation"] = registry._find_builtin( + torch.tanh + ) + torch.tanh = original + return result + + +class BuiltinCompatibilityTests(unittest.TestCase): + def fake_torch(self, operator): + table, events = {}, [] + + def find(function): + events.append(("find", function)) + return operator + + def register(function, mapped_operator): + events.append(("register", function, mapped_operator)) + table[id(function)] = mapped_operator + + builtins = SimpleNamespace(_find_builtin=find, _register_builtin=register) + return SimpleNamespace(jit=SimpleNamespace(_builtins=builtins)), table, events + + def test_copies_exact_original_mapping(self): + fake, table, events = self.fake_torch("aten::tanh") + + def original(): + pass + + def wrapped(): + pass + + result = preserve_torchscript_builtin(original, wrapped, fake) + self.assertEqual(result, {"status": "preserved", "operator": "aten::tanh"}) + self.assertEqual( + events, [("find", original), ("register", wrapped, "aten::tanh")] + ) + self.assertEqual(table[id(wrapped)], "aten::tanh") + + def test_unregistered_function_is_unchanged(self): + fake, table, events = self.fake_torch(None) + original = lambda: None + wrapped = lambda: None + self.assertEqual( + preserve_torchscript_builtin(original, wrapped, fake)["status"], + "not-registered", + ) + self.assertEqual(table, {}) + self.assertEqual(events, [("find", original)]) + self.assertNotIn(wrapped, _REGISTERED_WRAPPERS) + + def test_unavailable_private_api_is_noop(self): + original = lambda: None + wrapped = lambda: None + missing_register = SimpleNamespace( + jit=SimpleNamespace( + _builtins=SimpleNamespace(_find_builtin=lambda f: "aten::tanh") + ) + ) + for fake in [SimpleNamespace(), missing_register]: + self.assertEqual( + preserve_torchscript_builtin(original, wrapped, fake)["status"], + "unsupported-private-api", + ) + self.assertNotIn(wrapped, _REGISTERED_WRAPPERS) + + def test_registered_wrapper_cannot_leave_stale_reusable_id(self): + fake, table, events = self.fake_torch("aten::tanh") + original = lambda: None + wrapped = lambda: None + weak = weakref.ref(wrapped) + address = id(wrapped) + preserve_torchscript_builtin(original, wrapped, fake) + events.clear() + del wrapped + gc.collect() + self.assertIsNotNone(weak()) + self.assertEqual(table[address], "aten::tanh") + self.assertEqual(_REGISTERED_WRAPPERS[weak()], "aten::tanh") + + +class TorchScriptWrapperTests(unittest.TestCase): + def test_fresh_process_bias_gelu_semantics_and_eager_traces(self): + registry = getattr(getattr(torch, "jit", None), "_builtins", None) + if not all( + callable(getattr(registry, name, None)) + for name in ["_find_builtin", "_register_builtin"] + ): + self.skipTest("PyTorch lacks the private builtin-registration APIs") + + source = Path(torchscript_compat.__file__).resolve().parents[2] + environment = os.environ | { + "PYTHONPATH": str(source), + "PYTHONDONTWRITEBYTECODE": "1", + "OMP_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "CUDA_VISIBLE_DEVICES": "", + } + modes = [ + "baseline", + "buggy-unprimed", + "buggy-warm", + "fixed-unprimed", + "fixed-warm", + "integrated-fixed-warm", + ] + results = {} + with tempfile.TemporaryDirectory(prefix="traincheck-torchscript-") as temporary: + for mode in modes: + output = Path(temporary) / (mode + ".json") + completed = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--probe", + mode, + "--output", + str(output), + ], + cwd=source, + env=environment, + capture_output=True, + text=True, + timeout=90, + ) + self.assertEqual( + completed.returncode, 0, completed.stdout + completed.stderr + ) + results[mode] = json.loads(output.read_text()) + + baseline = results["baseline"] + checks = 0 + for mode in modes[:-1]: + result = results[mode] + with self.subTest(mode=mode): + self.assertEqual(result["eager"], baseline["eager"]) + checks += 1 + if mode.startswith("buggy"): + self.assertEqual(result["script_compilation"], "failed") + self.assertEqual(result["exception_type"], "TypeError") + self.assertIn("builtin_function_or_method", result["exception"]) + self.assertIsNone(result["current_operator_after_compilation"]) + else: + self.assertEqual(result["script_compilation"], "passed") + self.assertTrue(result["graph_has_aten_tanh"]) + self.assertEqual(result["scripted"], baseline["scripted"]) + self.assertEqual( + result["current_operator_after_compilation"], "aten::tanh" + ) + self.assertEqual(result["scripted_records"], []) + checks += 1 + if mode != "baseline": + self.assert_paired_eager_record(result) + checks += 1 + integrated = results["integrated-fixed-warm"] + self.assertEqual(integrated["script_compilation"], "passed") + self.assertTrue(integrated["graph_has_aten_tanh"]) + self.assertEqual(integrated["scripted"], baseline["scripted"]) + self.assertEqual(integrated["eager"], baseline["eager"]) + self.assertEqual(integrated["scripted_records"], []) + self.assertEqual(integrated["current_operator_after_compilation"], "aten::tanh") + self.assert_paired_eager_record(integrated) + checks += 1 + self.assertEqual(checks, 13) + + def assert_paired_eager_record(self, result): + records = result["eager_records"] + self.assertEqual(len(records), 2) + self.assertEqual(records[0]["type"], "function_call (pre)") + self.assertEqual(records[1]["type"], "function_call (post)") + self.assertEqual(records[0]["func_call_id"], records[1]["func_call_id"]) + self.assertEqual(records[0]["function"], records[1]["function"]) + self.assertIn("tanh", records[0]["function"]) + self.assertTrue(records[0]["args"]) + self.assertTrue(records[1]["return_values"]) + + +if __name__ == "__main__": + if "--probe" in sys.argv: + parser = argparse.ArgumentParser() + parser.add_argument("--probe", required=True) + parser.add_argument("--output", required=True, type=Path) + arguments = parser.parse_args() + arguments.output.write_text(json.dumps(probe(arguments.probe)) + "\n") + else: + unittest.main() diff --git a/traincheck/instrumentor/torchscript_compat.py b/traincheck/instrumentor/torchscript_compat.py new file mode 100644 index 00000000..03087d84 --- /dev/null +++ b/traincheck/instrumentor/torchscript_compat.py @@ -0,0 +1,27 @@ +"""Preserve TorchScript operator identity for instrumented builtin functions.""" + +# TorchScript indexes by id only. Keep registered wrappers alive, like PyTorch's +# own _builtin_ops list, so a later unrelated object cannot reuse a stale ID. +_REGISTERED_WRAPPERS = {} + + +def preserve_torchscript_builtin(original, wrapped, torch_module): + """Call before publishing wrapped into its module, while original is visible. + + TorchScript compiles the existing operator directly and does not run Python + recorder hooks inside its graph. Eager calls still execute the wrapper. + Private API feature detection leaves unsupported versions unchanged. + """ + jit = getattr(torch_module, "jit", None) + builtins = getattr(jit, "_builtins", None) + find_builtin = getattr(builtins, "_find_builtin", None) + register_builtin = getattr(builtins, "_register_builtin", None) + if not callable(find_builtin) or not callable(register_builtin): + return {"status": "unsupported-private-api", "operator": None} + # _find_builtin also initializes the lazy table before the first replacement. + operator = find_builtin(original) + if operator is None: + return {"status": "not-registered", "operator": None} + register_builtin(wrapped, operator) + _REGISTERED_WRAPPERS[wrapped] = operator + return {"status": "preserved", "operator": operator} diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 33c8f05e..dc1cbc7c 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -39,6 +39,7 @@ funcs_to_be_replaced, is_funcs_to_be_unproxied, ) +from traincheck.instrumentor.torchscript_compat import preserve_torchscript_builtin from traincheck.utils import get_timestamp_ns, get_unique_id, typename _instancemethod_t = type(torch._C._distributed_c10d.ProcessGroup.broadcast) @@ -486,6 +487,7 @@ def wrapped(*args, **kwargs): wrapped._traincheck_original_function = original_function wrapped._traincheck_instrumented = True wrapped._traincheck_dump_disabled = disable_dump + preserve_torchscript_builtin(original_function, wrapped, torch) return wrapped From d4ebb12e3c5fd7531080d5623447f2af0a1ab3e9 Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:32:52 +0800 Subject: [PATCH 2/3] Enforce TorchScript compatibility metadata contracts --- tests/test_torchscript_compat.py | 10 +- tests/test_torchscript_contracts.py | 161 ++++++++++++++++++ traincheck/instrumentor/torchscript_compat.py | 80 +++++++-- traincheck/instrumentor/tracer.py | 58 ++++--- 4 files changed, 270 insertions(+), 39 deletions(-) create mode 100644 tests/test_torchscript_contracts.py diff --git a/tests/test_torchscript_compat.py b/tests/test_torchscript_compat.py index 05f2b8bb..8d6c7937 100644 --- a/tests/test_torchscript_compat.py +++ b/tests/test_torchscript_compat.py @@ -47,16 +47,16 @@ def evaluate(function): } -def probe(mode): +def probe(mode: str) -> dict[str, object]: """Run one control in its own interpreter so JIT caches cannot leak.""" from traincheck.instrumentor import tracer torch.set_num_threads(1) torch.set_num_interop_threads(1) - records = [] + records: list[dict[str, object]] = [] registry = torch.jit._builtins original = torch.tanh - result = { + result: dict[str, object] = { "mode": mode, "torch": torch.__version__, "initial_registry_initialized": registry._builtin_table is not None, @@ -87,7 +87,9 @@ def probe(mode): # Disable only registration to reproduce the former wrapper; # the helper controls then apply the production helper directly. with patch.object( - tracer, "preserve_torchscript_builtin", return_value=None + tracer, + "preserve_torchscript_builtin", + return_value={"status": "not-registered", "operator": None}, ): wrapped = tracer.wrapper( original, diff --git a/tests/test_torchscript_contracts.py b/tests/test_torchscript_contracts.py new file mode 100644 index 00000000..fa841ed1 --- /dev/null +++ b/tests/test_torchscript_contracts.py @@ -0,0 +1,161 @@ +"""Explicit unsupported, malformed, and failed TorchScript boundary contracts.""" + +import threading +import unittest +import warnings +from types import SimpleNamespace +from typing import Callable +from unittest.mock import patch + +from traincheck.instrumentor import torchscript_compat, tracer +from traincheck.instrumentor.torchscript_compat import preserve_torchscript_builtin + + +def original(value: int = 1) -> int: + return value + 1 + + +def candidate(value: int = 1) -> int: + return original(value) + + +def framework(find: object, register: object) -> SimpleNamespace: + return SimpleNamespace( + jit=SimpleNamespace( + _builtins=SimpleNamespace(_find_builtin=find, _register_builtin=register) + ) + ) + + +class TorchScriptContractTests(unittest.TestCase): + def test_present_noncallable_callback_is_rejected(self) -> None: + def find(function: Callable[..., object]) -> str: + return "aten::tanh" + + def register(function: Callable[..., object], operator: str) -> None: + self.fail("Invalid capabilities must not register a wrapper") + + malformed_callbacks: tuple[object, ...] = (None, 17, "callable", []) + for malformed in malformed_callbacks: + for name, fake in ( + ("_find_builtin", framework(malformed, register)), + ("_register_builtin", framework(find, malformed)), + ): + with self.subTest(name=name, malformed=malformed): + with self.assertRaisesRegex(TypeError, name): + preserve_torchscript_builtin(original, candidate, fake) + + def test_missing_callback_does_not_hide_an_invalid_present_callback(self) -> None: + fake = SimpleNamespace( + jit=SimpleNamespace(_builtins=SimpleNamespace(_register_builtin=None)) + ) + with self.assertRaisesRegex(TypeError, "_register_builtin"): + preserve_torchscript_builtin(original, candidate, fake) + + def test_malformed_operator_is_rejected_before_registration(self) -> None: + registered: list[str] = [] + + def register(function: Callable[..., object], operator: str) -> None: + registered.append(operator) + + malformed_operators: tuple[object, ...] = ("", 1, False, [], {}) + for malformed in malformed_operators: + + def find(function: Callable[..., object]) -> object: + return malformed + + with self.subTest(operator=malformed): + with self.assertRaisesRegex(TypeError, "nonempty string or None"): + preserve_torchscript_builtin( + original, candidate, framework(find, register) + ) + self.assertEqual(registered, []) + self.assertNotIn(candidate, torchscript_compat._REGISTERED_WRAPPERS) + + def test_find_failure_propagates_unchanged(self) -> None: + error = RuntimeError("registry lookup failed") + + def find(function: Callable[..., object]) -> str | None: + raise error + + def register(function: Callable[..., object], operator: str) -> None: + self.fail("Failed lookup must not register a wrapper") + + with self.assertRaises(RuntimeError) as caught: + preserve_torchscript_builtin(original, candidate, framework(find, register)) + self.assertIs(caught.exception, error) + self.assertNotIn(candidate, torchscript_compat._REGISTERED_WRAPPERS) + + def test_registration_failure_propagates_without_retention(self) -> None: + error = RuntimeError("registry update failed") + + def find(function: Callable[..., object]) -> str: + return "aten::tanh" + + def register(function: Callable[..., object], operator: str) -> None: + raise error + + with self.assertRaises(RuntimeError) as caught: + preserve_torchscript_builtin(original, candidate, framework(find, register)) + self.assertIs(caught.exception, error) + self.assertNotIn(candidate, torchscript_compat._REGISTERED_WRAPPERS) + + def test_wrapper_exposes_unsupported_and_warns_once(self) -> None: + with ( + patch.object(tracer, "torch", SimpleNamespace()), + patch.object(torchscript_compat, "_UNSUPPORTED_WARNING_EMITTED", False), + warnings.catch_warnings(record=True) as captured, + ): + warnings.simplefilter("always") + first = tracer.wrapper( + original, False, False, False, disable_dump=True, handle_proxy=False + ) + second = tracer.wrapper( + original, False, False, False, disable_dump=True, handle_proxy=False + ) + self.assertEqual(first(3), 4) + self.assertEqual(second(8), 9) + self.assertEqual(len(captured), 1) + self.assertIn( + "Eager instrumentation remains available", str(captured[0].message) + ) + for wrapped in (first, second): + # Function metadata is deliberately public to instrumentation users. + self.assertEqual( + getattr(wrapped, "_traincheck_torchscript_builtin"), + {"status": "unsupported-private-api", "operator": None}, + ) + + def test_unsupported_warning_is_once_across_threads(self) -> None: + with ( + patch.object(torchscript_compat, "_UNSUPPORTED_WARNING_EMITTED", False), + warnings.catch_warnings(record=True) as captured, + ): + warnings.simplefilter("always") + workers = [ + threading.Thread(target=torchscript_compat.warn_unsupported_torchscript) + for _ in range(8) + ] + for worker in workers: + worker.start() + for worker in workers: + worker.join(5) + self.assertFalse(worker.is_alive()) + self.assertEqual(len(captured), 1) + + def test_wrapper_propagates_metadata_failure(self) -> None: + def find(function: Callable[..., object]) -> object: + return 3 + + def register(function: Callable[..., object], operator: str) -> None: + self.fail("Malformed operator must not be registered") + + with patch.object(tracer, "torch", framework(find, register)): + with self.assertRaisesRegex(TypeError, "nonempty string or None"): + tracer.wrapper( + original, False, False, False, disable_dump=True, handle_proxy=False + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/traincheck/instrumentor/torchscript_compat.py b/traincheck/instrumentor/torchscript_compat.py index 03087d84..4f35f8c8 100644 --- a/traincheck/instrumentor/torchscript_compat.py +++ b/traincheck/instrumentor/torchscript_compat.py @@ -1,27 +1,87 @@ """Preserve TorchScript operator identity for instrumented builtin functions.""" +import threading +import warnings +from typing import Callable, Literal, TypedDict + + +class PreservedBuiltin(TypedDict): + status: Literal["preserved"] + operator: str + + +class UnregisteredBuiltin(TypedDict): + status: Literal["not-registered"] + operator: None + + +class UnsupportedBuiltinAPI(TypedDict): + status: Literal["unsupported-private-api"] + operator: None + + +BuiltinPreservation = PreservedBuiltin | UnregisteredBuiltin | UnsupportedBuiltinAPI + # TorchScript indexes by id only. Keep registered wrappers alive, like PyTorch's # own _builtin_ops list, so a later unrelated object cannot reuse a stale ID. -_REGISTERED_WRAPPERS = {} +# These callables are retained/passed through, never invoked by this module. +_REGISTERED_WRAPPERS: dict[Callable[..., object], str] = {} +_MISSING = object() +_UNSUPPORTED_WARNING_LOCK = threading.Lock() +_UNSUPPORTED_WARNING_EMITTED = False + + +def warn_unsupported_torchscript() -> None: + """Report missing private support once per process, including always filters.""" + global _UNSUPPORTED_WARNING_EMITTED + with _UNSUPPORTED_WARNING_LOCK: + if not _UNSUPPORTED_WARNING_EMITTED: + warnings.warn( + "TrainCheck cannot preserve TorchScript builtin identity: this " + "PyTorch version lacks the private builtin-registration API. " + "Eager instrumentation remains available.", + RuntimeWarning, + stacklevel=2, + ) + _UNSUPPORTED_WARNING_EMITTED = True -def preserve_torchscript_builtin(original, wrapped, torch_module): - """Call before publishing wrapped into its module, while original is visible. +def preserve_torchscript_builtin( + original: Callable[..., object], + wrapped: Callable[..., object], + torch_module: object, +) -> BuiltinPreservation: + """Preserve a builtin mapping before publishing wrapped into its module. TorchScript compiles the existing operator directly and does not run Python recorder hooks inside its graph. Eager calls still execute the wrapper. - Private API feature detection leaves unsupported versions unchanged. + Missing private APIs are unsupported; present malformed APIs or metadata + are errors. Exceptions from framework callbacks propagate unchanged. """ - jit = getattr(torch_module, "jit", None) - builtins = getattr(jit, "_builtins", None) - find_builtin = getattr(builtins, "_find_builtin", None) - register_builtin = getattr(builtins, "_register_builtin", None) - if not callable(find_builtin) or not callable(register_builtin): + # Missing capabilities are expected across PyTorch versions. A sentinel + # distinguishes absence from present-but-invalid None callback attributes. + jit: object = getattr(torch_module, "jit", _MISSING) + builtins: object = getattr(jit, "_builtins", _MISSING) + find_builtin: object = getattr(builtins, "_find_builtin", _MISSING) + register_builtin: object = getattr(builtins, "_register_builtin", _MISSING) + for name, callback in ( + ("_find_builtin", find_builtin), + ("_register_builtin", register_builtin), + ): + if callback is not _MISSING and not callable(callback): + raise TypeError(f"torch.jit._builtins.{name} must be callable") + if find_builtin is _MISSING or register_builtin is _MISSING: return {"status": "unsupported-private-api", "operator": None} + # These checks narrow dynamic framework attributes after the missing case. + assert callable(find_builtin) and callable(register_builtin) # _find_builtin also initializes the lazy table before the first replacement. - operator = find_builtin(original) + operator: object = find_builtin(original) if operator is None: return {"status": "not-registered", "operator": None} + if not isinstance(operator, str) or not operator: + raise TypeError( + "TorchScript builtin operator must be a nonempty string or None" + ) register_builtin(wrapped, operator) _REGISTERED_WRAPPERS[wrapped] = operator return {"status": "preserved", "operator": operator} diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index d8eaa9bf..64b5d87d 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -39,7 +39,10 @@ funcs_to_be_replaced, is_funcs_to_be_unproxied, ) -from traincheck.instrumentor.torchscript_compat import preserve_torchscript_builtin +from traincheck.instrumentor.torchscript_compat import ( + preserve_torchscript_builtin, + warn_unsupported_torchscript, +) from traincheck.utils import get_timestamp_ns, get_unique_id, typename _instancemethod_t = type(torch._C._distributed_c10d.ProcessGroup.broadcast) @@ -200,7 +203,7 @@ def function_wrapper( dump_ret: bool, dump_ret_config, trigger_var_dump: bool, - var_dump_config: dict, + var_dump_config: dict[str, object] | None, need_unproxy_args_kwargs: bool, *args, **kwargs, @@ -445,19 +448,19 @@ def core_wrapper_proxy(original_function, *args, **kwargs): def wrapper( - original_function, - is_bound_method, - scan_proxy_in_args, - dump_stack_trace, - disable_dump=False, - dump_args=True, - dump_args_config=None, - dump_ret=True, - dump_ret_config=None, - handle_proxy=True, - trigger_var_dump=False, - var_dump_config=None, -): + original_function: Callable[..., object], + is_bound_method: bool, + scan_proxy_in_args: bool, + dump_stack_trace: bool, + disable_dump: bool = False, + dump_args: bool = True, + dump_args_config: dict[str, object] | None = None, + dump_ret: bool = True, + dump_ret_config: dict[str, object] | None = None, + handle_proxy: bool = True, + trigger_var_dump: bool = False, + var_dump_config: dict[str, object] | None = None, +) -> Callable[..., object]: is_builtin = is_c_level_function(original_function) need_unproxy_args_kwargs = handle_proxy and ( is_builtin or is_funcs_to_be_unproxied(original_function) @@ -468,7 +471,7 @@ def wrapper( METRIC_INSTRUMENTED_FUNC_LIST["dump"].append(original_function_name) @functools.wraps(original_function) - def wrapped(*args, **kwargs): + def wrapped(*args: object, **kwargs: object) -> object: return function_wrapper( original_function, original_function_name, @@ -491,19 +494,24 @@ def wrapped(*args, **kwargs): if need_unproxy_args_kwargs: @functools.wraps(original_function) - def wrapped(*args, **kwargs): + def wrapped(*args: object, **kwargs: object) -> object: return core_wrapper_proxy(original_function, *args, **kwargs) else: @functools.wraps(original_function) - def wrapped(*args, **kwargs): + def wrapped(*args: object, **kwargs: object) -> object: return original_function(*args, **kwargs) - wrapped._traincheck_original_function = original_function - wrapped._traincheck_instrumented = True - wrapped._traincheck_dump_disabled = disable_dump - preserve_torchscript_builtin(original_function, wrapped, torch) + setattr(wrapped, "_traincheck_original_function", original_function) + setattr(wrapped, "_traincheck_instrumented", True) + setattr(wrapped, "_traincheck_dump_disabled", disable_dump) + builtin_preservation = preserve_torchscript_builtin( + original_function, wrapped, torch + ) + setattr(wrapped, "_traincheck_torchscript_builtin", builtin_preservation) + if builtin_preservation["status"] == "unsupported-private-api": + warn_unsupported_torchscript() return wrapped @@ -955,9 +963,9 @@ def get_wrapped_function(self, func_obj: Callable) -> Callable: if func_name not in self.instr_opts.funcs_instr_opts: return wrapper( func_obj, - is_bound_method=None, - scan_proxy_in_args=None, - dump_stack_trace=None, + is_bound_method=False, + scan_proxy_in_args=False, + dump_stack_trace=False, disable_dump=True, handle_proxy=used_proxy, ) From 7bba4038577c243d835f0f90f7ff3f479499901e Mon Sep 17 00:00:00 2001 From: Yuxuan Jiang Date: Sun, 25 Oct 2026 21:35:26 +0800 Subject: [PATCH 3/3] Document TorchScript contract boundaries and validation --- .../torchscript-builtin-validation.md | 41 +++++++++++++++++++ traincheck/instrumentor/tracer.py | 2 + 2 files changed, 43 insertions(+) diff --git a/docs/experiments/torchscript-builtin-validation.md b/docs/experiments/torchscript-builtin-validation.md index 71c9f4b3..b87e8442 100644 --- a/docs/experiments/torchscript-builtin-validation.md +++ b/docs/experiments/torchscript-builtin-validation.md @@ -36,3 +36,44 @@ TrainCheck imports had already initialized the registry in both tested environme Scripted graphs call `aten::tanh` directly and produce no inner Python tanh records. Eager tracing remains intact; complete compiled-graph tracing is outside this change. Full instrument-before-import Megatron startup and checkpoint-shape compatibility were not validated by this focused fixture. Registered wrappers remain alive for the process lifetime; repeatedly creating new wrappers retains additional objects. The helper depends on feature-detected private APIs and does not establish compatibility with every future PyTorch version or wrapped operator. + +## Contract hardening (2026-10-25) + +Runtime checkpoint `d4ebb12e3c5fd7531080d5623447f2af0a1ab3e9` replaces the +untyped cache and status dictionaries with a typed retention registry and a +required, discriminated result: `preserved` has a nonempty string operator, +`not-registered` has no mapping, and `unsupported-private-api` reports missing +private support. Present noncallable callbacks and nonstring/empty operator +metadata raise `TypeError` before registration; lookup/registration exceptions +propagate unchanged. Failed registration does not retain the wrapper. + +The wrapper publishes `_traincheck_torchscript_builtin` metadata and warns once +per process when private support is unavailable, including under an `always` +warning filter and concurrent callers. Unsupported support still permits eager +instrumentation. Its factory and nested callables now have explicit annotations; +`var_dump_config=None` is passed unchanged to the existing registry. Disabled +tracing passes explicit false flags for its unused tracing options. + +Validation used Python 3.10.16/PyTorch 2.2.2, one OpenMP/MKL/OpenBLAS thread, +and the repository-pinned hooks (including mypy 1.9): + +```sh +python -m unittest discover -s tests -p 'test_torchscript*.py' +python -m unittest discover -s tests -p 'test_trac*.py' +pre-commit run --files traincheck/instrumentor/tracer.py traincheck/instrumentor/torchscript_compat.py tests/test_torchscript_compat.py tests/test_torchscript_contracts.py +mypy --follow-imports=silent --ignore-missing-imports --check-untyped-defs --disallow-untyped-defs --disallow-any-generics --warn-return-any traincheck/instrumentor/torchscript_compat.py tests/test_torchscript_contracts.py tests/test_tracing_contracts.py +``` + +All 13 TorchScript tests passed, including 8 new negative/capability contracts and +the existing six fresh-process controls for outputs, gradients, compiled graphs, +and eager records. All 15 tracing tests and all pinned hooks passed. The stronger +check above passed without adding `Any`, casts, or ignore directives. Callable +ellipsis parameters describe arbitrary framework functions retained or passed +through without invocation by the compatibility module. + +A separate `--check-untyped-defs` check of the legacy tracer/dumper files found +7 errors at original head `9ac61d1140cf8679c24442a174ec12f4de069980` and +4 after hardening: existing logger-key type inference, two unannotated queue +locals, and owner-class traversal assignment. Legacy serialization/proxy helpers +remain partially annotated; this is not a repository-wide strictness claim. +No timing or benchmark improvement is claimed for this contract-only change. diff --git a/traincheck/instrumentor/tracer.py b/traincheck/instrumentor/tracer.py index 64b5d87d..10bb81ab 100644 --- a/traincheck/instrumentor/tracer.py +++ b/traincheck/instrumentor/tracer.py @@ -503,6 +503,8 @@ def wrapped(*args: object, **kwargs: object) -> object: def wrapped(*args: object, **kwargs: object) -> object: return original_function(*args, **kwargs) + # Function attributes are runtime instrumentation metadata, attached after + # functools.wraps copies the original function's attributes. setattr(wrapped, "_traincheck_original_function", original_function) setattr(wrapped, "_traincheck_instrumented", True) setattr(wrapped, "_traincheck_dump_disabled", disable_dump)