trainscope 0.1.0__tar.gz

trainscope-0.1.0/.envrc

@@ -0,0 +1 @@
+use flake

trainscope-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,52 @@
+name: CI
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install dependencies
+        run: pip install -e ".[dev]"
+
+      - name: Lint (ruff)
+        run: ruff check trainscope/ tests/
+
+      - name: Test
+        run: pytest -q
+
+  frontend:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: frontend
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Node
+        uses: actions/setup-node@v4
+        with:
+          node-version: "20"
+
+      - name: Install
+        run: npm install
+
+      - name: Build
+        run: npm run build

trainscope-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,38 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write  # required for trusted publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+
+      - name: Install build tools
+        run: pip install hatch
+
+      - name: Build frontend
+        working-directory: frontend
+        run: |
+          npm install
+          npm run build
+        # vite outDir is ../trainscope/ui/static — build lands directly in the package
+
+      - name: Build wheel + sdist
+        run: hatch build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

trainscope-0.1.0/.gitignore

@@ -0,0 +1,64 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+.venv/
+venv/
+ENV/
+env/
+.mypy_cache/
+.ruff_cache/
+.pytest_cache/
+.coverage
+.coverage.*
+htmlcov/
+.tox/
+.nox/
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Frontend (Vite build lands in trainscope/ui/static)
+frontend/node_modules/
+trainscope/ui/static/
+
+# Runtime / local data
+trainscope_runs/
+*.parquet
+
+# IDE / editor
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Nix
+result
+result-*
+.direnv/
+flake.lock.bak
+
+# Local tooling
+.claude/settings.local.json

trainscope-0.1.0/PKG-INFO

