Metadata-Version: 2.5
Name: sherlorch
Version: 0.1.0
Summary: PyTorch-first tensor provenance debugger: find the exact op that introduced a NaN/Inf.
Project-URL: Homepage, https://github.com/SankaVaas/sherlorch
Project-URL: Issues, https://github.com/SankaVaas/sherlorch/issues
Project-URL: Changelog, https://github.com/SankaVaas/sherlorch/blob/main/CHANGELOG.md
Author: Sanka Vaas
License: MIT
License-File: LICENSE
Keywords: autograd,debugging,deep-learning,nan,provenance,pytorch
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Debuggers
Classifier: Typing :: Typed
Requires-Python: >=3.9
Requires-Dist: torch>=2.0
Provides-Extra: benchmark
Requires-Dist: matplotlib>=3.5; extra == 'benchmark'
Provides-Extra: dev
Requires-Dist: build; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: twine; extra == 'dev'
Description-Content-Type: text/markdown

![sherlorch name logo](<docs/images/sherlorch name logo.png>)

# sherlorch

[![CI](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml/badge.svg)](https://github.com/SankaVaas/sherlorch/actions/workflows/ci.yml)
[![PyPI version](https://img.shields.io/pypi/v/sherlorch.svg)](https://pypi.org/project/sherlorch/)
[![Python versions](https://img.shields.io/pypi/pyversions/sherlorch.svg)](https://pypi.org/project/sherlorch/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Code style](https://img.shields.io/badge/types-checked-brightgreen.svg)](https://peps.python.org/pep-0561/)

**Find the exact op that introduced a NaN or Inf into your PyTorch model — not just the line where you noticed it.**

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.

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.

```
sherlorch diagnosis for [14] layer_norm (shape=(4, 4, 16, 32)):
  Found 1 culprit op(s) out of 15 traced ops.

  Culprit #1: [10] log  shape=(4, 4, 16, 16) dtype=torch.float32  [model.layer1]  (train.py:27)  <-- Inf
    Path to target:
      -> [10] log  ...  [model.layer1]  <-- Inf
      -> [11] mean  ...  [model]  <-- Inf
      -> [12] unsqueeze  ...  [model]  <-- Inf
      -> [13] add  ...  [model]  <-- Inf
      -> [14] layer_norm  ...  [model.norm]  <-- NaN
```

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.

## Contents

- [sherlorch](#sherlorch)
  - [Contents](#contents)
  - [Install](#install)
  - [Quickstart](#quickstart)
    - [Options](#options)
  - [Module tagging](#module-tagging)
  - [Perf mode](#perf-mode)
  - [Benchmarks](#benchmarks)
  - [How it works](#how-it-works)
  - [Costs and limitations](#costs-and-limitations)
  - [Try it](#try-it)
  - [Contributing](#contributing)
  - [License](#license)

## Install

```bash
pip install sherlorch
```

Or from a clone, for local development:

```bash
git clone https://github.com/SankaVaas/sherlorch.git
cd sherlorch
pip install -e ".[dev]"
```

Requires `torch>=2.0`, Python `>=3.9`.

## Quickstart

```python
import torch
import sherlorch

with sherlorch.trace() as tracker:
    out = model(x)
    loss = criterion(out, y)
    loss.backward()

if tracker.has_issue(loss):
    print(tracker.diagnose(loss))
```

Example output:

```
sherlorch diagnosis for [812] sum (shape=()):
  Found 1 culprit op(s) out of 811 traced ops.

  Culprit #1: [340] div  shape=(32, 128) dtype=torch.float32  (train.py:57)  <-- Inf
    Path to target:
      -> [340] div  shape=(32, 128) dtype=torch.float32  (train.py:57)  <-- Inf
      -> [341] mul  shape=(32, 128) dtype=torch.float32  <-- Inf
      -> [812] sum  shape=()  <-- Inf
```

### Options

| Option | Default | What it does |
|---|---|---|
| `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. |
| `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). |
| `tracker.watch(model, name="model")` | — | Tag every recorded op with the submodule that produced it. See [Module tagging](#module-tagging). |
| `tracker.disable()` / `enable()` | — | Pause/resume recording inside the `with` block, e.g. to skip a warmup loop. |
| `tracker.reset()` | — | Drop all recorded provenance, e.g. between training steps, to bound memory on long runs. |
| `tracker.has_issue(tensor)` | — | Cheap NaN/Inf check on any live tensor, tracked or not. |

## Module tagging

Wrap `model(x)` with `tracker.watch(model)` and every op's diagnosis includes which submodule produced it:

```python
with sherlorch.trace() as tracker:
    tracker.watch(model, name="model")
    out = model(x)

print(tracker.diagnose(out))
```

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.

## Perf mode

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).

## Benchmarks

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`):

<p align="center">
  <img src="assets/benchmark_overhead.png" width="600" alt="sherlorch tracing overhead vs model size">
</p>

<p align="center">
  <img src="assets/benchmark_speedup.png" width="600" alt="sherlorch overhead multiplier, fast_shape_ops on vs off">
</p>

Takeaways from the actual measured numbers:

- **`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.
- **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.
- **`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.

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.

## How it works

`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.

`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.

## Costs and limitations

- This is a **debugging tool**, not something to leave on during real training — see [Benchmarks](#benchmarks) for the actual measured overhead.
- Backward-pass ops are captured on a best-effort basis; some fused/kernel-level backward computation may not surface individual sub-ops.
- 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()`.
- Module tagging via `watch()` assumes normally-nested forward calls; activation checkpointing or re-entrant forward passes can produce imprecise tags.

## Try it

- [`examples/find_nan_in_transformer.py`](examples/find_nan_in_transformer.py) — a runnable script reproducing the example at the top of this README.
- [`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.

## Contributing

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).

## License

MIT — see [LICENSE](LICENSE).
