Konsti/olmo sft example - #1
Merged
Merged
Conversation
The self-pair loop cached each eval batch's flat parameter-space gradient unconditionally, but only cross pairs consume it. For large models that kept a P-sized fp32 vector (~29 GB for a 7B model) resident while subsequent batches ran their own backward — enough to OOM a 94 GB H100 even when no cross pair was requested (the second batch's ~84 GB compute peak coincided with the first batch's retained gradient). Cache a batch's gradient only if some cross pair references it, and drop the local reference otherwise so it is freed before the next batch's backward. Self-pair-only runs now peak at a single batch's footprint. Unit suite green; ruff + mypy --strict clean.
A realistic, production-scale vatis example: load 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, chat-templated with completion-only loss masking (the actual SFT objective). Two-stage offline workflow for the internet-less compute nodes: prefetch.py (login node, online) selects samples -> samples.json and warms the model into the HF cache (needs 'datasets', a documented manual venv install, deliberately not a repo dep); olmo_sft_analyze.py (GPU node, offline) runs analyze() with resource tracking; analyze_results.py turns parquet+resources into plots/tables with zero vatis imports. resource_tracker.py records wall time, torch peak memory, and GPU engine activity via DCGM dmon (nvidia-smi fallback). mem_profile.py is the GPU memory profiler used to localize the 7B memory cost. Adds tokenize_chat_sft_sample / stack_lm_batch to examples/_helpers.py. First run (1x H100): chi_pos ~1e-9 — the SFT loss gradient sits deep in the spectral tail, as the dpo_spectral_filter hypothesis predicts for a late-SFT checkpoint; peak 84 GB, 195 s, autograd/memory-bound (tensor-core ~0%). Cross-pair alignment is left off (needs two gradients resident); CPU-offload is the planned follow-up.
…t RAM The cross-pair dot needs both batches' full P-sized gradients; holding the second on the GPU during a later batch's backward OOMs large models (~104 GB for a 7B). Add cross_grad_storage='auto'|'gpu'|'cpu' (default auto): auto offloads the cached gradients to host RAM when an on-GPU cache wouldn't leave room for a self-pair backward (the device-agnostic chunked _dot_fp64 then runs on CPU), and keeps small models on-device (no transfer, no numeric change — preserves the perspic cross-validation numerics). One-time ~29 GB host copy, negligible vs the analyze phase. Adds gpu-vs-cpu equivalence + validation tests; 54 unit tests pass, ruff + mypy --strict clean.
Expose from_pretrained's attn_implementation (e.g. 'sdpa', 'flash_attention_2', 'eager'). Default None keeps the model's own default. Useful where the attention backend changes the activation-memory class at long sequence lengths.
…sh outputs run.sbatch now computes the (english, code) cross pair with --cross-grad-storage auto (offloads the cached gradients to host RAM, dot on CPU; GPU peak stays ~84 GB) and exposes --attn-implementation (a no-op for OLMo3, which is O(S^2) under sdpa regardless). Refreshed results.parquet / resources.json / plots from the cross-enabled run. README trimmed to the running version: how-to-run (two-stage offline workflow), the analyzer config + rationale for adapting to your own model, and a resources-only findings section (wall ~210 s, peak ~84 GB, memory-bandwidth-bound). Documents the ~3k seq_len ceiling (O(S^2) attention) with gradient checkpointing as the future long-context lever. Drops the diagnostic mem_profile.md (keeps mem_profile.py as a pre-flight utility).
Move the auto offload size-arithmetic out of _resolve_cross_grad_storage into a pure, device-free _should_offload_cross_grads() so it is unit-testable. Add a near-threshold boundary test and a CUDA-gated test that exercises the real GPU->CPU offload (skipped in CI).
M1: default --chi-net-method to hutchinson (matches README/sbatch; per_sequence_cv OOMs at 7B). M4: tokenize_chat_sft_sample raises if truncation leaves zero completion tokens. m4: resource_tracker probes nvidia-smi before selecting it (else sampler='none'). Strip the dpo_spectral_filter research ref from the example README. Untrack samples.json/results.parquet/resources.json (regenerable; gitignored); keep the committed plots.
Bump 0.1.0 -> 0.2.0 (pyproject + _version.py); CHANGELOG [0.2.0] with date + compare link. README: document cross_grad_storage=auto and link the production-scale olmo_sft_analysis example. hf.py: note attn_implementation's O(S) win is model-dependent (OLMo3 stays O(S^2) under sdpa).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
vatis v0.2.0: cross-pair gradient offload + OLMo-3 SFT example
Adds
cross_grad_storage(auto/gpu/cpu) so cross-pair observables fit on large models — the cached gradients offload to host RAM and the dot runs on the CPU when two full gradients won't co-reside on the GPU; small models stay on-device unchanged. Also: anattn_implementationpassthrough inload_hf_model, and a fix so the analyzer only caches a batch's gradient when a cross pair needs it (the unconditional cache OOM'd large models).New example
examples/olmo_sft_analysis/: the LNP decomposition (incl. the code×English cross pair) onOlmo-3-7B-Think-SFTover real SFT chat data, with completion-only loss masking and wall-time / memory / GPU-activity tracking; a worked end-to-end demonstration on a real 7B checkpoint.Unit suite green (incl. a CUDA-gated offload-equivalence test); ruff + mypy clean.