@@ -0,0 +1,192 @@
+Metadata-Version: 2.4
+Name: trainscope
+Version: 0.1.0
+Summary: Post-mortem debugger for LLM training loss spikes
+Project-URL: Homepage, https://github.com/kaelvalen/trainscope
+Project-URL: Repository, https://github.com/kaelvalen/trainscope
+Project-URL: Bug Tracker, https://github.com/kaelvalen/trainscope/issues
+License: MIT
+Keywords: debugging,deep-learning,llm,loss-spike,pytorch,training
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: click>=8.1
+Requires-Dist: fastapi>=0.110
+Requires-Dist: numpy>=1.26
+Requires-Dist: pyarrow>=14.0
+Requires-Dist: scipy>=1.12
+Requires-Dist: torch>=2.0
+Requires-Dist: uvicorn[standard]>=0.27
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# trainscope
+
+Post-mortem debugger for LLM training loss spikes.
+
+When a spike hits, you usually know *that* it happened but not *why*. trainscope records per-layer gradients, weight distributions, and activation kurtosis at every step, then lets you scrub back through the event in a browser UI.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+Dependencies: `torch`, `pyarrow`, `fastapi`, `uvicorn`, `click`, `numpy`.
+
+## Quickstart
+
+```python
+from trainscope import TrainScope
+from trainscope.core.config import TrainScopeConfig
+
+scope = TrainScope(model, optimizer, config=TrainScopeConfig()).attach()
+
+for step, batch in enumerate(dataloader):
+    loss = forward_and_backward(batch)
+    optimizer.step()
+
+    spike = scope.step(loss.item(), batch_index=step)
+    if spike:
+        print(f"Spike at step {spike['step']}, z={spike['z_score']:.2f}")
+
+scope.writer.close()
+scope.detach()
+```
+
+Then open the UI:
+
+```bash
+trainscope ui --run ./trainscope_runs/<run-name>
+```
+
+## What gets recorded
+
+**Per step (global)**
+- Train loss, global grad norm (pre- and post-clip), learning rate
+- Adam second-moment (v) norm — stale momentum indicator
+- Step time, batch index
+
+**Per step, per layer**
+- Gradient L2 norm
+- Weight L2 norm
+- Activation mean / std / max-abs / kurtosis — kurtosis is the earliest spike signal
+- NaN/Inf ratio in gradients
+- 16-bin weight histogram
+
+**On spike**
+- Full snapshot of the surrounding window (configurable before/after)
+- Per-layer data for the same window
+- RNG state at the spike step (for exact replay)
+
+## Overhead
+
+Measured on CPU with a 2-layer GPT-2 (144 parameters). GPU overhead is ~3–8× lower.
+
+| Config | CPU overhead | GPU overhead |
+|--------|-------------|-------------|
+| Default (`hist/50`, `act/5`) | ~55% | ~4% |
+| + `activation_layer_filter=["attn","mlp"]` | ~38% | ~2% |
+| Minimal (`hist/50`, `act/50`, filter) | ~18% | ~1% |
+
+CPU measured on 2-layer mini-GPT (144 params), Apple M2. GPU measured on the same model with CUDA. Results will differ on larger models — histogram cost scales with parameter count, activation cost scales with layer count × sequence length.
+
+## UI
+
+Four views, one command:
+
+| View | What it shows |
+|------|---------------|
+| **Timeline** | Loss + grad norm, top-8 layers by grad variance |
+| **Layer Drill-down** | Kurtosis / grad norm / weight norm per layer; histogram scrubber |
+| **Diff View** | KL divergence of weight distributions between any two steps |
+| **Spike Inspector** | Per-spike window: loss+grad timeline and layer kurtosis/grad breakdown |
+
+The UI works immediately after `pip install` — a built-in fallback HTML with Plotly CDN is served when the React build is absent. For the full React build:
+
+```bash
+cd frontend && npm install && npm run build
+```
+
+## CLI
+
+```bash
+# Open UI for a completed or in-progress run
+trainscope ui --run ./trainscope_runs/run_20250516_143022 [--host 127.0.0.1] [--port 7007]
+
+# Generate replay_config.json (does NOT resume training automatically)
+trainscope replay --checkpoint ./checkpoints/step_4400.pt --skip-batches 4521,4522,4523 [--resume]
+```
+
+To actually skip batches, use `SkippingDataLoader` in your training script:
+
+```python
+from trainscope.replay import SkippingDataLoader
+import json
+
+with open("replay_config.json") as f:
+    cfg = json.load(f)
+
+loader = SkippingDataLoader(original_loader, skip_batches=cfg["skip_batches"])
+for batch in loader:
+    ...
+```
+
+## Configuration
+
+```python
+TrainScopeConfig(
+    run_dir="./trainscope_runs",                 # output root
+    spike_threshold=3.5,                     # z-score threshold (rolling window baseline)
+    full_resolution_window=500,              # last N steps at full resolution
+    decimation_factor=10,                    # older steps: keep every Nth
+    spike_window_before=50,                  # steps before spike to save (≤ full_resolution_window)
+    spike_window_after=10,                   # steps after spike to save
+    histogram_every_n_steps=50,             # weight histograms are expensive; sample them
+    activation_metrics_every_n_steps=5,     # kurtosis sampling; always captured at spike
+    activation_layer_filter=["attn", "mlp"],# None = all leaf layers
+    stop_on_spike=False,                     # raise StopTraining on detection
+    trace_every_n_steps=1,                   # subsample for very large models
+    rank=None,                               # DDP rank → adds _rank{N} suffix to run dir
+)
+```
+
+## Demo
+
+```bash
+python examples/gpt2_spike_demo.py
+```
+
+Trains a 2-layer mini-GPT, injects a ×50 loss spike at step 50, and shows trainscope detecting it. Run `trainscope ui` on the output directory to explore the event.
+
+## Storage layout
+
+```
+trainscope_runs/<run-name>/
+    meta.json                          model config + trainscope config
+    global.arrow                       step-level scalars (Arrow IPC)
+    layers/<param-name>.arrow          per-layer metrics
+    spikes/spike_step_<N>.arrow        global window around spike N
+    spikes/spike_step_<N>_layers/      per-layer data for that window
+    rng_states/step_<N>.pkl            RNG state for replay
+```
+
+Estimated storage: ~10 MB/step at full resolution. Rolling 500-step window → ~5 GB max for a 1B-param model. Spike windows are small.
+
+## Publishing
+
+CI runs on every push to `main` and every PR (`pytest` + `ruff`, Python 3.11 + 3.12, Vite build).
+
+To publish a release to PyPI:
+1. Set up [Trusted Publishing](https://docs.pypi.org/trusted-publishers/) on PyPI for this repo (environment name: `pypi`).
+2. Tag and push: `git tag v0.1.0 && git push origin v0.1.0`
+
+The publish workflow builds the React frontend, bundles it into the wheel, and uploads via OIDC — no API token needed.

trainscope-0.1.0/README.md

@@ -0,0 +1,162 @@
+# trainscope
+
+Post-mortem debugger for LLM training loss spikes.
+
+When a spike hits, you usually know *that* it happened but not *why*. trainscope records per-layer gradients, weight distributions, and activation kurtosis at every step, then lets you scrub back through the event in a browser UI.
+
+## Install
+
+```bash
+pip install -e .
+```
+
+Dependencies: `torch`, `pyarrow`, `fastapi`, `uvicorn`, `click`, `numpy`.
+
+## Quickstart
+
+```python
+from trainscope import TrainScope
+from trainscope.core.config import TrainScopeConfig
+
+scope = TrainScope(model, optimizer, config=TrainScopeConfig()).attach()
+
+for step, batch in enumerate(dataloader):
+    loss = forward_and_backward(batch)
+    optimizer.step()
+
+    spike = scope.step(loss.item(), batch_index=step)
+    if spike:
+        print(f"Spike at step {spike['step']}, z={spike['z_score']:.2f}")
+
+scope.writer.close()
+scope.detach()
+```
+
+Then open the UI:
+
+```bash
+trainscope ui --run ./trainscope_runs/<run-name>
+```
+
+## What gets recorded
+
+**Per step (global)**
+- Train loss, global grad norm (pre- and post-clip), learning rate
+- Adam second-moment (v) norm — stale momentum indicator
+- Step time, batch index
+
+**Per step, per layer**
+- Gradient L2 norm
+- Weight L2 norm
+- Activation mean / std / max-abs / kurtosis — kurtosis is the earliest spike signal
+- NaN/Inf ratio in gradients
+- 16-bin weight histogram
+
+**On spike**
+- Full snapshot of the surrounding window (configurable before/after)
+- Per-layer data for the same window
+- RNG state at the spike step (for exact replay)
+
+## Overhead
+
+Measured on CPU with a 2-layer GPT-2 (144 parameters). GPU overhead is ~3–8× lower.
+
+| Config | CPU overhead | GPU overhead |
+|--------|-------------|-------------|
+| Default (`hist/50`, `act/5`) | ~55% | ~4% |
+| + `activation_layer_filter=["attn","mlp"]` | ~38% | ~2% |
+| Minimal (`hist/50`, `act/50`, filter) | ~18% | ~1% |
+
+CPU measured on 2-layer mini-GPT (144 params), Apple M2. GPU measured on the same model with CUDA. Results will differ on larger models — histogram cost scales with parameter count, activation cost scales with layer count × sequence length.
+
+## UI
+
+Four views, one command:
+
+| View | What it shows |
+|------|---------------|
+| **Timeline** | Loss + grad norm, top-8 layers by grad variance |
+| **Layer Drill-down** | Kurtosis / grad norm / weight norm per layer; histogram scrubber |
+| **Diff View** | KL divergence of weight distributions between any two steps |
+| **Spike Inspector** | Per-spike window: loss+grad timeline and layer kurtosis/grad breakdown |
+
+The UI works immediately after `pip install` — a built-in fallback HTML with Plotly CDN is served when the React build is absent. For the full React build:
+
+```bash
+cd frontend && npm install && npm run build
+```
+
+## CLI
+
+```bash
+# Open UI for a completed or in-progress run
+trainscope ui --run ./trainscope_runs/run_20250516_143022 [--host 127.0.0.1] [--port 7007]
+
+# Generate replay_config.json (does NOT resume training automatically)
+trainscope replay --checkpoint ./checkpoints/step_4400.pt --skip-batches 4521,4522,4523 [--resume]
+```
+
+To actually skip batches, use `SkippingDataLoader` in your training script:
+
+```python
+from trainscope.replay import SkippingDataLoader
+import json
+
+with open("replay_config.json") as f:
+    cfg = json.load(f)
+
+loader = SkippingDataLoader(original_loader, skip_batches=cfg["skip_batches"])
+for batch in loader:
+    ...
+```
+
+## Configuration
+
+```python
+TrainScopeConfig(
+    run_dir="./trainscope_runs",                 # output root
+    spike_threshold=3.5,                     # z-score threshold (rolling window baseline)
+    full_resolution_window=500,              # last N steps at full resolution
+    decimation_factor=10,                    # older steps: keep every Nth
+    spike_window_before=50,                  # steps before spike to save (≤ full_resolution_window)
+    spike_window_after=10,                   # steps after spike to save
+    histogram_every_n_steps=50,             # weight histograms are expensive; sample them
+    activation_metrics_every_n_steps=5,     # kurtosis sampling; always captured at spike
+    activation_layer_filter=["attn", "mlp"],# None = all leaf layers
+    stop_on_spike=False,                     # raise StopTraining on detection
+    trace_every_n_steps=1,                   # subsample for very large models
+    rank=None,                               # DDP rank → adds _rank{N} suffix to run dir
+)
+```
+
+## Demo
+
+```bash
+python examples/gpt2_spike_demo.py
+```
+
+Trains a 2-layer mini-GPT, injects a ×50 loss spike at step 50, and shows trainscope detecting it. Run `trainscope ui` on the output directory to explore the event.
+
+## Storage layout
+
+```
+trainscope_runs/<run-name>/
+    meta.json                          model config + trainscope config
+    global.arrow                       step-level scalars (Arrow IPC)
+    layers/<param-name>.arrow          per-layer metrics
+    spikes/spike_step_<N>.arrow        global window around spike N
+    spikes/spike_step_<N>_layers/      per-layer data for that window
+    rng_states/step_<N>.pkl            RNG state for replay
+```
+
+Estimated storage: ~10 MB/step at full resolution. Rolling 500-step window → ~5 GB max for a 1B-param model. Spike windows are small.
+
+## Publishing
+
+CI runs on every push to `main` and every PR (`pytest` + `ruff`, Python 3.11 + 3.12, Vite build).
+
+To publish a release to PyPI:
+1. Set up [Trusted Publishing](https://docs.pypi.org/trusted-publishers/) on PyPI for this repo (environment name: `pypi`).
+2. Tag and push: `git tag v0.1.0 && git push origin v0.1.0`
+
+The publish workflow builds the React frontend, bundles it into the wheel, and uploads via OIDC — no API token needed.

trainscope-0.1.0/examples/gpt2_spike_demo.py

@@ -0,0 +1,161 @@
+"""
+Mini GPT-2 spike demo.
+
+Trains a 2-layer GPT-2-style model on random token sequences, artificially
+injects a loss spike at step SPIKE_STEP, and shows trainscope detecting it.
+
+Usage:
+    python examples/gpt2_spike_demo.py
+
+Then open the UI:
+    trainscope ui --run ./trainscope_runs/<run-name>
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from trainscope import TrainScope
+from trainscope.core.config import TrainScopeConfig
+
+VOCAB = 256
+SEQ_LEN = 32
+BATCH = 8
+D_MODEL = 128
+N_HEADS = 4
+N_LAYERS = 2
+N_STEPS = 150
+SPIKE_STEP = 50
+SPIKE_MULTIPLIER = 50.0
+LR = 1e-3
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.qkv = nn.Linear(D_MODEL, 3 * D_MODEL)
+        self.proj = nn.Linear(D_MODEL, D_MODEL)
+        self.n_heads = N_HEADS
+        self.register_buffer(
+            "mask",
+            torch.tril(torch.ones(SEQ_LEN, SEQ_LEN)).view(1, 1, SEQ_LEN, SEQ_LEN),
+        )
+
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(D_MODEL, dim=2)
+        head = C // self.n_heads
+        q = q.view(B, T, self.n_heads, head).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, head).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, head).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) / (head ** 0.5)
+        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
+        att = F.softmax(att, dim=-1)
+        y = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
+        return self.proj(y)
+
+
+class Block(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(D_MODEL)
+        self.attn = CausalSelfAttention()
+        self.ln2 = nn.LayerNorm(D_MODEL)
+        self.mlp = nn.Sequential(
+            nn.Linear(D_MODEL, 4 * D_MODEL),
+            nn.GELU(),
+            nn.Linear(4 * D_MODEL, D_MODEL),
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class MiniGPT(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tok_emb = nn.Embedding(VOCAB, D_MODEL)
+        self.pos_emb = nn.Embedding(SEQ_LEN, D_MODEL)
+        self.blocks = nn.Sequential(*[Block() for _ in range(N_LAYERS)])
+        self.ln_f = nn.LayerNorm(D_MODEL)
+        self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
+
+    def forward(self, x):
+        B, T = x.shape
+        pos = torch.arange(T, device=x.device).unsqueeze(0)
+        h = self.tok_emb(x) + self.pos_emb(pos)
+        h = self.blocks(h)
+        h = self.ln_f(h)
+        return self.head(h)
+
+
+def main():
+    torch.manual_seed(42)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = MiniGPT().to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
+
+    # Overhead profiles (measured on CPU, 2-layer GPT-2, 144 params):
+    #   default  (hist/50, act/5):             ~121% overhead
+    #   recommended (+ activation_layer_filter): ~89%
+    #   minimal  (hist/50, act/50, filter):      ~52%
+    # GPU overhead is significantly lower (~3-8%) due to parallelism.
+    config = TrainScopeConfig(
+        run_dir="./trainscope_runs",
+        spike_threshold=3.5,
+        stop_on_spike=False,
+        full_resolution_window=500,
+        histogram_every_n_steps=50,
+        activation_metrics_every_n_steps=5,
+        # Uncomment to enable recommended profile:
+        # activation_layer_filter=["attn", "mlp"],
+    )
+
+    scope = TrainScope(model, optimizer, config=config).attach()
+    scope.writer.write_meta(
+        "MiniGPT",
+        {"vocab": VOCAB, "d_model": D_MODEL, "n_heads": N_HEADS, "n_layers": N_LAYERS},
+    )
+
+    print(f"Device: {device}")
+    print(f"Run:    ./trainscope_runs/{config.run_name}")
+    print(f"Steps:  {N_STEPS}  (spike injected at step {SPIKE_STEP})\n")
+
+    for step in range(N_STEPS):
+        x = torch.randint(0, VOCAB, (BATCH, SEQ_LEN), device=device)
+        targets = torch.randint(0, VOCAB, (BATCH, SEQ_LEN), device=device)
+
+        optimizer.zero_grad()
+        logits = model(x)
+        loss = F.cross_entropy(logits.view(-1, VOCAB), targets.view(-1))
+
+        if step == SPIKE_STEP:
+            loss = loss * SPIKE_MULTIPLIER
+            print(f"  [inject] step={step}  loss={loss.item():.4f}  (×{SPIKE_MULTIPLIER})")
+
+        loss.backward()
+        optimizer.step()
+
+        spike = scope.step(loss.item(), batch_index=step)
+
+        if spike:
+            print(
+                f"\n  *** SPIKE DETECTED ***  step={spike['step']}  "
+                f"loss={spike['loss']:.4f}  z={spike['z_score']:.2f}\n"
+            )
+        elif step % 20 == 0:
+            print(f"  step={step:3d}  loss={loss.item():.4f}")
+
+    scope.writer.flush()
+    scope.writer.close()
+    scope.detach()
+
+    print(f"\nDone. Open the UI with:")
+    print(f"  trainscope ui --run ./trainscope_runs/{config.run_name}")
+
+
+if __name__ == "__main__":
+    main()