sherlorch 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,22 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ubuntu-latest
11
+ strategy:
12
+ matrix:
13
+ python-version: ["3.9", "3.10", "3.11", "3.12"]
14
+ steps:
15
+ - uses: actions/checkout@v4
16
+ - uses: actions/setup-python@v5
17
+ with:
18
+ python-version: ${{ matrix.python-version }}
19
+ - name: Install package
20
+ run: pip install -e ".[dev]"
21
+ - name: Run tests
22
+ run: pytest tests/ -v
@@ -0,0 +1,50 @@
1
+ name: Publish to PyPI
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+ workflow_dispatch:
7
+
8
+ jobs:
9
+ build:
10
+ runs-on: ubuntu-latest
11
+ steps:
12
+ - uses: actions/checkout@v4
13
+
14
+ - uses: actions/setup-python@v5
15
+ with:
16
+ python-version: "3.x"
17
+
18
+ - name: Install build tooling
19
+ run: python -m pip install --upgrade build twine
20
+
21
+ - name: Build sdist and wheel
22
+ run: python -m build
23
+
24
+ - name: Check distribution metadata
25
+ run: twine check dist/*
26
+
27
+ - name: Upload build artifacts
28
+ uses: actions/upload-artifact@v4
29
+ with:
30
+ name: dist
31
+ path: dist/
32
+
33
+ publish:
34
+ if: github.event_name == 'release'
35
+ needs: build
36
+ runs-on: ubuntu-latest
37
+ environment:
38
+ name: pypi
39
+ url: https://pypi.org/project/sherlorch/
40
+ steps:
41
+ - uses: actions/download-artifact@v4
42
+ with:
43
+ name: dist
44
+ path: dist/
45
+
46
+ - name: Publish to PyPI
47
+ uses: pypa/gh-action-pypi-publish@release/v1
48
+ with:
49
+ password: ${{ secrets.PYPI_API_TOKEN }}
50
+ skip-existing: true
@@ -0,0 +1,11 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ .pytest_cache/
5
+ .venv/
6
+ venv/
7
+ build/
8
+ dist/
9
+ *.whl
10
+ .DS_Store
11
+ .ipynb_checkpoints/
@@ -0,0 +1,26 @@
1
+ # Changelog
2
+
3
+ All notable changes to this project are documented here.
4
+ Format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this
5
+ project uses [Semantic Versioning](https://semver.org/).
6
+
7
+ ## [Unreleased]
8
+
9
+ ## [0.1.0] - 2026-08-23
10
+
11
+ ### Added
12
+ - Core `ProvenanceTracker` (`torch.overrides.TorchFunctionMode`-based) recording
13
+ an op-level provenance DAG for every tensor produced while tracing is active.
14
+ - `diagnose(tensor)`: backward search from a non-finite tensor to the op(s) that
15
+ first introduced NaN/Inf, with the path from culprit to target.
16
+ - `tracker.watch(model)` / `unwatch()`: tag recorded ops with the submodule that
17
+ produced them, via forward hooks. Hooks auto-removed on `with` block exit.
18
+ - `fast_shape_ops` (default `True`): shape-preserving ops (`view`, `permute`,
19
+ `cat`, etc.) inherit finiteness from parent nodes instead of re-scanning,
20
+ exact rather than approximate.
21
+ - `capture_stack`: optional `file:line` capture per recorded op.
22
+ - `has_issue()`, `enable()`/`disable()`, `reset()` utility methods.
23
+ - Benchmark suite (`benchmarks/bench_overhead.py`) and demo notebook
24
+ (`notebooks/sherlorch_demo.ipynb`).
25
+ - Test suite covering division-by-zero, log-of-negative, multi-source culprit
26
+ merging, module tagging, fast-path correctness, and lifecycle edge cases.
@@ -0,0 +1,19 @@
1
+ # Code of Conduct
2
+
3
+ This project follows the spirit of the [Contributor Covenant](https://www.contributor-covenant.org/):
4
+ be respectful, assume good faith, and keep disagreements focused on the technical merits.
5
+
6
+ In short:
7
+
8
+ - Be welcoming and patient, especially with people new to PyTorch internals or open source.
9
+ - Give and receive feedback about code, not people.
10
+ - Harassment, discriminatory language, and personal attacks are not tolerated, in issues,
11
+ PRs, or any other project space.
12
+ - Maintainers may edit, remove, or reject contributions and comments that don't follow this,
13
+ and may block repeat offenders.
14
+
15
+ ## Reporting
16
+
17
+ If you experience or witness unacceptable behavior, please open an issue (or, for sensitive
18
+ reports, contact the maintainers directly via the contact info on their GitHub profile).
19
+ Reports will be handled with discretion.
@@ -0,0 +1,72 @@
1
+ # Contributing to sherlorch
2
+
3
+ Thanks for considering a contribution — this project is intentionally small and
4
+ sharply scoped, so even a modest PR can meaningfully move it forward.
5
+
6
+ ## Setup
7
+
8
+ ```bash
9
+ git clone https://github.com/SankaVaas/sherlorch.git
10
+ cd sherlorch
11
+ pip install -e ".[dev]"
12
+ ```
13
+
14
+ ## Running tests
15
+
16
+ ```bash
17
+ pytest tests/ -v
18
+ ```
19
+
20
+ All new behavior needs a test. If you're fixing a bug, add a test that fails
21
+ without your fix and passes with it.
22
+
23
+ ## Running the benchmark
24
+
25
+ ```bash
26
+ python benchmarks/bench_overhead.py
27
+ ```
28
+
29
+ This regenerates `benchmarks/results/overhead_results.csv` and the two plots in
30
+ `assets/`. If your change affects performance (positively or negatively),
31
+ please re-run this and include the updated numbers in your PR description —
32
+ we'd rather see the real effect than guess at it.
33
+
34
+ ## Design principles
35
+
36
+ - **Never crash the user's workload.** All instrumentation in `_record()` is
37
+ wrapped in `try/except Exception: pass`. If you add new instrumentation,
38
+ keep that guarantee.
39
+ - **No strong references to traced tensors.** We track tensors by `id()` with
40
+ a weakref cleanup callback specifically so tracing doesn't extend a
41
+ tensor's lifetime or blow up memory on long runs. Don't reintroduce a
42
+ `WeakKeyDictionary` keyed by the tensor itself — `torch.Tensor.__eq__` is
43
+ elementwise, which breaks dict lookup machinery in ways that are easy to
44
+ miss in a quick test (see the history of `tracker.py` for the exact
45
+ failure mode).
46
+ - **Exactness over approximation for `fast_shape_ops`.** Any op added to
47
+ `_SHAPE_PRESERVING_OPS` in `tracker.py` must be provably incapable of
48
+ turning a finite input into a non-finite output. When in doubt, leave it
49
+ out — a missed optimization is a performance cost, but a wrong one is a
50
+ correctness bug in a debugging tool, which defeats the point.
51
+
52
+ ## Good first contributions
53
+
54
+ - More reliable backward-pass op-name capture
55
+ - An optional Graphviz/HTML export of the provenance DAG
56
+ - Expanding `_SHAPE_PRESERVING_OPS` to cover more of the dispatcher surface
57
+ (with a test proving each addition is safe)
58
+ - A sampling mode for very large tensors, if you can characterize the
59
+ correctness/performance tradeoff clearly
60
+
61
+ ## Pull requests
62
+
63
+ - Keep PRs focused — one behavior change per PR is easier to review and revert.
64
+ - Update `CHANGELOG.md` under `[Unreleased]`.
65
+ - Update the README if you're changing public API or documented behavior.
66
+
67
+ ## Reporting bugs
68
+
69
+ Open an issue with a minimal repro. If it's a case where `sherlorch` gives a
70
+ wrong diagnosis (rather than crashing), that's especially important — please
71
+ include the model/ops involved, since it likely means `_SHAPE_PRESERVING_OPS`
72
+ or the finiteness-checking logic needs a fix.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 tensortrace contributors
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,199 @@
1
+ Metadata-Version: 2.5
2
+ Name: sherlorch
3
+ Version: 0.1.0
4
+ Summary: PyTorch-first tensor provenance debugger: find the exact op that introduced a NaN/Inf.
5
+ Project-URL: Homepage, https://github.com/SankaVaas/sherlorch
6
+ Project-URL: Issues, https://github.com/SankaVaas/sherlorch/issues
7
+ Project-URL: Changelog, https://github.com/SankaVaas/sherlorch/blob/main/CHANGELOG.md
8
+ Author: Sanka Vaas
9
+ License: MIT
10
+ License-File: LICENSE
11
+ Keywords: autograd,debugging,deep-learning,nan,provenance,pytorch
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: License :: OSI Approved :: MIT License
16
+ Classifier: Operating System :: OS Independent
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Programming Language :: Python :: 3.9
19
+ Classifier: Programming Language :: Python :: 3.10
20
+ Classifier: Programming Language :: Python :: 3.11
21
+ Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
23
+ Classifier: Topic :: Software Development :: Debuggers
24
+ Classifier: Typing :: Typed
25
+ Requires-Python: >=3.9
26
+ Requires-Dist: torch>=2.0
27
+ Provides-Extra: benchmark
28
+ Requires-Dist: matplotlib>=3.5; extra == 'benchmark'
29
+ Provides-Extra: dev
30
+ Requires-Dist: build; extra == 'dev'
31
+ Requires-Dist: pytest>=7.0; extra == 'dev'
32
+ Requires-Dist: twine; extra == 'dev'
33
+ Description-Content-Type: text/markdown
34
+
35
+ ![sherlorch name logo](<docs/images/sherlorch name logo.png>)
36
+
37
+ # sherlorch
38
+
39
+ [![CI](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml/badge.svg)](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml)
40
+ [![PyPI version](https://img.shields.io/pypi/v/sherlorch.svg)](https://pypi.org/project/sherlorch/)
41
+ [![Python versions](https://img.shields.io/pypi/pyversions/sherlorch.svg)](https://pypi.org/project/sherlorch/)
42
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
43
+ [![Code style](https://img.shields.io/badge/types-checked-brightgreen.svg)](https://peps.python.org/pep-0561/)
44
+
45
+ **Find the exact op that introduced a NaN or Inf into your PyTorch model — not just the line where you noticed it.**
46
+
47
+ Every deep learning researcher knows this moment: `loss.backward()` returns `nan`, and now you're bisecting your forward pass by hand, printing `.isnan().any()` after every other line. `sherlorch` automates that entire process.
48
+
49
+ It works by transparently intercepting every tensor operation while active (via `torch.overrides.TorchFunctionMode` — no hooks to register, no model changes needed), building a lightweight op-level provenance graph. When you point it at a bad tensor, it walks that graph backward and tells you precisely which operation, on which inputs, first produced the non-finite value — plus the path from there to the tensor you noticed the problem in.
50
+
51
+ ```
52
+ sherlorch diagnosis for [14] layer_norm (shape=(4, 4, 16, 32)):
53
+ Found 1 culprit op(s) out of 15 traced ops.
54
+
55
+ Culprit #1: [10] log shape=(4, 4, 16, 16) dtype=torch.float32 [model.layer1] (train.py:27) <-- Inf
56
+ Path to target:
57
+ -> [10] log ... [model.layer1] <-- Inf
58
+ -> [11] mean ... [model] <-- Inf
59
+ -> [12] unsqueeze ... [model] <-- Inf
60
+ -> [13] add ... [model] <-- Inf
61
+ -> [14] layer_norm ... [model.norm] <-- NaN
62
+ ```
63
+
64
+ That's real output from [`examples/find_nan_in_transformer.py`](examples/find_nan_in_transformer.py) — a `log(relu(x))` bug buried two submodules deep, found and localized to `model.layer1` in one traced run, no bisection required.
65
+
66
+ ## Contents
67
+
68
+ - [sherlorch](#sherlorch)
69
+ - [Contents](#contents)
70
+ - [Install](#install)
71
+ - [Quickstart](#quickstart)
72
+ - [Options](#options)
73
+ - [Module tagging](#module-tagging)
74
+ - [Perf mode](#perf-mode)
75
+ - [Benchmarks](#benchmarks)
76
+ - [How it works](#how-it-works)
77
+ - [Costs and limitations](#costs-and-limitations)
78
+ - [Try it](#try-it)
79
+ - [Contributing](#contributing)
80
+ - [License](#license)
81
+
82
+ ## Install
83
+
84
+ ```bash
85
+ pip install sherlorch
86
+ ```
87
+
88
+ Or from a clone, for local development:
89
+
90
+ ```bash
91
+ git clone https://github.com/SankaVaas/sherlorch.git
92
+ cd sherlorch
93
+ pip install -e ".[dev]"
94
+ ```
95
+
96
+ Requires `torch>=2.0`, Python `>=3.9`.
97
+
98
+ ## Quickstart
99
+
100
+ ```python
101
+ import torch
102
+ import sherlorch
103
+
104
+ with sherlorch.trace() as tracker:
105
+ out = model(x)
106
+ loss = criterion(out, y)
107
+ loss.backward()
108
+
109
+ if tracker.has_issue(loss):
110
+ print(tracker.diagnose(loss))
111
+ ```
112
+
113
+ Example output:
114
+
115
+ ```
116
+ sherlorch diagnosis for [812] sum (shape=()):
117
+ Found 1 culprit op(s) out of 811 traced ops.
118
+
119
+ Culprit #1: [340] div shape=(32, 128) dtype=torch.float32 (train.py:57) <-- Inf
120
+ Path to target:
121
+ -> [340] div shape=(32, 128) dtype=torch.float32 (train.py:57) <-- Inf
122
+ -> [341] mul shape=(32, 128) dtype=torch.float32 <-- Inf
123
+ -> [812] sum shape=() <-- Inf
124
+ ```
125
+
126
+ ### Options
127
+
128
+ | Option | Default | What it does |
129
+ |---|---|---|
130
+ | `sherlorch.trace(capture_stack=True)` | `False` | Attach a `file:line` to each recorded op. Useful while actively hunting a bug; adds real overhead (see [Benchmarks](#benchmarks)), especially inside Jupyter/IPython. |
131
+ | `sherlorch.trace(fast_shape_ops=False)` | `True` | Force a full isnan/isinf scan on every op, including shape-only ops. See [Perf mode](#perf-mode). |
132
+ | `tracker.watch(model, name="model")` | — | Tag every recorded op with the submodule that produced it. See [Module tagging](#module-tagging). |
133
+ | `tracker.disable()` / `enable()` | — | Pause/resume recording inside the `with` block, e.g. to skip a warmup loop. |
134
+ | `tracker.reset()` | — | Drop all recorded provenance, e.g. between training steps, to bound memory on long runs. |
135
+ | `tracker.has_issue(tensor)` | — | Cheap NaN/Inf check on any live tensor, tracked or not. |
136
+
137
+ ## Module tagging
138
+
139
+ Wrap `model(x)` with `tracker.watch(model)` and every op's diagnosis includes which submodule produced it:
140
+
141
+ ```python
142
+ with sherlorch.trace() as tracker:
143
+ tracker.watch(model, name="model")
144
+ out = model(x)
145
+
146
+ print(tracker.diagnose(out))
147
+ ```
148
+
149
+ That `[model.layer1]` tag comes straight from `model.named_modules()`, so it matches whatever names you gave your submodules. Hooks are removed automatically when the `with` block exits, or manually via `tracker.unwatch()`. Best-effort: relies on forward hooks firing in normal LIFO order, which activation checkpointing or re-entrant forward calls can disrupt.
150
+
151
+ ## Perf mode
152
+
153
+ By default (`fast_shape_ops=True`), ops that can only rearrange, copy, index, or combine existing values — `view`, `reshape`, `permute`, `transpose`, `cat`, `stack`, `narrow`, and similar — skip the isnan/isinf tensor scan entirely and instead **inherit** finiteness from their parent node(s). This is exact, not an approximation: those ops provably cannot introduce a NaN/Inf that wasn't already in one of their inputs (dtype-narrowing ops like `.to()`/`.half()` are deliberately excluded, since precision loss can itself overflow to Inf). Since transformer-style models are full of reshapes and permutes, this cuts meaningful overhead — and the savings grow with tensor size, since the skipped scan cost is proportional to tensor size while inheritance is O(1).
154
+
155
+ ## Benchmarks
156
+
157
+ Measured on a small transformer attention block (`benchmarks/bench_overhead.py`, CPU, median of 15 runs after warmup — see [`benchmarks/results/overhead_results.csv`](benchmarks/results/overhead_results.csv) for raw numbers and reproduce with `python benchmarks/bench_overhead.py`):
158
+
159
+ <p align="center">
160
+ <img src="assets/benchmark_overhead.png" width="600" alt="sherlorch tracing overhead vs model size">
161
+ </p>
162
+
163
+ <p align="center">
164
+ <img src="assets/benchmark_speedup.png" width="600" alt="sherlorch overhead multiplier, fast_shape_ops on vs off">
165
+ </p>
166
+
167
+ Takeaways from the actual measured numbers:
168
+
169
+ - **`fast_shape_ops=True` (default) consistently cuts 35–47% off tracing overhead** compared to always doing a full scan, across model sizes from d_model=64 to d_model=1024.
170
+ - **The overhead multiplier shrinks as models get bigger** — at d_model=64 the default mode is ~3.9x baseline forward-pass time, dropping to ~1.5x at d_model=1024, because per-op Python bookkeeping is a fixed cost while actual tensor compute grows with size and dominates more.
171
+ - **`capture_stack=True` is the expensive option** — it's fine in a plain script, but calls `inspect.stack()`, which is dramatically more costly under Jupyter/IPython's deeper call stacks (see the walkthrough and real numbers in [`notebooks/sherlorch_demo.ipynb`](notebooks/sherlorch_demo.ipynb)). Turn it off, or use it sparingly, inside notebooks.
172
+
173
+ This is a **debugging tool**, not something to leave on during real training — the numbers above are the honest cost of that convenience, not a claim that tracing is free.
174
+
175
+ ## How it works
176
+
177
+ `ProvenanceTracker` is a `TorchFunctionMode`. While active, every `torch.*` call is intercepted: the op runs normally, and a small `OpNode` (op name, shape, dtype, finite/nan/inf flags, optional module tag and source location) is recorded and linked to the `OpNode`s of its input tensors. Tensors are tracked by `id()` in a plain dict, not by the tensor object itself — `torch.Tensor.__eq__` is elementwise, which breaks the equality checks a `WeakKeyDictionary` needs internally. A `weakref` callback removes each entry as soon as its tensor is garbage collected, so a later `id()` reuse can never collide with a stale entry, and traced tensors are never kept alive artificially.
178
+
179
+ `diagnose(tensor)` then does a backward BFS from the target tensor's node, following only the non-finite ancestors, until it finds nodes whose inputs were *all* finite — those are the culprits: the operations that actually introduced the corruption, as opposed to ones that merely inherited it.
180
+
181
+ ## Costs and limitations
182
+
183
+ - This is a **debugging tool**, not something to leave on during real training — see [Benchmarks](#benchmarks) for the actual measured overhead.
184
+ - Backward-pass ops are captured on a best-effort basis; some fused/kernel-level backward computation may not surface individual sub-ops.
185
+ - Only tensors actually produced *while tracing was active* have provenance; tensors created outside the `with` block, or moved off the tracked identity (e.g. via certain C++-side aliasing), will raise `KeyError` on `diagnose()`.
186
+ - Module tagging via `watch()` assumes normally-nested forward calls; activation checkpointing or re-entrant forward passes can produce imprecise tags.
187
+
188
+ ## Try it
189
+
190
+ - [`examples/find_nan_in_transformer.py`](examples/find_nan_in_transformer.py) — a runnable script reproducing the example at the top of this README.
191
+ - [`notebooks/sherlorch_demo.ipynb`](notebooks/sherlorch_demo.ipynb) — an executed walkthrough covering the minimal repro, module tagging on a small transformer, the full benchmark suite with plots, and a comparison with how you'd normally debug this.
192
+
193
+ ## Contributing
194
+
195
+ Issues and PRs welcome — this is intentionally a small, sharply-scoped tool. See [CONTRIBUTING.md](CONTRIBUTING.md) for setup, design principles, and good first contributions, and [CODE_OF_CONDUCT.md](CODE_OF_CONDUCT.md) for community expectations. Changes are tracked in [CHANGELOG.md](CHANGELOG.md).
196
+
197
+ ## License
198
+
199
+ MIT — see [LICENSE](LICENSE).
@@ -0,0 +1,165 @@
1
+ ![sherlorch name logo](<docs/images/sherlorch name logo.png>)
2
+
3
+ # sherlorch
4
+
5
+ [![CI](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml/badge.svg)](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml)
6
+ [![PyPI version](https://img.shields.io/pypi/v/sherlorch.svg)](https://pypi.org/project/sherlorch/)
7
+ [![Python versions](https://img.shields.io/pypi/pyversions/sherlorch.svg)](https://pypi.org/project/sherlorch/)
8
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
9
+ [![Code style](https://img.shields.io/badge/types-checked-brightgreen.svg)](https://peps.python.org/pep-0561/)
10
+
11
+ **Find the exact op that introduced a NaN or Inf into your PyTorch model — not just the line where you noticed it.**
12
+
13
+ Every deep learning researcher knows this moment: `loss.backward()` returns `nan`, and now you're bisecting your forward pass by hand, printing `.isnan().any()` after every other line. `sherlorch` automates that entire process.
14
+
15
+ It works by transparently intercepting every tensor operation while active (via `torch.overrides.TorchFunctionMode` — no hooks to register, no model changes needed), building a lightweight op-level provenance graph. When you point it at a bad tensor, it walks that graph backward and tells you precisely which operation, on which inputs, first produced the non-finite value — plus the path from there to the tensor you noticed the problem in.
16
+
17
+ ```
18
+ sherlorch diagnosis for [14] layer_norm (shape=(4, 4, 16, 32)):
19
+ Found 1 culprit op(s) out of 15 traced ops.
20
+
21
+ Culprit #1: [10] log shape=(4, 4, 16, 16) dtype=torch.float32 [model.layer1] (train.py:27) <-- Inf
22
+ Path to target:
23
+ -> [10] log ... [model.layer1] <-- Inf
24
+ -> [11] mean ... [model] <-- Inf
25
+ -> [12] unsqueeze ... [model] <-- Inf
26
+ -> [13] add ... [model] <-- Inf
27
+ -> [14] layer_norm ... [model.norm] <-- NaN
28
+ ```
29
+
30
+ That's real output from [`examples/find_nan_in_transformer.py`](examples/find_nan_in_transformer.py) — a `log(relu(x))` bug buried two submodules deep, found and localized to `model.layer1` in one traced run, no bisection required.
31
+
32
+ ## Contents
33
+
34
+ - [sherlorch](#sherlorch)
35
+ - [Contents](#contents)
36
+ - [Install](#install)
37
+ - [Quickstart](#quickstart)
38
+ - [Options](#options)
39
+ - [Module tagging](#module-tagging)
40
+ - [Perf mode](#perf-mode)
41
+ - [Benchmarks](#benchmarks)
42
+ - [How it works](#how-it-works)
43
+ - [Costs and limitations](#costs-and-limitations)
44
+ - [Try it](#try-it)
45
+ - [Contributing](#contributing)
46
+ - [License](#license)
47
+
48
+ ## Install
49
+
50
+ ```bash
51
+ pip install sherlorch
52
+ ```
53
+
54
+ Or from a clone, for local development:
55
+
56
+ ```bash
57
+ git clone https://github.com/SankaVaas/sherlorch.git
58
+ cd sherlorch
59
+ pip install -e ".[dev]"
60
+ ```
61
+
62
+ Requires `torch>=2.0`, Python `>=3.9`.
63
+
64
+ ## Quickstart
65
+
66
+ ```python
67
+ import torch
68
+ import sherlorch
69
+
70
+ with sherlorch.trace() as tracker:
71
+ out = model(x)
72
+ loss = criterion(out, y)
73
+ loss.backward()
74
+
75
+ if tracker.has_issue(loss):
76
+ print(tracker.diagnose(loss))
77
+ ```
78
+
79
+ Example output:
80
+
81
+ ```
82
+ sherlorch diagnosis for [812] sum (shape=()):
83
+ Found 1 culprit op(s) out of 811 traced ops.
84
+
85
+ Culprit #1: [340] div shape=(32, 128) dtype=torch.float32 (train.py:57) <-- Inf
86
+ Path to target:
87
+ -> [340] div shape=(32, 128) dtype=torch.float32 (train.py:57) <-- Inf
88
+ -> [341] mul shape=(32, 128) dtype=torch.float32 <-- Inf
89
+ -> [812] sum shape=() <-- Inf
90
+ ```
91
+
92
+ ### Options
93
+
94
+ | Option | Default | What it does |
95
+ |---|---|---|
96
+ | `sherlorch.trace(capture_stack=True)` | `False` | Attach a `file:line` to each recorded op. Useful while actively hunting a bug; adds real overhead (see [Benchmarks](#benchmarks)), especially inside Jupyter/IPython. |
97
+ | `sherlorch.trace(fast_shape_ops=False)` | `True` | Force a full isnan/isinf scan on every op, including shape-only ops. See [Perf mode](#perf-mode). |
98
+ | `tracker.watch(model, name="model")` | — | Tag every recorded op with the submodule that produced it. See [Module tagging](#module-tagging). |
99
+ | `tracker.disable()` / `enable()` | — | Pause/resume recording inside the `with` block, e.g. to skip a warmup loop. |
100
+ | `tracker.reset()` | — | Drop all recorded provenance, e.g. between training steps, to bound memory on long runs. |
101
+ | `tracker.has_issue(tensor)` | — | Cheap NaN/Inf check on any live tensor, tracked or not. |
102
+
103
+ ## Module tagging
104
+
105
+ Wrap `model(x)` with `tracker.watch(model)` and every op's diagnosis includes which submodule produced it:
106
+
107
+ ```python
108
+ with sherlorch.trace() as tracker:
109
+ tracker.watch(model, name="model")
110
+ out = model(x)
111
+
112
+ print(tracker.diagnose(out))
113
+ ```
114
+
115
+ That `[model.layer1]` tag comes straight from `model.named_modules()`, so it matches whatever names you gave your submodules. Hooks are removed automatically when the `with` block exits, or manually via `tracker.unwatch()`. Best-effort: relies on forward hooks firing in normal LIFO order, which activation checkpointing or re-entrant forward calls can disrupt.
116
+
117
+ ## Perf mode
118
+
119
+ By default (`fast_shape_ops=True`), ops that can only rearrange, copy, index, or combine existing values — `view`, `reshape`, `permute`, `transpose`, `cat`, `stack`, `narrow`, and similar — skip the isnan/isinf tensor scan entirely and instead **inherit** finiteness from their parent node(s). This is exact, not an approximation: those ops provably cannot introduce a NaN/Inf that wasn't already in one of their inputs (dtype-narrowing ops like `.to()`/`.half()` are deliberately excluded, since precision loss can itself overflow to Inf). Since transformer-style models are full of reshapes and permutes, this cuts meaningful overhead — and the savings grow with tensor size, since the skipped scan cost is proportional to tensor size while inheritance is O(1).
120
+
121
+ ## Benchmarks
122
+
123
+ Measured on a small transformer attention block (`benchmarks/bench_overhead.py`, CPU, median of 15 runs after warmup — see [`benchmarks/results/overhead_results.csv`](benchmarks/results/overhead_results.csv) for raw numbers and reproduce with `python benchmarks/bench_overhead.py`):
124
+
125
+ <p align="center">
126
+ <img src="assets/benchmark_overhead.png" width="600" alt="sherlorch tracing overhead vs model size">
127
+ </p>
128
+
129
+ <p align="center">
130
+ <img src="assets/benchmark_speedup.png" width="600" alt="sherlorch overhead multiplier, fast_shape_ops on vs off">
131
+ </p>
132
+
133
+ Takeaways from the actual measured numbers:
134
+
135
+ - **`fast_shape_ops=True` (default) consistently cuts 35–47% off tracing overhead** compared to always doing a full scan, across model sizes from d_model=64 to d_model=1024.
136
+ - **The overhead multiplier shrinks as models get bigger** — at d_model=64 the default mode is ~3.9x baseline forward-pass time, dropping to ~1.5x at d_model=1024, because per-op Python bookkeeping is a fixed cost while actual tensor compute grows with size and dominates more.
137
+ - **`capture_stack=True` is the expensive option** — it's fine in a plain script, but calls `inspect.stack()`, which is dramatically more costly under Jupyter/IPython's deeper call stacks (see the walkthrough and real numbers in [`notebooks/sherlorch_demo.ipynb`](notebooks/sherlorch_demo.ipynb)). Turn it off, or use it sparingly, inside notebooks.
138
+
139
+ This is a **debugging tool**, not something to leave on during real training — the numbers above are the honest cost of that convenience, not a claim that tracing is free.
140
+
141
+ ## How it works
142
+
143
+ `ProvenanceTracker` is a `TorchFunctionMode`. While active, every `torch.*` call is intercepted: the op runs normally, and a small `OpNode` (op name, shape, dtype, finite/nan/inf flags, optional module tag and source location) is recorded and linked to the `OpNode`s of its input tensors. Tensors are tracked by `id()` in a plain dict, not by the tensor object itself — `torch.Tensor.__eq__` is elementwise, which breaks the equality checks a `WeakKeyDictionary` needs internally. A `weakref` callback removes each entry as soon as its tensor is garbage collected, so a later `id()` reuse can never collide with a stale entry, and traced tensors are never kept alive artificially.
144
+
145
+ `diagnose(tensor)` then does a backward BFS from the target tensor's node, following only the non-finite ancestors, until it finds nodes whose inputs were *all* finite — those are the culprits: the operations that actually introduced the corruption, as opposed to ones that merely inherited it.
146
+
147
+ ## Costs and limitations
148
+
149
+ - This is a **debugging tool**, not something to leave on during real training — see [Benchmarks](#benchmarks) for the actual measured overhead.
150
+ - Backward-pass ops are captured on a best-effort basis; some fused/kernel-level backward computation may not surface individual sub-ops.
151
+ - Only tensors actually produced *while tracing was active* have provenance; tensors created outside the `with` block, or moved off the tracked identity (e.g. via certain C++-side aliasing), will raise `KeyError` on `diagnose()`.
152
+ - Module tagging via `watch()` assumes normally-nested forward calls; activation checkpointing or re-entrant forward passes can produce imprecise tags.
153
+
154
+ ## Try it
155
+
156
+ - [`examples/find_nan_in_transformer.py`](examples/find_nan_in_transformer.py) — a runnable script reproducing the example at the top of this README.
157
+ - [`notebooks/sherlorch_demo.ipynb`](notebooks/sherlorch_demo.ipynb) — an executed walkthrough covering the minimal repro, module tagging on a small transformer, the full benchmark suite with plots, and a comparison with how you'd normally debug this.
158
+
159
+ ## Contributing
160
+
161
+ Issues and PRs welcome — this is intentionally a small, sharply-scoped tool. See [CONTRIBUTING.md](CONTRIBUTING.md) for setup, design principles, and good first contributions, and [CODE_OF_CONDUCT.md](CODE_OF_CONDUCT.md) for community expectations. Changes are tracked in [CHANGELOG.md](CHANGELOG.md).
162
+
163
+ ## License
164
+
165
+ MIT — see [LICENSE](LICENSE).