Konsti/olmo sft example

.gitignore

@@ -44,3 +44,11 @@ BLOCKED.md
 # Benchmark run artifacts (logs + regenerable scratch output dirs)
 *.log
 examples/benchmark/*/out/
+
+# OLMo SFT example: regenerable run artifacts (recreate by running the scripts).
+# The committed plots (*.png) are kept on dev as reference figures; everything
+# below is selection/telemetry output and stays out of version control.
+examples/olmo_sft_analysis/samples.json
+examples/olmo_sft_analysis/results.parquet
+examples/olmo_sft_analysis/resources.json
+examples/olmo_sft_analysis/mem_profile.md

CHANGELOG.md

@@ -4,6 +4,28 @@ All notable changes to vatis are documented here. The format follows
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and the project aims
 to follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.0] — 2026-06-17
+
+### Added
+- `Analyzer` / `analyze()` gained **`cross_grad_storage`** (`"auto"` | `"gpu"` |
+  `"cpu"`, default `"auto"`): cross-pair gradients can be offloaded to host RAM
+  so cross-pair observables fit on large models whose two full gradients won't
+  co-reside on the GPU; small models stay on-device (no transfer, no numeric
+  change). `"auto"` offloads only when an on-GPU cache wouldn't leave room for a
+  self-pair backward.
+- `load_hf_model(..., attn_implementation=...)` — passthrough to
+  `from_pretrained` (e.g. `"sdpa"`, `"flash_attention_2"`, `"eager"`).
+- `examples/olmo_sft_analysis/` — a production-scale deployment example: the LNP
+  decomposition on `Olmo-3-7B-Think-SFT` over real SFT chat data with
+  completion-only loss masking, code-vs-English cross-pair analysis, and
+  wall-time / peak-memory / GPU-engine-activity tracking.
+
+### Fixed
+- The analyzer no longer caches a batch's full parameter-space gradient unless a
+  cross pair consumes it. Caching it unconditionally kept a P-sized vector
+  (~29 GB for a 7B model) resident across batches and could OOM large models
+  even when no cross pair was requested.
+
 ## [0.1.0] — initial public release
 
 First public release.
@@ -23,4 +45,5 @@ First public release.
 - Specification in `docs/SPEC.md`; LNP derivation in
   [arXiv:2605.31244](https://arxiv.org/abs/2605.31244).
 
+[0.2.0]: https://github.com/KonstiNik/vatis/releases/tag/v0.2.0
 [0.1.0]: https://github.com/KonstiNik/vatis/releases/tag/v0.1.0

README.md

@@ -1,9 +1,8 @@
 # vatis
 
-Compute **LNP observables** — `chi_loss`, `chi_net`, `chi_pos` — on **pretrained
-HF checkpoints** (Pythia, OLMo, GPT-NeoX), scaled across GPUs with DDP. Sister
-package to `perspic`, which computes the same quantities *during* Lightning
-training; vatis is for when you can only probe published weights.
+Compute  **Spectral Position** — and LNP observables `chi_loss`, `chi_net`, `chi_pos` — on **pretrained HF checkpoints** (Pythia, OLMo, GPT-NeoX), scaled across GPUs with DDP. Sister
+package to [perspic](https://github.com/zincware/perspic), which computes the same quantities *during* Lightning
+training; vatis is for when you want to probe existing checkpoints.
 
 The LNP decomposition factors the linearized loss change into three terms —
 **L**oss, **N**etwork, **P**osition:
@@ -57,14 +56,11 @@ uv pip install -e /path/to/perspic
 pytest tests/cross_validation -m cross_validation -q
 ```
 
-GPU-only tests are marked `@pytest.mark.gpu` and auto-skip when no CUDA device
-is present, so the unit tier above runs anywhere (and is the tier CI runs).
+The unit tier above is CPU-only, runs anywhere, and is the tier CI runs.
 
 To run the example/benchmark scripts (they need matplotlib + tqdm, which the
 package itself does not): `uv pip install ".[examples]"`.
 
-Optional logging backends: `uv pip install ".[wandb]"` or `".[tensorboard]"`.
-
 ## Quick start
 
 ```python
@@ -75,17 +71,28 @@ analyze(
     revisions=["step1000", "step8000", "step143000"],
     eval_batches={"prose": prose_batch, "code": code_batch},
     cross_pairs=[("prose", "code")],   # adds δL(prose,code) + chi_pos, no extra backward
+    cross_grad_storage="auto",         # offload cross-pair grads to host RAM on large models
     sink="results.parquet",
 )
 ```
 
+`cross_grad_storage="auto"` (the default) lets cross-pair analysis scale to
+models whose two full gradients won't co-reside on the GPU: it offloads the
+cached gradients to host RAM and runs the dot on the CPU, while small models
+stay fully on-device.
+
 Runnable end-to-end (downloads pythia-14m, ~20 s warm):
 
 ```bash
 python examples/pythia_sweep.py        # compute  → results.parquet + plots
 python examples/analyze_results.py     # analyze  → derived plots (zero vatis imports)
 ```
 
+For a **production-scale** demonstration — the full LNP decomposition on a real
+7B SFT checkpoint (`Olmo-3-7B-Think-SFT`) over real chat data, with
+code-vs-English cross-pair analysis and wall-time / memory / GPU-activity
+tracking — see [`examples/olmo_sft_analysis/`](examples/olmo_sft_analysis/README.md).
+
 Models download to the HuggingFace cache (`~/.cache/huggingface` by default).
 To relocate it, set `HF_HOME` in your shell or copy `.env.example` to `.env`
 and edit it — the example scripts read it via `run_config.py`. SLURM benchmark

examples/_helpers.py

@@ -5,6 +5,154 @@
 import torch
 from transformers import AutoTokenizer
 
+# A "chat sample" is a list of OpenAI-style message dicts, e.g.
+# ``[{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]``.
+ChatSample = list[dict[str, str]]
+
+
+def _common_prefix_len(a: list[int], b: list[int]) -> int:
+    """Length of the longest shared prefix of two token-id lists."""
+    n = 0
+    for x, y in zip(a, b, strict=False):
+        if x != y:
+            break
+        n += 1
+    return n
+
+
+def _render_full_and_prompt_len(
+    messages: ChatSample, tokenizer: AutoTokenizer
+) -> tuple[list[int], int]:
+    """Render the conversation and locate the prompt/completion boundary.
+
+    transformers 5.x returns a ``BatchEncoding`` from
+    ``apply_chat_template(tokenize=True)``, not a flat ``list[int]``. To stay
+    robust across versions we render to *text* first (``tokenize=False``) and
+    tokenize the string with ``add_special_tokens=False`` (the template already
+    emits the special / role tokens). This reproduces the tokenized template
+    exactly while giving plain id lists to diff.
+
+    Returns ``(full_ids, prompt_len)`` where ``prompt_len`` is the longest
+    common prefix of the prompt-only (``messages[:-1]``, ``add_generation_prompt
+    =True``) and full renderings — the token index at which the assistant's
+    completion content begins. The common-prefix rule is robust to templates
+    that re-emit the generation header slightly differently in the two passes.
+    """
+    full_text = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
+    prompt_text = tokenizer.apply_chat_template(
+        messages[:-1], add_generation_prompt=True, tokenize=False
+    )
+    full_ids = tokenizer(full_text, add_special_tokens=False)["input_ids"]
+    prompt_ids = tokenizer(prompt_text, add_special_tokens=False)["input_ids"]
+    return full_ids, _common_prefix_len(prompt_ids, full_ids)
+
+
+def sft_prompt_completion_lengths(
+    messages: ChatSample,
+    tokenizer: AutoTokenizer,
+) -> tuple[int, int]:
+    """Return ``(prompt_len, total_len)`` in tokens for an SFT chat sample.
+
+    ``prompt_len`` is the number of leading tokens that belong to the prompt
+    (everything up to and including the final assistant header, i.e. the point
+    where the assistant's *content* begins); ``total_len`` is the length of the
+    full rendered conversation. Used by ``prefetch.py`` to keep only samples
+    that will still have real completion tokens after truncation to ``seq_len``.
+
+    The boundary is computed as the longest common prefix between the
+    prompt-only and full renderings (see :func:`_render_full_and_prompt_len`).
+    """
+    full_ids, prompt_len = _render_full_and_prompt_len(messages, tokenizer)
+    return prompt_len, len(full_ids)
+
+
+def tokenize_chat_sft_sample(
+    messages: ChatSample,
+    tokenizer: AutoTokenizer,
+    *,
+    seq_len: int,
+) -> dict[str, torch.Tensor]:
+    """Render one chat sample into a fixed-length, completion-only LM example.
+
+    This is the *realistic* SFT objective: cross-entropy is computed only on
+    the assistant's response tokens; the prompt (system + user turns, plus the
+    assistant header) is masked out with ``-100``. Contrast with
+    :func:`tokenize_into_lm_batch`, which trains on every token of a flat text.
+
+    The model's own chat template defines the prompt/completion boundary. We
+    render the conversation twice — full and prompt-only — and take the common
+    prefix as the boundary (see :func:`sft_prompt_completion_lengths`).
+
+    Labels follow the **pre-shifted** HF-causal-LM convention that vatis expects
+    (``causal_lm_loss`` compares ``logits[i]`` against ``labels[i]`` with no
+    further shift): ``labels[i] = ids[i + 1]`` when ``ids[i + 1]`` is a
+    completion token, else ``-100``; ``labels[-1] = -100``.
+
+    The example is then right-truncated (keeping the prompt prefix and as much
+    of the completion as fits) or right-padded to exactly ``seq_len``. Padding
+    uses ``pad_token_id`` (falling back to ``eos_token_id``) with
+    ``attention_mask = 0`` and ``labels = -100`` so pads count toward neither
+    the loss nor the valid-token denominator.
+
+    Returns a dict of ``(seq_len,)`` tensors: ``input_ids``, ``attention_mask``,
+    ``labels``. Stack several with :func:`stack_lm_batch` to form a batch.
+    """
+    if seq_len < 2:
+        raise ValueError(f"seq_len must be >= 2, got {seq_len}")
+
+    full_ids, prompt_len = _render_full_and_prompt_len(messages, tokenizer)
+
+    ids = torch.tensor(full_ids, dtype=torch.long)
+    total = ids.shape[0]
+
+    # Pre-shifted next-token targets: labels[i] = ids[i + 1].
+    labels = torch.full((total,), fill_value=-100, dtype=torch.long)
+    labels[:-1] = ids[1:]
+    # Mask every target that lands inside the prompt. Position i predicts
+    # ids[i + 1]; that target is a completion token iff i + 1 >= prompt_len,
+    # i.e. i >= prompt_len - 1. So positions [0, prompt_len - 2] are masked.
+    if prompt_len >= 1:
+        labels[: prompt_len - 1] = -100
+
+    attention_mask = torch.ones(total, dtype=torch.long)
+
+    # Truncate (keep the prefix) or pad on the right to exactly seq_len.
+    pad_id = tokenizer.pad_token_id
+    if pad_id is None:
+        pad_id = tokenizer.eos_token_id
+    if pad_id is None:
+        pad_id = 0
+
+    if total >= seq_len:
+        ids = ids[:seq_len].contiguous()
+        labels = labels[:seq_len].contiguous()
+        attention_mask = attention_mask[:seq_len].contiguous()
+    else:
+        pad = seq_len - total
+        ids = torch.cat([ids, torch.full((pad,), pad_id, dtype=torch.long)])
+        labels = torch.cat([labels, torch.full((pad,), -100, dtype=torch.long)])
+        attention_mask = torch.cat([attention_mask, torch.zeros(pad, dtype=torch.long)])
+
+    # Fail loudly rather than silently contribute a sequence with no supervised
+    # tokens: if the prompt is at least seq_len long, truncation drops the whole
+    # completion and every label is -100. The caller must pick a seq_len larger
+    # than the prompt (prefetch.py enforces this via --max-prompt-frac, but a
+    # smaller --seq-len at analysis time can re-open the gap).
+    if int((labels != -100).sum()) == 0:
+        raise ValueError(
+            f"no completion tokens survive truncation to seq_len={seq_len} "
+            f"(prompt_len={prompt_len}); increase seq_len or drop this sample."
+        )
+
+    return {"input_ids": ids, "attention_mask": attention_mask, "labels": labels}
+
+
+def stack_lm_batch(samples: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+    """Stack a list of equal-length single-example dicts into a ``(B, S)`` batch."""
+    if not samples:
+        raise ValueError("cannot stack an empty list of samples")
+    return {key: torch.stack([s[key] for s in samples], dim=0) for key in samples[0]}
+
 
 def tokenize_into_lm_batch(
     text: str,

examples/olmo_sft_analysis/README.md

@@ -0,0 +1,123 @@
+# Realistic SFT analysis: `Olmo-3-7B-Think-SFT`
+
+Load Ai2's `Olmo-3-7B-Think-SFT` and compute the LNP decomposition —
+`chi_loss`, `chi_net`, `delta_loss`, and the headline **`chi_pos`** (spectral
+position) — on real `Dolci-Think-SFT` data: 10 code + 10 English samples,
+tokenized with the model's chat template and **completion-only loss masking**
+(cross-entropy on the assistant's response only — the actual SFT objective).
+Wall time, peak GPU memory, and GPU engine activity are tracked throughout.
+
+It doubles as a production-scale template: point it at your own causal-LM
+checkpoint + chat data and the same config should run (see
+[Adapting to your model](#adapting-to-your-model)).
+
+## What it computes
+
+Two eval batches (`code`, `english`) as self-pairs, plus the `(english, code)`
+cross pair:
+
+| observable | meaning |
+|---|---|
+| `chi_pos` (self) | spectral position of the batch's SFT loss gradient, in `[0, 1]`: →1 bulk, →0 tail. |
+| `chi_pos` (cross) | parameter-space cosine between the two batches' SFT gradients, in `[-1, 1]`. |
+| `delta_loss` | `‖∇θ L‖²` (self) / `⟨∇θ L_A, ∇θ L_B⟩` (cross). |
+| `chi_loss`, `chi_net` | loss-curvature and Jacobian-norm factors (and normalized forms). |
+
+On-distribution data matters: `chi_net` is the Jacobian norm *evaluated at the
+input*, so evaluate on data the model actually trained on (here, its own SFT
+mixture) — off-distribution tokens make the value meaningless.
+
+## How to run
+
+The compute nodes here are offline (`HF_HUB_OFFLINE=1` in the repo `.env`), so
+it runs in two stages: an online prefetch, then offline compute.
+
+**Prerequisites:** `HF_HOME` set (shell or `.env`, see
+[`../../run_config.py`](../../run_config.py)); `uv sync --extra examples` for
+matplotlib; and `datasets` for the prefetch *only* — deliberately not a vatis
+dependency, so install it into the venv:
+
+```bash
+uv pip install --python .venv/bin/python datasets
+```
+
+**1. Prefetch — login node, online.** Selects 10 code + 10 English samples into
+`samples.json` and warms the model into the HF cache:
+
+```bash
+.venv/bin/python examples/olmo_sft_analysis/prefetch.py
+```
+
+**2. Analyze — GPU node, offline.** Runs `analyze()` with resource tracking →
+`results.parquet` + `resources.json`:
+
+```bash
+examples/benchmark/submit.sh examples/olmo_sft_analysis/run.sbatch
+```
+
+**3. Plots & tables — anywhere, CPU.** Reads the parquet + resources (zero vatis
+imports) → `spectral_position.png`, `lnp_components.png`, `gpu_activity.png`:
+
+```bash
+.venv/bin/python examples/olmo_sft_analysis/analyze_results.py
+```
+
+> Optional pre-flight: `mem_profile.sbatch` measures peak GPU memory for your
+> model/config before a long run — useful when adapting to a new checkpoint.
+
+## The analyzer config, and why
+
+`run.sbatch` calls `analyze()` with the settings below. They are what let a 7B
+model fit a single ~93 GiB GPU; the rationale is what you adapt for your own model:
+
+- **`chi_net_method="hutchinson"`** — one backward per probe. `per_sequence_cv`
+  (the small-batch default) keeps extra gradients resident and overflows the
+  card at 7B; hutchinson holds one gradient at a time (~73 GB). Trade-off:
+  higher variance per probe — raise `--n-hutchinson` to tighten `chi_net`.
+- **`micro_batch_size=1`** — one sequence per backward, the safe default at 7B.
+  Peak is essentially sequence-length-independent here (it is param-bound).
+- **`cross_grad_storage="auto"`** — the cross-pair dot needs *both* batches'
+  parameter-space gradients; at 7B they won't co-reside on the GPU (~104 GB), so
+  `auto` offloads the cached gradients to host RAM and runs the dot on the CPU
+  (a one-time ~29 GB copy, negligible vs the analyze phase). Small models stay
+  on-GPU automatically. Force with `--cross-grad-storage gpu|cpu`.
+- **`dtype="bf16"`** — matches OLMo's training precision; gradients still
+  accumulate in fp32 (needed for `delta_loss` correctness).
+- **`seq_len=1024`** — this is the binding limit at 7B. Attention activations
+  grow ~**O(S²)** (per-layer scores retained across the 32 layers), so peak
+  climbs steeply: ~78 GB at S=2048, ~98 GB at S=4096, OOM by S=8192. Treat
+  **~3k as the ceiling** on a 93 GiB card. The attention backend does *not*
+  help (sdpa, the transformers default, still materializes the scores for this
+  model); covering the full ~16k-token code traces needs **activation/gradient
+  checkpointing** (recompute instead of retain) — a future addition.
+
+## Resources (measured: 1× H100, ~93 GiB)
+
+From `resources.json` / `gpu_activity.png` for the shipped config (code +
+English self-pairs **plus** the cross pair, `seq_len=1024`):
+
+- **wall:** ~210 s (model load ~18 s + analyze ~193 s).
+- **peak GPU memory:** ~84 GB allocated / ~95 GB reserved — the cross pair adds
+  no GPU peak because its gradients are offloaded to host RAM.
+- **GPU engine activity (DCGM):** GR-engine ~82%, SM-active ~73%, DRAM
+  (memory-bandwidth) ~51%, tensor-core ~14%, ~430 W (peak ~580 W) — a
+  memory-bandwidth-heavy profile (autograd over the 100k-vocab head), not
+  compute/tensor-core-bound. (DCGM's `FB_USED` under-reports the process; the
+  `torch.cuda` peak is the authoritative memory figure.)
+
+Tracking lives in [`resource_tracker.py`](resource_tracker.py): wall time +
+per-phase peak `torch.cuda` memory + GPU engine activity sampled via
+`dcgmi dmon`, with an automatic `nvidia-smi` fallback.
+
+## Adapting to your model
+
+- Point `prefetch.py --model / --dataset` at your causal-LM checkpoint + chat
+  dataset (keep them on-distribution). `--n-per-batch`, `--seq-len`,
+  `--max-prompt-frac` tune sample selection.
+- **Larger model:** keep `cross_grad_storage=auto` (offload) and `hutchinson`.
+  **Smaller model:** `auto` keeps everything on-GPU, and `per_sequence_cv`
+  becomes affordable (lower-variance `chi_net`).
+- `--no-cross-pairs` skips the cross observables; `--cross-grad-storage gpu|cpu`
+  overrides the auto memory decision.
+- The example uses the standard cross-entropy SFT loss — vatis v1 supports the
+  closed-form CE path for `chi_loss`; custom losses are future work.