silentbug 0.1.0__py3-none-any.whl

silentbug/__init__.py

@@ -0,0 +1,7 @@
+"""silentbug: find silent PyTorch training bugs and the line that caused them."""
+
+from .report import BUG, CHECK, Finding, Report
+from .session import Session, check
+
+__all__ = ["check", "Session", "Report", "Finding", "BUG", "CHECK"]
+__version__ = "0.1.0"

silentbug/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

silentbug/_where.py

@@ -0,0 +1,100 @@
+"""Find the line of user code that called the current torch op."""
+
+from __future__ import annotations
+
+import linecache
+import os
+import sys
+import sysconfig
+from dataclasses import dataclass
+
+import torch
+
+_TORCH_DIR = os.path.dirname(os.path.realpath(torch.__file__)) + os.sep
+_OWN_DIR = os.path.dirname(os.path.realpath(__file__)) + os.sep
+_LIB_DIRS = tuple(
+    {os.path.realpath(p) + os.sep for p in (sysconfig.get_paths().get("purelib"), sysconfig.get_paths().get("platlib")) if p}
+)
+_STDLIB = os.path.realpath(sysconfig.get_paths()["stdlib"]) + os.sep
+
+_real_cache: dict[str, str] = {}
+
+
+def _real(path: str) -> str:
+    r = _real_cache.get(path)
+    if r is None:
+        r = os.path.realpath(path) if not path.startswith("<") else path
+        _real_cache[path] = r
+    return r
+
+
+def _skip(path: str) -> bool:
+    if path.startswith("<"):
+        return True
+    p = _real(path)
+    return p.startswith(_TORCH_DIR) or p.startswith(_OWN_DIR) or p.startswith(_STDLIB) and not any(
+        p.startswith(d) for d in _LIB_DIRS
+    )
+
+
+def _is_lib(path: str) -> bool:
+    p = _real(path)
+    return any(p.startswith(d) for d in _LIB_DIRS)
+
+
+@dataclass(frozen=True)
+class Where:
+    file: str
+    line: int
+    func: str
+    code: str
+    lib: bool = False  # True if the line is inside an installed package
+
+    def short(self) -> str:
+        f = self.file
+        try:
+            rel = os.path.relpath(f)
+            if not rel.startswith(".."):
+                f = rel
+        except ValueError:
+            pass
+        return f"{f}:{self.line}"
+
+    def to_dict(self) -> dict:
+        return {"file": self.file, "line": self.line, "function": self.func, "code": self.code, "library": self.lib}
+
+
+# torch packages that copy, sync or update parameters as bookkeeping, not as part of the model
+_INFRA = tuple(_TORCH_DIR + d + os.sep for d in (
+    "distributed", os.path.join("nn", "parallel"), "optim", os.path.join("nn", "utils"),
+    "amp", os.path.join("cuda", "amp")))
+
+
+def torch_caller(skip: int = 2) -> str | None:
+    """'infra' or 'torch' when the nearest frame outside silentbug is PyTorch's own code."""
+    f = sys._getframe(skip)
+    while f is not None:
+        p = _real(f.f_code.co_filename)
+        if not p.startswith(_OWN_DIR):
+            if p.startswith(_INFRA) or p.startswith(_TORCH_DIR + "_tensor_str"):
+                return "infra"
+            return "torch" if p.startswith(_TORCH_DIR) else None
+        f = f.f_back
+    return None
+
+
+def where(skip: int = 2) -> tuple[Where | None, Where | None]:
+    """Return (first frame outside torch, first frame outside installed packages)."""
+    f = sys._getframe(skip)
+    first = None
+    while f is not None:
+        path = f.f_code.co_filename
+        if not _skip(path):
+            lib = _is_lib(path)
+            w = Where(path, f.f_lineno, f.f_code.co_name, linecache.getline(path, f.f_lineno).strip(), lib)
+            if first is None:
+                first = w
+            if not lib:
+                return first, (w if w != first else None)
+        f = f.f_back
+    return first, None

silentbug/cli.py

