diff --git a/.gitignore b/.gitignore index 73e86af..9ba102f 100644 --- a/.gitignore +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ce13d5..0d4e61a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 944b35d..c6277f0 100644 --- a/README.md +++ b/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,10 +71,16 @@ 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 @@ -86,6 +88,11 @@ 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 diff --git a/examples/_helpers.py b/examples/_helpers.py index 1d50209..f67f79a 100644 --- a/examples/_helpers.py +++ b/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, diff --git a/examples/olmo_sft_analysis/README.md b/examples/olmo_sft_analysis/README.md new file mode 100644 index 0000000..ac47e04 --- /dev/null +++ b/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. diff --git a/examples/olmo_sft_analysis/analyze_results.py b/examples/olmo_sft_analysis/analyze_results.py new file mode 100644 index 0000000..c5f4c80 --- /dev/null +++ b/examples/olmo_sft_analysis/analyze_results.py @@ -0,0 +1,264 @@ +"""Stage 3 (analysis, CPU): plot + tabulate the OLMo SFT LNP results. + +Reads ``results.parquet`` (the canonical vatis output) and ``resources.json`` +(time / memory / GPU activity) and produces: + +- ``spectral_position.png`` — ``chi_pos`` for the code and English self-pairs + (the spectral position of each batch's SFT loss gradient, in [0, 1]) plus the + cross-pair ``chi_pos(english, code)`` (the parameter-space gradient cosine, + in [-1, 1]). +- ``lnp_components.png`` — the full decomposition: ``chi_loss_normalized``, + ``chi_net_normalized``, ``delta_loss``, ``chi_pos`` as small multiples. +- ``gpu_activity.png`` — the DCGM activity timeline (SM / tensor / DRAM active, + framebuffer memory) with the load vs analyze phases shaded. + +It also prints a markdown observable table and a resource summary. + +Like ``examples/analyze_results.py``, this script has **zero vatis imports** — +only ``pyarrow`` and ``matplotlib`` (and stdlib ``json``). The parquet schema +plus the resources JSON are the entire contract with the compute step. + +Run (CPU is fine):: + + .venv/bin/python examples/olmo_sft_analysis/analyze_results.py +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") # headless + +import matplotlib.pyplot as plt # noqa: E402 +import pyarrow.parquet as pq # noqa: E402 + +_HERE = Path(__file__).resolve() +PARQUET_PATH = _HERE.parent / "results.parquet" +RESOURCES_PATH = _HERE.parent / "resources.json" + +CODE = "code" +ENGLISH = "english" +CROSS = (ENGLISH, CODE) + +# (key, human label) for the fraction-valued DCGM activity channels. +_ACTIVITY_CHANNELS = [ + ("sm_active", "SM active"), + ("tensor_active", "tensor core"), + ("dram_active", "DRAM (mem BW)"), + ("gr_engine_active", "GR engine"), +] + + +def pair_label(a: str, b: str) -> str: + return a if a == b else f"{a}×{b}" + + +def load_observables(path: Path) -> dict[str, dict[str, float]]: + """Return ``{observable: {pair_label: value}}`` from the parquet.""" + rows = pq.read_table(path).to_pylist() + out: dict[str, dict[str, float]] = {} + for r in rows: + out.setdefault(r["observable"], {})[pair_label(r["batch_a"], r["batch_b"])] = r["value"] + return out + + +def plot_spectral_position(obs: dict[str, dict[str, float]], path: Path) -> None: + chi_pos = obs.get("chi_pos", {}) + labels = [CODE, ENGLISH, pair_label(*CROSS)] + values = [chi_pos.get(lbl) for lbl in labels] + present = [(lbl, v) for lbl, v in zip(labels, values, strict=True) if v is not None] + if not present: + return + fig, ax = plt.subplots(figsize=(6, 4.5)) + xs = range(len(present)) + colors = ["#1f77b4", "#1f77b4", "#d62728"] + ax.bar(list(xs), [v for _, v in present], color=colors[: len(present)]) + for x, (_, v) in zip(xs, present, strict=True): + ax.text(x, v, f"{v:.2e}", ha="center", va="bottom" if v >= 0 else "top", fontsize=9) + ax.set_xticks(list(xs)) + ax.set_xticklabels([lbl for lbl, _ in present]) + ax.axhline(0.0, color="black", linewidth=0.7, alpha=0.5) + ax.set_ylabel("chi_pos (spectral position)") + ax.set_title( + "Spectral position of the SFT loss gradient\n" + "self pairs ∈ [0,1]: bulk→1, tail→0 | cross = gradient cosine ∈ [-1,1]" + ) + ax.grid(True, axis="y", alpha=0.3) + fig.tight_layout() + fig.savefig(path, dpi=120) + plt.close(fig) + print(f"wrote {path}") + + +def plot_lnp_components(obs: dict[str, dict[str, float]], path: Path) -> None: + panels = ["chi_loss_normalized", "chi_net_normalized", "delta_loss", "chi_pos"] + logy = {"chi_net_normalized"} + labels = [CODE, ENGLISH, pair_label(*CROSS)] + fig, axes = plt.subplots(1, 4, figsize=(15, 4)) + for ax, name in zip(axes, panels, strict=True): + series = obs.get(name, {}) + present = [(lbl, series[lbl]) for lbl in labels if lbl in series] + if not present: + ax.set_visible(False) + continue + xs = range(len(present)) + ax.bar(list(xs), [v for _, v in present], color="#4c72b0") + ax.set_xticks(list(xs)) + ax.set_xticklabels([lbl for lbl, _ in present], rotation=20, ha="right") + ax.set_title(name) + ax.axhline(0.0, color="black", linewidth=0.7, alpha=0.5) + if name in logy and all(v > 0 for _, v in present): + ax.set_yscale("log") + ax.grid(True, axis="y", alpha=0.3) + fig.suptitle("LNP decomposition — Olmo-3-7B-Think-SFT on Dolci-Think-SFT samples") + fig.tight_layout() + fig.savefig(path, dpi=120) + plt.close(fig) + print(f"wrote {path}") + + +def plot_gpu_activity(resources: dict, path: Path) -> None: + timeline = resources.get("timeline", []) + summary = resources.get("summary", {}) + if not timeline: + print("no GPU timeline to plot (sampler collected no samples).") + return + ts = [s["t"] for s in timeline] + fig, ax = plt.subplots(figsize=(9, 4.5)) + plotted = False + for key, label in _ACTIVITY_CHANNELS: + ys = [s.get(key) for s in timeline] + if any(y is not None for y in ys): + ax.plot( + ts, + [(y * 100.0 if y is not None else None) for y in ys], + marker=".", + ms=3, + label=label, + ) + plotted = True + # Coarse util% fallback if no profiling channels (nvidia-smi sampler). + if not plotted: + ys = [s.get("gpu_util_pct") for s in timeline] + if any(y is not None for y in ys): + ax.plot(ts, ys, marker=".", ms=3, label="GPU util (coarse)") + plotted = True + ax.set_xlabel("time since tracker start (s)") + ax.set_ylabel("engine active (%)") + ax.set_ylim(0, 105) + + # Framebuffer memory on a secondary axis. + fb = [s.get("fb_used_mb") for s in timeline] + if any(v is not None for v in fb): + ax2 = ax.twinx() + ax2.plot( + ts, + [(v / 1024.0 if v is not None else None) for v in fb], + color="gray", + linestyle="--", + label="FB mem (GB)", + ) + ax2.set_ylabel("framebuffer memory (GB)") + ax2.legend(loc="upper right") + + # Shade the named phases. + phase_colors = {"load_model": "#ffeda0", "analyze": "#c7e9c0"} + for p in summary.get("phases", []): + t0, t1 = p.get("t_start"), p.get("t_end") + if t0 is None or t1 is None: + continue + ax.axvspan(t0, t1, color=phase_colors.get(p["name"], "#eeeeee"), alpha=0.5, zorder=0) + ax.text((t0 + t1) / 2, 102, p["name"], ha="center", va="top", fontsize=9) + + sampler = summary.get("sampler", "?") + ax.set_title(f"GPU activity during analysis (sampler: {sampler})") + ax.legend(loc="upper left") + ax.grid(True, alpha=0.3) + fig.tight_layout() + fig.savefig(path, dpi=120) + plt.close(fig) + print(f"wrote {path}") + + +def print_observable_table(obs: dict[str, dict[str, float]]) -> None: + cols = [CODE, ENGLISH, pair_label(*CROSS)] + order = [ + "chi_pos", + "delta_loss", + "chi_loss_normalized", + "chi_net_normalized", + "chi_loss", + "chi_net", + ] + names = [o for o in order if o in obs] + [o for o in obs if o not in order] + print("\n### LNP observables\n") + print("| observable | " + " | ".join(cols) + " |") + print("|" + "---|" * (len(cols) + 1)) + for name in names: + cells = [] + for c in cols: + v = obs[name].get(c) + cells.append(f"{v:+.4e}" if v is not None else "—") + print(f"| {name} | " + " | ".join(cells) + " |") + + +def print_resource_summary(resources: dict) -> None: + s = resources.get("summary", {}) + if not s: + return + print("\n### Resources\n") + print( + f"- device: `{s.get('device')}`" + (f" ({s.get('gpu_name')})" if s.get("gpu_name") else "") + ) + print( + f"- GPU sampler: `{s.get('sampler')}` " + f"({s.get('n_gpu_samples')} samples @ {s.get('sampler_interval_s')}s)" + ) + print(f"- wall total: {s.get('wall_s_total', 0):.1f} s") + print( + f"- torch peak: {s.get('torch_peak_alloc_mb', 0):.0f} MB allocated, " + f"{s.get('torch_peak_reserved_mb', 0):.0f} MB reserved" + ) + gpu = s.get("gpu", {}) + + def stat(key: str, scale: float = 1.0, unit: str = "") -> str: + st = gpu.get(key) + return ( + f"mean {st['mean'] * scale:.1f}{unit}, max {st['max'] * scale:.1f}{unit}" if st else "—" + ) + + print(f"- SM active: {stat('sm_active', 100, '%')}") + print(f"- tensor-core active: {stat('tensor_active', 100, '%')}") + print(f"- DRAM (mem-BW) active: {stat('dram_active', 100, '%')}") + print(f"- power: {stat('power_w', 1, ' W')}") + print("\n| phase | wall_s | torch peak MB | SM active mean |") + print("|---|---|---|---|") + for p in s.get("phases", []): + sm = p.get("gpu", {}).get("sm_active") + smtxt = f"{sm['mean'] * 100:.0f}%" if sm else "—" + wall = f"{p['wall_s']:.1f}" if p.get("wall_s") is not None else "—" + print(f"| {p['name']} | {wall} | {p.get('torch_peak_alloc_mb', 0):.0f} | {smtxt} |") + + +def main() -> None: + if not PARQUET_PATH.exists(): + raise SystemExit(f"{PARQUET_PATH} not found — run olmo_sft_analyze.py first.") + obs = load_observables(PARQUET_PATH) + + plot_spectral_position(obs, _HERE.parent / "spectral_position.png") + plot_lnp_components(obs, _HERE.parent / "lnp_components.png") + print_observable_table(obs) + + if RESOURCES_PATH.exists(): + resources = json.loads(RESOURCES_PATH.read_text()) + plot_gpu_activity(resources, _HERE.parent / "gpu_activity.png") + print_resource_summary(resources) + else: + print(f"\n(no {RESOURCES_PATH.name} — skipping resource plots/table)") + + +if __name__ == "__main__": + main() diff --git a/examples/olmo_sft_analysis/gpu_activity.png b/examples/olmo_sft_analysis/gpu_activity.png new file mode 100644 index 0000000..c51e549 Binary files /dev/null and b/examples/olmo_sft_analysis/gpu_activity.png differ diff --git a/examples/olmo_sft_analysis/lnp_components.png b/examples/olmo_sft_analysis/lnp_components.png new file mode 100644 index 0000000..cc7a4d2 Binary files /dev/null and b/examples/olmo_sft_analysis/lnp_components.png differ diff --git a/examples/olmo_sft_analysis/mem_profile.py b/examples/olmo_sft_analysis/mem_profile.py new file mode 100644 index 0000000..43d1753 --- /dev/null +++ b/examples/olmo_sft_analysis/mem_profile.py @@ -0,0 +1,205 @@ +"""GPU memory profiler for the OLMo SFT analysis (diagnostic). + +Runs the **real vatis code paths** (no replica) on tiny synthetic batches and +reports peak GPU memory for each step, across a small `(B, S)` grid, so we can +see exactly where the memory goes and how it scales with batch/sequence: + +- model weights (resident after load), +- the analyzer's ``delta_loss`` gradient accumulation (forward + backward + + the fp32 flat gradient vector), +- ``chi_net`` via ``hutchinson`` and ``per_sequence_cv``. + +It also prints the actual gradient-tuple dtype/size — the number that decides +whether the param-bound floor is ~58 GB (bf16 grads) or ~73 GB (fp32 grads). + +Synthetic random tokens on purpose: this measures memory only, not observable +values (see CLAUDE.md). Run on a GPU node (offline, model from cache): + + .venv/bin/python examples/olmo_sft_analysis/mem_profile.py + # or via SLURM: + examples/benchmark/submit.sh examples/olmo_sft_analysis/mem_profile.sbatch +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve() +_REPO_ROOT = _HERE.parents[2] +sys.path.insert(0, str(_REPO_ROOT)) + +from run_config import configure_hf_cache # noqa: E402 + +configure_hf_cache() + +import torch # noqa: E402 + +from vatis.analyzer import _add_grads_into_flat, _zero_param_vector # noqa: E402 +from vatis.core.chi_net.hutchinson import HutchinsonEstimator # noqa: E402 +from vatis.core.chi_net.per_seq_cv import PerSequenceControlVariateEstimator # noqa: E402 +from vatis.data.collate import iter_micro_batches # noqa: E402 +from vatis.models.hf import load_hf_model # noqa: E402 + +_GB = 1e9 +SAMPLES_PATH = _HERE.parent / "samples.json" + + +def gpu_now_gb() -> float: + return torch.cuda.memory_allocated() / _GB + + +def reset_peak() -> None: + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + +def peak_gb() -> float: + torch.cuda.synchronize() + return torch.cuda.max_memory_allocated() / _GB + + +def make_batch(vocab: int, b: int, s: int, device: torch.device) -> dict[str, torch.Tensor]: + g = torch.Generator().manual_seed(0) + ids = torch.randint(0, vocab, (b, s), generator=g).to(device) + labels = torch.full_like(ids, -100) + labels[:, :-1] = ids[:, 1:] + return {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": labels} + + +def run_delta_loss(bundle, batch, micro_bs: int) -> tuple[float, str, float]: + """Replicate the analyzer's delta_loss gradient accumulation. Returns + (peak_gb, grad_dtype, grad_tuple_gb).""" + device = next(bundle.model.parameters()).device + params = bundle.params + reset_peak() + flat = _zero_param_vector(params, device) + grad_dtype = "?" + grad_bytes = 0 + n_micros = max(1, (batch["input_ids"].shape[0] + micro_bs - 1) // micro_bs) + for _start, _stop, micro in iter_micro_batches(batch, micro_bs): + logits = bundle.forward_fn(bundle.model, micro) + loss = bundle.loss_fn(logits, micro) + grads = torch.autograd.grad(loss / n_micros, params, retain_graph=False, allow_unused=True) + if grad_dtype == "?": + present = [g for g in grads if g is not None] + grad_dtype = str(present[0].dtype) if present else "none" + grad_bytes = sum(g.numel() * g.element_size() for g in present) + _add_grads_into_flat(flat, grads, params) + del logits, loss, grads + pk = peak_gb() + del flat + return pk, grad_dtype, grad_bytes / _GB + + +def run_chi_net(bundle, batch, micro_bs: int, method: str, n_h: int) -> float: + device = next(bundle.model.parameters()).device + if method == "hutchinson": + est = HutchinsonEstimator(n_hutchinson=n_h) + else: + est = PerSequenceControlVariateEstimator(n_hutchinson=n_h) + gen = torch.Generator(device=device).manual_seed(0) + reset_peak() + est.compute( + bundle.model, + batch, + bundle.forward_fn, + loss_fn=bundle.loss_fn, + valid_mask_fn=bundle.valid_mask_fn, + params=bundle.params, + micro_batch_size=micro_bs, + generator=gen, + ) + return peak_gb() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default=None, help="default: model from samples.json") + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[2, 10]) + parser.add_argument("--seq-lens", type=int, nargs="+", default=[64, 256, 512]) + parser.add_argument("--micro-batch-size", type=int, default=1) + parser.add_argument("--n-hutchinson", type=int, default=8) + parser.add_argument("--methods", nargs="+", default=["hutchinson", "per_sequence_cv"]) + parser.add_argument( + "--attn-implementation", + default=None, + help="attention-backend passthrough (None = model default). NOTE: for " + "OLMo3, sdpa still materializes scores (O(S^2)), same as eager.", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device — run on a GPU node.") + device = torch.device("cuda") + model_name = args.model or json.loads(SAMPLES_PATH.read_text())["meta"]["model"] + + total_gb = torch.cuda.get_device_properties(0).total_memory / _GB + print(f"# mem_profile — {model_name}") + print(f"device: {torch.cuda.get_device_name(0)} total={total_gb:.1f} GB") + print( + f"micro_batch_size={args.micro_batch_size}, n_hutchinson={args.n_hutchinson}, " + f"attn_implementation={args.attn_implementation}\n" + ) + + reset_peak() + bundle = load_hf_model( + model_name, dtype="bf16", device=device, attn_implementation=args.attn_implementation + ) + p = sum(x.numel() for x in bundle.params) + weights_gb = gpu_now_gb() + print(f"n_params={p:,}") + print(f"weights resident after load: {weights_gb:.1f} GB (bf16 theory {p * 2 / _GB:.1f})") + print(f"flat_grad fp32 (theory): {p * 4 / _GB:.1f} GB\n") + + vocab = bundle.model.config.vocab_size + header = "| B | S | delta_loss peak | grad dtype | grad tuple | hutch peak | pseqcv peak |" + print(header) + print("|---|---|---|---|---|---|---|") + rows = [] + for b in args.batch_sizes: + for s in args.seq_lens: + batch = make_batch(vocab, b, s, device) + try: + d_pk, gdt, g_gb = run_delta_loss(bundle, batch, args.micro_batch_size) + d_str = f"{d_pk:.1f} GB" + except torch.OutOfMemoryError: + d_pk, gdt, g_gb, d_str = float("nan"), "OOM", float("nan"), "OOM" + torch.cuda.empty_cache() + peaks = {} + for m in args.methods: + try: + peaks[m] = ( + f"{run_chi_net(bundle, batch, args.micro_batch_size, m, args.n_hutchinson):.1f} GB" + ) + except torch.OutOfMemoryError: + peaks[m] = "OOM" + torch.cuda.empty_cache() + row = ( + f"| {b} | {s} | {d_str} | {gdt} | {g_gb:.1f} GB | " + f"{peaks.get('hutchinson', '-')} | {peaks.get('per_sequence_cv', '-')} |" + ) + print(row, flush=True) + rows.append(row) + + out = _HERE.parent / "mem_profile.md" + with open(out, "w") as f: + f.write(f"# mem_profile — {model_name}\n\n") + f.write(f"- device: {torch.cuda.get_device_name(0)} ({total_gb:.1f} GB)\n") + f.write( + f"- n_params: {p:,}; weights {weights_gb:.1f} GB; flat_grad fp32 {p * 4 / _GB:.1f} GB\n" + ) + f.write( + f"- micro_batch_size={args.micro_batch_size}, n_hutchinson={args.n_hutchinson}, " + f"attn_implementation={args.attn_implementation}\n\n" + ) + f.write(header + "\n|---|---|---|---|---|---|---|\n") + f.write("\n".join(rows) + "\n") + print(f"\nwrote {out}") + + +if __name__ == "__main__": + main() diff --git a/examples/olmo_sft_analysis/mem_profile.sbatch b/examples/olmo_sft_analysis/mem_profile.sbatch new file mode 100644 index 0000000..37c47ba --- /dev/null +++ b/examples/olmo_sft_analysis/mem_profile.sbatch @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +#SBATCH --job-name=olmo-mem-profile +#SBATCH --nodes=1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=6 +#SBATCH --mem=64G +#SBATCH --time=00:20:00 +#SBATCH --output=olmo_mem_profile_%j.log + +set -euo pipefail + +# Submit from the repo root via the wrapper that pulls SBATCH_* from .env: +# examples/benchmark/submit.sh examples/olmo_sft_analysis/mem_profile.sbatch +REPO="${SLURM_SUBMIT_DIR:-$PWD}" +cd "$REPO" + +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 +# Measure the allocator's true peak, not fragmentation artefacts. +export PYTORCH_ALLOC_CONF=expandable_segments:True + +echo "node: $(hostname)" +nvidia-smi -L + +"$REPO"/.venv/bin/python examples/olmo_sft_analysis/mem_profile.py "$@" diff --git a/examples/olmo_sft_analysis/olmo_sft_analyze.py b/examples/olmo_sft_analysis/olmo_sft_analyze.py new file mode 100644 index 0000000..28a8314 --- /dev/null +++ b/examples/olmo_sft_analysis/olmo_sft_analyze.py @@ -0,0 +1,221 @@ +"""Stage 2 (OFFLINE, GPU node): LNP observables on Olmo-3-7B-Think-SFT. + +A realistic single-checkpoint SFT analysis. Loads ``Olmo-3-7B-Think-SFT`` from +the (pre-populated, offline) HF cache, builds two real eval batches from the +samples ``prefetch.py`` saved — 10 *code* SFT examples and 10 *English* SFT +examples, each tokenized with the model's chat template and **completion-only +loss masking** (the actual SFT objective: cross-entropy on assistant tokens +only) — and computes the LNP decomposition with vatis: + +- ``chi_pos`` — the **spectral position** of the SFT loss gradient, in [0, 1]. + Near 1 = the gradient exploits the bulk (well-resolved, dominant eigenmodes); + near 0 = it lives in the tail (weak, fine-grained directions). +- ``chi_loss`` / ``chi_net`` (and their normalized forms), ``delta_loss``. +- the **cross pair** ``(english, code)``: ``delta_loss`` and ``chi_pos`` between + the two batches' gradients — do an SFT step on English and one on code help or + fight each other at this checkpoint? (positive ``delta_loss`` = transfer, + negative = interference.) + +Resource usage (wall time, peak GPU memory, GPU engine activity via DCGM) is +tracked throughout and written to ``resources.json``; see ``resource_tracker``. + +Run offline on a GPU node (the cache must already hold the model — run +``prefetch.py`` on a login node first). Via SLURM:: + + examples/benchmark/submit.sh examples/olmo_sft_analysis/run.sbatch + +or interactively on a GPU:: + + .venv/bin/python examples/olmo_sft_analysis/olmo_sft_analyze.py + +Outputs (next to this script): + +- ``results.parquet`` — canonical long-format LNP table +- ``resources.json`` — time / memory / GPU-activity summary + timeline +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve() +_REPO_ROOT = _HERE.parents[2] +_EXAMPLES_DIR = _HERE.parents[1] +sys.path.insert(0, str(_REPO_ROOT)) # run_config +sys.path.insert(0, str(_EXAMPLES_DIR)) # _helpers + +from run_config import configure_hf_cache # noqa: E402 + +# On the compute node this pins HF_HOME and (via .env) HF_HUB_OFFLINE=1 so HF +# only ever touches the pre-populated cache. Call before importing torch/transformers. +configure_hf_cache() + +import time # noqa: E402 + +import pyarrow.parquet as pq # noqa: E402 +import torch # noqa: E402 +from _helpers import stack_lm_batch, tokenize_chat_sft_sample # noqa: E402 +from resource_tracker import ResourceTracker # noqa: E402 +from transformers import AutoTokenizer # noqa: E402 + +from vatis import analyze # noqa: E402 +from vatis.models.hf import load_hf_model # noqa: E402 + +SAMPLES_PATH = _HERE.parent / "samples.json" +PARQUET_PATH = _HERE.parent / "results.parquet" +RESOURCES_PATH = _HERE.parent / "resources.json" + +CODE_NAME = "code" +ENGLISH_NAME = "english" +CROSS_PAIRS = [(ENGLISH_NAME, CODE_NAME)] + + +def build_batch(samples: list[dict], tokenizer, seq_len: int) -> dict[str, torch.Tensor]: + """Tokenize a list of chat samples into one stacked completion-only LM batch.""" + return stack_lm_batch( + [tokenize_chat_sft_sample(s["messages"], tokenizer, seq_len=seq_len) for s in samples] + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--samples", default=str(SAMPLES_PATH)) + parser.add_argument( + "--seq-len", type=int, default=None, help="override the seq_len recorded in samples.json" + ) + parser.add_argument("--n-hutchinson", type=int, default=32) + parser.add_argument("--micro-batch-size", type=int, default=1) + parser.add_argument( + "--chi-net-method", + default="hutchinson", + help="chi_net estimator. Default 'hutchinson' (one backward per probe) is " + "what fits a 7B model; 'per_sequence_cv' is lower-variance but OOMs at 7B " + "(use it only for smaller models). See the README config section.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--sampler-interval-s", type=float, default=0.25) + parser.add_argument( + "--no-cross-pairs", + action="store_true", + help="skip the (english, code) cross pair. With --cross-grad-storage " + "auto the cross pair fits even on a 7B model (the cached gradients " + "offload to host RAM and the dot runs on CPU), so this is opt-out.", + ) + parser.add_argument( + "--cross-grad-storage", + default="auto", + choices=["auto", "gpu", "cpu"], + help="where cached cross-pair gradients live. 'auto' (default) offloads " + "to host RAM for large models so the cross dot fits, keeps small models " + "on-device; 'gpu'/'cpu' force it. See vatis.Analyzer.", + ) + parser.add_argument( + "--attn-implementation", + default=None, + help="optional attention-backend passthrough to from_pretrained " + "(e.g. 'sdpa', 'flash_attention_2', 'eager'). Default None = the model's " + "own default (sdpa for OLMo3). NOTE: for OLMo3 sdpa does NOT reduce the " + "O(S^2) attention memory, so it does not raise the seq_len ceiling here.", + ) + args = parser.parse_args() + + samples_path = Path(args.samples) + if not samples_path.exists(): + raise SystemExit( + f"{samples_path} not found. Run prefetch.py on a login node first " + "(see this example's README)." + ) + payload = json.loads(samples_path.read_text()) + meta = payload["meta"] + model_name = meta["model"] + seq_len = args.seq_len or meta["seq_len_target"] + + if torch.cuda.is_available(): + device = torch.device("cuda") + dtype = "bf16" # matches OLMo training precision; grads accumulate in fp32 + else: + device = torch.device("cpu") + dtype = "fp32" + print("WARNING: no CUDA device — a 7B model will not realistically run on CPU.") + + cross_pairs = [] if args.no_cross_pairs else CROSS_PAIRS + + print(f"olmo SFT analysis — model={model_name}") + print(f" device={device}, dtype={dtype}, seq_len={seq_len}, attn={args.attn_implementation}") + print( + f" chi_net_method={args.chi_net_method}, n_hutchinson={args.n_hutchinson}, " + f"micro_batch_size={args.micro_batch_size}, cross_pairs={cross_pairs}, " + f"cross_grad_storage={args.cross_grad_storage}" + ) + + tokenizer = AutoTokenizer.from_pretrained(model_name) + eval_batches = { + CODE_NAME: build_batch(payload["batches"][CODE_NAME], tokenizer, seq_len), + ENGLISH_NAME: build_batch(payload["batches"][ENGLISH_NAME], tokenizer, seq_len), + } + for name, batch in eval_batches.items(): + n_tokens = int(batch["input_ids"].numel()) + n_valid = int((batch["labels"] != -100).sum().item()) + b, s = tuple(batch["input_ids"].shape) + print( + f" {name}: B={b} S={s}, {n_valid} valid completion tokens " + f"({n_valid / n_tokens:.1%} of {n_tokens})" + ) + + if PARQUET_PATH.exists(): + PARQUET_PATH.unlink() # fresh write; ParquetSink truncates on re-open anyway + + tracker = ResourceTracker(device=device, interval_s=args.sampler_interval_s) + with tracker: + with tracker.measure("load_model"): + t0 = time.perf_counter() + bundle = load_hf_model( + model_name, + dtype=dtype, + device=device, + attn_implementation=args.attn_implementation, + ) + n_params = sum(p.numel() for p in bundle.params) + print(f" loaded {n_params:,} params in {time.perf_counter() - t0:.1f}s") + + with tracker.measure("analyze"): + analyze( + model=bundle, + revisions=["sft"], # single checkpoint; label only + eval_batches=eval_batches, + cross_pairs=cross_pairs, + cross_grad_storage=args.cross_grad_storage, + chi_net_method=args.chi_net_method, + n_hutchinson=args.n_hutchinson, + micro_batch_size=args.micro_batch_size, + seed=args.seed, + device=device, + dtype=dtype, + sink=str(PARQUET_PATH), + ) + + tracker.save(str(RESOURCES_PATH)) + print(f"\nwrote {PARQUET_PATH}") + print(f"wrote {RESOURCES_PATH}") + + # ---- compact console summary of the spectral-position story ---- + table = pq.read_table(PARQUET_PATH).to_pylist() + print("\n=== LNP observables ===") + wanted = ("chi_pos", "delta_loss", "chi_loss_normalized", "chi_net_normalized") + for obs in wanted: + for row in table: + if row["observable"] != obs: + continue + a, b = row["batch_a"], row["batch_b"] + pair = a if a == b else f"{a}×{b}" + print(f" {obs:<20} {pair:<16} {row['value']:+.4e}") + + print("\n=== resources ===") + print(tracker.format_summary()) + + +if __name__ == "__main__": + main() diff --git a/examples/olmo_sft_analysis/prefetch.py b/examples/olmo_sft_analysis/prefetch.py new file mode 100644 index 0000000..85a42a5 --- /dev/null +++ b/examples/olmo_sft_analysis/prefetch.py @@ -0,0 +1,211 @@ +"""Stage 1 (ONLINE, login node): prefetch the model + select the eval samples. + +The compute nodes on this cluster have **no internet** (``HF_HUB_OFFLINE=1`` in +the repo ``.env``), so everything the offline analysis needs must be pulled into +the HF cache and written to a local file first. This script does both: + +1. **Selects the eval data** — streams ``allenai/Dolci-Think-SFT`` (the dataset + ``Olmo-3-7B-Think-SFT`` was trained on, so the observables are evaluated + on-distribution) and picks the first ``--n-per-batch`` *code* samples and + ``--n-per-batch`` *English* samples whose prompt is short enough to leave + real completion tokens after truncation to ``--seq-len``. The raw chat + ``messages`` are written to ``samples.json`` so the offline step needs no + dataset dependency at all. + +2. **Warms the model into the HF cache** — downloads the tokenizer and (unless + ``--skip-model-download``) the model weights so ``olmo_sft_analyze.py`` can + load them offline. + +This is the only part of the example that needs the ``datasets`` package and a +network connection. ``datasets`` is intentionally **not** a repo dependency +(vatis is dataset-agnostic); install it into the venv just for this step:: + + uv pip install --python .venv/bin/python datasets + +Run on a login node (which has internet):: + + .venv/bin/python examples/olmo_sft_analysis/prefetch.py + +Set ``HF_HOME`` (shell or repo-root ``.env``) to relocate the cache. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve() +_REPO_ROOT = _HERE.parents[2] +_EXAMPLES_DIR = _HERE.parents[1] +sys.path.insert(0, str(_REPO_ROOT)) # run_config +sys.path.insert(0, str(_EXAMPLES_DIR)) # _helpers + +# Force ONLINE before run_config loads .env (which pins offline=1 for compute +# nodes). configure_hf_cache() uses os.environ.setdefault, so pre-setting these +# to "0" wins while HF_HOME is still taken from .env. +os.environ["HF_HUB_OFFLINE"] = "0" +os.environ["TRANSFORMERS_OFFLINE"] = "0" + +from run_config import configure_hf_cache # noqa: E402 + +configure_hf_cache() + +from _helpers import sft_prompt_completion_lengths # noqa: E402 + +# ---------------------------------------------------------------- config + +MODEL = "allenai/Olmo-3-7B-Think-SFT" +DATASET = "allenai/Dolci-Think-SFT" +SPLIT = "train" +OUT_PATH = _HERE.parent / "samples.json" + +# Code samples: source name names a programming/algorithms split. +CODE_SOURCE_RE = re.compile(r"(Code|Python|Algorithms|Nemotron)", re.IGNORECASE) +# English natural-language chat: curated non-code, non-math, English sources. +# (Aya is excluded — it is multilingual; math/STEM splits are excluded — they +# are English text but not natural-*language* in the prose sense we want as the +# contrast with code.) +ENGLISH_SOURCES = { + "Persona Precise IF", + "OpenAssistant", + "CoCoNot", + "WildChat", + "Dolci Think Precise IF", +} + + +def classify(source: str) -> str | None: + if CODE_SOURCE_RE.search(source or ""): + return "code" + if source in ENGLISH_SOURCES: + return "english" + return None + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default=MODEL) + parser.add_argument("--dataset", default=DATASET) + parser.add_argument("--n-per-batch", type=int, default=10) + parser.add_argument("--seq-len", type=int, default=1024) + parser.add_argument( + "--max-prompt-frac", + type=float, + default=0.5, + help="reject a sample if its prompt exceeds this fraction of seq_len " + "(guarantees real completion tokens survive truncation).", + ) + parser.add_argument( + "--scan-limit", + type=int, + default=20000, + help="max dataset rows to stream before giving up on filling a batch.", + ) + parser.add_argument("--out", default=str(OUT_PATH)) + parser.add_argument( + "--skip-model-download", + action="store_true", + help="select data and fetch the tokenizer only; skip the ~14GB weights.", + ) + args = parser.parse_args() + + from datasets import load_dataset + from transformers import AutoTokenizer + + print(f"prefetch — model={args.model} dataset={args.dataset}") + print(f" HF_HOME={os.environ.get('HF_HOME', '(default)')}") + print( + f" selecting {args.n_per_batch} code + {args.n_per_batch} english, " + f"seq_len={args.seq_len}, max_prompt_frac={args.max_prompt_frac}" + ) + + tokenizer = AutoTokenizer.from_pretrained(args.model) + max_prompt = int(args.max_prompt_frac * args.seq_len) + + ds = load_dataset(args.dataset, split=SPLIT, streaming=True) + batches: dict[str, list[dict]] = {"code": [], "english": []} + selected_meta: dict[str, list[dict]] = {"code": [], "english": []} + n_scanned = 0 + + for ex in ds: + n_scanned += 1 + if n_scanned > args.scan_limit: + break + kind = classify(ex.get("source", "")) + if kind is None or len(batches[kind]) >= args.n_per_batch: + continue + messages = ex["messages"] + # Need a final assistant turn to have a completion to learn. + if not messages or messages[-1].get("role") != "assistant": + continue + try: + prompt_len, total_len = sft_prompt_completion_lengths(messages, tokenizer) + except Exception: # noqa: BLE001 — skip samples the template can't render + continue + if prompt_len > max_prompt or total_len <= prompt_len: + continue # no usable completion tokens at this seq_len + + batches[kind].append({"id": ex.get("id"), "source": ex.get("source"), "messages": messages}) + selected_meta[kind].append( + { + "id": ex.get("id"), + "source": ex.get("source"), + "prompt_len": prompt_len, + "total_len": total_len, + "completion_len_at_seq_len": max(0, min(total_len, args.seq_len) - prompt_len), + } + ) + if all(len(batches[k]) >= args.n_per_batch for k in batches): + break + + for k in batches: + got = len(batches[k]) + if got < args.n_per_batch: + print( + f" WARNING: only found {got}/{args.n_per_batch} '{k}' samples " + f"in {n_scanned} rows. Raise --scan-limit or relax the filters." + ) + print(f" {k}: {got} samples; sources={sorted({m['source'] for m in selected_meta[k]})}") + + payload = { + "meta": { + "model": args.model, + "dataset": args.dataset, + "split": SPLIT, + "seq_len_target": args.seq_len, + "n_per_batch": args.n_per_batch, + "max_prompt_frac": args.max_prompt_frac, + "code_source_regex": CODE_SOURCE_RE.pattern, + "english_sources": sorted(ENGLISH_SOURCES), + "n_rows_scanned": n_scanned, + "selected": selected_meta, + }, + "batches": batches, + } + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w") as f: + json.dump(payload, f, indent=2) + print(f"wrote {out_path} ({out_path.stat().st_size / 1024:.0f} KB)") + + if args.skip_model_download: + print("skipped model weight download (--skip-model-download); tokenizer cached.") + return + + print( + f"downloading model weights for {args.model} into the HF cache (this is the ~14GB step) ..." + ) + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained(args.model, dtype="auto") + n_params = sum(p.numel() for p in model.parameters()) + print(f" cached model: {n_params:,} parameters. Ready for offline analysis.") + del model + + +if __name__ == "__main__": + main() diff --git a/examples/olmo_sft_analysis/resource_tracker.py b/examples/olmo_sft_analysis/resource_tracker.py new file mode 100644 index 0000000..e5a0cfa --- /dev/null +++ b/examples/olmo_sft_analysis/resource_tracker.py @@ -0,0 +1,404 @@ +"""Resource tracking for the OLMo SFT example: time, memory, and GPU activity. + +Three signals, recorded while a computation runs: + +- **wall time** — overall and per named phase (``perf_counter``); +- **peak GPU memory** — ``torch.cuda.max_memory_allocated`` / + ``max_memory_reserved``, per phase and overall (always on, zero deps); +- **GPU activity** — sampled in a background thread. The primary sampler is + **NVIDIA DCGM** (``dcgmi dmon``), which exposes the *meaningful* engine + counters on A100-class GPUs: SM-active, tensor-pipe-active, DRAM-active and + GR-engine-active fractions, plus power / framebuffer / temperature. These say + how hard the GPU is actually working — unlike the coarse "GPU utilization %" + (percent of wall time at least one kernel ran), which on a single-stream + backward-heavy job just pins near 100%. + +If ``dcgmi`` is missing or its profiling fields are unavailable (e.g. on a MIG +slice, or without the DCGM host engine), the tracker **automatically falls +back** to sampling ``nvidia-smi`` (coarse util% / memory / power). The summary +records which sampler was used. + +Usage:: + + tracker = ResourceTracker(device) + with tracker: + with tracker.measure("load_model"): + bundle = load_hf_model(...) + with tracker.measure("analyze"): + analyze(...) + tracker.save("resources.json") + print(tracker.format_summary()) + +The chosen sampler needs no Python package (``dcgmi`` / ``nvidia-smi`` are +system CLIs). See this example's README for the DCGM note. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import IO + +import torch + +_MB = 1024.0 * 1024.0 + +# DCGM field IDs we stream, in order. The (id, key, kind) tuples drive both the +# `-e` argument and the row parser, so they cannot drift apart. +# kind "frac": 0..1 engine-active fraction (shown as % in the text summary) +# kind "pct": 0..100 coarse utilization +# kind "mb"/"w"/"c": framebuffer MB / power W / temperature C +_DCGM_FIELDS: list[tuple[int, str, str]] = [ + (1001, "gr_engine_active", "frac"), # DCGM_FI_PROF_GR_ENGINE_ACTIVE + (1002, "sm_active", "frac"), # DCGM_FI_PROF_SM_ACTIVE + (1004, "tensor_active", "frac"), # DCGM_FI_PROF_PIPE_TENSOR_ACTIVE + (1005, "dram_active", "frac"), # DCGM_FI_PROF_DRAM_ACTIVE + (203, "gpu_util_pct", "pct"), # DCGM_FI_DEV_GPU_UTIL (coarse, for reference) + (252, "fb_used_mb", "mb"), # DCGM_FI_DEV_FB_USED + (155, "power_w", "w"), # DCGM_FI_DEV_POWER_USAGE + (150, "gpu_temp_c", "c"), # DCGM_FI_DEV_GPU_TEMP +] + +# nvidia-smi fallback query fields, mapped onto the same keys where they exist. +_SMI_FIELDS: list[tuple[str, str]] = [ + ("utilization.gpu", "gpu_util_pct"), + ("memory.used", "fb_used_mb"), + ("power.draw", "power_w"), + ("temperature.gpu", "gpu_temp_c"), +] + +# Keys whose values are 0..1 fractions — rendered as percentages in text. +_FRACTION_KEYS = {"gr_engine_active", "sm_active", "tensor_active", "dram_active"} + + +@dataclass +class _Sample: + t: float # seconds since tracker entry + vals: dict[str, float | None] + + +@dataclass +class _Phase: + name: str + t_start: float + t_end: float | None = None + torch_peak_alloc_mb: float = 0.0 + torch_peak_reserved_mb: float = 0.0 + + +def _to_float(x: str) -> float | None: + try: + return float(x) + except ValueError: + return None # "N/A" / "[N/A]" + + +@dataclass +class ResourceTracker: + """Track wall time, peak GPU memory, and sampled GPU activity (DCGM).""" + + device: torch.device | str = "cuda" + interval_s: float = 0.25 + gpu_index: int | str | None = None + + _samples: list[_Sample] = field(default_factory=list, init=False) + _phases: list[_Phase] = field(default_factory=list, init=False) + _stop: threading.Event = field(default_factory=threading.Event, init=False) + _thread: threading.Thread | None = field(default=None, init=False) + _proc: subprocess.Popen | None = field(default=None, init=False) + _t0: float | None = field(default=None, init=False) + sampler: str = field(default="none", init=False) + + def __post_init__(self) -> None: + self.device = torch.device(self.device) + self.is_cuda = self.device.type == "cuda" and torch.cuda.is_available() + self._gpu_id = self._resolve_gpu_id() if self.is_cuda else None + + # ---- GPU selection ------------------------------------------------- + + def _resolve_gpu_id(self) -> str: + """Return the GPU id (for ``-i``) of the device torch will use.""" + if self.gpu_index is not None: + return str(self.gpu_index) + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + torch_idx = self.device.index if self.device.index is not None else 0 + if visible: + tokens = [t.strip() for t in visible.split(",") if t.strip()] + if 0 <= torch_idx < len(tokens): + return tokens[torch_idx] # physical index or UUID + return str(torch_idx) + + # ---- DCGM sampler (primary) ---------------------------------------- + + def _dcgm_cmd(self, count: int | None) -> list[str]: + fields = ",".join(str(f[0]) for f in _DCGM_FIELDS) + cmd = [ + "dcgmi", + "dmon", + "-e", + fields, + "-d", + str(max(1, int(self.interval_s * 1000))), + "-i", + str(self._gpu_id), + ] + if count is not None: + cmd += ["-c", str(count)] + return cmd + + @staticmethod + def _parse_dcgm_row(line: str) -> dict[str, float | None] | None: + """Parse one ``dcgmi dmon`` data row into the field keys, or None.""" + toks = line.split() + # Data rows look like: "GPU 0 0.50 0.45 ...". Header lines start "#". + if len(toks) < 2 + len(_DCGM_FIELDS) or toks[0] != "GPU": + return None + if not toks[1].lstrip("-").isdigit(): + return None + values = toks[2 : 2 + len(_DCGM_FIELDS)] + return {key: _to_float(v) for (_, key, _kind), v in zip(_DCGM_FIELDS, values, strict=True)} + + def _dcgm_available(self) -> bool: + """One-shot probe: does ``dcgmi dmon`` yield a parseable row?""" + try: + out = subprocess.run( + self._dcgm_cmd(count=1), + capture_output=True, + text=True, + timeout=10.0, + ) + except (FileNotFoundError, OSError, subprocess.SubprocessError): + return False + if out.returncode != 0: + return False + return any(self._parse_dcgm_row(ln) is not None for ln in out.stdout.splitlines()) + + def _read_dcgm_stream(self, stream: IO[str]) -> None: + assert self._t0 is not None + for line in stream: + if self._stop.is_set(): + break + vals = self._parse_dcgm_row(line) + if vals is not None: + self._samples.append(_Sample(t=time.perf_counter() - self._t0, vals=vals)) + + def _start_dcgm(self) -> bool: + try: + self._proc = subprocess.Popen( # noqa: S603 — fixed argv, no shell + self._dcgm_cmd(count=None), + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + except (FileNotFoundError, OSError): + return False + assert self._proc.stdout is not None + self._thread = threading.Thread( + target=self._read_dcgm_stream, args=(self._proc.stdout,), daemon=True + ) + self._thread.start() + return True + + # ---- nvidia-smi sampler (fallback) --------------------------------- + + def _poll_smi(self) -> _Sample | None: + assert self._t0 is not None + query = ",".join(f for f, _ in _SMI_FIELDS) + cmd = [ + "nvidia-smi", + f"--query-gpu={query}", + "--format=csv,noheader,nounits", + "-i", + str(self._gpu_id), + ] + try: + out = subprocess.run( + cmd, capture_output=True, text=True, timeout=5.0, check=True + ).stdout.strip() + except (subprocess.SubprocessError, FileNotFoundError, OSError): + return None + line = out.splitlines()[0] if out else "" + parts = [p.strip() for p in line.split(",")] + if len(parts) < len(_SMI_FIELDS): + return None + vals = {key: _to_float(parts[i]) for i, (_, key) in enumerate(_SMI_FIELDS)} + return _Sample(t=time.perf_counter() - self._t0, vals=vals) + + def _run_smi_sampler(self) -> None: + while not self._stop.is_set(): + sample = self._poll_smi() + if sample is not None: + self._samples.append(sample) + self._stop.wait(self.interval_s) + + # ---- context management -------------------------------------------- + + def __enter__(self) -> ResourceTracker: + self._t0 = time.perf_counter() + if not self.is_cuda: + return self + torch.cuda.reset_peak_memory_stats(self.device) + torch.cuda.synchronize(self.device) + if self._dcgm_available() and self._start_dcgm(): + self.sampler = "dcgm" + elif self._poll_smi() is not None: + # nvidia-smi answered a probe → use it. (Without this probe a node + # with neither dcgmi nor nvidia-smi would silently report + # sampler="nvidia-smi" with zero samples.) + self.sampler = "nvidia-smi" + self._thread = threading.Thread(target=self._run_smi_sampler, daemon=True) + self._thread.start() + else: + self.sampler = "none" # no GPU sampler available; torch peak still tracked + return self + + def __exit__(self, *exc: object) -> None: + self._stop.set() + if self._proc is not None: + self._proc.terminate() + try: + self._proc.wait(timeout=5.0) + except subprocess.TimeoutExpired: + self._proc.kill() + if self._thread is not None: + self._thread.join(timeout=5.0) + if self.is_cuda: + torch.cuda.synchronize(self.device) + + @contextmanager + def measure(self, name: str) -> Iterator[None]: + """Time a named phase and record its peak torch memory.""" + assert self._t0 is not None, "use `with tracker:` before measure()" + if self.is_cuda: + torch.cuda.synchronize(self.device) + torch.cuda.reset_peak_memory_stats(self.device) + phase = _Phase(name=name, t_start=time.perf_counter() - self._t0) + self._phases.append(phase) + try: + yield + finally: + if self.is_cuda: + torch.cuda.synchronize(self.device) + phase.torch_peak_alloc_mb = torch.cuda.max_memory_allocated(self.device) / _MB + phase.torch_peak_reserved_mb = torch.cuda.max_memory_reserved(self.device) / _MB + phase.t_end = time.perf_counter() - self._t0 + + # ---- summary ------------------------------------------------------- + + @staticmethod + def _stats(values: list[float | None]) -> dict[str, float] | None: + vals = [v for v in values if v is not None] + if not vals: + return None + return {"mean": sum(vals) / len(vals), "max": max(vals), "min": min(vals)} + + def _all_keys(self) -> list[str]: + keys: list[str] = [] + for s in self._samples: + for k in s.vals: + if k not in keys: + keys.append(k) + return keys + + def _agg(self, samples: list[_Sample]) -> dict[str, dict[str, float] | None]: + return {k: self._stats([s.vals.get(k) for s in samples]) for k in self._all_keys()} + + def _slice(self, t_start: float, t_end: float | None) -> list[_Sample]: + end = t_end if t_end is not None else float("inf") + return [s for s in self._samples if t_start <= s.t <= end] + + def summary(self) -> dict: + total_wall = ( + max((p.t_end for p in self._phases if p.t_end is not None), default=0.0) + if self._phases + else (time.perf_counter() - self._t0 if self._t0 is not None else 0.0) + ) + return { + "device": str(self.device), + "gpu_name": torch.cuda.get_device_name(self.device) if self.is_cuda else None, + "gpu_id": self._gpu_id, + "sampler": self.sampler, + "sampler_interval_s": self.interval_s, + "n_gpu_samples": len(self._samples), + "wall_s_total": total_wall, + "torch_peak_alloc_mb": max((p.torch_peak_alloc_mb for p in self._phases), default=0.0), + "torch_peak_reserved_mb": max( + (p.torch_peak_reserved_mb for p in self._phases), default=0.0 + ), + "gpu": self._agg(self._samples), + "phases": [ + { + "name": p.name, + "t_start": p.t_start, + "t_end": p.t_end, + "wall_s": (p.t_end - p.t_start) if p.t_end is not None else None, + "torch_peak_alloc_mb": p.torch_peak_alloc_mb, + "torch_peak_reserved_mb": p.torch_peak_reserved_mb, + "gpu": self._agg(self._slice(p.t_start, p.t_end)), + } + for p in self._phases + ], + } + + def timeline(self) -> list[dict]: + """Raw per-sample GPU telemetry (for plotting an activity trace).""" + return [{"t": s.t, **s.vals} for s in self._samples] + + def save(self, path: str) -> None: + with open(path, "w") as f: + json.dump({"summary": self.summary(), "timeline": self.timeline()}, f, indent=2) + + @staticmethod + def _fmt_pct(key: str, stat: dict[str, float] | None) -> str | None: + if stat is None: + return None + scale = 100.0 if key in _FRACTION_KEYS else 1.0 + unit = "%" if key in _FRACTION_KEYS or key.endswith("_pct") else "" + return f"mean {stat['mean'] * scale:.0f}{unit}, max {stat['max'] * scale:.0f}{unit}" + + def format_summary(self) -> str: + s = self.summary() + lines = [ + f"device: {s['device']}" + (f" ({s['gpu_name']})" if s["gpu_name"] else ""), + f"GPU sampler: {s['sampler']} ({s['n_gpu_samples']} samples @ {s['sampler_interval_s']}s)", + f"wall (total): {s['wall_s_total']:.1f} s", + f"torch peak: {s['torch_peak_alloc_mb']:.0f} MB allocated, " + f"{s['torch_peak_reserved_mb']:.0f} MB reserved", + ] + gpu = s["gpu"] + labels = [ + ("gr_engine_active", "GR engine"), + ("sm_active", "SM active"), + ("tensor_active", "tensor core"), + ("dram_active", "DRAM (mem BW)"), + ("gpu_util_pct", "GPU util"), + ("fb_used_mb", "FB mem (MB)"), + ("power_w", "power (W)"), + ("gpu_temp_c", "temp (C)"), + ] + for key, label in labels: + stat = gpu.get(key) + if stat is None: + continue + if key in _FRACTION_KEYS or key.endswith("_pct"): + txt = self._fmt_pct(key, stat) + else: + txt = f"mean {stat['mean']:.0f}, max {stat['max']:.0f}" + lines.append(f" {label:<14} {txt}") + lines.append("phases:") + for p in s["phases"]: + busy = p["gpu"].get("sm_active") or p["gpu"].get("gpu_util_pct") + scale = 100.0 if (p["gpu"].get("sm_active")) else 1.0 + bstr = f", busy mean {busy['mean'] * scale:.0f}%" if busy else "" + wall = f"{p['wall_s']:.1f}s" if p["wall_s"] is not None else "—" + lines.append( + f" {p['name']:<14} {wall:>8} peak {p['torch_peak_alloc_mb']:.0f} MB{bstr}" + ) + return "\n".join(lines) diff --git a/examples/olmo_sft_analysis/run.sbatch b/examples/olmo_sft_analysis/run.sbatch new file mode 100644 index 0000000..1283ad0 --- /dev/null +++ b/examples/olmo_sft_analysis/run.sbatch @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +#SBATCH --job-name=olmo-sft-lnp +# Partition / account come from .env via examples/benchmark/submit.sh +# (SBATCH_PARTITION / SBATCH_ACCOUNT); override on the CLI to change per-submit. +#SBATCH --nodes=1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=6 +#SBATCH --mem=64G +#SBATCH --time=00:30:00 +#SBATCH --output=olmo_sft_lnp_%j.log + +set -euo pipefail + +# IMPORTANT: run prefetch.py on a LOGIN node first (it needs internet) so the +# model weights are in the HF cache and samples.json exists. This job runs +# OFFLINE — see this example's README. +# +# Submit from the repo root via the wrapper that pulls SBATCH_* from .env: +# examples/benchmark/submit.sh examples/olmo_sft_analysis/run.sbatch +REPO="${SLURM_SUBMIT_DIR:-$PWD}" +cd "$REPO" + +# Cache-only on the compute node (no internet). run_config also reads HF_HOME +# from .env; these exports make the offline intent explicit and independent of it. +export HF_HUB_OFFLINE=1 +export TRANSFORMERS_OFFLINE=1 + +# Reclaim reserved-but-unallocated fragmentation. A 7B chi_net run sits near the +# 7B-on-94GB ceiling (the analyzer holds one batch's full gradient for the +# pending cross-pair while computing the next), so this measurably lowers peak. +export PYTORCH_ALLOC_CONF=expandable_segments:True + +echo "node: $(hostname)" +nvidia-smi -L +dcgmi discovery -l 2>/dev/null | head -5 || echo "(dcgmi discovery unavailable; tracker will fall back to nvidia-smi)" + +# Memory (MEASURED via mem_profile.py on a 93 GiB H100, this 7B model): +# per-batch peak ~75-84 GB, essentially FLAT in seq_len and batch size — it is +# param-bound (bf16 weights 14.6 + fp32 flat-grad 29.2 + bf16 grad tuple 14.6 +# + workspace). Activations are negligible, so seq_len is ~free. +# The cross pair needs BOTH batches' gradients for the dot; holding the second +# ~29 GB gradient on-GPU during the next backward would OOM (~104 GB). So we let +# cross_grad_storage=auto offload the cached gradients to host RAM (the dot then +# runs on CPU) -> GPU peak stays at one batch's ~84 GB, the cross story is back. +# hutchinson keeps the chi_net step (~73 GB with the live flat-grad) under the +# delta_loss peak; per_sequence_cv would push it to ~94 GB (too tight). +# seq_len=1024 captures more of the long traces at no memory cost; raise +# it (and --time) for more. +"$REPO"/.venv/bin/python examples/olmo_sft_analysis/olmo_sft_analyze.py \ + --chi-net-method hutchinson --seq-len 1024 --cross-grad-storage auto "$@" + +# Plots + tables (CPU-only; safe to run here or later on a login node): +"$REPO"/.venv/bin/python examples/olmo_sft_analysis/analyze_results.py diff --git a/examples/olmo_sft_analysis/spectral_position.png b/examples/olmo_sft_analysis/spectral_position.png new file mode 100644 index 0000000..089baf4 Binary files /dev/null and b/examples/olmo_sft_analysis/spectral_position.png differ diff --git a/pyproject.toml b/pyproject.toml index 1684ddb..ca25357 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "vatis" -version = "0.1.0" +version = "0.2.0" description = "Compute LNP observables (chi_loss, chi_net, chi_pos) on pretrained model checkpoints, scaled across multiple GPUs via DDP." readme = "README.md" requires-python = ">=3.11" diff --git a/tests/unit/test_analyzer.py b/tests/unit/test_analyzer.py index 777211b..6384f85 100644 --- a/tests/unit/test_analyzer.py +++ b/tests/unit/test_analyzer.py @@ -20,7 +20,7 @@ mlp_valid_mask, ) from vatis import analyze -from vatis.analyzer import ALL_OBSERVABLES, Analyzer +from vatis.analyzer import ALL_OBSERVABLES, Analyzer, _should_offload_cross_grads from vatis.models.hf import ModelBundle from vatis.sinks.parquet import ParquetSink @@ -347,3 +347,125 @@ def test_parquet_sink_streaming_flush(tmp_path: Path) -> None: ) table = pq.read_table(out) assert table.num_rows == 12 + + +def test_cross_grad_storage_gpu_cpu_equivalent() -> None: + """Cross-pair observables must match whether the cached gradients are held + on-device (``"gpu"``) or offloaded to host RAM (``"cpu"``). + + The offload is a value-preserving device copy plus a device-agnostic + ``_dot_fp64``, so forcing either storage on the same seeded run must agree. + (On a CPU model both keep the gradient on the CPU, so this guards the + plumbing — the resolver honouring the override, the ``.to('cpu')`` cache + write, and the dot — rather than an actual cross-device transfer.) + """ + batch_a = make_tiny_lm_batch(batch_size=4, seed=0) + batch_b = make_tiny_lm_batch(batch_size=4, seed=1) + + def run(storage: str) -> dict[tuple[str, str, str], float]: + results = analyze( + model=_build_lm_bundle(), # fresh model, fixed seed → identical weights + revisions=["step0"], + eval_batches={"a": batch_a, "b": batch_b}, + cross_pairs=[("a", "b")], + cross_grad_storage=storage, + n_hutchinson=8, + micro_batch_size=2, + seed=0, + sink=None, + ) + return {(r.batch_a, r.batch_b, r.observable): r.value for r in results[0].rows} + + gpu = run("gpu") + cpu = run("cpu") + assert gpu.keys() == cpu.keys() + for key in gpu: + assert gpu[key] == pytest.approx(cpu[key], rel=1e-12, abs=1e-12), key + # And the cross pair was actually emitted (not silently skipped). + assert ("a", "b", "delta_loss") in gpu + + +def test_invalid_cross_grad_storage_raises() -> None: + with pytest.raises(ValueError, match=r"cross_grad_storage must be"): + Analyzer(eval_batches={"v": make_tiny_lm_batch(batch_size=2)}, cross_grad_storage="disk") + + +def test_should_offload_cross_grads_heuristic() -> None: + """The auto offload decision: large models offload, small stay on-GPU.""" + gib93 = 99_900_000_000 # ~93 GiB H100 + # 7B bf16, two cross batches → one self-pair (~weights + 8·P) plus the other + # batch's fp32 gradient (4·P) overflows the card → offload. + assert _should_offload_cross_grads( + n_params=7_300_000_000, + weights_bytes=14_600_000_000, + n_cross_batches=2, + total_device_bytes=gib93, + ) + # A 160M model on the same card fits comfortably → stay on-GPU. + assert not _should_offload_cross_grads( + n_params=160_000_000, + weights_bytes=320_000_000, + n_cross_batches=2, + total_device_bytes=gib93, + ) + # Fewer than two cross batches → nothing to offload, regardless of size. + assert not _should_offload_cross_grads( + n_params=7_300_000_000, + weights_bytes=14_600_000_000, + n_cross_batches=1, + total_device_bytes=gib93, + ) + # Boundary: with P=1e9, weights=2e9, n=2 the footprint is weights + 12·P = + # 14e9 bytes, compared against 0.85·total. The decision must flip right + # around total = 14e9 / 0.85 ≈ 16.47e9 — this catches a wrong factor that + # the 45× large-vs-small gap above would not. + common = dict(n_params=1_000_000_000, weights_bytes=2_000_000_000, n_cross_batches=2) + assert _should_offload_cross_grads(**common, total_device_bytes=16_000_000_000) # 14 > 13.6 + assert not _should_offload_cross_grads( + **common, total_device_bytes=17_000_000_000 + ) # 14 < 14.45 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="real GPU→CPU offload needs CUDA") +def test_cross_grad_storage_offload_equivalent_on_cuda() -> None: + """On a CUDA device, the cross observables must be identical whether the + cached gradients stay on the GPU (``"gpu"``) or are offloaded to host RAM + (``"cpu"``) — a genuine device-boundary copy plus a CPU-side dot. The + CPU-model test above cannot exercise this (both paths keep the gradient on + the CPU); this is the one that actually validates the offload. + """ + batch_a = make_tiny_lm_batch(batch_size=4, seed=0) + batch_b = make_tiny_lm_batch(batch_size=4, seed=1) + + def run(storage: str) -> dict[tuple[str, str, str], float]: + torch.manual_seed(0) # identical weights each call + model = TinyTransformer().eval().cuda() + for p in model.parameters(): + p.requires_grad_(True) + bundle = ModelBundle( + model=model, + params=list(model.parameters()), + forward_fn=lambda m, b: m(b["input_ids"]), + loss_fn=causal_lm_loss, + valid_mask_fn=causal_lm_valid_mask, + identifier="toy@step0", + ) + results = analyze( + model=bundle, + revisions=["step0"], + eval_batches={"a": batch_a, "b": batch_b}, + cross_pairs=[("a", "b")], + cross_grad_storage=storage, + n_hutchinson=8, + micro_batch_size=2, + seed=0, + device="cuda", + sink=None, + ) + return {(r.batch_a, r.batch_b, r.observable): r.value for r in results[0].rows} + + gpu = run("gpu") + cpu = run("cpu") + assert gpu.keys() == cpu.keys() + for key in gpu: + assert gpu[key] == pytest.approx(cpu[key], rel=1e-6, abs=1e-6), key diff --git a/vatis/_version.py b/vatis/_version.py index 3dc1f76..d3ec452 100644 --- a/vatis/_version.py +++ b/vatis/_version.py @@ -1 +1 @@ -__version__ = "0.1.0" +__version__ = "0.2.0" diff --git a/vatis/analyzer.py b/vatis/analyzer.py index 5355cf1..3dccc71 100644 --- a/vatis/analyzer.py +++ b/vatis/analyzer.py @@ -119,6 +119,15 @@ class Analyzer: cross_pairs: optional list of ``(name_a, name_b)`` tuples to compute ``delta_loss(A, B)`` (and the corresponding ``chi_pos``) for. Self pairs ``(name, name)`` are always computed. + cross_grad_storage: where a batch's cached full gradient lives while the + cross-pair dot product is computed — ``"gpu"``, ``"cpu"``, or + ``"auto"`` (default). The cross dot needs both batches' P-sized + gradients; for large models keeping them on the GPU alongside a + later batch's backward can OOM. ``"cpu"`` offloads them to host RAM + (the dot then runs on the CPU); ``"gpu"`` keeps them on-device + (fast, no transfer); ``"auto"`` offloads only when an on-GPU cache + would not leave room for a self-pair backward. Only matters when + ``cross_pairs`` is non-empty. sink: a :class:`ResultSink`, a path str (auto-wraps to :class:`ParquetSink`), or a list of sinks. seed: base seed used to derive Hutchinson probes; the analyzer mixes @@ -142,6 +151,7 @@ def __init__( hutchinson_distribution: str = "rademacher", micro_batch_size: int = 1, cross_pairs: list[tuple[str, str]] | None = None, + cross_grad_storage: str = "auto", sink: ResultSink | str | list[ResultSink] | None = None, seed: int = 0, device: torch.device | str = "cpu", @@ -161,6 +171,11 @@ def __init__( self.hutchinson_distribution = hutchinson_distribution self.micro_batch_size = micro_batch_size self.cross_pairs = list(cross_pairs or []) + if cross_grad_storage not in ("auto", "gpu", "cpu"): + raise ValueError( + f"cross_grad_storage must be 'auto', 'gpu', or 'cpu'; got {cross_grad_storage!r}" + ) + self.cross_grad_storage = cross_grad_storage self.seed = seed self.device = torch.device(device) self.dtype = dtype @@ -199,6 +214,38 @@ def _normalize_sinks( return out raise TypeError(f"unsupported sink type: {type(sink).__name__}") + def _resolve_cross_grad_storage(self, bundle: ModelBundle) -> str: + """Decide where cached cross-pair gradients live: ``"gpu"`` or ``"cpu"``. + + Explicit ``"gpu"``/``"cpu"`` pass through. ``"auto"`` offloads to host + RAM only when keeping the cached fp32 gradient(s) on the GPU would not + leave room for a self-pair's backward. For a 7B model the cached + gradient (~29 GB) plus one self-pair's footprint (~weights + 8·P) + overflows an 80–90 GB card, so ``"auto"`` offloads; small models stay + GPU-resident (no transfer, no numeric change — important for the + perspic cross-validation tests). + + The estimate is intentionally conservative (a 0.85 fraction over an + optimistic ``8·P`` self-pair floor); if ``"gpu"`` still OOMs on a + borderline model, pass ``cross_grad_storage="cpu"`` explicitly. + """ + if self.cross_grad_storage in ("gpu", "cpu"): + return self.cross_grad_storage + device = self.device + if device.type != "cuda" or not torch.cuda.is_available(): + return "gpu" # no host-offload target that would help + cross_names = {n for pair in self.cross_pairs for n in pair if pair[0] != pair[1]} + n_params = sum(p.numel() for p in bundle.params) + weights_bytes = sum(p.numel() * p.element_size() for p in bundle.params) + total = torch.cuda.get_device_properties(device).total_memory + offload = _should_offload_cross_grads( + n_params=n_params, + weights_bytes=weights_bytes, + n_cross_batches=len(cross_names), + total_device_bytes=total, + ) + return "cpu" if offload else "gpu" + # ------------------------------------------------------------------ runs def run( @@ -250,18 +297,31 @@ def _run_one_checkpoint( for p in bundle.params: p.requires_grad_(True) - # Self pairs first; we cache the per-batch grads (parameter-vector flat) - # so the cross pairs can re-use them via dot product. + # Self pairs first. We cache a batch's full parameter-space gradient + # ONLY if some cross pair will consume it via the dot product. Caching it + # unconditionally keeps a P-sized vector resident (e.g. ~29 GB fp32 for a + # 7B model) while subsequent batches run their own backward — for large + # models that retained gradient is enough to OOM even when no cross pair + # was requested. So we cache selectively and drop the local reference + # otherwise, freeing it before the next batch's backward. + cross_names = {n for pair in self.cross_pairs for n in pair if pair[0] != pair[1]} + grad_storage = self._resolve_cross_grad_storage(bundle) if cross_names else "gpu" per_batch_grad_cache: dict[str, torch.Tensor] = {} per_batch_n_valid: dict[str, int] = {} for name, spec in self.eval_batches.items(): self_result, g_full = self._compute_self_pair(bundle, ckpt_id, revision, name, spec) result.rows.extend(self_result) - # The cached gradient lives on the model device; sized P (params). - if g_full is not None: - per_batch_grad_cache[name] = g_full per_batch_n_valid[name] = self._last_n_valid + if g_full is not None and name in cross_names: + # Cache for the cross-pair dot. Offload to host RAM when keeping + # it on the GPU would crowd out a later batch's backward (large + # models); otherwise keep it on-device. The cross dot + # (``_dot_fp64``) runs on whichever device the gradients live on. + per_batch_grad_cache[name] = g_full.to("cpu") if grad_storage == "cpu" else g_full + # Drop the local (GPU) reference; if it wasn't cached on-device above + # it is freed now, so gradients can't pile up across batches. + g_full = None # Cross pairs (if any) — only delta_loss + chi_pos. for a, b in self.cross_pairs: @@ -617,6 +677,35 @@ def _emit_rows( # ------------------------------------------------------------- helpers +def _should_offload_cross_grads( + *, + n_params: int, + weights_bytes: int, + n_cross_batches: int, + total_device_bytes: int, + fraction: float = 0.85, +) -> bool: + """Whether cached cross-pair gradients should be offloaded to host RAM. + + Pure size arithmetic (no device access) so it is unit-testable. The cross + dot needs ``n_cross_batches`` gradients; while one batch runs its backward + the others' fp32 gradients (``4·P`` bytes each) sit resident alongside that + backward's footprint — weights + the fp32 flat-grad result + a transient + grad tuple + autograd/cuBLAS workspace, estimated as ``weights + 8·P`` + (measured ~``weights + 9.5·P`` for a 7B bf16 model; ``8·P`` is a deliberately + optimistic floor, paired with the conservative ``fraction``). Offload when + that total would exceed ``fraction`` of device memory. + + Returns ``False`` for ``n_cross_batches < 2`` — no cross dot holds two + gradients at once, so there is nothing to offload. + """ + if n_cross_batches < 2: + return False + self_pair_bytes = weights_bytes + 8 * n_params + coexisting_cached = (n_cross_batches - 1) * 4 * n_params + return (self_pair_bytes + coexisting_cached) > fraction * total_device_bytes + + def _zero_param_vector(params: list[torch.nn.Parameter], device: torch.device) -> torch.Tensor: sizes = [p.numel() for p in params] if not sizes: @@ -721,6 +810,7 @@ def analyze( hutchinson_distribution: str = "rademacher", micro_batch_size: int = 1, cross_pairs: list[tuple[str, str]] | None = None, + cross_grad_storage: str = "auto", sink: ResultSink | str | list[ResultSink] | None = None, seed: int = 0, device: torch.device | str = "cpu", @@ -746,6 +836,7 @@ def analyze( hutchinson_distribution=hutchinson_distribution, micro_batch_size=micro_batch_size, cross_pairs=cross_pairs, + cross_grad_storage=cross_grad_storage, sink=sink, seed=seed, device=device, diff --git a/vatis/models/hf.py b/vatis/models/hf.py index 81244b3..31f8ca7 100644 --- a/vatis/models/hf.py +++ b/vatis/models/hf.py @@ -160,6 +160,7 @@ def load_hf_model( dtype: str = "bf16", device: torch.device | str = "cpu", trust_remote_code: bool = False, + attn_implementation: str | None = None, ) -> ModelBundle: """Load an HF causal LM by name + revision and wrap it in a ModelBundle. @@ -171,6 +172,18 @@ def load_hf_model( Gradients are always accumulated in fp32 by the analyzer. device: target device for the model. trust_remote_code: passed through to ``from_pretrained``. + attn_implementation: optional ``from_pretrained`` attention backend, + e.g. ``"sdpa"``, ``"flash_attention_2"``, or ``"eager"``. ``None`` + (default) keeps the model's own default. A memory-efficient backend + *can* turn the attention activation memory from ``"eager"``'s O(S²) + (materialized scores) into O(S), which is the difference between + fitting and OOMing at long sequence lengths — **but only if the + model's attention module actually routes through it**. Some + architectures still materialize scores under ``"sdpa"`` (e.g. OLMo3 + in this project's testing), so don't assume it lowers the seq-len + ceiling without measuring. Note: flash/sdpa kernels use a different + reduction order than eager, so observables can shift at the ~1e-5 + level (below the Hutchinson noise floor). Returns: A ModelBundle ready to feed to ``Analyzer``. @@ -188,6 +201,7 @@ def load_hf_model( revision=revision, torch_dtype=torch_dtype, trust_remote_code=trust_remote_code, + attn_implementation=attn_implementation, ) model.to(device=device) model.eval()