sherlorch

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 — a log(relu(x)) bug buried two submodules deep, found and localized to model.layer1 in one traced run, no bisection required.

Contents

• sherlorch
  • Contents
  • Install
  • Quickstart
    • Options
  • Module tagging
  • Perf mode
  • Benchmarks
  • How it works
  • Costs and limitations
  • Try it
  • Contributing
  • License

Install

pip install sherlorch

Or from a clone, for local development:

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

Requires torch>=2.0, Python >=3.9.

Quickstart

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

Module tagging

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

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 for raw numbers and reproduce with python benchmarks/bench_overhead.py):

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). 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 OpNodes 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 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 — a runnable script reproducing the example at the top of this README.
• 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 for setup, design principles, and good first contributions, and CODE_OF_CONDUCT.md for community expectations. Changes are tracked in CHANGELOG.md.

License

MIT — see LICENSE.