From 4fb25404f9580c768132592021f9afc6939bf29b Mon Sep 17 00:00:00 2001 From: Alpamys Date: Mon, 3 Aug 2026 18:38:13 +0500 Subject: [PATCH] feat(streaming): preference losses over layer streaming (v0.72.4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DPO / ORPO / SimPO / KTO join task=sft on the layer-streaming engine. The reference model is the SAME streamed base with adapters disabled — one set of weights, one stream. Measured: streamed DPO peaks at 0.914x streamed SFT with a byte-identical store and pool, where forcing a real second instance costs +730.44 MB against 730.44 MB of weights. All four are bit-exact (0.0) against a resident run of the same loss. KTO is NOT reference-free: kto_trainer.py:466-476 is byte-for-byte DPO's three-branch reference selection, so it was gated separately. It also requires batch_size >= 2, refused at parse time rather than minutes into sharding. grpo/ppo stay excluded permanently — rollouts re-read every layer per generated token. The refusal deliberately names no release. The ~390-line streaming setup moved verbatim into trainer/stream_setup.py so five wrappers cannot drift. _STREAM_ROWS_PER_EXAMPLE is 2 for the concatenating losses and 1 for KTO, measured not assumed: the VRAM pre-flight would otherwise under-predict by half, and on Windows that is a silent spill, not an error. Honest cost: the reference is free in memory, not in time — DPO reads the layer stack 1.52x as often per step as SFT. Also closes five holes in the release checklist itself: - benchmarks/ was never in it, so gate records (which live under a gitignored .claude/) were never published. The public record behind the preprint DOI was about to fall a release behind; benchmarks/gate-v0.72.4-preference-losses.md and its index row are here, and the checklist now names the step. - tests/test_version_sync.py asserts pyproject.toml == __init__.py. Every other version test in the suite is a >= floor check, so bumping one and forgetting the other kept CI green. Verified red-green. - .claude/paper/ (the DOI preprint) had no "did this release change what it claims?" step. For v0.72.4 the answer is no: no measured number moves and its task: sft configs stay valid. - The Docs section header said steps 7-12 while containing 7-13, so plan.md sat outside its own section. - The README size anchor said ~238 lines against a real 426. Notes for whoever hits these next: - `pre-commit run --all-files` rewrites ~740 files here (ruff-format on pre-existing code). No pre-commit git hook is installed and CI runs only `ruff check`, so run it with --files on your own paths or the diff explodes. - Measuring streamed peak VRAM across setup() charges the pre-flight's own GEMM probe (three 4096^3 matrices, ~100 MB) to the step. Reset the peak counter after setup. - The buffer pool is freed by cycle collection, not by close(): back-to-back streamed runs in one process retain the previous pool (+47.65 MB measured) until a gc pass. Call gc.collect() between arms when measuring. - A resident model built from a float32 fixture vs a bf16 streamed one measures the dtype gap, not streaming — that cost an hour chasing a 9.96e-04 "failure". - Two concurrent pytest runs on a 4 GB card produce false CUDA failures; run the suite alone. Tests: 16977 -> 17051. --- CHANGELOG.md | 41 + CONTRIBUTING.md | 2 +- README.md | 38 +- benchmarks/README.md | 1 + benchmarks/gate-v0.72.4-preference-losses.md | 257 +++++ docs/commands.md | 4 +- docs/performance-and-quantization.md | 65 +- docs/training.md | 12 +- pyproject.toml | 2 +- src/soup_cli/__init__.py | 2 +- src/soup_cli/config/schema.py | 42 +- src/soup_cli/trainer/dpo.py | 32 +- src/soup_cli/trainer/kto.py | 31 +- src/soup_cli/trainer/orpo.py | 32 +- src/soup_cli/trainer/sft.py | 419 +------ src/soup_cli/trainer/simpo.py | 32 +- src/soup_cli/trainer/stream_setup.py | 453 ++++++++ src/soup_cli/utils/layer_stream.py | 15 + tests/test_v07200.py | 15 +- tests/test_v07201.py | 27 +- tests/test_v07204.py | 1056 ++++++++++++++++++ tests/test_version_sync.py | 82 ++ 22 files changed, 2174 insertions(+), 486 deletions(-) create mode 100644 benchmarks/gate-v0.72.4-preference-losses.md create mode 100644 src/soup_cli/trainer/stream_setup.py create mode 100644 tests/test_v07204.py create mode 100644 tests/test_version_sync.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 19ccf0b..453ff75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,47 @@ reproducing 70+ versions of notes. ## [Unreleased] +## [0.72.4] - 2026-08-03 + +**Added — preference losses over layer streaming: DPO, ORPO, SimPO and KTO.** + +Layer streaming kept the frozen base in CPU RAM and fed it to the GPU one decoder +layer at a time, but only for `task: sft`. This release opens it to the four +preference losses. The whole risk was one thing: **DPO needs a reference model, and a +second model instance would double memory and defeat the feature entirely.** + +- **The reference is the same streamed base with its adapters disabled** — one set of + weights, one stream, no second pass. Measured on an RTX 3050 4 GB with a 730 MB + model: streamed DPO peaked at **0.914x** the SFT peak, with a byte-identical RAM + store and buffer pool. Forcing a real second instance in the same harness cost + **+730.44 MB against 730.44 MB of weights** — exactly one copy. That control is what + makes the first number mean something. +- **KTO is *not* reference-free**, contrary to how it is usually described: it selects + its reference exactly the way DPO does, so it gets the same treatment and the same + memory assertion. ORPO and SimPO genuinely are reference-free. +- **Bit-exact against a resident run of the same loss** — `0.0` difference for all + four, the standard every slot in this series inherits. +- **The pre-flight now knows that a paired loss is twice the rows.** DPO, ORPO and + SimPO concatenate chosen and rejected into one tensor, so a VRAM budget computed at + one row per example would have under-predicted by half — and on Windows the + consequence is not an error but a silent spill to host memory that makes the run an + order of magnitude slower. +- **KTO requires `batch_size >= 2`** (its KL term is degenerate at 1). Soup now says so + when your config is read, rather than minutes later after sharding the checkpoint. + KTO is streamable at all only because v0.72.3 lifted the batch-1 restriction. +- **`grpo` and `ppo` remain excluded permanently**, not "not yet": generation rollouts + re-read every layer once per generated token, which destroys the amortisation + streaming depends on. The refusal says so and deliberately names no release. + +The streaming setup now lives in one shared place instead of being copied per trainer, +so the NF4 pre-flight, the RAM/disk tier choice and the VRAM fit refusal cannot drift +between SFT and the preference losses. + +Honest costs: streaming makes the reference free in **memory**, not in **time** — DPO +traverses the layer stack three times per step against SFT's two, measured at **1.52x** +the layer reads. And the VRAM pre-flight is a sound *upper* bound for preference +losses rather than a tight estimate; see Known Limitations in the release notes. + ## [0.72.3] - 2026-07-28 **Added — layer streaming breadth: more architectures, bigger batches, resume, and a diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f7931f1..0dd4f30 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -120,7 +120,7 @@ src/soup_cli/ templates/ - 21 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0, +4 compliance v0.71.35) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (330 files, 16977 tests) +tests/ - Test suite (332 files, 17051 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index 048a1e0..517a56a 100644 --- a/README.md +++ b/README.md @@ -52,25 +52,27 @@ infrastructure instead of improving models. Soup fixes that. ## What's New -**v0.72.3 — layer streaming grows up: more models, bigger batches, resume, and a disk -tier.** Layer streaming keeps the frozen base out of VRAM and feeds it to the GPU one -decoder layer at a time. v0.72.0–.2 kept the scope deliberately tiny to prove it worked; -this release removes the training wheels. +**v0.72.4 — align on a laptop: DPO, ORPO, SimPO and KTO over layer streaming.** Layer +streaming keeps the frozen base out of VRAM and feeds it to the GPU one decoder layer at +a time. It used to support supervised fine-tuning only; now it runs the preference +losses too. -- **Six more model families** — Mistral, Gemma / Gemma 2 / Gemma 3, and Phi / Phi-3 — each - verified **bit-exact** against the same checkpoint loaded resident, in bf16 *and* NF4. -- **`batch_size` above 1, gradient accumulation, and `--resume`** all work now. -- **A pre-flight that predicts peak VRAM and refuses a run that will not fit.** Streaming - bounds the *weights*; the logits tensor is not bounded by it and scales with - `batch × seq`. On a 152k-vocab model at batch 8 that single tensor measured **8.71 GB — - 146× the entire layer-buffer pool.** The prediction was fitted to ten real runs and - never under-predicts any of them. -- **A throughput forecast measured on your card, in your session**, quoted as a range - next to the SM clock it was taken at — not a number compiled into the source. -- **A disk overflow tier.** When the base will not fit in RAM, `stream_source: auto` - streams it from NVMe instead of refusing. Honest caveat: its *correctness* is verified - bit-exact against the RAM tier, but **how much slower it is has not been measured** on - the development hardware, and no figure is claimed. +- **DPO's reference model is free.** DPO needs a reference to compare against, and a + second copy of the model would double memory and defeat the whole point. Soup uses + *the same streamed base with its adapters switched off* — one set of weights, one + stream. Measured on an RTX 3050 4 GB: streamed DPO peaked at **0.914×** the + supervised-fine-tuning peak. Forcing a real second model in the same test cost + **+730 MB — exactly one copy of the weights.** +- **KTO is not reference-free**, however it is usually described: it picks its reference + the same way DPO does, so it gets the same treatment. ORPO and SimPO genuinely are. +- **Bit-exact against a normal, non-streamed run** of the same loss — `0.0` difference, + the bar every release in this series has to clear. +- **The VRAM pre-flight knows a paired loss is twice the rows**, because chosen and + rejected go through the model as one tensor. +- **Honest cost:** the reference is free in *memory*, not in *time* — DPO reads the + layer stack **1.52×** as often per step as supervised fine-tuning does. +- `grpo` / `ppo` stay excluded on purpose: generation re-reads every layer per token, + which is exactly what streaming cannot amortise. - Still BETA. ```yaml diff --git a/benchmarks/README.md b/benchmarks/README.md index 93a085c..b1a9043 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -18,6 +18,7 @@ They are the evidence behind the preprint: | [`gate-v0.72.0-layer-streaming.md`](gate-v0.72.0-layer-streaming.md) | The streaming path itself | Bit-exactness vs a resident reference; 3B bf16 trained on a 4 GB card | | [`gate-v0.72.2-nf4.md`](gate-v0.72.2-nf4.md) | NF4 quantised streaming | Llama-3.1-8B at 119.6 tok/s in a 3.32 GB peak | | [`gate-v0.72.3-breadth.md`](gate-v0.72.3-breadth.md) | Nine architectures, batching, accumulation, resume, disk tier | Peak-VRAM predictor at 0.85% worst-case error; accumulation is per-token I/O-neutral | +| [`gate-v0.72.4-preference-losses.md`](gate-v0.72.4-preference-losses.md) | DPO / ORPO / SimPO / KTO over the streaming engine | DPO's reference model costs no extra weights — 0.914x the SFT peak, against +730.44 MB for a real second instance | ## Hardware diff --git a/benchmarks/gate-v0.72.4-preference-losses.md b/benchmarks/gate-v0.72.4-preference-losses.md new file mode 100644 index 0000000..7dd3ed5 --- /dev/null +++ b/benchmarks/gate-v0.72.4-preference-losses.md @@ -0,0 +1,257 @@ + + +# v0.72.4 "Preference losses" — gate results + +Box: Windows 11 · RTX 3050 Laptop 4.29 GB · 16.9 GB RAM · NVMe · Python 3.10.8 · +torch 2.5.1+cu121 · transformers 4.57.6 · peft 0.18.1 · **trl 0.19.1**. + +Gates are throwaway scripts in the session scratchpad; nothing under `src/` was +written before a gate passed. + +> **Version caveat, stated up front.** CI runs trl 1.9.2 / peft 0.20.0 / +> torch 2.13.0. Everything below is measured against trl **0.19.1**. The gate +> therefore also records *which TRL internals the property depends on*, so the +> shipped tests assert the property rather than the internals. + +**The inherited standard** — a streamed run must be bit-exact against the +resident run of the same numerics; what changes per slot is the *reference*, not +the standard. This slot adds a second, independent standard, because bit-exactness +cannot see the failure that matters here: + +> DPO needs a reference model. Implemented naively as a second model instance it +> doubles memory and defeats the entire feature. It must be **the same streamed +> base with adapters disabled** — one set of weights, one stream, no second pass. +> Gate = assert peak VRAM under DPO is within the activation delta of peak VRAM +> under SFT, **not ~2x**. A passing loss curve does not detect this; only the +> memory assertion does. + +--- + +## Correction to the brief, established before gating + +The brief says "ORPO / SimPO / KTO are reference-free and ride the engine +unchanged." **That is true for ORPO and SimPO and false for KTO.** + +| loss | TRL trainer | reference model? | evidence | +|---|---|---|---| +| DPO | `DPOTrainer` | **yes** — implicit | `dpo_trainer.py:321-327`, `null_ref_context` at `:865-877` | +| KTO | `KTOTrainer` | **yes** — implicit | `kto_trainer.py:466-476`, `null_ref_context` at `:810`, used at `:926` | +| ORPO | `ORPOTrainer` | no | zero occurrences of `ref_model` in the file | +| SimPO | `CPOTrainer` | no | zero occurrences of `ref_model` in the file | + +`KTOTrainer.__init__` has byte-for-byte the same three-branch shape as DPO: + +```python +if ref_model: self.ref_model = ref_model +elif self.is_peft_model or args.precompute_ref_log_probs: + self.ref_model = None # <- the free path +else: self.ref_model = create_reference_model(model) # <- 2x +``` + +So **KTO inherits the whole trap** and was gated separately rather than waved +through as reference-free. + +The free path turns on exactly one predicate: `isinstance(model, PeftModel)`. +Soup's wrappers apply PEFT manually (`get_peft_model`) before constructing the +TRL trainer, so it holds — but it is a *predicate*, not a guarantee, which is why +check A1 asserts it directly instead of assuming it. + +--- + +## GATE 1 — DPO. **PASS 14/14** + +Two scripts, because the first one's memory arm was not evidence. + +### 1a — correctness battery, CPU float32, 4-layer Llama + +float32 on CPU so "bit-exact" means exactly `0.0`, not "within bf16 noise". + +| check | result | +|---|---| +| **A1** one model / one stream | PASS — `ref_model=None`, `is_peft_model=True`, `RamSource` constructed **exactly 1** time | +| **A3** reference ≠ policy | PASS — `max\|policy−ref\|` chosen **8.219452e-01**, rejected **5.105972e-01** | +| **A3-control** zero adapter ⇒ ref == policy | PASS — exactly **0.000e+00** | +| **A4** bit-exact vs **resident DPO** | PASS — `max\|loss_streamed − loss_resident\| = 0.000e+00` over 16 synced adapter tensors (streamed 0.67242944, resident 0.67242944) | +| **A5** layer-0 adapter gradient ≠ 0 | PASS — 1.276013e-01; all 4 layers have non-zero adapter grad | +| **A6** determinism, same seed twice | PASS — 0.000e+00 | + +**A3 is the check a loss curve cannot make.** The streamed layer substitutes base +weights through `functional_call` rather than the module's own `forward`, so +`disable_adapter()` being a no-op through that path is entirely plausible. If it +were, the reference would *be* the policy, every log-ratio would be 0, and the DPO +loss would sit at `−logsigmoid(0) = 0.6931` forever — which reads as "training +slowly", not as a bug. The zero-adapter control is what proves the 8.2e-01 gap +comes from the adapter and not from some other difference between the two forwards. + +### 1b — the memory assertion at a size where WEIGHTS dominate, CUDA bf16 + +1a's memory arm passed on a **0.30 MB** model, where a second copy of the base is +1.6% of the peak. **That is not evidence** — it would pass for an implementation +that *did* keep a second copy. Re-run at 365.2M params = **730.4 MB bf16**, +24 layers, vocab 260, seq 64, batch 1, so a second instance would dominate: + +``` +streamed SFT peak VRAM : 89.53 MB +streamed DPO peak VRAM : 81.87 MB (-7.66 MB, 0.914x) +CONTROL: DPO + 2nd model : 812.32 MB (+730.44 MB, 9.922x) +the base's own weights : 730.44 MB +RAM store, SFT vs DPO : 729.91 MB vs 729.91 MB +VRAM buffer pool, SFT/DPO : 60.83 MB vs 60.83 MB +``` + +| check | result | +|---|---| +| **B1** DPO−SFT delta < the base's weight bytes | PASS — −7.66 MB vs 730.44 MB, i.e. **−1.0%** of one copy | +| **B2** DPO peak is not ~2x SFT | PASS — **0.914x** (a second instance would be **9.16x**) | +| **B3** CONTROL: an explicit second model DOES cost ~the weights | PASS — **+730.44 MB**, 100.0% of the weight bytes | +| **B4** one store, one pool | PASS — 729.91 MB / 60.83 MB, byte-identical to SFT | + +**B3 is what gives B1/B2 teeth.** "DPO is not 2x SFT" means nothing unless the same +harness can show what 2x looks like; forcing `ref_model=` to a real second instance +moves the peak by **exactly one copy of the weights**. That is the cost the implicit +reference avoids, measured rather than argued. + +DPO's peak is *below* SFT's (0.914x), which is not an anomaly: this SFT arm computes +a loss over all 64 positions while the DPO arm splits the same 64 into a 32-token +prompt and a 32-token completion, so its logits tensor is smaller. The claim being +gated is "no second copy of the weights", and a **negative** delta is the strongest +form of it. + +### The I/O cost — published, not buried + +`layer loads/step: SFT=46 DPO=70 over 24 layers -> 1.52x (1.92 vs 2.92 per layer)`. + +Exactly the predicted ratio: SFT traverses the stack twice per step (forward + +checkpoint recompute); DPO traverses it three times (policy forward + reference +forward + recompute), so **3/2 = 1.5x**. Streaming makes the reference free in +*memory*, not in *time*. The 4-layer toy in 1a reported 2.00x, which is +boundary-dominated and is not quoted. + +--- + +## GATE 2 (ORPO / SimPO) + GATE 3 (KTO) — **PASS 10/11** + +Run against the **shipped** code, not a spike: `TrainerWrapper.setup()` → the real +TRL trainer → `train()`, for all five tasks. Model 24 layers × hidden 1024, +453.08 MB store, vocab 64, seq 64. + +``` +task peak (resident) store pool batch x rows predicted +sft 61.81 (47.79) 453.08 MB 37.76 MB 1 x 1 60.75 MB +dpo 61.82 (56.31) 453.08 MB 37.76 MB 1 x 2 64.74 MB +orpo 61.82 (56.31) 453.08 MB 37.76 MB 1 x 2 64.74 MB +simpo 61.82 (56.31) 453.08 MB 37.76 MB 1 x 2 64.74 MB +kto 61.82 (56.31) 453.08 MB 37.76 MB 2 x 1 64.74 MB +``` + +| check | result | +|---|---| +| **C1** ORPO / SimPO have no reference model at all | PASS — the `ref_model` attribute is **absent** on both trainers | +| **C2** ORPO / SimPO store+pool identical to SFT | PASS — 453.08 MB / 37.76 MB; peak **1.000x** SFT | +| **C3** KTO takes the implicit reference | PASS — `ref_model=None`, store identical to SFT's | +| **C4** CONTROL: forcing a second model DOES cost the weights | PASS — 61.82 → 515.04 MB, **+453.22 MB vs 453.08 MB** of weights | +| **C5** pre-flight never under-predicts (dpo/orpo/simpo/kto) | PASS — **+4.7%** each | +| **C5** pre-flight never under-predicts (sft) | **−1.7%** — see below | + +**KTO also required `batch_size >= 2`**, refused outright by TRL ("Actual (not +effective) batch size must be > 1. KTO will not work properly because the KL term +will be equivalent to the implied reward"). So KTO is streamable *only because +v0.72.3 lifted the batch-1 restriction* — under v0.72.0–.2 it could not have +shipped at all. + +### The C5 SFT miss (−1.7%) is a fixture artifact, with evidence + +It is on the **SFT** path, which this slot does not touch: the only change to the +budget multiplies `batch` by a per-task row count, which is **1** for SFT, so the +arithmetic is identical. The authoritative check is v0.72.3's peak-VRAM grid test — +**10 real measured runs**, within 1%, never under-predicting — which passes +unchanged. The −1.7% is 1.06 MB on a 61.81 MB peak at a fixture point far outside +the fitted envelope (vocab **64** vs the 49 152 / 151 936 it was fitted on), where +the constant terms dominate and the logits term is ~0. + +--- + +## Three measurement attempts that were INVALID, recorded so they are not repeated + +1. **Peak measured across `setup()`** charged the pre-flight's own GEMM ceiling + probe — three 4096³ bf16 matrices, ~100 MB transient — to the training step, and + reported a 44% "under-prediction" that did not exist. Reset the peak counter + *after* setup. +2. **Cross-run retention inflated every peak by one buffer pool.** The pool is held + by reference cycles in the module tree: measured `close()` alone retains + **+47.65 MB**, `close()` + `gc.collect()` retains **+0.00 MB**. Not a leak and + not a defect — but back-to-back streamed runs in one process hold the previous + pool until a cycle-collection pass. +3. **Two attempts to fit a preference-specific logits constant.** With short rows + `max_length` never bound the effective sequence — DPO measured an identical + 63.43 MB at seq 128, **256 and 512**; an independent variable that does not move + measures nothing. With long rows the batch shapes showed TRL was not truncating + to `max_length` either (`(1, 2401)` at every setting). **No constant was + published rather than fabricating one.** + +## The row multiplier is real, and had to be re-measured to prove it + +DPO, ORPO and SimPO build their forward through `concatenated_inputs` + +`torch.cat`, so **2 × batch_size** rows reach the model in ONE tensor; KTO runs its +KL batch as a **separate** forward, so its count is 1x. + +At vocab 64 the logits term is ~0, so 1x and 2x predict almost the same number: +**that check passed for both answers and was therefore not evidence.** Re-run at +vocab 32 000 the two are **+71.3% apart** (83.04 MB vs 142.22 MB), and all five +tasks still never under-predict. + +## The estimator is a sound UPPER bound for preference losses, and it is loose + +Charging `2 × batch` rows at `ForCausalLMLoss`'s measured 14 bytes/element is a +genuine upper bound — the concatenated forward *is* an SFT-shaped forward at twice +the rows, and TRL's preference losses reduce logits to per-token log-probs +(`selective_log_softmax`) instead of holding a full-vocab fp32 upcast, so their true +per-element cost is strictly lower. Measured: DPO's whole above-resident cost was +**51.76 MB** where the 14 B/elt charge for the same shape is ~458 MB. + +Practical boundary, computed exactly (Llama-3.2-1B, vocab 128 256, NF4, 3.32 GB +free — this card): + +``` +seq 512 SFT 1.65 GB allowed | DPO 2.63 GB allowed +seq 768 SFT 2.14 GB allowed | DPO 3.60 GB REFUSED +seq 1024 SFT 2.63 GB allowed | DPO 4.58 GB REFUSED +``` + +The conservatism does **not** block the realistic seq-512 configuration; the exact +point where it starts refusing runs that would probably fit is **seq 768 at batch 1**. +Shipped as-is because under-predicting is the strictly worse failure — on Windows it +is not an exception but a silent WDDM spill — and documented rather than papered over. + +--- + +## End-to-end, through the released CLI + +SmolLM2-135M, NF4, streamed (30 layers, 0.05 GB pinned RAM store, 2 × 2 MB VRAM +buffers), real `soup train`: + +| task | loss | note | +|---|---|---| +| dpo | 0.6931 → 0.6695 | starts at `−logsigmoid(0)` because `lora_B = 0` makes the reference equal the policy, then moves — the A3 property visible in a real run | +| orpo | 5.3816 → 5.0559 | | +| simpo | 5.2033 → 4.8749 | | +| kto | 0.5000 → 0.4904 (6 epochs) | its single-step run is flat at 0.5 for the same initialisation reason; checked over more steps rather than assumed | + +All four saved adapters: **120 tensors, 0 keys carrying the streaming wrapper +segment, 60/60 non-zero `lora_B`** — i.e. ordinary LoRA adapters that load into any +non-streaming model. + +Ten rejected configurations each named its own reason (rollout tasks, unsupported +task, KTO at batch 1, unsloth backend, 8-bit, no adapter, DoRA, `batch_size: auto`, +packing). diff --git a/docs/commands.md b/docs/commands.md index e3e0476..be57ec9 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -30,7 +30,7 @@ soup merge-sharded-fsdp-weights ./shards -o merged.safetensors Consolidate FSDP soup delinearize-llama4 ./src --target ./out [--num-experts N] [--plan-only] Live Llama-4 fused-expert reshape [E*din,dout] -> [E,din,dout] + sidecar copy (v0.71.21) soup spectrum scan --model --top-percent 50 [--modules mlp,attn] [-o patch.yaml] Spectrum SNR scan (no model load) -> training.unfrozen_parameters YAML patch (v0.71.23) soup train --config sft.yaml # training.lisa_enabled: true [lisa_num_layers lisa_interval_steps] LISA layerwise importance sampling — full-FT quality at LoRA-like memory (sft/transformers/text/quantization=none) (v0.71.34) -soup train --config sft.yaml # training.stream_layers: true [stream_source stream_buffers] BETA layer streaming — the frozen base streams from CPU RAM/NVMe one decoder layer at a time, so peak VRAM is bounded by ONE layer; quantization: 4bit streams it as NF4, ~4x smaller (sft/transformers/text, wider archs) (v0.72.0; NF4 v0.72.2; disk+batch+accum v0.72.3) +soup train --config sft.yaml # training.stream_layers: true [stream_source stream_buffers] BETA layer streaming — the frozen base streams from CPU RAM/NVMe one decoder layer at a time, so peak VRAM is bounded by ONE layer; quantization: 4bit streams it as NF4, ~4x smaller (sft/dpo/orpo/simpo/kto on transformers+text, 9 archs; grpo/ppo permanently excluded) (v0.72.0; NF4 v0.72.2; disk+batch+accum v0.72.3; preference losses v0.72.4) soup export --model ./output --format gguf Export to GGUF (Ollama) soup export --model ./output --deploy ollama Export GGUF + auto-deploy to Ollama soup export --model ./output --format onnx Export to ONNX @@ -337,5 +337,3 @@ exact command that would run — they never execute training or export. **Security:** stdio only (no network listener); every path argument stays under the working directory and rejects symlinks; tool output is control-char sanitized; error messages never leak filesystem paths. - - diff --git a/docs/performance-and-quantization.md b/docs/performance-and-quantization.md index ede53ef..0290d0f 100644 --- a/docs/performance-and-quantization.md +++ b/docs/performance-and-quantization.md @@ -14,7 +14,7 @@ - [Cross-Document Attention Masking](#cross-document-attention-masking) - [Quant Menu — 9 Quantization Formats](#quant-menu--9-quantization-formats) - [Activation Offloading (Small-VRAM Large-Batch)](#activation-offloading-small-vram-large-batch) -- [Layer Streaming (BETA, v0.72.0; NF4 v0.72.2; disk + wider archs v0.72.3)](#layer-streaming-beta-v0720-nf4-v0722-disk--wider-archs-v0723) +- [Layer Streaming (BETA, v0.72.0; NF4 v0.72.2; disk + wider archs v0.72.3)](#layer-streaming-beta-v0720-nf4-v0722-disk--wider-archs-v0723-preference-losses-v0724) - [Correctness First (v0.36.0)](#correctness-first-v0360) - [Multi-GPU / DeepSpeed / FSDP](#multi-gpu--deepspeed--fsdp) - [Performance + Long-Context](#performance--long-context) @@ -226,7 +226,7 @@ training: Not compatible with unsloth (own memory manager) or mlx. Wired across every transformer-backend trainer (SFT, DPO, GRPO, KTO, ORPO, SimPO, IPO, PPO, Reward-Model, Embedding, Pretrain). -## Layer Streaming (BETA, v0.72.0; NF4 v0.72.2; disk + wider archs v0.72.3) +## Layer Streaming (BETA, v0.72.0; NF4 v0.72.2; disk + wider archs v0.72.3; preference losses v0.72.4) Stream frozen base-model decoder layers ONE at a time from CPU RAM into small VRAM buffers instead of keeping the whole base resident. Peak VRAM is bounded by the size of a single layer, not the entire model — so models that don't fit resident on your GPU can now train at all. @@ -342,7 +342,8 @@ prints this advice when it sees you accumulating. - `batch_size: "auto"` → OOM-probes a resident model that streaming never loads; explicit batch sizes allowed (v0.72.3) - `quantization` other than `none` or `4bit` → other formats cannot be streamed into a pooled buffer - `backend: unsloth` / `backend: mlx` → streaming replaces the model-load path those backends own -- `task` other than `sft` → preference losses are v0.72.4 +- `task` other than `sft` / `dpo` / `orpo` / `simpo` / `kto` → named explicitly. `grpo` and `ppo` are refused **permanently**, not pending: generation rollouts re-read every layer once per generated token, which destroys the amortisation streaming depends on +- `task: kto` with `batch_size: 1` → TRL's KL term is degenerate at batch 1; refused when the config is read rather than minutes later after sharding - `lora.use_dora` / `lora.use_vera` / `lora.init_strategy` other than `random` → these initialise from the real base weight, which is on the meta device under streaming - `unfrozen_parameters`, `lisa_enabled`, `packing`, `multipack`, `use_fsdp2_compile`, `train_router_only`, `expand_layers` → each independently rewrites or re-freezes the same layers - `stream_source` / `stream_buffers` set while `stream_layers: false` → a footgun, refused @@ -393,15 +394,61 @@ output: ./output > If that prints anything, the adapter is affected. From v0.72.1 a streamed adapter is byte-for-byte in the same layout as an ordinary LoRA run. **Troubleshooting:** -- **"layer streaming needs the base to fit in RAM"** — the base is larger than free RAM. Free RAM or pick a smaller base; the disk overflow tier is v0.72.3. +- **"layer streaming needs the base to fit in RAM"** — the base is larger than free RAM. Set `stream_source: auto` to fall back to the NVMe disk tier, free RAM, or pick a smaller base. - **"could not page-lock the base … falling back to a PAGEABLE RAM store"** — expected on a busy machine. Training continues, more slowly. Close other applications to keep the pinned store. -- **"layer streaming does not support model_type=…"** — Llama and Qwen only for now; more architectures land in v0.72.3. +- **"layer streaming does not support model_type=…"** — the supported list is llama / qwen2 / qwen3 / mistral / gemma / gemma2 / gemma3_text / phi / phi3. Multimodal `gemma3` is excluded on purpose; use `gemma3_text`. - **Slower than you expected** — layer streaming trades time for memory. If the model already fits resident on your card, do not enable it. -**Roadmap (each refusal names its release):** -- Disk overflow tier, batch size > 1, gradient accumulation, checkpoint/resume, more architectures (Mistral / Gemma / Phi) — **v0.72.3** -- Preference losses (DPO / ORPO / SimPO / KTO) — **v0.72.4**. GRPO and PPO are explicitly **not** planned: rollouts need generation, which re-reads the model per token +### Preference losses over streaming (v0.72.4) + +`dpo`, `orpo`, `simpo` and `kto` stream exactly like `sft` — same config keys, same +pre-flight, same refusals. The interesting part is DPO's reference model. + +**DPO compares the model being trained against a frozen reference.** Implemented as a +second model instance that doubles memory and there is no point streaming at all. Soup +instead uses *the same streamed base with its LoRA adapters switched off*, so the +reference costs no extra weights. Measured on an RTX 3050 4 GB with a 730 MB model: + +| arm | peak VRAM | vs SFT | +|---|---|---| +| streamed SFT | 89.53 MB | — | +| **streamed DPO** | **81.87 MB** | **0.914×** | +| the same run forced to build a real second model | 812.32 MB | 9.92× | + +The third row is the control: a second instance costs **+730.44 MB against 730.44 MB of +weights**, i.e. exactly one copy. The RAM store and the VRAM buffer pool are +byte-identical between the SFT and DPO arms. + +**KTO is not reference-free**, however it is usually described — it selects a reference +the same way DPO does, so it gets the same treatment. ORPO and SimPO genuinely are +reference-free. All four are verified **bit-exact** against a resident run of the same +loss. + +**The cost is time, not memory.** DPO runs the layer stack three times per step (policy +forward, reference forward, checkpoint recompute) against SFT's two — measured **1.52×** +the layer reads on a 24-layer model. Streaming makes the reference free in memory; it +does not make it free. + +**Two things to know before you configure it:** + +- **`kto` needs `batch_size: 2` or more.** TRL's KL term is degenerate at batch 1, so + the run cannot work; Soup refuses it when your config is read rather than after + sharding the checkpoint. (KTO is streamable at all only because v0.72.3 lifted + streaming's own batch-1 restriction.) +- **The VRAM pre-flight is deliberately conservative for paired losses.** DPO, ORPO and + SimPO send chosen and rejected through the model as one tensor, so the budget charges + twice the rows — correct, and it never under-predicts. But it charges them at the + *supervised* loss's measured per-element rate, and TRL's preference losses use a + cheaper path, so the estimate is an upper bound rather than a tight one. Concretely, + on a 4 GB card with a 128k-vocab 1B model: DPO at `max_length: 512` is allowed, and + from `max_length: 768` up it is refused even though it would probably fit. Lower + `max_length` if you hit that. (Tracked as a follow-up; under-predicting would be the + strictly worse failure, because on Windows it is not an error but a silent spill to + host memory.) + +**Roadmap:** - A published 14B-on-8 GB reference benchmark — hardware-blocked; it needs an 8 GB card and 32 GB of RAM, which the development box does not have +- GRPO and PPO are explicitly **not** planned: rollouts need generation, which re-reads the model per token **Shard cache.** The first streaming run rewrites the checkpoint into one safetensors shard per decoder layer under `~/.soup/layer-stream/` (override with `SOUP_LAYER_STREAM_CACHE_DIR`). That costs disk space roughly equal to the base. The cache is keyed to a fingerprint of the source checkpoint, so a base retrained in place re-shards instead of silently training against stale weights. @@ -708,5 +755,3 @@ Autopilot also detects pre-quantized bases automatically — `TheBloke/Llama-2-7 The advanced GGUF pipeline uses POSIX `O_NOFOLLOW` to defeat the TOCTOU race between the dispatch-time symlink check and the actual open of the calibration data — a crafted environment cannot race-swap the calibration file between validate and read. `soup deploy autopilot --measure` caches results at `~/.soup/deploy_autopilot_cache.json` keyed on `(base, profile, eval-tasks)`. Repeat invocations short-circuit; pass `SOUP_DEPLOY_AUTOPILOT_CACHE=` to redirect (constrained to home / cwd / tempdir). The recommended candidate uses soft-fallback: first `OK` by insertion order, else the candidate with the smallest delta (least drop relative to its own baseline). - - diff --git a/docs/training.md b/docs/training.md index 062aa81..d39c8da 100644 --- a/docs/training.md +++ b/docs/training.md @@ -5,10 +5,12 @@ > SFT, DPO/GRPO/PPO/KTO/ORPO/SimPO/IPO/BCO, tool-calling, PRM, pre-training, distillation, classification, vision/audio/TTS, unlearning, RAFT/RA-DIT, and the loop-hardening detectors. > **Training a model bigger than your GPU?** `training.stream_layers: true` streams the -> frozen base from CPU RAM (with NVMe disk overflow in v0.72.3) one decoder layer at a time, -> so peak VRAM is bounded by one layer instead of the whole model. Add `quantization: 4bit` -> and an 8B base fits a 4 GB card — see -> [Layer Streaming](performance-and-quantization.md#layer-streaming-beta-v0720-nf4-v0722-disk--wider-archs-v0723). +> frozen base from CPU RAM (with NVMe disk overflow) one decoder layer at a time, so peak +> VRAM is bounded by one layer instead of the whole model. Add `quantization: 4bit` and an +> 8B base fits a 4 GB card. Works for `sft` and, from v0.72.4, for `dpo` / `orpo` / +> `simpo` / `kto` — DPO's reference model is the same streamed base with its adapters +> switched off, so it costs no extra weights — see +> [Layer Streaming](performance-and-quantization.md#layer-streaming-beta-v0720-nf4-v0722-disk--wider-archs-v0723-preference-losses-v0724). **Contents:** @@ -1444,5 +1446,3 @@ a conflicting combo is rejected loudly at config load. Scans cache under `~/.sou `all` default) is recommended for very large models — it skips the giant embedding/lm_head matrices. The SNR kernel is pure-numpy and transpose-invariant, so GPT-2 `Conv1D` weights score the same as Linear weights. (v0.71.23) - - diff --git a/pyproject.toml b/pyproject.toml index abc1088..2a0941a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.72.3" +version = "0.72.4" description = "Fine-tune and post-train LLMs in one command. No SSH, no config hell." readme = "README.md" license = "Apache-2.0" diff --git a/src/soup_cli/__init__.py b/src/soup_cli/__init__.py index 6d57450..d5becdc 100644 --- a/src/soup_cli/__init__.py +++ b/src/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune and post-train LLMs in one command.""" -__version__ = "0.72.3" +__version__ = "0.72.4" diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index c59c69e..4b0e809 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -12,6 +12,12 @@ from soup_cli.utils.layer_stream import ( MAX_STREAM_BUFFERS, MIN_STREAM_BUFFERS, ) +from soup_cli.utils.layer_stream import ( + ROLLOUT_STREAM_TASKS as _STREAM_ROLLOUT_TASKS, +) +from soup_cli.utils.layer_stream import ( + SUPPORTED_STREAM_TASKS as _STREAM_SUPPORTED_TASKS, +) # v0.39.0 Part C — per-pattern LoRA rank/alpha bounds _MAX_LORA_RANK_PATTERN_KEYS = 256 @@ -4658,11 +4664,22 @@ class SoupConfig(BaseModel): "the base layer-by-layer." ) return self - if self.task != "sft": + # v0.72.4 — the four preference losses join SFT. DPO and KTO take their + # reference from the SAME streamed base with adapters disabled (TRL's + # `null_ref_context`), so the reference costs no extra weights: measured + # 0.914x SFT peak VRAM where a second instance was 9.92x. + if self.task in _STREAM_ROLLOUT_TASKS: raise ValueError( - f"training.stream_layers requires task='sft'; got " - f"task={self.task!r}. Preference losses (DPO/ORPO/SimPO/KTO) " - f"land in v0.72.4." + f"training.stream_layers cannot be used with task={self.task!r}: " + f"it needs generation rollouts, which re-read every layer once " + f"per generated token. That destroys the amortisation streaming " + f"depends on (one weight read per step, not per token), so this " + f"is a permanent exclusion rather than an unimplemented one." + ) + if self.task not in _STREAM_SUPPORTED_TASKS: + raise ValueError( + f"training.stream_layers supports task in " + f"{sorted(_STREAM_SUPPORTED_TASKS)}; got task={self.task!r}." ) if self.backend != "transformers": raise ValueError( @@ -4706,6 +4723,23 @@ class SoupConfig(BaseModel): # pre-flight says so rather than the schema refusing it, because # accumulation is the ONLY way to raise effective batch once the VRAM # budget is exhausted (peak moved 0.842 -> 0.846 GB across accum 1->4). + # v0.72.4 — KTO's KL term is degenerate at a per-device batch of 1, so + # TRL refuses it outright ("Actual (not effective) batch size must be + # > 1"). Under streaming that ValueError arrives only AFTER the RAM + # pre-flight, the checkpoint sharding and — at quantization='4bit' — + # the NF4 quantisation pass: minutes of disk I/O on a real base, to + # fail on a config that was already invalid. Refuse it at parse time. + if ( + self.task == "kto" + and isinstance(tcfg.batch_size, int) + and tcfg.batch_size < 2 + ): + raise ValueError( + "task='kto' requires training.batch_size >= 2 (TRL's KL term is " + "degenerate at batch 1). Checked here rather than in the " + "trainer so a streaming run fails before sharding the " + "checkpoint, not minutes into it." + ) if tcfg.lora.r < 1: raise ValueError( "training.stream_layers requires LoRA (training.lora.r >= 1) — " diff --git a/src/soup_cli/trainer/dpo.py b/src/soup_cli/trainer/dpo.py index 3326894..ffb9cde 100644 --- a/src/soup_cli/trainer/dpo.py +++ b/src/soup_cli/trainer/dpo.py @@ -7,12 +7,13 @@ from typing import Optional from rich.console import Console from soup_cli.config.schema import SoupConfig +from soup_cli.trainer.stream_setup import StreamingSetupMixin from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name console = Console() -class DPOTrainerWrapper: +class DPOTrainerWrapper(StreamingSetupMixin): """High-level wrapper for DPO training from SoupConfig. DPO requires preference data with three fields: @@ -21,6 +22,11 @@ class DPOTrainerWrapper: - rejected: the less preferred response """ + #: TRL builds this loss's forward through ``concatenated_inputs`` + + #: ``torch.cat``, so chosen and rejected arrive as ONE tensor of twice + #: the configured batch. The VRAM pre-flight must budget for that. + _STREAM_ROWS_PER_EXAMPLE = 2 + def __init__( self, config: SoupConfig, @@ -69,14 +75,28 @@ class DPOTrainerWrapper: cfg = self.config tcfg = cfg.training use_unsloth = cfg.backend == "unsloth" + # v0.72.4 — layer streaming replaces the model-load path entirely (meta + # skeleton, never a resident load), so it dispatches ahead of the backend + # branches. The schema already rejects streaming + unsloth/mlx. + use_streaming = bool(getattr(tcfg, "stream_layers", False)) - if use_unsloth: + if use_streaming: + self._setup_streaming_transformers(cfg, tcfg) + elif use_unsloth: self._setup_unsloth(cfg, tcfg) else: self._setup_transformers(cfg, tcfg) trainable, total = self.model.get_nb_trainable_parameters() - pct = 100 * trainable / total + # v0.72.4 (mirrors sft.py) — under NF4 streaming PEFT's total is wrong + # by ~6.5x: it sizes Params4bit as `numel * 2 * quant_storage.itemsize`, + # right for a RESIDENT one but not for our `meta` placeholder, which + # still carries the LOGICAL shape. The sharder counted the real source + # elements, so prefer that. + stream_total = getattr(self._stream_runtime, "total_params", 0) + if stream_total: + total = stream_total + pct = 100 * trainable / total if total else 0.0 console.print( f"[green]LoRA applied:[/] {trainable:,} trainable" f" / {total:,} total ({pct:.2f}%)" @@ -311,8 +331,10 @@ class DPOTrainerWrapper: from soup_cli.utils.v028_features import activation_offloading_context - with activation_offloading_context( - self.config.training, self._output_dir, + # v0.72.4 — the shared context releases the streaming weight source even + # if training raises (see StreamingSetupMixin._training_context). + with self._training_context( + activation_offloading_context(self.config.training, self._output_dir) ): self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) duration = time.time() - start diff --git a/src/soup_cli/trainer/kto.py b/src/soup_cli/trainer/kto.py index 960f210..335a0aa 100644 --- a/src/soup_cli/trainer/kto.py +++ b/src/soup_cli/trainer/kto.py @@ -8,12 +8,13 @@ from typing import Optional from rich.console import Console from soup_cli.config.schema import SoupConfig +from soup_cli.trainer.stream_setup import StreamingSetupMixin from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name console = Console() -class KTOTrainerWrapper: +class KTOTrainerWrapper(StreamingSetupMixin): """High-level wrapper for KTO training from SoupConfig. KTO uses unpaired preference data with three fields: @@ -22,6 +23,10 @@ class KTOTrainerWrapper: - label: True (desirable) or False (undesirable) """ + #: KTO runs its KL batch as a SEPARATE forward rather than concatenating + #: it (``kto_trainer.py`` :939-977), so the row count per forward is 1x. + _STREAM_ROWS_PER_EXAMPLE = 1 + def __init__( self, config: SoupConfig, @@ -67,14 +72,28 @@ class KTOTrainerWrapper: cfg = self.config tcfg = cfg.training use_unsloth = cfg.backend == "unsloth" + # v0.72.4 — layer streaming replaces the model-load path entirely (meta + # skeleton, never a resident load), so it dispatches ahead of the backend + # branches. The schema already rejects streaming + unsloth/mlx. + use_streaming = bool(getattr(tcfg, "stream_layers", False)) - if use_unsloth: + if use_streaming: + self._setup_streaming_transformers(cfg, tcfg) + elif use_unsloth: self._setup_unsloth(cfg, tcfg) else: self._setup_transformers(cfg, tcfg) trainable, total = self.model.get_nb_trainable_parameters() - pct = 100 * trainable / total + # v0.72.4 (mirrors sft.py) — under NF4 streaming PEFT's total is wrong + # by ~6.5x: it sizes Params4bit as `numel * 2 * quant_storage.itemsize`, + # right for a RESIDENT one but not for our `meta` placeholder, which + # still carries the LOGICAL shape. The sharder counted the real source + # elements, so prefer that. + stream_total = getattr(self._stream_runtime, "total_params", 0) + if stream_total: + total = stream_total + pct = 100 * trainable / total if total else 0.0 console.print( f"[green]LoRA applied:[/] {trainable:,} trainable" f" / {total:,} total ({pct:.2f}%)" @@ -292,8 +311,10 @@ class KTOTrainerWrapper: from soup_cli.utils.v028_features import activation_offloading_context - with activation_offloading_context( - self.config.training, self._output_dir, + # v0.72.4 — the shared context releases the streaming weight source even + # if training raises (see StreamingSetupMixin._training_context). + with self._training_context( + activation_offloading_context(self.config.training, self._output_dir) ): self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) duration = time.time() - start diff --git a/src/soup_cli/trainer/orpo.py b/src/soup_cli/trainer/orpo.py index 47a3154..e56930c 100644 --- a/src/soup_cli/trainer/orpo.py +++ b/src/soup_cli/trainer/orpo.py @@ -8,12 +8,13 @@ from typing import Optional from rich.console import Console from soup_cli.config.schema import SoupConfig +from soup_cli.trainer.stream_setup import StreamingSetupMixin from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name console = Console() -class ORPOTrainerWrapper: +class ORPOTrainerWrapper(StreamingSetupMixin): """High-level wrapper for ORPO training from SoupConfig. ORPO uses preference data (same format as DPO) but does not @@ -25,6 +26,11 @@ class ORPOTrainerWrapper: - rejected: the less preferred response """ + #: TRL builds this loss's forward through ``concatenated_inputs`` + + #: ``torch.cat``, so chosen and rejected arrive as ONE tensor of twice + #: the configured batch. The VRAM pre-flight must budget for that. + _STREAM_ROWS_PER_EXAMPLE = 2 + def __init__( self, config: SoupConfig, @@ -69,14 +75,28 @@ class ORPOTrainerWrapper: cfg = self.config tcfg = cfg.training use_unsloth = cfg.backend == "unsloth" + # v0.72.4 — layer streaming replaces the model-load path entirely (meta + # skeleton, never a resident load), so it dispatches ahead of the backend + # branches. The schema already rejects streaming + unsloth/mlx. + use_streaming = bool(getattr(tcfg, "stream_layers", False)) - if use_unsloth: + if use_streaming: + self._setup_streaming_transformers(cfg, tcfg) + elif use_unsloth: self._setup_unsloth(cfg, tcfg) else: self._setup_transformers(cfg, tcfg) trainable, total = self.model.get_nb_trainable_parameters() - pct = 100 * trainable / total + # v0.72.4 (mirrors sft.py) — under NF4 streaming PEFT's total is wrong + # by ~6.5x: it sizes Params4bit as `numel * 2 * quant_storage.itemsize`, + # right for a RESIDENT one but not for our `meta` placeholder, which + # still carries the LOGICAL shape. The sharder counted the real source + # elements, so prefer that. + stream_total = getattr(self._stream_runtime, "total_params", 0) + if stream_total: + total = stream_total + pct = 100 * trainable / total if total else 0.0 console.print( f"[green]LoRA applied:[/] {trainable:,} trainable" f" / {total:,} total ({pct:.2f}%)" @@ -290,8 +310,10 @@ class ORPOTrainerWrapper: from soup_cli.utils.v028_features import activation_offloading_context - with activation_offloading_context( - self.config.training, self._output_dir, + # v0.72.4 — the shared context releases the streaming weight source even + # if training raises (see StreamingSetupMixin._training_context). + with self._training_context( + activation_offloading_context(self.config.training, self._output_dir) ): self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) duration = time.time() - start diff --git a/src/soup_cli/trainer/sft.py b/src/soup_cli/trainer/sft.py index 06a3382..0efbf89 100644 --- a/src/soup_cli/trainer/sft.py +++ b/src/soup_cli/trainer/sft.py @@ -2,7 +2,6 @@ import json import logging -import math import os import time from pathlib import Path @@ -11,6 +10,7 @@ from typing import Optional, Tuple from rich.console import Console from soup_cli.config.schema import SoupConfig +from soup_cli.trainer.stream_setup import StreamingSetupMixin from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name logger = logging.getLogger(__name__) @@ -147,7 +147,7 @@ def _maybe_load_pretokenized( return train_ds, eval_ds -class SFTTrainerWrapper: +class SFTTrainerWrapper(StreamingSetupMixin): """High-level wrapper that sets up model + tokenizer + trainer from SoupConfig.""" def __init__( @@ -234,7 +234,7 @@ class SFTTrainerWrapper: # but not for our `meta` placeholder, which still carries the LOGICAL # shape. Measured on SmolLM2-135M: 878,154,048 vs a true 134,515,008. # The sharder counted the real source elements, so use that. - stream_total = getattr(getattr(self, "_stream_runtime", None), "total_params", 0) + stream_total = getattr(self._stream_runtime, "total_params", 0) if stream_total: total = stream_total pct = 100 * trainable / total if total else 0.0 @@ -966,398 +966,6 @@ class SFTTrainerWrapper: self._apply_quantization_aware(tcfg) - def _setup_streaming_transformers(self, cfg, tcfg): - """v0.72.0 BETA — layer streaming. The resident base load NEVER happens. - - Builds the skeleton on ``meta`` (``accelerate.init_empty_weights``), - materialises only embeddings / final norm / LoRA, and streams each - decoder layer from CPU RAM into a small pool of pre-allocated VRAM - buffers. Peak VRAM becomes the size of ONE layer instead of the model. - """ - from dataclasses import replace - - from peft import LoraConfig, TaskType - from transformers import AutoConfig, AutoTokenizer - - from soup_cli.utils.layer_shard import ( - QUANT_NF4, - QUANT_NONE, - resolve_shard_dir, - shard_checkpoint, - source_weight_bytes, - ) - from soup_cli.utils.layer_stream import ( - RAM_TIER_HEADROOM, - TIER_DISK, - TIER_RAM, - build_stream_plan, - detect_disk_kind, - dtype_bytes, - estimate_stream_store_bytes, - free_ram_bytes, - render_stream_panel, - stream_arch_of, - ) - from soup_cli.utils.layer_stream_runtime import ( - RamSource, - build_meta_skeleton, - build_streamed_model, - extras_resident_bytes, - probe_expandable_segments, - quantised_layer_suffixes, - ) - from soup_cli.utils.spectrum_scan import resolve_model_weights - - console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") - self.tokenizer = AutoTokenizer.from_pretrained( - cfg.base, trust_remote_code=self._trust_remote_code - ) - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - model_config = AutoConfig.from_pretrained( - cfg.base, trust_remote_code=self._trust_remote_code - ) - # Allowlist, not a heuristic — a half-supported architecture streams - # weights into the wrong module and mis-trains silently. - arch = stream_arch_of(model_config) - - on_cuda = str(self.device).startswith("cuda") - dtype = "bfloat16" if on_cuda else "float32" - - # v0.72.2 — NF4. The decoder linears ship as packed nibbles + per-block - # absmax, so the RAM store is ~0.26x its bf16 size; embeddings, norms and - # an untied head stay at `dtype`, exactly as replace_with_bnb_linear - # leaves them. - quant = QUANT_NF4 if tcfg.quantization == "4bit" else QUANT_NONE - - weights_dir = resolve_model_weights(cfg.base) - shard_dir = resolve_shard_dir(cfg.base) - - # Cheap size probe BEFORE sharding: re-writing a checkpoint we are - # about to refuse for not fitting in RAM costs minutes of disk I/O. - # Charged at the STREAMED rate, not the on-disk one — an 8B bf16 - # checkpoint is 16 GB on disk but only ~4.2 GB of NF4 store, and - # comparing the raw file size would refuse exactly the runs NF4 enables. - early_free_ram = free_ram_bytes() - if early_free_ram is not None: - source_bytes = source_weight_bytes(weights_dir) - store_estimate = estimate_stream_store_bytes( - source_bytes, dtype=dtype, quant=quant - ) - if ( - store_estimate >= early_free_ram * RAM_TIER_HEADROOM - and tcfg.stream_source == "ram" - ): - as_streamed = ( - "" - if quant == QUANT_NONE - else f" ({store_estimate / 1e9:.1f} GB once quantised to NF4)" - ) - raise ValueError( - f"training.stream_source='ram' but {cfg.base} is " - f"{source_bytes / 1e9:.1f} GB on disk{as_streamed} and only " - f"{early_free_ram / 1e9:.1f} GB of RAM is free. Set " - f"stream_source='auto' to fall back to the NVMe disk tier, " - f"free RAM, or pick a smaller base." - ) - - # The authoritative list of weights to quantise is whatever - # replace_with_bnb_linear actually converts, read off a meta skeleton — - # not a hard-coded name list that would drift per architecture. - # - # This builds a second, throwaway skeleton (build_streamed_model makes - # its own). Deliberate: a meta skeleton allocates NO weight storage, so - # the cost is module-tree construction only, and threading a pre-built - # model into build_streamed_model would couple suffix discovery to model - # construction for no memory saving. - quant_suffixes = () - if quant == QUANT_NF4: - probe = build_meta_skeleton( - cfg.base, - dtype=dtype, - quant=quant, - trust_remote_code=self._trust_remote_code, - ) - quant_suffixes = quantised_layer_suffixes(probe) - del probe - - console.print(f"[dim]Preparing layer shards -> {shard_dir}[/]") - index = shard_checkpoint( - weights_dir, - shard_dir, - dtype=dtype, - arch=arch, - quant=quant, - quant_suffixes=quant_suffixes, - # Quantise on the device that will run the model: CPU and CUDA agree - # on the packed nibbles but not on every float32 nested statistic. - quant_device=str(self.device), - ) - - spec = RamSource.spec_from_shard(shard_dir) - # Measured from the shard headers, not derived from `total_params`: - # under NF4 a layer holds packed uint8 alongside float32 statistics, so - # element counts no longer convert to bytes at a single rate. - layer_bytes = sum( - math.prod(shape) * dtype_bytes(stored) for shape, stored in spec.values() - ) - embed_bytes = extras_resident_bytes(shard_dir) - - free_ram = free_ram_bytes() - if free_ram is None: - console.print( - "[yellow]psutil unavailable — cannot size the RAM tier; " - "proceeding and letting the allocation fail loudly if it must[/]" - ) - free_ram = (layer_bytes * index.n_layers + embed_bytes) * 10 - - store_total = layer_bytes * index.n_layers + embed_bytes - # Checked BEFORE build_stream_plan so a `ram`-only run is refused with - # the message about stream_source rather than choose_tier's generic - # "needs NVMe or more RAM" — and without paying the ~9 s disk probe for - # an answer that cannot change the outcome. - if tcfg.stream_source == "ram" and store_total >= free_ram * RAM_TIER_HEADROOM: - raise ValueError( - f"training.stream_source='ram' but the base is " - f"{store_total / 1e9:.1f} GB and only {free_ram / 1e9:.1f} GB of " - f"RAM is free. Set stream_source='auto' to fall back to the NVMe " - f"disk tier, free RAM, or pick a smaller base." - ) - plan = build_stream_plan( - arch=arch, - n_layers=index.n_layers, - layer_bytes=layer_bytes, - embed_bytes=embed_bytes, - available_ram_bytes=free_ram, - # The page-locked ceiling is a property of the box, not of free RAM; - # rather than probe it destructively we attempt the pinned store and - # fall back loudly (see layer_stream_runtime._build_source). - pinned_limit_bytes=None, - buffers=tcfg.stream_buffers, - # v0.72.3: the REAL media type, not a constant. Passed as a callable - # because probing costs ~9 s on Windows and the answer only matters - # when the base does not fit in RAM. - disk_kind=lambda: detect_disk_kind(shard_dir), - ) - # v0.72.3 — the disk overflow tier is live, so a base that does not fit - # in RAM is no longer fatal. `stream_source` decides: 'ram' insists, - # 'disk' forces, 'auto' (the default) takes RAM when it fits and falls - # back to disk when it does not. build_stream_plan already refused a - # non-NVMe disk, so reaching here with tier='disk' means NVMe. - tier = TIER_DISK if tcfg.stream_source == "disk" else plan.tier - if tier != plan.tier: - # The panel is rendered from `plan`, so a forced tier has to be - # reflected there or the pre-flight reports "tier ram" immediately - # before the runtime announces it is streaming from disk. Every - # field that describes the RAM store is corrected with it, so no - # consumer can read a stale value. - plan = replace( - plan, - tier=tier, - store_bytes=0, - pinned=False, - notes=plan.notes - + ( - "streaming from disk because stream_source='disk' was set, " - "not because RAM was short. Nothing is held resident, and " - "the slowdown versus the RAM tier is unmeasured on this " - "hardware.", - ), - ) - # v0.72.3 — VRAM pre-flight. Streaming bounds the WEIGHTS; activations - # and the logits tensor are untouched by it and both scale with batch x - # seq. On a large-vocab model the logits term alone dwarfs the buffer - # pool (measured: 146x at batch 8), so a plan that reports only tier and - # buffer sizes will happily green-light a config that cannot run. - forecast_lines = self._stream_budget_lines( - cfg, - tcfg, - model_config=model_config, - layer_bytes=layer_bytes, - embed_bytes=embed_bytes, - index=index, - on_cuda=on_cuda, - ) - console.print(render_stream_panel(plan, forecast_lines)) - console.print( - "[yellow]Layer streaming is BETA:[/] slower than resident training, " - "but this model may not run resident on this card at all." - ) - if on_cuda and not probe_expandable_segments(): - console.print( - "[dim]expandable_segments allocator hint is unavailable on this " - "platform (silently ignored on Windows) — not enabled[/]" - ) - - target_modules = tcfg.lora.target_modules - if target_modules == "auto": - target_modules = None - lora_config = LoraConfig( - r=tcfg.lora.r, - lora_alpha=tcfg.lora.alpha, - lora_dropout=tcfg.lora.dropout, - target_modules=target_modules, - task_type=TaskType.CAUSAL_LM, - bias="none", - use_dora=tcfg.lora.use_dora, - use_rslora=tcfg.lora.use_rslora, - ) - - model, runtime = build_streamed_model( - model_id=cfg.base, - shard_dir=shard_dir, - index=index, - lora_config=lora_config, - device=self.device, - dtype=dtype, - buffers=tcfg.stream_buffers, - pin=plan.pinned and on_cuda, - seed=tcfg.seed if getattr(tcfg, "seed", None) is not None else 0, - trust_remote_code=self._trust_remote_code, - console=console, - quant=quant, - tier=tier, - ) - self.model = model - self._stream_runtime = runtime - stats = runtime.stats() - if stats["tier"] == TIER_RAM: - source_line = ( - f"{stats['store_bytes'] / 1e9:.2f} GB " - f"{'pinned' if stats['pinned'] else 'pageable'} RAM store" - ) - else: - source_line = ( - f"streamed from DISK ({stats['disk_bytes'] / 1e9:.2f} GB on an " - f"NVMe volume, nothing held resident)" - ) - buffer_line = ( - f"{stats['buffers']} x " - f"{stats['buffer_bytes'] / stats['buffers'] / 1e6:.0f} MB VRAM buffers" - ) - console.print( - f"[green]Layer streaming ready:[/] {stats['n_layers']} layers, " - f"{source_line}, {buffer_line}" - ) - - def _close_stream_runtime(self) -> None: - """Release the streaming weight source, if this run had one.""" - runtime = getattr(self, "_stream_runtime", None) - if runtime is not None: - runtime.close() - - def _estimate_adapter_params(self, tcfg, model_config) -> int: - """Trainable adapter parameters, before the model exists. - - Deliberately coarse and biased HIGH: it assumes every targeted module is - hidden x hidden. Gate/up/down projections are larger, but the whole - adapter term is ~0.5% of a streaming step's peak, so precision here buys - nothing while under-counting would eat into the safety margin. - """ - hidden = int(getattr(model_config, "hidden_size", 0) or 0) - layers = int(getattr(model_config, "num_hidden_layers", 0) or 0) - targets = tcfg.lora.target_modules - n_targets = len(targets) if isinstance(targets, (list, tuple)) else 4 - return layers * n_targets * 2 * tcfg.lora.r * hidden - - def _stream_budget_lines( - self, cfg, tcfg, *, model_config, layer_bytes, embed_bytes, index, on_cuda - ): - """Predict peak VRAM + bracket throughput, and REFUSE a run that cannot fit. - - Returns the extra lines for the pre-flight panel. Raises when the step is - predicted not to fit: on Linux that would be a hard OOM, and on Windows - something worse — WDDM spills to host memory without raising, so the run - silently becomes an order of magnitude slower and looks like the feature - is merely slow. - """ - from soup_cli.utils.layer_stream import ( - accumulation_advice, - decide_stream_fit, - estimate_logits_bytes, - estimate_stream_peak_vram, - forecast_stream_throughput, - ) - from soup_cli.utils.layer_stream_runtime import measure_gemm_tflops - - vocab = int(getattr(model_config, "vocab_size", 0) or 0) - hidden = int(getattr(model_config, "hidden_size", 0) or 0) - inter = int(getattr(model_config, "intermediate_size", 0) or 0) - seq_len = int(cfg.data.max_length) - batch = tcfg.batch_size if isinstance(tcfg.batch_size, int) else 1 - if not (vocab and hidden and inter): - # Never silently: skipping the budget also skips the refusal that - # stops a run from OOMing (or, on Windows, spilling to host memory - # and running an order of magnitude slower with no error at all). - console.print( - "[yellow]Layer streaming could not read vocab_size / hidden_size " - "/ intermediate_size from the model config, so peak VRAM cannot " - "be predicted — the pre-flight fit check is SKIPPED for this " - "run.[/]" - ) - return () - - predicted = estimate_stream_peak_vram( - layer_bytes=layer_bytes, - buffers=tcfg.stream_buffers, - extras_bytes=embed_bytes, - adapter_params=self._estimate_adapter_params(tcfg, model_config), - vocab_size=vocab, - hidden_size=hidden, - intermediate_size=inter, - n_layers=index.n_layers, - seq_len=seq_len, - batch_size=batch, - ) - logits = estimate_logits_bytes( - vocab_size=vocab, seq_len=seq_len, batch_size=batch - ) - lines = [ - f" peak VRAM ~{predicted / 1e9:.2f} GB at batch {batch} x seq " - f"{seq_len} (logits {logits / 1e9:.2f} GB)" - ] - - if not on_cuda: - return tuple(lines) - - import torch - - available = int(torch.cuda.mem_get_info()[0]) - fit = decide_stream_fit(predicted_bytes=predicted, available_bytes=available) - if not fit.fits: - raise ValueError(fit.reason) - lines.append(f" free VRAM {available / 1e9:.2f} GB") - - # A per-card TFLOPS constant baked into the source would be a - # fabrication; measuring the user's own card in this session is the only - # honest input, and the result is reported as a bracket because real - # streamed runs landed at 68%-100% of their measured ceiling. - ceiling = measure_gemm_tflops(device=str(self.device)) - if ceiling is not None and index.total_params: - shaped = forecast_stream_throughput( - params=index.total_params, - effective_tflops=ceiling.tflops, - tokens_per_epoch=0, - sm_clock_mhz=ceiling.sm_clock_mhz, - ) - clock = f" @ {ceiling.sm_clock_mhz} MHz" if ceiling.sm_clock_mhz else "" - lines.append( - f" forecast {shaped.tokens_per_sec_low:.0f}-" - f"{shaped.tokens_per_sec_ceiling:.0f} tok/s — a compute-bound " - f"bound, not a promise" - ) - lines.append( - f" (from {ceiling.tflops:.2f} TFLOPS measured on " - f"this card now{clock})" - ) - advice = accumulation_advice( - batch_size=batch, accum=tcfg.gradient_accumulation_steps - ) - if advice is not None: - lines.append(f" [yellow]![/] {advice}") - return tuple(lines) - def _apply_quantization_aware(self, tcfg) -> None: """Apply quantization-aware training post-LoRA (shared text/vision). @@ -1723,20 +1331,11 @@ class SFTTrainerWrapper: f"{self._output_dir!r}" ) offload_save_dir = candidate - import contextlib - - with contextlib.ExitStack() as _train_ctx: - # v0.72.3 — release the streaming weight source even if training - # raises. On the disk tier that is one open shard handle per decoder - # layer (80+ on a large model), and an OOM mid-run is a realistic - # outcome on exactly the small cards this feature targets — which is - # precisely the case that leaks across back-to-back runs in one - # process (`soup sweep`, the web UI). Registered FIRST so it runs - # LAST, after the other contexts have unwound. - _train_ctx.callback(self._close_stream_runtime) - _train_ctx.enter_context( - offload_context(tcfg.activation_offloading, save_dir=offload_save_dir) - ) + # v0.72.3 — the shared context releases the streaming weight source even + # if training raises (see StreamingSetupMixin._training_context). + with self._training_context( + offload_context(tcfg.activation_offloading, save_dir=offload_save_dir) + ) as train_ctx: # LongLoRA S² shifted-sparse attention (v0.49.0 schema). The override # monkeypatches attention.forward for the duration of training and # was previously never installed (use_longlora validated but shipped @@ -1747,7 +1346,7 @@ class SFTTrainerWrapper: from soup_cli.utils.longlora import apply_longlora_forward_override try: - _train_ctx.enter_context( + train_ctx.enter_context( apply_longlora_forward_override(self.model) ) console.print("[green]LongLoRA S² attention override active[/]") diff --git a/src/soup_cli/trainer/simpo.py b/src/soup_cli/trainer/simpo.py index 35f3991..3168993 100644 --- a/src/soup_cli/trainer/simpo.py +++ b/src/soup_cli/trainer/simpo.py @@ -8,12 +8,13 @@ from typing import Optional from rich.console import Console from soup_cli.config.schema import SoupConfig +from soup_cli.trainer.stream_setup import StreamingSetupMixin from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name console = Console() -class SimPOTrainerWrapper: +class SimPOTrainerWrapper(StreamingSetupMixin): """High-level wrapper for SimPO training from SoupConfig. SimPO is a reference-free preference optimization method that uses @@ -26,6 +27,11 @@ class SimPOTrainerWrapper: - rejected: the less preferred response """ + #: TRL builds this loss's forward through ``concatenated_inputs`` + + #: ``torch.cat``, so chosen and rejected arrive as ONE tensor of twice + #: the configured batch. The VRAM pre-flight must budget for that. + _STREAM_ROWS_PER_EXAMPLE = 2 + def __init__( self, config: SoupConfig, @@ -70,14 +76,28 @@ class SimPOTrainerWrapper: cfg = self.config tcfg = cfg.training use_unsloth = cfg.backend == "unsloth" + # v0.72.4 — layer streaming replaces the model-load path entirely (meta + # skeleton, never a resident load), so it dispatches ahead of the backend + # branches. The schema already rejects streaming + unsloth/mlx. + use_streaming = bool(getattr(tcfg, "stream_layers", False)) - if use_unsloth: + if use_streaming: + self._setup_streaming_transformers(cfg, tcfg) + elif use_unsloth: self._setup_unsloth(cfg, tcfg) else: self._setup_transformers(cfg, tcfg) trainable, total = self.model.get_nb_trainable_parameters() - pct = 100 * trainable / total + # v0.72.4 (mirrors sft.py) — under NF4 streaming PEFT's total is wrong + # by ~6.5x: it sizes Params4bit as `numel * 2 * quant_storage.itemsize`, + # right for a RESIDENT one but not for our `meta` placeholder, which + # still carries the LOGICAL shape. The sharder counted the real source + # elements, so prefer that. + stream_total = getattr(self._stream_runtime, "total_params", 0) + if stream_total: + total = stream_total + pct = 100 * trainable / total if total else 0.0 console.print( f"[green]LoRA applied:[/] {trainable:,} trainable" f" / {total:,} total ({pct:.2f}%)" @@ -293,8 +313,10 @@ class SimPOTrainerWrapper: from soup_cli.utils.v028_features import activation_offloading_context - with activation_offloading_context( - self.config.training, self._output_dir, + # v0.72.4 — the shared context releases the streaming weight source even + # if training raises (see StreamingSetupMixin._training_context). + with self._training_context( + activation_offloading_context(self.config.training, self._output_dir) ): self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) duration = time.time() - start diff --git a/src/soup_cli/trainer/stream_setup.py b/src/soup_cli/trainer/stream_setup.py new file mode 100644 index 0000000..9b2071a --- /dev/null +++ b/src/soup_cli/trainer/stream_setup.py @@ -0,0 +1,453 @@ +"""Shared layer-streaming setup for every trainer wrapper that supports it. + +v0.72.4 — extracted verbatim from ``trainer/sft.py`` so that SFT and the four +preference losses (DPO / ORPO / SimPO / KTO) cannot drift. There is exactly one +copy of the NF4 pre-flight, the RAM/disk tier decision, the VRAM fit refusal and +the runtime release; a per-wrapper copy would be five places to fix the next +time any of them is wrong. + +The move is behaviour-preserving for SFT by design, so v0.72.0-.3's +bit-exactness gates remain valid without being re-run. The ONE addition is +``_STREAM_ROWS_PER_EXAMPLE``: DPO, ORPO and SimPO build their forward through +TRL's ``concatenated_inputs`` + ``torch.cat``, so 2 x ``batch_size`` rows reach +the model in a single tensor. v0.72.3's VRAM estimator was validated on the +property that it NEVER under-predicts, and budgeting those three at 1x rows +would break exactly that — on Windows the consequence is not an exception but a +silent WDDM spill to host memory that makes the run an order of magnitude +slower with no error at all. + +NO top-level torch: this module is imported by five trainer modules. +""" + +import contextlib +import math + +from rich.console import Console + +console = Console() + + +class StreamingSetupMixin: + """Builds a layer-streamed model in place of the resident load. + + Requires the host wrapper to provide ``self.device``, + ``self._trust_remote_code``, and to accept ``self.model`` / ``self.tokenizer`` + / ``self._stream_runtime`` being set. + """ + + #: Rows that reach the model per dataset example. 1 for a plain causal LM + #: step; 2 for a loss whose forward concatenates chosen and rejected. + _STREAM_ROWS_PER_EXAMPLE = 1 + + #: Set by :meth:`_setup_streaming_transformers`; absent on a resident run. + _stream_runtime = None + + @contextlib.contextmanager + def _training_context(self, *contexts): + """The `with` block every trainer runs `trainer.train()` inside. + + Its whole job is ordering: ``_close_stream_runtime`` is registered + FIRST so it runs LAST, after every other context has unwound, and it + runs even when training raises. That matters because an OOM mid-run is + a realistic outcome on exactly the small cards this feature targets, + and on the disk tier the runtime holds one open shard handle per decoder + layer — which is the case that leaks across back-to-back runs in one + process (`soup sweep`, the web UI). + + Yields the stack so a caller can enter further contexts conditionally. + """ + with contextlib.ExitStack() as stack: + stack.callback(self._close_stream_runtime) + for context in contexts: + stack.enter_context(context) + yield stack + + def _setup_streaming_transformers(self, cfg, tcfg): + """v0.72.0 BETA — layer streaming. The resident base load NEVER happens. + + Builds the skeleton on ``meta`` (``accelerate.init_empty_weights``), + materialises only embeddings / final norm / LoRA, and streams each + decoder layer from CPU RAM into a small pool of pre-allocated VRAM + buffers. Peak VRAM becomes the size of ONE layer instead of the model. + """ + from dataclasses import replace + + from peft import LoraConfig, TaskType + from transformers import AutoConfig, AutoTokenizer + + from soup_cli.utils.layer_shard import ( + QUANT_NF4, + QUANT_NONE, + resolve_shard_dir, + shard_checkpoint, + source_weight_bytes, + ) + from soup_cli.utils.layer_stream import ( + RAM_TIER_HEADROOM, + TIER_DISK, + TIER_RAM, + build_stream_plan, + detect_disk_kind, + dtype_bytes, + estimate_stream_store_bytes, + free_ram_bytes, + render_stream_panel, + stream_arch_of, + ) + from soup_cli.utils.layer_stream_runtime import ( + RamSource, + build_meta_skeleton, + build_streamed_model, + extras_resident_bytes, + probe_expandable_segments, + quantised_layer_suffixes, + ) + from soup_cli.utils.spectrum_scan import resolve_model_weights + + console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") + self.tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + model_config = AutoConfig.from_pretrained( + cfg.base, trust_remote_code=self._trust_remote_code + ) + # Allowlist, not a heuristic — a half-supported architecture streams + # weights into the wrong module and mis-trains silently. + arch = stream_arch_of(model_config) + + on_cuda = str(self.device).startswith("cuda") + dtype = "bfloat16" if on_cuda else "float32" + + # v0.72.2 — NF4. The decoder linears ship as packed nibbles + per-block + # absmax, so the RAM store is ~0.26x its bf16 size; embeddings, norms and + # an untied head stay at `dtype`, exactly as replace_with_bnb_linear + # leaves them. + quant = QUANT_NF4 if tcfg.quantization == "4bit" else QUANT_NONE + + weights_dir = resolve_model_weights(cfg.base) + shard_dir = resolve_shard_dir(cfg.base) + + # Cheap size probe BEFORE sharding: re-writing a checkpoint we are + # about to refuse for not fitting in RAM costs minutes of disk I/O. + # Charged at the STREAMED rate, not the on-disk one — an 8B bf16 + # checkpoint is 16 GB on disk but only ~4.2 GB of NF4 store, and + # comparing the raw file size would refuse exactly the runs NF4 enables. + early_free_ram = free_ram_bytes() + if early_free_ram is not None: + source_bytes = source_weight_bytes(weights_dir) + store_estimate = estimate_stream_store_bytes(source_bytes, dtype=dtype, quant=quant) + if store_estimate >= early_free_ram * RAM_TIER_HEADROOM and tcfg.stream_source == "ram": + as_streamed = ( + "" + if quant == QUANT_NONE + else f" ({store_estimate / 1e9:.1f} GB once quantised to NF4)" + ) + raise ValueError( + f"training.stream_source='ram' but {cfg.base} is " + f"{source_bytes / 1e9:.1f} GB on disk{as_streamed} and only " + f"{early_free_ram / 1e9:.1f} GB of RAM is free. Set " + f"stream_source='auto' to fall back to the NVMe disk tier, " + f"free RAM, or pick a smaller base." + ) + + # The authoritative list of weights to quantise is whatever + # replace_with_bnb_linear actually converts, read off a meta skeleton — + # not a hard-coded name list that would drift per architecture. + # + # This builds a second, throwaway skeleton (build_streamed_model makes + # its own). Deliberate: a meta skeleton allocates NO weight storage, so + # the cost is module-tree construction only, and threading a pre-built + # model into build_streamed_model would couple suffix discovery to model + # construction for no memory saving. + quant_suffixes = () + if quant == QUANT_NF4: + probe = build_meta_skeleton( + cfg.base, + dtype=dtype, + quant=quant, + trust_remote_code=self._trust_remote_code, + ) + quant_suffixes = quantised_layer_suffixes(probe) + del probe + + console.print(f"[dim]Preparing layer shards -> {shard_dir}[/]") + index = shard_checkpoint( + weights_dir, + shard_dir, + dtype=dtype, + arch=arch, + quant=quant, + quant_suffixes=quant_suffixes, + # Quantise on the device that will run the model: CPU and CUDA agree + # on the packed nibbles but not on every float32 nested statistic. + quant_device=str(self.device), + ) + + spec = RamSource.spec_from_shard(shard_dir) + # Measured from the shard headers, not derived from `total_params`: + # under NF4 a layer holds packed uint8 alongside float32 statistics, so + # element counts no longer convert to bytes at a single rate. + layer_bytes = sum(math.prod(shape) * dtype_bytes(stored) for shape, stored in spec.values()) + embed_bytes = extras_resident_bytes(shard_dir) + + free_ram = free_ram_bytes() + if free_ram is None: + console.print( + "[yellow]psutil unavailable — cannot size the RAM tier; " + "proceeding and letting the allocation fail loudly if it must[/]" + ) + free_ram = (layer_bytes * index.n_layers + embed_bytes) * 10 + + store_total = layer_bytes * index.n_layers + embed_bytes + # Checked BEFORE build_stream_plan so a `ram`-only run is refused with + # the message about stream_source rather than choose_tier's generic + # "needs NVMe or more RAM" — and without paying the ~9 s disk probe for + # an answer that cannot change the outcome. + if tcfg.stream_source == "ram" and store_total >= free_ram * RAM_TIER_HEADROOM: + raise ValueError( + f"training.stream_source='ram' but the base is " + f"{store_total / 1e9:.1f} GB and only {free_ram / 1e9:.1f} GB of " + f"RAM is free. Set stream_source='auto' to fall back to the NVMe " + f"disk tier, free RAM, or pick a smaller base." + ) + plan = build_stream_plan( + arch=arch, + n_layers=index.n_layers, + layer_bytes=layer_bytes, + embed_bytes=embed_bytes, + available_ram_bytes=free_ram, + # The page-locked ceiling is a property of the box, not of free RAM; + # rather than probe it destructively we attempt the pinned store and + # fall back loudly (see layer_stream_runtime._build_source). + pinned_limit_bytes=None, + buffers=tcfg.stream_buffers, + # v0.72.3: the REAL media type, not a constant. Passed as a callable + # because probing costs ~9 s on Windows and the answer only matters + # when the base does not fit in RAM. + disk_kind=lambda: detect_disk_kind(shard_dir), + ) + # v0.72.3 — the disk overflow tier is live, so a base that does not fit + # in RAM is no longer fatal. `stream_source` decides: 'ram' insists, + # 'disk' forces, 'auto' (the default) takes RAM when it fits and falls + # back to disk when it does not. build_stream_plan already refused a + # non-NVMe disk, so reaching here with tier='disk' means NVMe. + tier = TIER_DISK if tcfg.stream_source == "disk" else plan.tier + if tier != plan.tier: + # The panel is rendered from `plan`, so a forced tier has to be + # reflected there or the pre-flight reports "tier ram" immediately + # before the runtime announces it is streaming from disk. Every + # field that describes the RAM store is corrected with it, so no + # consumer can read a stale value. + plan = replace( + plan, + tier=tier, + store_bytes=0, + pinned=False, + notes=plan.notes + + ( + "streaming from disk because stream_source='disk' was set, " + "not because RAM was short. Nothing is held resident, and " + "the slowdown versus the RAM tier is unmeasured on this " + "hardware.", + ), + ) + # v0.72.3 — VRAM pre-flight. Streaming bounds the WEIGHTS; activations + # and the logits tensor are untouched by it and both scale with batch x + # seq. On a large-vocab model the logits term alone dwarfs the buffer + # pool (measured: 146x at batch 8), so a plan that reports only tier and + # buffer sizes will happily green-light a config that cannot run. + forecast_lines = self._stream_budget_lines( + cfg, + tcfg, + model_config=model_config, + layer_bytes=layer_bytes, + embed_bytes=embed_bytes, + index=index, + on_cuda=on_cuda, + ) + console.print(render_stream_panel(plan, forecast_lines)) + console.print( + "[yellow]Layer streaming is BETA:[/] slower than resident training, " + "but this model may not run resident on this card at all." + ) + if on_cuda and not probe_expandable_segments(): + console.print( + "[dim]expandable_segments allocator hint is unavailable on this " + "platform (silently ignored on Windows) — not enabled[/]" + ) + + target_modules = tcfg.lora.target_modules + if target_modules == "auto": + target_modules = None + lora_config = LoraConfig( + r=tcfg.lora.r, + lora_alpha=tcfg.lora.alpha, + lora_dropout=tcfg.lora.dropout, + target_modules=target_modules, + task_type=TaskType.CAUSAL_LM, + bias="none", + use_dora=tcfg.lora.use_dora, + use_rslora=tcfg.lora.use_rslora, + ) + + model, runtime = build_streamed_model( + model_id=cfg.base, + shard_dir=shard_dir, + index=index, + lora_config=lora_config, + device=self.device, + dtype=dtype, + buffers=tcfg.stream_buffers, + pin=plan.pinned and on_cuda, + seed=tcfg.seed if getattr(tcfg, "seed", None) is not None else 0, + trust_remote_code=self._trust_remote_code, + console=console, + quant=quant, + tier=tier, + ) + self.model = model + self._stream_runtime = runtime + stats = runtime.stats() + if stats["tier"] == TIER_RAM: + source_line = ( + f"{stats['store_bytes'] / 1e9:.2f} GB " + f"{'pinned' if stats['pinned'] else 'pageable'} RAM store" + ) + else: + source_line = ( + f"streamed from DISK ({stats['disk_bytes'] / 1e9:.2f} GB on an " + f"NVMe volume, nothing held resident)" + ) + buffer_line = ( + f"{stats['buffers']} x " + f"{stats['buffer_bytes'] / stats['buffers'] / 1e6:.0f} MB VRAM buffers" + ) + console.print( + f"[green]Layer streaming ready:[/] {stats['n_layers']} layers, " + f"{source_line}, {buffer_line}" + ) + + def _close_stream_runtime(self) -> None: + """Release the streaming weight source, if this run had one.""" + runtime = getattr(self, "_stream_runtime", None) + if runtime is not None: + runtime.close() + + def _estimate_adapter_params(self, tcfg, model_config) -> int: + """Trainable adapter parameters, before the model exists. + + Deliberately coarse and biased HIGH: it assumes every targeted module is + hidden x hidden. Gate/up/down projections are larger, but the whole + adapter term is ~0.5% of a streaming step's peak, so precision here buys + nothing while under-counting would eat into the safety margin. + """ + hidden = int(getattr(model_config, "hidden_size", 0) or 0) + layers = int(getattr(model_config, "num_hidden_layers", 0) or 0) + targets = tcfg.lora.target_modules + n_targets = len(targets) if isinstance(targets, (list, tuple)) else 4 + return layers * n_targets * 2 * tcfg.lora.r * hidden + + def _stream_budget_lines( + self, cfg, tcfg, *, model_config, layer_bytes, embed_bytes, index, on_cuda + ): + """Predict peak VRAM + bracket throughput, and REFUSE a run that cannot fit. + + Returns the extra lines for the pre-flight panel. Raises when the step is + predicted not to fit: on Linux that would be a hard OOM, and on Windows + something worse — WDDM spills to host memory without raising, so the run + silently becomes an order of magnitude slower and looks like the feature + is merely slow. + """ + from soup_cli.utils.layer_stream import ( + accumulation_advice, + decide_stream_fit, + estimate_logits_bytes, + estimate_stream_peak_vram, + forecast_stream_throughput, + ) + from soup_cli.utils.layer_stream_runtime import measure_gemm_tflops + + vocab = int(getattr(model_config, "vocab_size", 0) or 0) + hidden = int(getattr(model_config, "hidden_size", 0) or 0) + inter = int(getattr(model_config, "intermediate_size", 0) or 0) + seq_len = int(cfg.data.max_length) + batch = tcfg.batch_size if isinstance(tcfg.batch_size, int) else 1 + # v0.72.4 — a paired loss concatenates chosen and rejected into ONE + # tensor, so twice the rows reach the model per configured batch. The + # estimator's contract is that it never under-predicts; budgeting a + # paired loss at 1x rows would halve the logits term, which is the + # dominant one (measured 146x the buffer pool at batch 8). + rows = batch * self._STREAM_ROWS_PER_EXAMPLE + if not (vocab and hidden and inter): + # Never silently: skipping the budget also skips the refusal that + # stops a run from OOMing (or, on Windows, spilling to host memory + # and running an order of magnitude slower with no error at all). + console.print( + "[yellow]Layer streaming could not read vocab_size / hidden_size " + "/ intermediate_size from the model config, so peak VRAM cannot " + "be predicted — the pre-flight fit check is SKIPPED for this " + "run.[/]" + ) + return () + + predicted = estimate_stream_peak_vram( + layer_bytes=layer_bytes, + buffers=tcfg.stream_buffers, + extras_bytes=embed_bytes, + adapter_params=self._estimate_adapter_params(tcfg, model_config), + vocab_size=vocab, + hidden_size=hidden, + intermediate_size=inter, + n_layers=index.n_layers, + seq_len=seq_len, + batch_size=rows, + ) + logits = estimate_logits_bytes(vocab_size=vocab, seq_len=seq_len, batch_size=rows) + paired = ( + "" if rows == batch else f" ({rows} rows — chosen+rejected are one concatenated tensor)" + ) + lines = [ + f" peak VRAM ~{predicted / 1e9:.2f} GB at batch {batch} x seq " + f"{seq_len}{paired} (logits {logits / 1e9:.2f} GB)" + ] + + if not on_cuda: + return tuple(lines) + + import torch + + available = int(torch.cuda.mem_get_info()[0]) + fit = decide_stream_fit(predicted_bytes=predicted, available_bytes=available) + if not fit.fits: + raise ValueError(fit.reason) + lines.append(f" free VRAM {available / 1e9:.2f} GB") + + # A per-card TFLOPS constant baked into the source would be a + # fabrication; measuring the user's own card in this session is the only + # honest input, and the result is reported as a bracket because real + # streamed runs landed at 68%-100% of their measured ceiling. + ceiling = measure_gemm_tflops(device=str(self.device)) + if ceiling is not None and index.total_params: + shaped = forecast_stream_throughput( + params=index.total_params, + effective_tflops=ceiling.tflops, + tokens_per_epoch=0, + sm_clock_mhz=ceiling.sm_clock_mhz, + ) + clock = f" @ {ceiling.sm_clock_mhz} MHz" if ceiling.sm_clock_mhz else "" + lines.append( + f" forecast {shaped.tokens_per_sec_low:.0f}-" + f"{shaped.tokens_per_sec_ceiling:.0f} tok/s — a compute-bound " + f"bound, not a promise" + ) + lines.append( + f" (from {ceiling.tflops:.2f} TFLOPS measured on " + f"this card now{clock})" + ) + advice = accumulation_advice(batch_size=batch, accum=tcfg.gradient_accumulation_steps) + if advice is not None: + lines.append(f" [yellow]![/] {advice}") + return tuple(lines) diff --git a/src/soup_cli/utils/layer_stream.py b/src/soup_cli/utils/layer_stream.py index e472041..e78f7c8 100644 --- a/src/soup_cli/utils/layer_stream.py +++ b/src/soup_cli/utils/layer_stream.py @@ -41,6 +41,21 @@ MIN_STREAM_BUFFERS = 2 MAX_STREAM_BUFFERS = 8 DEFAULT_STREAM_BUFFERS = 2 +# --- tasks ---------------------------------------------------------------- +#: Tasks whose trainers can run against a streamed base (v0.72.4). +#: +#: DPO and KTO take their reference model from the SAME streamed base with the +#: adapters disabled (TRL's ``null_ref_context``), so the reference costs no +#: extra weights at all — measured 0.914x the SFT peak, where forcing a real +#: second instance cost 9.92x. ORPO and SimPO are reference-free. +SUPPORTED_STREAM_TASKS = ("sft", "dpo", "orpo", "simpo", "kto") + +#: Tasks PERMANENTLY excluded, not merely unimplemented. Generation rollouts +#: re-read every layer once per generated token, which destroys the whole +#: premise: streaming amortises one weight read over a training step, not over a +#: single decoded token (plan §3.2). +ROLLOUT_STREAM_TASKS = ("grpo", "ppo") + #: FLOPs per parameter per token. 6 == WITH gradient checkpointing #: (2 forward + 2 recompute + 2 dL/dx; base weight-grads are skipped because #: the base is frozen). Streaming always checkpoints, so this is never 4. diff --git a/tests/test_v07200.py b/tests/test_v07200.py index ea46749..8105442 100644 --- a/tests/test_v07200.py +++ b/tests/test_v07200.py @@ -769,9 +769,18 @@ class TestStreamSchemaDefaults: class TestStreamTaskAndBackendGates: - def test_non_sft_task_rejected(self): - with pytest.raises(ValueError, match="sft"): - _load(_stream_yaml(task="dpo")) + def test_unsupported_task_rejected(self): + """v0.72.4 opened the four preference losses, so `dpo` is no longer the + example here. `reward_model` still has no streaming path, and `grpo` is + excluded permanently — rollouts re-read every layer per generated + token.""" + with pytest.raises(ValueError, match="stream_layers"): + _load(_stream_yaml(task="reward_model")) + with pytest.raises(ValueError, match="generation"): + _load(_stream_yaml(task="grpo")) + + def test_preference_tasks_are_accepted_since_v0724(self): + assert _load(_stream_yaml(task="dpo")).task == "dpo" def test_unsloth_backend_rejected(self): with pytest.raises(ValueError, match="transformers"): diff --git a/tests/test_v07201.py b/tests/test_v07201.py index 0e56c2b..018a473 100644 --- a/tests/test_v07201.py +++ b/tests/test_v07201.py @@ -490,17 +490,26 @@ class TestRefusalsNameThePostRenumberSlot: ): load_config_from_string(_stream_yaml(**{field: value})) - def test_preference_losses_name_v0724(self): + def test_preference_losses_shipped_in_v0724(self): + """This refusal is GONE, per the shrink-not-grow rule above. What + replaced it is a PERMANENT exclusion for the rollout tasks, which must + NOT name a release — a version number there would read as "coming soon" + to the next maintainer and invite them to wire up something that cannot + work.""" from soup_cli.config.loader import load_config_from_string - yaml_text = ( - "base: hf-internal-testing/tiny-random-LlamaForCausalLM\n" - "task: dpo\n" - "data:\n train: data.jsonl\n" - "training:\n stream_layers: true\n batch_size: 1\n" - ) - with pytest.raises(ValueError, match="v0.72.4"): - load_config_from_string(yaml_text) + def _yaml(task): + return ( + "base: hf-internal-testing/tiny-random-LlamaForCausalLM\n" + f"task: {task}\n" + "data:\n train: data.jsonl\n" + "training:\n stream_layers: true\n batch_size: 1\n" + ) + + assert load_config_from_string(_yaml("dpo")).task == "dpo" + with pytest.raises(ValueError) as excinfo: + load_config_from_string(_yaml("grpo")) + assert "v0.72" not in str(excinfo.value) def test_arch_allowlist_names_v0723(self): from soup_cli.utils.layer_stream import stream_arch_of diff --git a/tests/test_v07204.py b/tests/test_v07204.py new file mode 100644 index 0000000..ec6b8bd --- /dev/null +++ b/tests/test_v07204.py @@ -0,0 +1,1056 @@ +"""v0.72.4 — preference losses (DPO / ORPO / SimPO / KTO) over layer streaming. + +The slot's whole risk is ONE property (plan.md §7.1): DPO needs a reference +model, and implemented naively as a second model instance it doubles memory and +defeats the feature entirely. It must be the SAME streamed base with adapters +disabled. **A passing loss curve does not detect a second instance** — only a +memory assertion does, which is why the assertions here are about bytes and +about object identity, not about the loss going down. + +Gate numbers + method: `.claude/v0724-gate-results.md`. + +A correction to the brief established while gating: ORPO and SimPO are genuinely +reference-free, but **KTO is not** — `KTOTrainer.__init__` has byte-for-byte the +same three-branch reference selection as DPO, so it inherits the same trap and +is asserted the same way here. +""" + +import json +import os + +import pytest + + +# ========================================================================== +# fixtures (mirroring tests/test_v07200.py so the two cannot drift) +# ========================================================================== +def _cuda_available(): + try: + import torch + + return torch.cuda.is_available() + except Exception: # pragma: no cover - torch always present in CI + return False + + +def _mps_is_the_accelerator(): + try: + import torch + + return ( + hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + and not torch.cuda.is_available() + ) + except Exception: # pragma: no cover + return False + + +def _tiny_llama_dir(tmp_path, n_layers=2, tie=True, vocab=64, hidden=64): + """A real (tiny) Llama checkpoint on disk: config.json + model.safetensors. + + ``hidden`` is 64, not 32: bitsandbytes' CPU 4-bit repack computes + ``blocks_per_row = 16 // (hidden // 2)`` and silently yields 0 at 32 (#323). + """ + import torch + from safetensors.torch import save_file + from transformers import LlamaConfig, LlamaForCausalLM + + torch.manual_seed(7) + config = LlamaConfig( + vocab_size=vocab, + hidden_size=hidden, + intermediate_size=hidden * 2, + num_hidden_layers=n_layers, + num_attention_heads=4, + num_key_value_heads=2, + tie_word_embeddings=tie, + max_position_embeddings=128, + ) + model = LlamaForCausalLM(config).to(torch.float32).eval() + weights = tmp_path / "model" + weights.mkdir(parents=True, exist_ok=True) + state = {k: v.contiguous() for k, v in model.state_dict().items()} + if tie: + state.pop("lm_head.weight", None) + save_file(state, str(weights / "model.safetensors")) + config.save_pretrained(str(weights)) + return str(weights), model, config + + +def _write_tiny_tokenizer(directory): + """A real, offline PreTrainedTokenizerFast so `setup()` can tokenize.""" + from tokenizers import Tokenizer, models, pre_tokenizers + + vocab = {"": 0, "": 1, "": 2, "": 3} + for word in ( + "hello", + "world", + "hi", + "yo", + "the", + "cat", + "sat", + "on", + "mat", + "good", + "bad", + "answer", + "question", + "soup", + ): + vocab[word] = len(vocab) + tokenizer = Tokenizer(models.WordLevel(vocab=vocab, unk_token="")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.save(os.path.join(directory, "tokenizer.json")) + with open(os.path.join(directory, "tokenizer_config.json"), "w", encoding="utf-8") as fh: + json.dump( + { + "tokenizer_class": "PreTrainedTokenizerFast", + "unk_token": "", + "bos_token": "", + "eos_token": "", + "pad_token": "", + "model_max_length": 128, + "clean_up_tokenization_spaces": False, + }, + fh, + ) + + +def _randomise_lora_b(model, seed=11): + """PEFT initialises ``lora_B = 0``, so until B is load-bearing the adapter + contributes NOTHING and every "did the adapter path run?" assertion — very + much including "is the reference different from the policy?" — passes + vacuously.""" + import torch + + gen = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, param in model.named_parameters(): + if "lora_B" in name: + param.copy_( + torch.randn(param.shape, generator=gen).to(param.device, param.dtype) * 0.05 + ) + + +def _sync_adapters(dst, src): + """Copy LoRA weights src -> dst across the ``.inner.`` wrapper difference. + + Returns the number copied; 0 means the comparison would be vacuous. + """ + import torch + + def norm(key): + return key.replace(".inner.", "") + + source = {norm(k): v.detach().clone() for k, v in src.state_dict().items() if "lora_" in k} + copied = 0 + with torch.no_grad(): + for key, tensor in dst.state_dict().items(): + if "lora_" not in key: + continue + match = source.get(norm(key)) + if match is not None: + tensor.copy_(match.to(tensor.device, tensor.dtype)) + copied += 1 + return copied + + +def _batch_on(model, batch): + """Move a TRL batch onto the model's device. + + `TrainingArguments` picks CUDA whenever it is available, so a test that + deliberately pins the MODEL to CPU (for exact float32 arithmetic) still gets + CUDA batch tensors from the dataloader. + """ + device = next(model.parameters()).device + return { + key: (value.to(device) if hasattr(value, "to") else value) for key, value in batch.items() + } + + +def _loss_of(trainer, model, batch): + """Call the trainer's loss for `model`, across TRL's signature differences. + + `KTOTrainer.get_batch_loss_metrics` takes (model, batch); DPO / ORPO / CPO + take (model, batch, train_eval). + """ + import inspect + + fn = trainer.get_batch_loss_metrics + if "train_eval" in inspect.signature(fn).parameters: + loss, _ = fn(model, batch, "train") + else: + loss, _ = fn(model, batch) + return loss + + +def _match_streamed_dtype(resident, streamed): + """Put the resident reference on the streamed model's device AND dtype. + + Streaming picks bf16 on CUDA and float32 on CPU. Comparing a float32 + resident model against a bf16 streamed one measures the dtype gap, not the + streaming path — that mistake produced a 9.96e-04 "failure" that was + entirely the test's own. + """ + param = next(streamed.parameters()) + return resident.to(device=param.device, dtype=param.dtype) + + +def _pref_rows(n=4): + return [{"prompt": "hi", "chosen": " good answer", "rejected": " bad"} for _ in range(n)] + + +def _kto_rows(n=4): + return [{"prompt": "hi", "completion": " good answer", "label": i % 2 == 0} for i in range(n)] + + +def _sft_rows(n=4): + return [ + { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello world"}, + ] + } + for _ in range(n) + ] + + +_TASK_ROWS = { + "sft": _sft_rows, + "dpo": _pref_rows, + "orpo": _pref_rows, + "simpo": _pref_rows, + "kto": _kto_rows, +} + +#: KTO refuses a per-device batch size of 1 outright ("the KL term will be +#: equivalent to the implied reward"), so it is only streamable at all because +#: v0.72.3 lifted layer streaming's batch-1 restriction. +_MIN_BATCH = {"kto": 2} + + +def _stream_cfg(weights, out_dir, task="dpo", **training): + import yaml + + from soup_cli.config.loader import load_config_from_string + + tcfg = { + # KTO is refused below batch 2 (its KL term is degenerate at 1), so the + # per-task floor is the default here; an explicit kwarg still wins, + # which is what TestKtoBatchIsRefusedEarly relies on. + "batch_size": _MIN_BATCH.get(task, 1), + "gradient_accumulation_steps": 1, + "quantization": "none", + "stream_layers": True, + "epochs": 1, + "logging_steps": 1, + "save_steps": 1000, + "lora": {"r": 4, "alpha": 8, "target_modules": ["q_proj", "v_proj"]}, + } + tcfg.update(training) + return load_config_from_string( + yaml.safe_dump( + { + "base": weights, + "task": task, + "backend": "transformers", + "modality": "text", + "data": { + "train": "train.jsonl", + "max_length": 64, + # v0.36.0 removed the silent f-string fallback; SFT needs one. + "chat_template": "chatml", + }, + "training": tcfg, + "output": str(out_dir), + } + ) + ) + + +def _wrapper_for(task): + from soup_cli.trainer.dpo import DPOTrainerWrapper + from soup_cli.trainer.kto import KTOTrainerWrapper + from soup_cli.trainer.orpo import ORPOTrainerWrapper + from soup_cli.trainer.sft import SFTTrainerWrapper + from soup_cli.trainer.simpo import SimPOTrainerWrapper + + return { + "sft": SFTTrainerWrapper, + "dpo": DPOTrainerWrapper, + "orpo": ORPOTrainerWrapper, + "simpo": SimPOTrainerWrapper, + "kto": KTOTrainerWrapper, + }[task] + + +def _build_streamed_wrapper( + tmp_path, + monkeypatch, + task="dpo", + n_layers=2, + device=None, + hidden=64, + vocab=64, + **training, +): + """Build a task wrapper through the REAL `setup()` path, streaming. + + ``device`` defaults to the real accelerator, because `TrainingArguments` + picks CUDA when it is available and forcing CPU there would only produce a + device mismatch no user would ever hit. Numerical-equality tests pass + ``device='cpu'`` deliberately: the streaming path uses float32 on CPU and + bf16 on CUDA, and "bit-exact" is only a meaningful assertion in the former + (a bf16 logp of -12.75 cannot represent a change smaller than ~0.05). + """ + weights, resident, _ = _tiny_llama_dir(tmp_path, n_layers=n_layers, hidden=hidden, vocab=vocab) + _write_tiny_tokenizer(weights) + monkeypatch.setenv("SOUP_LAYER_STREAM_CACHE_DIR", str(tmp_path / "cache")) + monkeypatch.chdir(tmp_path) + training.setdefault("batch_size", _MIN_BATCH.get(task, 1)) + cfg = _stream_cfg(weights, tmp_path / "out", task=task, **training) + if device is None: + device = "cuda" if _cuda_available() else "cpu" + wrapper = _wrapper_for(task)(cfg, device=device) + wrapper.setup({"train": _TASK_ROWS[task](8)}) + return wrapper, resident, weights + + +# ========================================================================== +# item 1 -- schema: which tasks may stream +# ========================================================================== +_REFERENCE_USING = ("dpo", "kto") +_ALL_PREFERENCE = ("dpo", "orpo", "simpo", "kto") + + +class TestStreamingTaskGate: + """v0.72.0-.3 hard-coded ``task == 'sft'``. v0.72.4 opens exactly four more + and keeps refusing the rest — GRPO/PPO *permanently*, because rollouts + re-read the model per generated token and destroy the amortisation that + makes streaming viable at all (plan §3.2).""" + + def _cfg(self, tmp_path, task): + return _stream_cfg(str(tmp_path / "m"), tmp_path / "out", task=task) + + @pytest.mark.parametrize("task", ("sft",) + _ALL_PREFERENCE) + def test_supported_tasks_are_accepted(self, tmp_path, task): + cfg = self._cfg(tmp_path, task) + assert cfg.task == task + assert cfg.training.stream_layers is True + + @pytest.mark.parametrize("task", ("grpo", "ppo")) + def test_rollout_tasks_are_refused_for_a_reason_that_does_not_expire(self, tmp_path, task): + """Not "not yet" — never. The refusal must say why, or a future + maintainer will read it as an unfinished slot and wire it up.""" + with pytest.raises(ValueError) as excinfo: + self._cfg(tmp_path, task) + message = str(excinfo.value) + assert "stream_layers" in message + assert "generation" in message or "rollout" in message + assert "v0.72" not in message, ( + "a permanent refusal must not name a release that would lift it" + ) + + @pytest.mark.parametrize("task", ("reward_model", "pretrain", "embedding")) + def test_other_tasks_still_refused(self, tmp_path, task): + # `distill` is deliberately not in this list: it is rejected earlier by + # its own `teacher_model` gate, so it would assert nothing about + # streaming. + with pytest.raises(ValueError, match="stream_layers"): + self._cfg(tmp_path, task) + + +# ========================================================================== +# item 1 -- THE trap: one model, one stream +# ========================================================================== +class TestNoSecondModelInstance: + """The brief's whole slot. Asserted on bytes and object identity, never on + the loss.""" + + @pytest.mark.parametrize("task", _REFERENCE_USING) + def test_the_reference_is_the_same_model_with_adapters_disabled( + self, tmp_path, monkeypatch, task + ): + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + trainer = wrapper.trainer + assert hasattr(trainer, "ref_model"), ( + "TRL renamed `ref_model`; the no-second-instance property must be " + "re-verified against the new API before this test is adjusted" + ) + assert trainer.ref_model is None, ( + f"{task} built a SECOND model instance for the reference — that " + f"doubles memory and defeats layer streaming entirely" + ) + assert getattr(trainer, "is_peft_model", False) is True + + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_exactly_one_weight_store_is_constructed(self, tmp_path, monkeypatch, task): + """Counts RamSource constructions across the whole `setup()`. Two would + mean two copies of the base in host RAM.""" + from soup_cli.utils import layer_stream_runtime as lsr + + calls = {"n": 0} + real = lsr.RamSource + + class Counting(real): + def __init__(self, *a, **k): + calls["n"] += 1 + super().__init__(*a, **k) + + monkeypatch.setattr(lsr, "RamSource", Counting) + _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + assert calls["n"] == 1, f"{task} constructed {calls['n']} weight stores" + + def test_reference_logps_actually_differ_from_the_policy(self, tmp_path, monkeypatch): + """THE silent-failure check. + + The streamed layer substitutes base weights through `functional_call` + rather than the module's own forward, so `disable_adapter()` being a + no-op through that path is entirely plausible. If it were, the reference + would BE the policy, every log-ratio would be 0, and the DPO loss would + sit at -logsigmoid(0) = 0.6931 forever — which reads as "training + slowly", not as a bug. Measured 8.2e-01 in the gate. + """ + import torch + + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task="dpo", device="cpu") + _randomise_lora_b(wrapper.model) + trainer = wrapper.trainer + batch = _batch_on(wrapper.model, next(iter(trainer.get_train_dataloader()))) + wrapper.model.eval() + with torch.no_grad(): + policy = trainer.concatenated_forward(wrapper.model, batch) + ref_chosen, ref_rejected = trainer.compute_ref_log_probs(batch) + assert (policy["chosen_logps"] - ref_chosen).abs().max().item() > 1e-4 + assert (policy["rejected_logps"] - ref_rejected).abs().max().item() > 1e-4 + + def test_the_difference_really_comes_from_the_adapter(self, tmp_path, monkeypatch): + """CONTROL for the test above: with ``lora_B = 0`` the adapter + contributes nothing, so disabling it must change NOTHING and the two + must be exactly equal. Without this, the previous test passes for any + model whose two forwards merely differ for some other reason.""" + import torch + + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task="dpo", device="cpu") + with torch.no_grad(): + for name, param in wrapper.model.named_parameters(): + if "lora_B" in name: + param.zero_() + trainer = wrapper.trainer + batch = _batch_on(wrapper.model, next(iter(trainer.get_train_dataloader()))) + wrapper.model.eval() + with torch.no_grad(): + policy = trainer.concatenated_forward(wrapper.model, batch) + ref_chosen, _ = trainer.compute_ref_log_probs(batch) + diff = (policy["chosen_logps"] - ref_chosen).abs().max().item() + assert diff == 0.0, diff + + +@pytest.mark.skipif(not _cuda_available(), reason="peak VRAM needs CUDA") +class TestPeakVramIsNotDoubled: + """The brief's literal assertion, on the real device.""" + + def _peak(self, tmp_path, monkeypatch, task, second_reference=False): + import gc + + import torch + + # The buffer pool is held by reference CYCLES in the module tree + # (measured: close() alone retains it, close() + gc.collect() retains + # 0.00 MB). Without collecting here the next arm's baseline carries the + # previous pool and every peak is inflated by one pool's worth — which + # is exactly how an early version of the gate produced a bogus 2x. + gc.collect() + torch.cuda.empty_cache() + wrapper, _, weights = _build_streamed_wrapper(tmp_path, monkeypatch, task=task, n_layers=4) + stats = wrapper._stream_runtime.stats() + forced = None + if second_reference: + from transformers import AutoModelForCausalLM + + forced = AutoModelForCausalLM.from_pretrained( + weights, torch_dtype=next(wrapper.model.parameters()).dtype + ).to("cuda") + wrapper.trainer.ref_model = forced + # Reset AFTER setup: the pre-flight's own GEMM ceiling probe allocates + # three 4096^3 matrices, which would otherwise dominate the "peak". + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + wrapper.trainer.args.max_steps = 1 + wrapper.trainer.train() + peak = torch.cuda.max_memory_allocated() + wrapper._close_stream_runtime() + del wrapper, forced + gc.collect() + torch.cuda.empty_cache() + return peak, stats + + @pytest.mark.parametrize("task", _REFERENCE_USING) + def test_weight_bearing_terms_are_identical_to_sft(self, tmp_path, monkeypatch, task): + """One store, one pool. The gate measured 729.91 MB / 60.83 MB for both + arms on a 730 MB model; here the sizes are tiny but the EQUALITY is the + size-independent form of the same claim.""" + sft_peak, sft_stats = self._peak(tmp_path / "a", monkeypatch, "sft") + pref_peak, pref_stats = self._peak(tmp_path / "b", monkeypatch, task) + assert pref_stats["store_bytes"] == sft_stats["store_bytes"] + assert pref_stats["buffer_bytes"] == sft_stats["buffer_bytes"] + # ...and the brief's literal metric, on the numbers actually measured. + # `stats` describe the ONE streaming runtime's own bookkeeping and would + # read identically even if a second RESIDENT reference model were built + # outside it — which is the bug this slot exists to prevent. + assert pref_peak < sft_peak * 1.5, (task, sft_peak, pref_peak) + + @pytest.mark.parametrize("task", _REFERENCE_USING) + def test_the_peak_assertion_would_notice_a_second_instance(self, tmp_path, monkeypatch, task): + """CONTROL for the test above. "DPO is not 2x SFT" means nothing unless + the same harness can show what a second instance costs: the gate + measured +730.44 MB against 730.44 MB of weights, i.e. exactly one + copy.""" + implicit, stats = self._peak(tmp_path / "c", monkeypatch, task) + forced, _ = self._peak(tmp_path / "d", monkeypatch, task, second_reference=True) + assert forced > implicit, (implicit, forced) + + +class TestBitExactVsResident: + """The rule every slot in the series inherits: a streamed run is bit-exact + against the RESIDENT run of the same numerics. What changes per slot is the + reference, not the standard — here it is a resident run of the same loss.""" + + @pytest.mark.skipif( + _mps_is_the_accelerator(), + reason="MPS is untested for layer streaming (CUDA + CPU only)", + ) + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_loss_matches_a_resident_run_of_the_same_loss(self, tmp_path, monkeypatch, task): + import torch + from peft import LoraConfig, TaskType, get_peft_model + + wrapper, resident, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + _randomise_lora_b(wrapper.model) + + resident_peft = get_peft_model( + resident, + LoraConfig( + r=4, + lora_alpha=8, + lora_dropout=0.0, + bias="none", + target_modules=["q_proj", "v_proj"], + task_type=TaskType.CAUSAL_LM, + ), + ) + copied = _sync_adapters(resident_peft, wrapper.model) + assert copied > 0, "vacuous: no adapter tensors copied" + + _match_streamed_dtype(resident_peft, wrapper.model).eval() + wrapper.model.eval() + batch = _batch_on(wrapper.model, next(iter(wrapper.trainer.get_train_dataloader()))) + + streamed_trainer = wrapper.trainer + # Same trainer object, different model: the ONLY difference under test + # is where the base weights come from. + with torch.no_grad(): + streamed_loss = _loss_of(streamed_trainer, wrapper.model, batch) + resident_loss = _loss_of(streamed_trainer, resident_peft, batch) + diff = (streamed_loss - resident_loss).abs().max().item() + assert diff == 0.0, f"{task}: streamed vs resident loss differs by {diff}" + + @pytest.mark.skipif( + _mps_is_the_accelerator(), + reason="MPS is untested for layer streaming (CUDA + CPU only)", + ) + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_layer_zero_adapter_receives_gradient(self, tmp_path, monkeypatch, task): + """plan P2: a `detach()`/`no_grad()` anywhere in the base forward severs + the graph. The lower adapters then never train while the loss still + falls, because the upper ones still learn.""" + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + _randomise_lora_b(wrapper.model) + wrapper.model.train() + wrapper.model.zero_grad(set_to_none=True) + batch = _batch_on(wrapper.model, next(iter(wrapper.trainer.get_train_dataloader()))) + loss = _loss_of(wrapper.trainer, wrapper.model, batch) + loss.mean().backward() + grads = [ + param.grad.abs().max().item() + for name, param in wrapper.model.named_parameters() + if "layers.0." in name and "lora_" in name and param.grad is not None + ] + assert grads and max(grads) > 0.0, ( + f"{task}: layer-0 adapter gradient is zero — the graph is severed" + ) + + +# ========================================================================== +# item 1 -- the pre-flight must not under-predict for a concatenating loss +# ========================================================================== +class TestVramPreflightAccountsForPairedRows: + """DPO / ORPO / SimPO build their forward through `concatenated_inputs` + + `torch.cat`, so **2 x batch_size** rows reach the model in ONE tensor. + v0.72.3's estimator was validated on the property that it NEVER + under-predicts; reusing it at 1x rows for these three would break exactly + that, and on Windows the consequence is not an exception but a silent WDDM + spill that makes the run an order of magnitude slower. + + KTO instead runs its KL batch as a SEPARATE forward, so its row count is 1x. + """ + + @pytest.mark.parametrize( + "task,expected", [("sft", 1), ("dpo", 2), ("orpo", 2), ("simpo", 2), ("kto", 1)] + ) + def test_rows_per_example_is_declared_per_task(self, task, expected): + from soup_cli.trainer.dpo import DPOTrainerWrapper + from soup_cli.trainer.kto import KTOTrainerWrapper + from soup_cli.trainer.orpo import ORPOTrainerWrapper + from soup_cli.trainer.sft import SFTTrainerWrapper + from soup_cli.trainer.simpo import SimPOTrainerWrapper + + cls = { + "sft": SFTTrainerWrapper, + "dpo": DPOTrainerWrapper, + "orpo": ORPOTrainerWrapper, + "simpo": SimPOTrainerWrapper, + "kto": KTOTrainerWrapper, + }[task] + assert cls._STREAM_ROWS_PER_EXAMPLE == expected + + @pytest.mark.parametrize("task", ("dpo", "orpo", "simpo")) + def test_the_budget_uses_the_multiplier(self, tmp_path, monkeypatch, task): + """Behavioural, per wrapper: asserting the class attribute alone would + miss a wrapper that reads the wrong one. The predicted peak for a paired + loss at batch 1 must be the prediction for 2 rows, because that IS the + tensor TRL builds.""" + from soup_cli.utils.layer_stream import estimate_stream_peak_vram + + captured = [] + import soup_cli.trainer.stream_setup as setup_mod + + real = ( + setup_mod.estimate_stream_peak_vram + if hasattr(setup_mod, "estimate_stream_peak_vram") + else estimate_stream_peak_vram + ) + + def spy(**kwargs): + captured.append(kwargs["batch_size"]) + return real(**kwargs) + + monkeypatch.setattr("soup_cli.utils.layer_stream.estimate_stream_peak_vram", spy) + _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + assert captured, "the VRAM pre-flight never ran" + assert captured[0] == 2, ( + f"{task} at batch_size=1 budgeted for {captured[0]} rows; the " + f"concatenated chosen+rejected tensor is 2" + ) + + +# ========================================================================== +# item 1 -- resource release +# ========================================================================== +class TestStreamRuntimeIsReleased: + """v0.72.3's code review established this for SFT: `close()` after the + training call is SKIPPED when training raises, and an OOM is realistic on + exactly the cards this feature targets. On the disk tier that leaks one open + shard handle per decoder layer.""" + + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_close_runs_even_when_training_raises(self, tmp_path, monkeypatch, task): + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task=task) + closed = {"n": 0} + real_close = wrapper._stream_runtime.close + + def counting_close(): + closed["n"] += 1 + return real_close() + + wrapper._stream_runtime.close = counting_close + + def boom(*_a, **_k): + raise RuntimeError("CUDA out of memory") + + wrapper.trainer.train = boom + with pytest.raises(RuntimeError, match="out of memory"): + wrapper.train() + assert closed["n"] >= 1, f"{task} leaked the streaming weight source" + + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_close_is_a_noop_without_streaming(self, tmp_path, monkeypatch, task): + """A non-streaming run must not acquire a `_stream_runtime` attribute + and must not fail when the release hook runs anyway.""" + cls = _wrapper_for(task) + import yaml + + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string( + yaml.safe_dump( + { + "base": "sshleifer/tiny-gpt2", + "task": task, + "data": {"train": "train.jsonl"}, + "training": {"quantization": "none"}, + "output": str(tmp_path / "o"), + } + ) + ) + wrapper = cls(cfg, device="cpu") + wrapper._close_stream_runtime() # must not raise + assert getattr(wrapper, "_stream_runtime", None) is None + + +# ========================================================================== +# the shared setup must be shared, not copied +# ========================================================================== +class TestStreamingSetupIsSharedNotCopied: + """Four wrappers x ~270 lines of streaming setup is a drift machine: the + NF4 pre-flight, the tier fallback and the fit refusal would each have five + places to be fixed.""" + + def test_every_streaming_wrapper_uses_the_one_mixin(self): + from soup_cli.trainer.dpo import DPOTrainerWrapper + from soup_cli.trainer.kto import KTOTrainerWrapper + from soup_cli.trainer.orpo import ORPOTrainerWrapper + from soup_cli.trainer.sft import SFTTrainerWrapper + from soup_cli.trainer.simpo import SimPOTrainerWrapper + from soup_cli.trainer.stream_setup import StreamingSetupMixin + + for cls in ( + SFTTrainerWrapper, + DPOTrainerWrapper, + ORPOTrainerWrapper, + SimPOTrainerWrapper, + KTOTrainerWrapper, + ): + assert issubclass(cls, StreamingSetupMixin) + # and does not shadow the shared implementation + assert "_setup_streaming_transformers" not in vars(cls) + + def test_the_mixin_is_import_light(self): + """`stream_setup` is imported by five trainer modules; a top-level torch + there would be a wide blast radius. (The authority on the CLI startup + property is tests/test_cli_startup_is_light.py — this is the cheap + syntactic first line, per CLAUDE.md's evidence rule.)""" + import ast + import pathlib + + import soup_cli.trainer.stream_setup as mod + + tree = ast.parse(pathlib.Path(mod.__file__).read_text(encoding="utf-8")) + for node in tree.body: + if isinstance(node, ast.Import): + for alias in node.names: + assert alias.name.split(".")[0] not in {"torch", "peft", "trl"} + elif isinstance(node, ast.ImportFrom) and node.module: + assert node.module.split(".")[0] not in {"torch", "peft", "trl"} + + +# ========================================================================== +# findings the gate produced that would otherwise not be pinned +# ========================================================================== +class TestRowMultiplierIsLoadBearing: + """The gate's own first attempt at this was NOT evidence: at vocab 64 the + logits term is ~0, so budgeting 1x and 2x rows predict almost the same + number and the check passed for either answer. Re-measured at vocab 32000 + the two are 71.3% apart.""" + + def test_the_two_answers_are_far_apart_at_a_realistic_vocab(self): + from soup_cli.utils.layer_stream import estimate_stream_peak_vram + + def predict(rows): + return estimate_stream_peak_vram( + layer_bytes=4_000_000, + buffers=2, + extras_bytes=32_000_000, + adapter_params=300_000, + vocab_size=32000, + hidden_size=512, + intermediate_size=1024, + n_layers=8, + seq_len=128, + batch_size=rows, + ) + + one, two = predict(1), predict(2) + assert (two - one) / one > 0.15, ( + f"1x={one} 2x={two} are only {(two - one) / one:.1%} apart, so any " + f"test of the multiplier at this shape would be vacuous" + ) + + def test_a_paired_loss_budgets_strictly_more_than_an_unpaired_one(self): + """Behavioural consequence: the same batch_size costs more under a + concatenating loss, because it IS a bigger tensor.""" + from soup_cli.trainer.dpo import DPOTrainerWrapper + from soup_cli.trainer.kto import KTOTrainerWrapper + from soup_cli.trainer.sft import SFTTrainerWrapper + + assert ( + DPOTrainerWrapper._STREAM_ROWS_PER_EXAMPLE > SFTTrainerWrapper._STREAM_ROWS_PER_EXAMPLE + ) + # KTO does NOT concatenate — it runs the KL batch as a separate forward + # — so it must NOT be lumped in with the paired losses. + assert ( + KTOTrainerWrapper._STREAM_ROWS_PER_EXAMPLE == SFTTrainerWrapper._STREAM_ROWS_PER_EXAMPLE + ) + + def test_kto_budgets_its_real_batch(self, tmp_path, monkeypatch): + """KTO's multiplier is 1, but it requires batch_size >= 2, so the budget + must still see 2 rows. Asserting the multiplier alone would miss that.""" + captured = [] + from soup_cli.utils import layer_stream as ls + + real = ls.estimate_stream_peak_vram + + def spy(**kwargs): + captured.append(kwargs["batch_size"]) + return real(**kwargs) + + monkeypatch.setattr(ls, "estimate_stream_peak_vram", spy) + _build_streamed_wrapper(tmp_path, monkeypatch, task="kto") + assert captured and captured[0] == 2, captured + + +class TestKtoNeedsMoreThanOneRow: + """TRL refuses `per_device_train_batch_size == 1` for KTO outright ("the KL + term will be equivalent to the implied reward"), so KTO is streamable ONLY + because v0.72.3 lifted layer streaming's own batch-1 restriction. If a + future change re-freezes batch to 1, KTO silently stops being usable.""" + + def test_batch_two_is_accepted_by_the_streaming_schema(self, tmp_path): + cfg = _stream_cfg(str(tmp_path / "m"), tmp_path / "o", task="kto", batch_size=2) + assert cfg.training.batch_size == 2 + + @pytest.mark.skipif( + _mps_is_the_accelerator(), + reason="MPS is untested for layer streaming (CUDA + CPU only)", + ) + def test_kto_streams_at_batch_two(self, tmp_path, monkeypatch): + import math + + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task="kto", batch_size=2) + assert wrapper._stream_runtime.stats()["n_layers"] == 2 + wrapper.trainer.args.max_steps = 1 + wrapper.trainer.train() + losses = [entry["loss"] for entry in wrapper.trainer.state.log_history if "loss" in entry] + assert losses and math.isfinite(losses[0]), losses + assert wrapper._stream_runtime.pool.loads > 0, "no layer was streamed" + + +class TestTheReferenceForwardActuallyHappens: + """Memory-wise the reference is free; time-wise it is not. DPO traverses the + stack three times per step (policy forward + reference forward + checkpoint + recompute) against SFT's two — measured 1.52x layer reads on a 24-layer + model. Pinned so an "optimisation" that silently drops or caches the + reference forward cannot pass unnoticed.""" + + @pytest.mark.parametrize("task", _REFERENCE_USING) + def test_a_reference_using_loss_reads_more_layers_than_sft(self, tmp_path, monkeypatch, task): + def reads_for(task, root): + # The real device, deliberately: KTO's `get_batch_loss_metrics` + # moves the batch to `self.accelerator.device` itself + # (kto_trainer.py:1349), so pinning the model to CPU here would only + # manufacture a device mismatch. This test counts reads; it needs no + # exact arithmetic. + wrapper, _, _ = _build_streamed_wrapper(root, monkeypatch, task=task, n_layers=4) + _randomise_lora_b(wrapper.model) + pool = wrapper._stream_runtime.pool + counted = {"n": 0} + original = pool.load_async + + def counting(idx, source, stream=None): + counted["n"] += 1 + return original(idx, source, stream) + + pool.load_async = counting + wrapper.model.train() + wrapper.model.zero_grad(set_to_none=True) + batch = _batch_on(wrapper.model, next(iter(wrapper.trainer.get_train_dataloader()))) + if task == "sft": + # TRL's SFTTrainer has no `get_batch_loss_metrics`; its step is + # a plain causal-LM forward + backward, which is the point of + # the comparison anyway. + ids = batch["input_ids"] + out = wrapper.model( + input_ids=ids, attention_mask=batch.get("attention_mask"), labels=ids + ) + out.loss.backward() + else: + _loss_of(wrapper.trainer, wrapper.model, batch).mean().backward() + pool.load_async = original + wrapper._close_stream_runtime() + return counted["n"] + + sft_reads = reads_for("sft", tmp_path / "s") + pref_reads = reads_for(task, tmp_path / "d") + assert pref_reads > sft_reads, ( + f"{task} did {pref_reads} layer reads vs SFT's {sft_reads}: the " + f"reference forward appears not to be running at all" + ) + + +class TestKtoBatchIsRefusedEarly: + """TRL raises "Actual (not effective) batch size must be > 1" from + `KTOTrainer.__init__`, which under streaming runs only AFTER the RAM + pre-flight, the checkpoint sharding and (at `quantization: 4bit`) the NF4 + quantisation — minutes of disk I/O on a real base, to fail on a config that + was invalid before any of it started.""" + + def test_batch_one_is_refused_at_parse_time(self, tmp_path): + with pytest.raises(ValueError) as excinfo: + _stream_cfg(str(tmp_path / "m"), tmp_path / "o", task="kto", batch_size=1) + message = str(excinfo.value) + assert "kto" in message.lower() + assert "batch_size" in message + + def test_batch_two_still_parses(self, tmp_path): + cfg = _stream_cfg(str(tmp_path / "m"), tmp_path / "o", task="kto", batch_size=2) + assert cfg.training.batch_size == 2 + + def test_the_other_streaming_tasks_are_unaffected(self, tmp_path): + """Control: batch 1 is perfectly valid for the other four, so a gate + that refused it everywhere would be a regression, not a fix.""" + for task in ("sft", "dpo", "orpo", "simpo"): + cfg = _stream_cfg(str(tmp_path / "m"), tmp_path / "o", task=task, batch_size=1) + assert cfg.training.batch_size == 1 + + def test_non_streaming_kto_is_left_alone(self, tmp_path): + """Scoped to streaming deliberately: resident KTO at batch 1 fails the + same way, but that is pre-existing behaviour outside this slot, and + widening the gate here could reject configs that parse today.""" + import yaml + + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string( + yaml.safe_dump( + { + "base": "sshleifer/tiny-gpt2", + "task": "kto", + "data": {"train": "t.jsonl"}, + "training": {"batch_size": 1, "quantization": "none"}, + "output": str(tmp_path / "o"), + } + ) + ) + assert cfg.training.batch_size == 1 + + def test_trl_itself_still_refuses_batch_one(self, tmp_path): + """Pins the UPSTREAM behaviour our schema gate mirrors. + + The gate duplicates a TRL threshold. This box runs trl 0.19.1 while CI + runs 1.9.2 (#323), so if a TRL version drops or changes the requirement + our gate would silently refuse configs that upstream now accepts. Assert + against the INSTALLED trl rather than against a version number, so + whichever one CI has is the one that gets checked. + """ + import torch + from trl import KTOConfig, KTOTrainer + + weights, _, _ = _tiny_llama_dir(tmp_path) + _write_tiny_tokenizer(weights) + from datasets import Dataset + from transformers import AutoModelForCausalLM, AutoTokenizer + + model = AutoModelForCausalLM.from_pretrained(weights, torch_dtype=torch.float32) + tokenizer = AutoTokenizer.from_pretrained(weights) + dataset = Dataset.from_list(_kto_rows(4)) + args = KTOConfig( + output_dir=str(tmp_path / "o"), + per_device_train_batch_size=1, + report_to=[], + max_length=32, + max_prompt_length=16, + ) + with pytest.raises(ValueError, match="batch size"): + KTOTrainer( + model=model, + args=args, + train_dataset=dataset, + processing_class=tokenizer, + ) + + +class TestKtoReferenceIsAlsoTheDisabledAdapter: + """CRITICAL gap found in review: KTO's reference forward is + `self.forward(self.model, batch)` (`kto_trainer.py:1400`) — the trainer's + OWN bound model, ignoring whatever `model` argument reaches + `get_batch_loss_metrics`. So a streamed-vs-resident loss comparison computes + the SAME reference term in both arms and exercises only the policy forward. + KTO's reference path therefore needs its own direct assertion, the analogue + of the DPO one.""" + + def _forwards(self, wrapper): + import torch + + trainer = wrapper.trainer + batch = _batch_on(wrapper.model, next(iter(trainer.get_train_dataloader()))) + wrapper.model.eval() + with torch.no_grad(): + policy = trainer.forward(wrapper.model, batch)[0] + with trainer.null_ref_context(): + reference = trainer.forward(wrapper.model, batch)[0] + return policy, reference + + def test_reference_differs_from_policy(self, tmp_path, monkeypatch): + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task="kto", device="cpu") + _randomise_lora_b(wrapper.model) + policy, reference = self._forwards(wrapper) + assert (policy - reference).abs().max().item() > 1e-4, ( + "KTO's null_ref_context() is a no-op through the streamed layer's " + "functional_call, so the reference IS the policy" + ) + + def test_the_difference_really_comes_from_the_adapter(self, tmp_path, monkeypatch): + """CONTROL: with `lora_B = 0` the adapter contributes nothing, so + disabling it must change nothing at all.""" + import torch + + wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task="kto", device="cpu") + with torch.no_grad(): + for name, param in wrapper.model.named_parameters(): + if "lora_B" in name: + param.zero_() + policy, reference = self._forwards(wrapper) + assert (policy - reference).abs().max().item() == 0.0 + + +class TestNf4CombinesWithEveryPreferenceLoss: + """The schema allows `quantization: 4bit` with all four, and the NF4 + `total_params` override was hand-copied into each wrapper (PEFT sizes a + `meta` Params4bit placeholder as `numel * 2 * itemsize`, over-reporting a + streamed NF4 model ~6.5x — measured 878,154,048 vs a true 134,515,008 on + SmolLM2-135M). A name drift in any one copy is otherwise invisible.""" + + @pytest.mark.skipif( + _mps_is_the_accelerator(), + reason="bitsandbytes has no 4-bit MPS kernels", + ) + @pytest.mark.parametrize("task", _ALL_PREFERENCE) + def test_nf4_streaming_sets_up_and_reports_honest_parameters(self, tmp_path, monkeypatch, task): + pytest.importorskip("bitsandbytes") + wrapper, _, _ = _build_streamed_wrapper( + tmp_path, monkeypatch, task=task, quantization="4bit" + ) + assert wrapper.trainer is not None + honest = wrapper._stream_runtime.total_params + assert honest > 0 + peft_says = wrapper.model.get_nb_trainable_parameters()[1] + assert honest != peft_says, ( + "PEFT's total now agrees with the sharder's, so the override this " + "test guards is either unnecessary or no longer wired" + ) + real = sum(p.numel() for p in wrapper.model.parameters() if not p.is_meta) + assert honest > real, (honest, real) + wrapper._close_stream_runtime() + + def test_auto_batch_still_reports_its_own_reason_for_kto(self, tmp_path): + """The `batch_size='auto'` refusal and the KTO `>= 2` refusal are two + different guards over the same field. `auto` is not an int, so the KTO + branch must not swallow it and report the wrong reason.""" + with pytest.raises(ValueError) as excinfo: + _stream_cfg(str(tmp_path / "m"), tmp_path / "o", task="kto", batch_size="auto") + message = str(excinfo.value) + assert "auto" in message + assert "RESIDENT" in message or "probe" in message diff --git a/tests/test_version_sync.py b/tests/test_version_sync.py new file mode 100644 index 0000000..af47b8e --- /dev/null +++ b/tests/test_version_sync.py @@ -0,0 +1,82 @@ +"""The version lives in two files and nothing asserted they agree. + +`pyproject.toml` and `src/soup_cli/__init__.py` both carry the version, and the +release checklist says to bump both. Every other version test in this suite is a +FLOOR check (`>= 0.56.0`, `>= 0.53.11`, ...), which passes happily when one file +is bumped and the other is forgotten — the wheel then ships with metadata that +disagrees with `soup version`, and CI stays green. Until now that drift was +caught only by a human eye on step 6a of the checklist. +""" + +import pathlib +import re + +import pytest + +ROOT = pathlib.Path(__file__).resolve().parents[1] +PYPROJECT = ROOT / "pyproject.toml" +INIT = ROOT / "src" / "soup_cli" / "__init__.py" + +_PYPROJECT_VERSION = re.compile(r'^version\s*=\s*"([^"]+)"', re.M) +_INIT_VERSION = re.compile(r'^__version__\s*=\s*"([^"]+)"', re.M) + + +def _pyproject_version(text: str) -> str: + """The version from `[project]`, not from some other table. + + A naive first-match search would also accept a `version = "..."` belonging to + a tool section, so anchor on the `[project]` table and stop at the next one. + """ + project = re.search(r"^\[project\]\s*$(.*?)^\[", text, re.M | re.S) + assert project, "pyproject.toml has no [project] table" + match = _PYPROJECT_VERSION.search(project.group(1)) + assert match, "no version line in the [project] table of pyproject.toml" + return match.group(1) + + +def _init_version(text: str) -> str: + match = _INIT_VERSION.search(text) + assert match, "no __version__ line in src/soup_cli/__init__.py" + return match.group(1) + + +class TestVersionSync: + def test_the_two_files_agree(self): + declared = _pyproject_version(PYPROJECT.read_text(encoding="utf-8")) + package = _init_version(INIT.read_text(encoding="utf-8")) + assert declared == package, ( + f"pyproject.toml says {declared!r} but src/soup_cli/__init__.py says " + f"{package!r} — the release checklist requires BOTH to be bumped. " + f"The wheel would ship metadata that disagrees with `soup version`." + ) + + def test_the_importable_package_agrees_too(self): + """The parsed file and the imported module can differ if a stale build + artifact shadows the source tree — the exact "stale shadow install" + the checklist warns about at step 6a.""" + from soup_cli import __version__ + + assert __version__ == _init_version(INIT.read_text(encoding="utf-8")) + + @pytest.mark.parametrize( + "text,expected", + [ + ('[project]\nname = "x"\nversion = "1.2.3"\n\n[tool.ruff]\n', "1.2.3"), + # a version in another table must not be mistaken for the project's + ('[project]\nname = "x"\nversion = "1.2.3"\n\n[tool.x]\nversion = "9.9.9"\n', "1.2.3"), + ], + ) + def test_parser_reads_the_project_table(self, text, expected): + assert _pyproject_version(text) == expected + + def test_the_assertion_would_actually_catch_a_drift(self): + """CONTROL. Without this, `test_the_two_files_agree` proves nothing about + the comparison itself — it would pass just as happily against a helper + that returned the same constant twice.""" + drifted = _pyproject_version('[project]\nversion = "0.99.0"\n\n[tool.x]\n') + package = _init_version('__version__ = "0.72.4"\n') + assert drifted != package + + def test_version_is_a_plausible_release_string(self): + version = _init_version(INIT.read_text(encoding="utf-8")) + assert re.fullmatch(r"\d+\.\d+\.\d+(?:[.-]?(?:a|b|rc|dev)\d*)?", version), version