@@ -0,0 +1,80 @@
+"""Command line: `silentbug run train.py [args...]`."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import runpy
+import sys
+
+from .session import Session, _Done
+
+USAGE = """silentbug run [--steps N] [--keep-going] [--json FILE] (script.py | -m module) [args...]
+
+Runs your training script with silentbug watching the first N optimizer steps
+(default 3), prints what it found, and stops the script. Exit code: 0 nothing
+found, 1 likely bugs found, 2 only things to check."""
+
+
+def _parse(argv: list[str]):
+    ap = argparse.ArgumentParser(prog="silentbug run", usage=USAGE)
+    ap.add_argument("--steps", type=int, default=3)
+    ap.add_argument("--keep-going", action="store_true", help="let the script keep running after the report")
+    ap.add_argument("--json", metavar="FILE", help="also write the report as JSON ('-' for stdout only JSON)")
+    ap.add_argument("-m", dest="module", action="store_true", help="run a module, like python -m")
+    ap.add_argument("target")
+    ap.add_argument("args", nargs=argparse.REMAINDER)
+    return ap.parse_args(argv)
+
+
+def run(argv: list[str]) -> int:
+    a = _parse(argv)
+    s = Session(max_steps=a.steps, stop_at_limit=not a.keep_going)
+    old_argv = sys.argv[:]
+    sys.argv = [a.target, *a.args]
+    if not a.module:
+        sys.path.insert(0, os.path.dirname(os.path.abspath(a.target)))
+    error = None
+    s.start()
+    try:
+        if a.module:
+            runpy.run_module(a.target, run_name="__main__", alter_sys=True)
+        else:
+            runpy.run_path(a.target, run_name="__main__")
+    except _Done:
+        pass
+    except SystemExit:
+        pass
+    except BaseException as e:  # the script crashed: still report what we saw
+        error = e
+    finally:
+        sys.argv = old_argv
+        rep = s.stop()
+    rank = os.environ.get("RANK", "0")  # under torchrun every rank sees the same code
+    if rank != "0":
+        pass
+    elif a.json == "-":
+        print(rep.to_json())
+    else:
+        print(rep, file=sys.stderr if a.keep_going else sys.stdout)
+        if a.json:
+            with open(a.json, "w") as f:
+                f.write(rep.to_json())
+    if error is not None:
+        print(f"\n(your script raised {type(error).__name__}: {error})", file=sys.stderr)
+    return 1 if rep.bugs else (2 if rep.findings else 0)
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = sys.argv[1:] if argv is None else argv
+    if not argv or argv[0] in ("-h", "--help"):
+        print(USAGE)
+        return 0
+    if argv[0] == "run":
+        return run(argv[1:])
+    if argv[0] in ("--version", "version"):
+        from . import __version__
+        print(__version__)
+        return 0
+    print(f"unknown command {argv[0]!r}\n\n{USAGE}", file=sys.stderr)
+    return 64

silentbug/explain.py

@@ -0,0 +1,225 @@
+"""Plain-language findings: what went wrong, why it matters, how to fix it."""
+
+from __future__ import annotations
+
+import math
+
+from ._where import Where
+from .report import BUG, CHECK, Finding
+
+
+def _s(shape) -> str:
+    return "(" + ", ".join(str(x) for x in shape) + ("," if len(shape) == 1 else "") + ")"
+
+
+# ---------------------------------------------------------------- shapes
+
+def broadcast(op: str, sym: str | None, sa, sb, out, same: bool, w: Where | None, c: Where | None) -> Finding:
+    expr = f"a {sym} b" if sym else f"{op}(a, b)"
+    big = math.prod(out)
+    small = max(math.prod(sa), math.prod(sb))
+    lo, hi = (sa, sb) if len(sa) < len(sb) else (sb, sa)
+    title = f"{expr} broadcasts {_s(sa)} with {_s(sb)} into {_s(out)}"
+    if same:
+        why = (f"the shapes have different ranks, so PyTorch puts a new leading axis on {_s(lo)} and stretches both "
+               f"tensors: you get {big} values instead of {small}. This is the classic pred (N,1) vs target (N,) "
+               "mistake, where every row is paired with every other row. Nothing errors; the loss is just wrong.")
+        level = BUG
+    else:
+        why = (f"both tensors get stretched ({_s(lo)} gets a new leading axis), giving {big} values from inputs of at "
+               f"most {small}. That is right for an outer product, and wrong anywhere you meant to combine matching "
+               "elements.")
+        level = CHECK
+    fix = (f"make the shapes match, e.g. `.squeeze(-1)` on the {_s(hi)} tensor or `.unsqueeze(-1)` on the {_s(lo)} one. "
+           "If you do want all pairs, write the new axis yourself (`x[:, None]`, `unsqueeze`) so the intent is visible; "
+           "silentbug then stays quiet.")
+    return Finding("broadcast", level, title, why, fix, w, c, detail={"a": list(sa), "b": list(sb), "out": list(out)})
+
+
+def loss_shape(op: str, sx, sy, so, w, c) -> Finding:
+    grows = math.prod(so) > max(math.prod(sx), math.prod(sy))
+    title = f"{op} got input {_s(sx)} and target {_s(sy)}"
+    if grows:
+        why = (f"they broadcast to {_s(so)}, so each prediction is compared with many targets instead of its own. "
+               "The loss still goes down a little, which hides the bug. PyTorch only prints a UserWarning.")
+    else:
+        why = (f"they broadcast to {_s(so)}: one tensor is reused across a dimension of the other "
+               "(for example a batch of 1 compared with every sample).")
+    fix = "reshape one side so input and target have the same shape (often `.squeeze(-1)` or `.view_as(input)`)"
+    return Finding("broadcast", BUG, title, why, fix, w, c, detail={"input": list(sx), "target": list(sy), "out": list(so)})
+
+
+# ---------------------------------------------------------------- activations into losses
+
+def activation_into_loss(loss: str, act: str, act_w: Where | None, w: Where | None, c: Where | None) -> Finding | None:
+    at = f" (made at {act_w.short()})" if act_w else ""
+    if loss == "cross_entropy" and act in ("softmax", "softmax-mix"):
+        what = "softmax probabilities" if act == "softmax" else "values built from softmax probabilities"
+        return Finding("double-activation", BUG, f"cross_entropy got {what}{at}",
+                       "cross_entropy applies log_softmax itself, so softmax runs twice. Outputs squash toward uniform, "
+                       "gradients shrink, and training stalls at a poor accuracy without any error.",
+                       "pass raw logits to the loss; apply softmax only for predictions and metrics", w, c)
+    if loss == "cross_entropy" and act == "sigmoid":
+        return Finding("double-activation", BUG, f"cross_entropy got sigmoid outputs{at}",
+                       "cross_entropy expects raw logits; squashing them into (0, 1) first caps how confident the "
+                       "model can get and weakens the gradient.",
+                       "pass the logits before sigmoid", w, c)
+    if loss == "nll_loss" and act in ("softmax", "sigmoid", "softmax-mix"):
+        return Finding("double-activation", BUG, f"nll_loss got {act} probabilities{at}, not log-probabilities",
+                       "nll_loss just picks -input[target]. With probabilities the loss becomes -p instead of -log p, "
+                       "which goes negative and gives tiny gradients.",
+                       "use log_softmax instead of softmax, or call cross_entropy on the logits", w, c)
+    if loss == "binary_cross_entropy_with_logits" and act == "sigmoid":
+        return Finding("double-activation", BUG, f"binary_cross_entropy_with_logits got sigmoid outputs{at}",
+                       "the loss applies sigmoid itself, so it runs twice: predictions are squeezed into (0.5, 0.73) "
+                       "and the loss can never reach zero.",
+                       "pass the logits, or switch to binary_cross_entropy if you must pass probabilities", w, c)
+    if loss == "kl_div" and act in ("softmax", "sigmoid"):
+        return Finding("double-activation", BUG, f"kl_div got {act} probabilities{at} as input",
+                       "F.kl_div expects log-probabilities as its first argument; with probabilities the value "
+                       "and its gradient are wrong.",
+                       "use log_softmax for the input (the target stays probabilities unless log_target=True)", w, c)
+    if loss in ("softmax", "log_softmax") and act == "softmax":
+        return Finding("double-activation", BUG, f"{loss} applied to an output that is already softmax{at}",
+                       "softmax runs twice, which squashes the distribution toward uniform and shrinks gradients.",
+                       "drop the first softmax and work from the logits", w, c)
+    return None
+
+
+# ---------------------------------------------------------------- gradients
+
+_CUT_TEXT = {
+    "detach": ("`.detach()` here cuts the graph, so no gradient flows back to these parameters",
+               "remove `.detach()` if these parameters should learn; keep it only for targets, logging, or a frozen branch"),
+    "data": ("`.data` gives a tensor with no gradient history, so everything computed from it is cut off",
+             "use the tensor itself (or `.detach()` where you really mean to stop gradients)"),
+    "value": ("`.{op}()` turns the tensor into a Python number or NumPy array; anything rebuilt from it has no gradient",
+              "keep the math in torch ops on tensors instead of going through Python numbers or NumPy"),
+    "new_tensor": ("`{op}(...)` makes a brand-new tensor with no gradient history",
+                   "use the existing tensor (or torch.stack/torch.cat to combine tensors) instead of building a new one"),
+    "no_grad": ("this line runs with gradients turned off (inside torch.no_grad() or a @torch.no_grad() function)",
+                "move this computation out of the no_grad block if its parameters should learn"),
+    "inference_mode": ("this line runs inside torch.inference_mode(), which records no gradients",
+                       "move this computation out of inference_mode if its parameters should learn"),
+    "requires_grad_false": ("`requires_grad_(False)` is called on a tensor computed from these parameters",
+                            "remove the requires_grad_(False) call"),
+    "int": ("`{op}` returns integer or boolean values (indices, comparisons, casts), which carry no gradient",
+            "use a differentiable path: return the scores as well as the indices and add them to the loss, use soft "
+            "weights (softmax) instead of argmax/topk indices, or a straight-through estimator"),
+    "other": ("`{op}` returns a tensor with no gradient history",
+              "check whether this op is differentiable; if not, route the loss through a differentiable alternative"),
+}
+
+
+def cut(params: list[str], ev) -> Finding:
+    why, fix = _CUT_TEXT.get(ev.kind, _CUT_TEXT["other"])
+    op = ev.op.replace("attr:", "")
+    why = why.format(op=op)
+    n = len(params)
+    title = f"{n} parameter{'s' if n != 1 else ''} never get a gradient: cut by {op}"
+    return Finding("dead-param", BUG, title, why + ". They will never train, and nothing errors (unless you use DDP).",
+                   fix, ev.where, ev.caller, sorted(params), {"cause": ev.kind, "op": op})
+
+
+def unused(params: list[str], module: str, called: bool, has_opt: bool,
+           w: Where | None = None, c: Where | None = None) -> Finding:
+    n = len(params)
+    mod = f"`{module}`" if module else "the model"
+    if module and not called:
+        title = f"{n} parameter{'s' if n != 1 else ''} never get a gradient: {mod} is never called"
+        if w is not None:
+            why = (f"{mod} is defined, but forward() never calls it: its weights are read directly at the line "
+                   "shown, which skips its forward() and anything it wraps (such as a LoRA adapter)")
+        else:
+            why = (f"{mod} is defined, but forward() never calls it (it may read its weights directly, like "
+                   "F.linear(x, layer.weight), which skips any wrapper such as a LoRA adapter)")
+        fix = f"call {mod} in forward(), or delete it / set requires_grad=False so it stops looking trainable"
+    else:
+        title = f"{n} parameter{'s' if n != 1 else ''} never get a gradient: never used in forward"
+        why = f"{mod} runs, but no op ever reads these parameters"
+        fix = "use them in forward(), or delete them / set requires_grad=False"
+    level = BUG if has_opt else CHECK
+    why += ". They sit in the model and the optimizer but never learn; DDP would fail with 'did not receive grad'."
+    return Finding("dead-param", level, title, why, fix, w, c, sorted(params), {"cause": "unused", "module": module})
+
+
+def orphan(params: list[str], w: Where | None, c: Where | None) -> Finding:
+    n = len(params)
+    title = f"{n} parameter{'s' if n != 1 else ''} never get a gradient: their output never reaches the loss"
+    why = ("they are used (first at the line shown), but nothing computed from them ends up in the value you call "
+           "backward() on: for example an extra output that is returned but never put into the loss, or a "
+           "feature map the next stage skips")
+    fix = "add that output to the loss, or remove the unused branch (or set requires_grad=False on it)"
+    return Finding("dead-param", BUG, title, why, fix, w, c, sorted(params), {"cause": "orphan"})
+
+
+def zero_grad(params: list[str], ev) -> Finding:
+    n = len(params)
+    if ev is not None:
+        title = f"{n} parameter{'s' if n != 1 else ''} only ever get a zero gradient: through {ev.op}()"
+        why = f"{ev.op}() has a zero derivative almost everywhere, so the gradient that reaches these parameters is exactly 0"
+        fix = f"use a straight-through estimator, e.g. `x + (x.{ev.op}() - x).detach()`"
+        return Finding("zero-grad", BUG, title, why, fix, ev.where, ev.caller, sorted(params), {"op": ev.op})
+    title = f"{n} parameter{'s' if n != 1 else ''} only ever get a zero gradient"
+    why = "every gradient element was exactly 0 in every step: they are masked out, behind a dead ReLU, or multiplied by zero"
+    fix = "check masks and activations on the path from these parameters to the loss"
+    return Finding("zero-grad", CHECK, title, why, fix, None, None, sorted(params))
+
+
+# ---------------------------------------------------------------- optimizer and modes
+
+def not_in_optimizer(params: list[str], numel: int) -> Finding:
+    n = len(params)
+    return Finding(
+        "not-in-optimizer", BUG, f"{n} trainable parameter{'s' if n != 1 else ''} ({numel:,} values) are in no optimizer",
+        "they get gradients but optimizer.step() never updates them. This usually means a layer was created or "
+        "replaced after the optimizer was built (for example in a subclass __init__ after super().__init__()).",
+        "build the optimizer after the whole model exists, or add these parameters with optimizer.add_param_group()",
+        params=sorted(params))
+
+
+def stale_optimizer(count: int, numel: int, replacements: list[str] | None = None) -> Finding:
+    why = ("these tensors are not part of any model that ran: a layer or parameter was probably replaced after the "
+           "optimizer was built, so the optimizer keeps updating the old copy.")
+    if replacements:
+        why += " The model now uses new parameters that no optimizer updates, so they never train."
+    return Finding(
+        "stale-optimizer", BUG, f"the optimizer updates {count} tensor{'s' if count != 1 else ''} ({numel:,} values) that the model no longer uses",
+        why, "rebuild the optimizer after replacing layers, or update parameters in place (param.data.copy_(new))",
+        params=sorted(replacements or []))
+
+
+def eval_mode(mods: list[str], whole: bool, w, c) -> Finding:
+    kinds = sorted(set(mods))
+    if whole:
+        title = "the model is in eval() mode while it trains"
+        why = (f"gradients flow, but the model's training flag is off, so {', '.join(kinds)} behave as at inference "
+               "(dropout does nothing, batch norm uses frozen running stats). Usually model.eval() was called for "
+               "validation and model.train() was never called again.")
+        level = BUG
+    else:
+        title = f"{len(mods)} layer{'s' if len(mods) != 1 else ''} ({', '.join(kinds)}) run in eval mode during training"
+        why = "part of the model is in eval mode while gradients flow. This is fine if you froze them on purpose."
+        level = CHECK
+    return Finding("eval-mode", level, title, why, "call model.train() before the training step (e.g. at the top of each epoch)", w, c)
+
+
+def batchnorm_batch1(shape, w, c) -> Finding:
+    return Finding("batchnorm-batch1", CHECK, f"batch norm in training mode got a batch of 1 {_s(shape)}",
+                   "with one sample, batch statistics are just that sample's own values, which makes the running "
+                   "stats noisy and the output nearly constant per channel",
+                   "use a larger batch, drop_last=True in the DataLoader, or GroupNorm/LayerNorm", w, c)
+
+
+def grad_not_cleared(w, c, manual: bool = False) -> Finding:
+    if manual:
+        return Finding("grad-not-cleared", BUG, "gradients from the previous step were not cleared",
+                       "this weight update used gradients that still held the previous step's values: backward() adds onto "
+                       ".grad, so every update uses the sum of all past gradients and the effective learning rate keeps growing",
+                       "clear the gradients once per step, e.g. model.zero_grad() or p.grad = None for each parameter, "
+                       "before backward()", w, c)
+    return Finding("grad-not-cleared", BUG, "gradients from the previous step were not cleared",
+                   "this optimizer.step() used gradients that still held the previous step's values: backward() adds onto .grad, so every step uses the "
+                   "sum of all past gradients and the effective learning rate keeps growing",
+                   "call optimizer.zero_grad() once per step (before backward() or after step()). For gradient "
+                   "accumulation, zero right after each optimizer.step().", w, c)

silentbug/report.py

@@ -0,0 +1,105 @@
+"""Findings and the printed report."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+
+from ._where import Where
+
+# How sure we are that a finding is a real bug.
+BUG = "bug"  # almost certainly wrong
+CHECK = "check"  # often wrong, sometimes intended
+
+
+@dataclass
+class Finding:
+    kind: str  # broadcast | dead-param | zero-grad | not-in-optimizer | stale-optimizer | double-activation | eval-mode | batchnorm-batch1 | grad-not-cleared
+    level: str  # BUG or CHECK
+    title: str
+    why: str
+    fix: str
+    where: Where | None = None
+    caller: Where | None = None  # your own line, when `where` is inside a library
+    params: list[str] = field(default_factory=list)
+    detail: dict = field(default_factory=dict)
+
+    def key(self) -> tuple:
+        loc = (self.where.file, self.where.line) if self.where else None
+        return (self.kind, loc, self.title)
+
+    def to_dict(self) -> dict:
+        return {
+            "kind": self.kind,
+            "level": self.level,
+            "title": self.title,
+            "why": self.why,
+            "fix": self.fix,
+            "where": self.where.to_dict() if self.where else None,
+            "caller": self.caller.to_dict() if self.caller else None,
+            "params": self.params,
+            "detail": self.detail,
+        }
+
+
+def _plist(names: list[str], limit: int = 6) -> str:
+    if len(names) <= limit:
+        return ", ".join(names)
+    return ", ".join(names[:limit]) + f" and {len(names) - limit} more"
+
+
+@dataclass
+class Report:
+    findings: list[Finding] = field(default_factory=list)
+    steps: int = 0  # optimizer steps seen
+    backwards: int = 0  # backward passes seen
+    models: list[str] = field(default_factory=list)
+
+    @property
+    def bugs(self) -> list[Finding]:
+        return [f for f in self.findings if f.level == BUG]
+
+    @property
+    def checks(self) -> list[Finding]:
+        return [f for f in self.findings if f.level == CHECK]
+
+    def kinds(self) -> set[str]:
+        return {f.kind for f in self.findings}
+
+    def __bool__(self) -> bool:  # truthy when something was found
+        return bool(self.findings)
+
+    def to_dict(self) -> dict:
+        return {
+            "findings": [f.to_dict() for f in self.findings],
+            "steps": self.steps,
+            "backwards": self.backwards,
+            "models": self.models,
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict(), indent=2)
+
+    def __str__(self) -> str:
+        head = f"silentbug: checked {self.backwards} backward pass(es), {self.steps} optimizer step(s)"
+        if self.models:
+            head += f" on {', '.join(self.models)}"
+        if not self.findings:
+            msg = head + "\n  no silent bugs found"
+            if self.backwards == 0:
+                msg += " (but no backward pass ran, so gradient checks were skipped)"
+            return msg
+        out = [head, f"  {len(self.bugs)} likely bug(s), {len(self.checks)} thing(s) to check", ""]
+        for i, f in enumerate(sorted(self.findings, key=lambda f: f.level != BUG), 1):
+            tag = "BUG  " if f.level == BUG else "CHECK"
+            out.append(f"[{tag}] {i}. {f.title}")
+            if f.where:
+                out.append(f"    at {f.where.short()}: {f.where.code}")
+            if f.caller:
+                out.append(f"    called from {f.caller.short()}: {f.caller.code}")
+            if f.params:
+                out.append(f"    parameters: {_plist(f.params)}")
+            out.append(f"    why: {f.why}")
+            out.append(f"    fix: {f.fix}")
+            out.append("")
+        return "\n".join(out).rstrip()