diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d707c77..032868e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -111,7 +111,7 @@ soup_cli/ templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (184 files, 7722 tests) +tests/ - Test suite (185 files, 7842 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index dcbb286..f7b7858 100644 --- a/README.md +++ b/README.md @@ -43,14 +43,14 @@ soup train Latest highlights only. Full history: [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). -**v0.53.1 — Quant Menu II + Export pipeline live**: Six v0.53.0 deferred stubs lifted — autopilot pre-quantized detection, single-stage BNB-4bit merge, TorchAO PTQ export, Unsloth Dynamic 2.0 GGUF ladder via llama.cpp `imatrix`, and `soup deploy autopilot --measure` Quant-Lobotomy scorecard. +**v0.53.2 — Modality II live trainers**: Four v0.52.0 deferred stubs lifted into real, end-to-end-trainable wrappers — knowledge distillation, sequence classification, EBFT / GDPO loss kernels, and gpt-oss-style `reasoning_effort` system-prompt injection. -- **Autopilot detects pre-quantized bases.** `TheBloke/Llama-2-7B-Chat-GPTQ` now recommends `gptq` instead of stacking BNB-4bit on top. Name-regex + `config.json` `quantization_config.quant_method` probe with cwd-containment + symlink rejection on the on-disk path. BNB aliases (`bitsandbytes_4bit` / `nf4` / `bnb_8bit`) canonicalise to `4bit` / `8bit`. -- **`soup merge --save-format 4bit | 4bit_forced`.** Single BNB-4bit-quantized merged checkpoint without the dequant→merge→requant cycle. `4bit_forced` quantizes every Linear (including `lm_head`). Output path is cwd-contained + symlink-rejected at CLI dispatch. -- **`soup export --format torchao --quant-config `.** Live `torchao.quantize_` + `save_pretrained`. Closed per-scheme kwarg allowlist (Int4WeightOnly accepts `{group_size, inner_k_tiles}`, NVFP4 accepts nothing extra) defeats kwarg-injection through the YAML. -- **`soup export --format gguf-ud --gguf-flavour `.** Three-stage pipeline: HF → f16 GGUF → optional importance-matrix (UD ladder + low-bit IQ) → quantize. All subprocess calls use argv-list form + 30-min timeout. Calibration JSONL is sanitised (null-byte stripped, newlines collapsed, 8 KB per-line + 50 MB total cap). POSIX `O_NOFOLLOW` defeats the TOCTOU race between the dispatch-time symlink check and the actual open. -- **`soup deploy autopilot --measure --tasks `.** Loops every candidate quant through the v0.26.0 `eval/quant_check` scorer, renders OK/MINOR/MAJOR table, picks the best-by-delta candidate. Results cached at `~/.soup/deploy_autopilot_cache.json` (atomic write, 0o600 perms on POSIX, symlink-rejected on both load and save). Cache key is SHA-256 of `(base, profile, eval-tasks)`. -- **+112 net new tests** (7610 → 7722) across `test_v0531_82.py`, `test_v0531_109.py`, `test_v0531_139.py`, `test_v0531_142.py`. Four review agents (python / code / security / tdd) ran; every CRITICAL / HIGH / MEDIUM / LOW finding fixed: per-scheme TorchAO kwarg allowlist (rejects dunders + unknown keys), corrected BNB 4-bit skip-modules kwarg name, shared `enforce_under_cwd_and_no_symlink` in `utils/paths.py` (single source of truth), `pick_best` switches from `max(after)` to `max(delta)` matching the v0.33.0 #54 design intent, `_run_convert_to_f16` verifies the convert script stays inside `llama_cpp_dir` via realpath + commonpath, `_safe_stderr` Rich-escapes subprocess stderr before exception propagation. +- **`soup train` with `task: distill`.** New `DistillTrainerWrapper`: student + frozen teacher both load via `AutoModelForCausalLM` (separate `trust_remote_code` resolution for each), KL / forward_KL / reverse_KL / JS divergence kernels scaled by `temperature**2` per the Hinton paper. Device-bridge: teacher inputs auto-move to the teacher's device, teacher logits move back onto the student's device before the KL kernel — survives HF Trainer's auto-CUDA promotion on a CPU-tagged run. `DataCollatorForSeq2Seq(label_pad_token_id=-100)` handles variable-length pre-tokenised loss-masked rows correctly. +- **`soup train` with `task: classifier | reranker | cross_encoder`.** New `ClassifierTrainerWrapper`: `AutoModelForSequenceClassification` with `num_labels` and `label_names`, auto-routes `single_label_classification` / `multi_label_classification` from `tcfg.classifier_kind`. Multi-label string labels resolved via the `label_names` map with a 1024-entry cap + dedup. Training Setup Panel renders `Head: num_labels=N, kind=...` instead of LoRA r/alpha for the classifier family. +- **EBFT structured / strided + GDPO standard / length_normalized / margin loss kernels.** `apply_ebft_loss` and `apply_gdpo_loss` exit the v0.52.0 `NotImplementedError` stubs with finite-only-input guards and bool-rejected numeric params. `attach_ebft_compute_loss(trainer, tcfg)` (SFT) and `attach_gdpo_compute_loss(trainer, tcfg)` (DPO) wrap `Trainer.compute_loss` idempotently — re-attach is a no-op via a marker attribute on the wrapped method. Auto-attached when the corresponding `*_variant` field is set on `TrainingConfig`. +- **gpt-oss `reasoning_effort` + `train_on_eot`.** `apply_reasoning_effort_prefix(messages, level)` injects `<|reasoning_effort|>{low,medium,high}<|/reasoning_effort|>` into the system turn (creates one if absent), returning a new list (caller's messages immutable). `build_assistant_only_labels(train_on_eot=True)` keeps the EOT/EOS token unmasked at the assistant-turn boundary so the model learns when to stop. Both gated to the SFT-family at config-load. +- **+120 net new tests** (7722 → 7842) across `test_v0532.py`. Four review agents (python / code / security / tdd) ran; every CRITICAL / HIGH / MEDIUM / LOW finding fixed — separate `trust_remote_code` resolution for student vs teacher, idempotent attach hooks with regression tests, 1024-entry multi-label cap, `dpo_margin` defaults to `None` (not `0.0`) so missing values raise rather than silently zero, source-grep regression guards on the trainer-routing call sites use the full instantiation expression (no comment-only false-positives), Panel renders the classifier head instead of LoRA r/alpha. +- **Local end-to-end CPU smoke** confirms both new wrappers train 2 steps with finite loss on `hf-internal-testing/tiny-random-gpt2`. Two real bugs surfaced and were fixed during the smoke (collator label padding + teacher / student device mismatch) — both have source-level regression guards in the test suite. ONNX export QA: pipeline integrity proven on tiny-gpt2; TinyLlama-1.1B full export is host-RAM-bound (documented in `tests/qa/v053_qa.md`). ## Why Soup? @@ -275,6 +275,97 @@ soup init --template pretrain soup train ``` +## Knowledge Distillation + +Train a small student model to match a larger teacher's output distribution. + +```yaml +base: HuggingFaceTB/SmolLM2-135M +task: distill +modality: text +backend: transformers + +data: + train: ./data/chat.jsonl + max_length: 2048 + chat_template: chatml + +training: + teacher_model: meta-llama/Llama-3.1-8B + distill_divergence: forward_kl # kl | forward_kl | reverse_kl | js + distill_temperature: 2.0 + epochs: 3 + lr: 5e-5 + quantization: 4bit # quantizes student only +``` + +Loss = student CE + (T**2) × KL(teacher_logits / T || student_logits / T). +Teacher is loaded once, frozen via `requires_grad_(False)` + `.eval()`, and its +inputs / logits are auto-bridged across CPU / CUDA devices. + +## Sequence Classification + +Train a classifier head on top of any base model — supports single-label, +multi-label, and cross-encoder reranking. + +```yaml +base: BAAI/bge-base-en-v1.5 +task: classifier # or `reranker`, `cross_encoder` +modality: text +backend: transformers + +data: + train: ./data/labelled.jsonl # rows: {"text": "...", "label": "spam"} or {"text": "...", "label": [0, 1, 0]} + max_length: 256 + +training: + num_labels: 3 + classifier_kind: single_label # or `multi_label` + label_names: [ham, spam, promo] # required when labels are strings + epochs: 5 + lr: 2e-5 + batch_size: 32 +``` + +Routes `classifier` / `reranker` / `cross_encoder` through +`AutoModelForSequenceClassification`. Multi-label heads cap at 1024 entries per +row, dedup via set conversion, and reject null bytes in label strings. + +## Reasoning Effort + EOT Control + +gpt-oss-style reasoning-effort control for instruction tuning. + +```yaml +training: + reasoning_effort: high # low | medium | high + train_on_eot: true # do NOT mask the EOT/EOS token in the loss +``` + +`reasoning_effort` injects `<|reasoning_effort|>high<|/reasoning_effort|>` into +the system turn (creating one if absent). `train_on_eot=True` makes the model +learn when to stop generating by training on the trailing EOS token instead of +masking it out. Both are gated to the SFT-family of tasks. + +## EBFT / GDPO Loss Variants + +Entropy-regularised SFT (`ebft_variant: structured | strided`) and generalised +DPO (`gdpo_variant: standard | length_normalized | margin`) — both attach +idempotently via `compute_loss` wrappers and auto-fire when the corresponding +variant field is set on `TrainingConfig`. + +```yaml +# SFT with EBFT structured +training: + ebft_variant: structured + ebft_temperature: 1.0 + +# DPO with GDPO length_normalized +task: dpo +training: + gdpo_variant: length_normalized + dpo_beta: 0.1 +``` + ## MoE Model Support Fine-tune Mixture of Experts models (Mixtral, Qwen3-30B-A3B, DeepSeek V3) with ScatterMoE LoRA — applies LoRA to both attention layers and expert FFN layers: diff --git a/SECURITY.md b/SECURITY.md index e7b42bb..7f08c88 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -9,7 +9,8 @@ We provide security updates for the following versions: - **Versions older than 3 minor versions:** No support Example: -- v0.53.1 -- Full support (latest) +- v0.53.2 -- Full support (latest) +- v0.53.1 -- Full support - v0.53.0 -- Full support - v0.52.0 -- Full support - v0.51.0 -- Full support @@ -147,6 +148,8 @@ No known critical vulnerabilities in current releases. - **v0.32.0 — Training Stability & Auto-Tuning**: `--find-lr-output` containment via shared `utils/paths.is_under_cwd` (prevents writes outside cwd); `save_lr_finder_report` rejects NaN / Infinity floats in `lrs` / `losses` and serialises with `allow_nan=False` (keeps the report parser-safe); `compute_lr_schedule` rejects non-positive `start_lr`, inverted ranges, and `num_steps` outside `[2, 10_000]`; `pick_mixed_precision` rejects empty / null-byte / >200-char model names and resolves multi-version quirks (`qwen2.5` vs `qwen2`, `phi-3.5` vs `phi-3`) by longest-substring-first iteration so an added family can never accidentally make a more-specific entry dead code; `compute_warmup_steps` clamps to `[10, 1000]` with a `ratio==0.0` short-circuit matching HF Trainer's "no warmup" convention; `SpikeRecoveryStrategy` is `@dataclass(frozen=True)` (post-construction mutation cannot bypass validation), `max_attempts ∈ [1, 10]`, `lr_decay ∈ (0, 1)`, `min_lr > 0`; cross-validator `_validate_spike_recovery_requires_watchdog` rejects `loss_spike_recovery=true, loss_watchdog=false` at config-load (fails fast instead of never triggering); `convergence_window ∈ [5, 10_000]`, `convergence_rel_tol ∈ (0, 1]`, `recommend_action` reuses `detect_plateau` so plateau heuristic stays single-source-of-truth; `GradAccumMonitor.recommend()` caps doubled `accum` at `MAX_ACCUM=1024` so a runaway advisory loop cannot blow up DataLoader prefetch; `generate_config` validates BOTH the YAML output path AND the embedded `decisions["output"]` field via `is_under_cwd` (closes the gap where a crafted `decisions["output"]="../../etc"` would have silently propagated into the rendered YAML) - **v0.34.0 — Observability & Dev UX**: `.crash` bundle generator (`utils/crash.py`) recursively redacts `hf_*` / `sk-*` / `Bearer …` token-shaped strings in any captured `config` and metric tail before serialisation, so a `.crash` file shared on a public GitHub issue cannot leak credentials; `output_dir` is reduced to `os.path.basename` so `$HOME` doesn't leak; `write_crash_bundle` uses `os.path.realpath + commonpath` for cwd containment (Windows-safe; raises `ValueError` not `PermissionError` so callers cannot silently swallow with `except OSError`); filename appends `secrets.token_hex(4)` so two crashes in the same UTC second don't collide; bundle truncated to `MAX_BUNDLE_BYTES=1_000_000`. `train.py` crash-write surfaces failures to the user (no silent missing-bundle). `profiling.py` `resolve_trace_path` rejects empty / `.` / `..` / `/` / `\\` / null-byte `run_id` (closes the `output_dir/profiles/../trace.json` escape) and uses `os.path.realpath + is_under_cwd`; profiles dir is created only on successful torch import (no stale empty dirs on torch-less CI). `tracker.get_run` LIKE-prefix match escapes `%` / `_` / `\\` and uses `ESCAPE '\\'` so a crafted `run_id` cannot widen the match (mirrors v0.26.0 registry policy). Lazy schema migration (`_ensure_schema`) tolerates the "duplicate column" race when two CLI processes start simultaneously on a fresh DB (fork-based multi-GPU training, TUI auto-refresh). `runs.py show/replay/clean` switched user `run_id` rendering to `markup_escape` and switched `clean` containment from broken `Path.resolve() + relative_to()` to project-standard `os.path.realpath + is_under_cwd`. `tui_app.py` lazy-imports `ExperimentTracker` and `markup_escape`s every DB-sourced string before passing into Textual widgets so a crafted base_model / experiment_name cannot inject `[bold red]…[/]` markup. `run_cost.estimate_run_cost_usd` rejects `bool` in `num_gpus` (bool is a subclass of int — same defence as v0.30.0 `Candidate.__post_init__`); duration clamped to `[0, 1 year]`; unknown GPU returns `None` so callers render `—` instead of fabricating `$0.00`. `log_level.parse_log_level` rejects non-string + null-byte input. - **v0.33.0 — Live Wire**: RLVR `code_exec_reward` adds OS-level isolation (Linux best-effort `os.unshare(CLONE_NEWUSER|CLONE_NEWNET|CLONE_NEWPID)`, macOS `sandbox-exec` with default-deny `MACOS_SANDBOX_PROFILE` narrowed to a 3-name `mach-lookup` allowlist to prevent DNS / NSURLSession bypass of `(deny network*)`); `prune_checkpoints` switches to TOCTOU-safe `os.lstat + S_ISLNK` + `shutil.rmtree(onerror=_abort_on_symlink)` so a symlink encountered mid-walk aborts rather than escapes; `run_gate` wraps each task scorer in a typed `try/except` so backend failures produce `score=None, error=str(exc)` (never silent `score=1.0`); `_parse_judge_url` removes the bare `http://` catch-all (defence-in-depth after the Pydantic GateTask validator); `soup can run` requires `--yes` or explicit consent callback and raises `ValueError` (not `PermissionError`, which is an `OSError` subclass that broad `except` blocks would swallow); GGUF `rglob` result for ollama deploy is `realpath+commonpath` checked against extract_dir (prevents symlink escape from a crafted can); `DeployTarget.path` validator normalises mixed `\\`/`/` separators before splitting (closes a Windows `..` bypass); `CAN_FORMAT_VERSION` 1→2 (additive — v1 still loads); `soup can publish` validates `repo_id` via `utils/hf.validate_repo_id`, resolves token via `resolve_token`, sanitises commit messages (first-line, 200-char cap), uses HTTPS-only HfApi; `_write_spike_recovery_hint` adds `is_under_cwd` containment check on `args.output_dir` from raw HF `TrainingArguments`; `lookup_entry_by_output_dir` emits `ResourceWarning` when 1000-row scan limit is hit (no silent miss); `CrossDocCollator` no longer mutates input feature dicts (HF Dataset rows are cached and reused — mutation broke subsequent batches); `Candidate` rejects `bool` in `score`/`latency_ms` (was sneaking past `int` isinstance check); `evaluate_candidate` latency mean now divides by *completed* prompts (excludes crashed) so a broken candidate isn't artificially fast; `auto_quant.run_auto_quant_picker` soft-falls-back to highest-scored candidate when no candidate clears `min_score` (server still binds); `build_logits_processors` returns `[]` when neither `outlines` nor `lm-format-enforcer` is installed (server degrades to free-form rather than 500); MII server uses loopback-only CORS, max_tokens cap [1, 16384], stream rejection, generic 500 with no stack-trace leak; `os.execvp` auto-reexec uses list args (no shell), all forwarded flags pre-validated; `cleanup_extract_dir` uses `os.path.commonpath` (Windows-safe) instead of `startswith`; `_run_subprocess` catches `TimeoutExpired` and returns rc=124 (coreutils convention) instead of an unhandled traceback; new `eval_results` and `tensorrt` artifact kinds in `RegistryStore._VALID_KINDS` +- **v0.53.2 — Modality II live trainers**: lifts four v0.52.0 deferred stubs (#137, #135, #133, #132) into real trainer wrappers while keeping the project's hardening invariants. (#137 reasoning_effort + train_on_eot) `apply_reasoning_effort_prefix` follows v0.41.0 / v0.51.0 validator policy (bool-first, null-byte / empty / oversize / case-insensitive normalisation); messages list is treated as immutable (returns a new list — matches v0.33.0 #47 `CrossDocCollator` policy). `build_assistant_only_labels(train_on_eot=True)` reuses the existing v0.36.0 mask infrastructure — same null-byte / max_length / bool guards. (#135 EBFT / GDPO) `apply_ebft_loss` and `apply_gdpo_loss` enforce **finite-only inputs** (`torch.isfinite` guard on tensor inputs + `math.isfinite` on scalar params) — NaN / Inf would silently corrupt training otherwise. `dpo_margin` defaults to `None` (not `0.0`) per security-review M3 fix: silent zeroing in the `margin` variant when the operator forgot to set the margin would have looked like training success but produced a meaningless gradient. Both attach hooks (`attach_ebft_compute_loss`, `attach_gdpo_compute_loss`) are **idempotent** via a marker attribute on the wrapped method — re-attach is a no-op and a dedicated test class verifies the invariant (code-review M2 fix). (#133 DistillTrainerWrapper) **Separate trust_remote_code resolution for student and teacher** (security-review L2 fix): `model_requires_trust_remote_code(teacher)` runs independently of the student probe, otherwise a malicious teacher could piggy-back on the student's opt-in. Teacher is loaded with `device_map="cpu" if device == "cpu" else "auto"`, frozen via `requires_grad_(False)` + `.eval()` immediately after load — never participates in gradient computation. `_DistillTrainer.compute_loss` device-bridge: `teacher_device = next(teacher_ref.parameters()).device`, `teacher_inputs.to(teacher_device)` before teacher forward, `teacher_logits.to(student_logits.device)` before KL kernel — defends against HF Trainer's auto-CUDA promotion silently producing cross-device `index_select` crashes. **DataCollator correctness fix** (surfaced during Wave 3 CPU smoke): `DataCollatorForLanguageModeling` does NOT pad pre-tokenised `labels` — switched to `DataCollatorForSeq2Seq(label_pad_token_id=-100, padding=True)` so variable-length loss-masked rows batch correctly without runtime crash. (#132 ClassifierTrainerWrapper) `_normalise_label` caps multi-label entries at **1024 per row** (matches v0.52.0 schema cap; security-review HIGH fix — unbounded would allow OOM via crafted JSONL), dedups via set conversion, validates `label_names` map entries reject null bytes + empty strings. `problem_type` is set explicitly from `tcfg.classifier_kind` (not silently inferred from labels) so a multi-label-shaped row in a single-label config raises rather than mis-trains. Training Setup Panel renders `Head: num_labels=N, kind=...` for classifier-family tasks instead of meaningless LoRA r/alpha lines (code-review L3 cosmetic fix — Panel no longer mis-represents what the wrapper is doing). (Cross-cutting) `commands/train.py` task routing branches added for `distill` and `classifier` / `reranker` / `cross_encoder` — source-grep regression guards in the test suite use the **full instantiation expression** `DistillTrainerWrapper(cfg, **trainer_kwargs)` so comment-only mentions of the class name cannot satisfy the regression check (TDD-review hardening). Both new factories (`build_distill_trainer`, `build_classifier_trainer`) reject unknown kwargs via Python signature contract — dedicated `pytest.raises(TypeError)` tests cover the path (TDD-review L1 fix). Test surface: 1 new test file (`test_v0532.py`) carrying 120 new tests across 14 classes. Known limitations: (1) `#71` TinyLlama-1.1B-LoRA full ONNX export is host-RAM-bound (≥16 GB free RAM needed for the `onnx.load(load_external_data=True)` post-process step); tiny-gpt2 smoke proves pipeline integrity — recorded in `tests/qa/v053_qa.md`. (2) Distillation supports same-tokenizer pairs only — cross-tokenizer (Llama → Qwen) needs a projection or sequence-level loss, out of scope. (3) Classifier wrapper has no LoRA path — full head + base training; LoRA classifier finetuning is a follow-up. (4) EBFT / GDPO auto-attach only fires when the corresponding `*_variant` field is set; manual `attach_*` invocation from custom training loops is supported and idempotent. (5) `reasoning_effort` injection happens at data-prep time inside `build_format_row`; changing the level between runs requires re-rendering the dataset. (v0.53.2) + - **v0.53.1 — Quant Menu II + Export pipeline live**: lifts six v0.53.0 deferred stubs to live wiring while keeping the project's hardening invariants. New shared helper `soup_cli/utils/paths.enforce_under_cwd_and_no_symlink` consolidates the v0.33.0 #22 TOCTOU pattern (cwd containment via `os.path.realpath + os.path.commonpath` + `os.lstat + S_ISLNK` rejection) — used by `commands/merge.py`, `commands/export.py`, `utils/save_formats.py`, and `utils/gguf_quant.py` so the same boundary check fires at every CLI dispatch point. `merge_4bit` and `export_torchao` (`utils/save_formats.py`): cwd containment + symlink rejection on `merged_dir` / `model_dir` / `output_dir`; `load_quant_config` enforces `yaml.safe_load` only + 256 KB cap + extension allowlist (`.yaml`/`.yml`); **per-scheme closed kwarg allowlist** rejects dunder keys + unknown params before the splat into `torchao.Config(**kwargs)` (security-review HIGH fix — `Int4WeightOnly` accepts `{group_size, inner_k_tiles}`, `NVFP4` accepts nothing extra). Corrected BNB-4bit skip-modules kwarg name from `llm_int8_skip_modules` to `bnb_4bit_skip_modules`. `export_advanced_gguf` (`utils/gguf_quant.py`): all three subprocess invocations (`convert_hf_to_gguf.py`, `llama-imatrix`, `llama-quantize`) use argv-list form with no shell, 30-min timeout, `sys.executable` for the convert script; `_run_convert_to_f16` realpath-verifies that `convert_hf_to_gguf.py` stays inside the `llama_cpp_dir` after resolution (security-review HIGH M5 fix — defends against a symlinked script escape). `_prepare_calibration_text` strips null bytes, collapses newlines to spaces, caps per-line at 8 KB + total at 50 MB (security-review M1), uses POSIX `O_NOFOLLOW` to refuse symlinks at the kernel level (security-review M3 — closes the TOCTOU window between the dispatch-time check and the actual `open()`); requires ≥ 1 usable row before invoking imatrix. `_safe_stderr` Rich-markup-escapes subprocess stderr before embedding in `RuntimeError` (security-review L4) so a crafted llama.cpp error cannot inject `[red]...[/]` into the operator-facing panel. UD-prefix stripped from flavour arg before passing to llama-quantize (`UD-Q4_K_XL` → `Q4_K_XL`). Calibration data path containment + symlink rejection fires at CLI dispatch in `commands/export.py::_export_gguf_advanced`. `detect_prequantized_format_from_path` (`autopilot/decisions.py`): cwd containment + `os.lstat + S_ISLNK` on `/config.json` (security-review HIGH H2 — out-of-cwd model paths silently return `None` to preserve soft-probe semantics so HF Hub repo IDs aren't rejected); null-byte rejection on `model_dir`. `commands/merge.py`: early `is_under_cwd(output)` check at CLI boundary (security-review M4) — consistent with the v0.20.0 / v0.40.2 containment-at-the-boundary policy. `deploy_measure.py`: cache file written atomically via `tempfile.mkstemp` + `os.replace` with `os.lstat + S_ISLNK` rejection on BOTH `load_cache` and `save_cache` (security-review M2 — was missing on the load side); env override `SOUP_DEPLOY_AUTOPILOT_CACHE` rejects null bytes + control chars before any path resolution and confines the override to home / cwd / tempdir; cache file gets best-effort 0o600 perms on POSIX (matches v0.26.0 registry.db policy); 1 MB cache-file cap. `_DEPLOY_MEASURE_BEFORE_GEN` / `_AFTER_FACTORY` module-level callables are documented as a non-public escape hatch (deferred until v0.46.1 live model-loader). Test surface: 4 new test files (`test_v0531_82.py` / `test_v0531_109.py` / `test_v0531_139.py` / `test_v0531_142.py`) carrying 112 new tests covering happy paths + failure modes + every security guard (POSIX symlink rejection, per-scheme kwarg allowlist, TOCTOU defences, `_MAX_CANDIDATES` cap, MINOR-verdict band, mxfp4 word boundary, BNB-alias detection, render-table markup escape). Known limitations: (1) `_DEPLOY_MEASURE_BEFORE_GEN` / `_AFTER_FACTORY` are a stop-gap until v0.46.1 ships first-party transformers / vLLM generator factories. (2) `#70` GGUF and `#72` AWQ/GPTQ manual QA smokes remain pending — require CUDA + llama.cpp build; recipes scripted in `tests/qa/v053_qa.md`. (3) BNB-4bit merge + TorchAO PTQ live happy-path is mock-covered only — CPU-only CI cannot execute the real BNB / torchao kernels. (4) `_prepare_calibration_text` accepts JSONL with `text` / `prompt` / `content` aliases + raw text fallback; other formats (parquet / markdown) are out of scope. (5) Cache key truncates `base_sha` to 16 hex chars at the call site (collision probability ≈ 1-in-2³² across ~4 billion entries). (6) Pre-quantized detection is heuristic — name regex + local `config.json` probe; HF Hub repo IDs without local download fall back to name-only matching. (7) `enforce_under_cwd_and_no_symlink` checks only the leaf path; deeper traversal relies on the per-file leaf check at each site. (v0.53.1) - **v0.53.0 — Quant Menu II (UD GGUFs + KV cache + NVFP4 + LF parity + save formats)**: 6 schema-only Parts; live wiring deferred to v0.53.1. Every new validator follows the project's established hardening policy: closed allowlists (`UD_GGUF_FORMATS`, `IQ_GGUF_FORMATS`, `APPLE_ARM_GGUF_FORMATS`, `KV_CACHE_TYPES`, `MERGE_SAVE_FORMATS`, `TORCHAO_PTQ_SCHEMES`) as `frozenset` so registries cannot be mutated; `_GGUF_METADATA` / `_KV_CACHE_METADATA` / `_MERGE_METADATA` / `_TORCHAO_METADATA` wrapped in `MappingProxyType`; `_LOWER_INDEX` for GGUF lookup is also `MappingProxyType`-wrapped (replaces O(N) walk with O(1) lookup — code-review MEDIUM fix). All string validators reject non-string / bool / empty / null-byte / oversize with case-insensitive normalisation (matches v0.41.0 `validate_optimizer_name` / v0.51.0 `validate_hub_name` policy); `validate_torchao_scheme` is INTENTIONALLY case-sensitive (PyTorch class names — `torchao.quantize_` looks them up by exact name) with the asymmetry documented at both validators (security-review LOW fix). `validate_calibration_data_path` + `validate_quant_config_path` are shape-only at this release; their docstrings name the exact controls a v0.53.1 CLI dispatch contributor MUST add (`os.path.realpath` + `os.path.commonpath` cwd containment, `os.lstat` + `stat.S_ISLNK` symlink rejection before `open()`, existence check, `yaml.safe_load`-only for quant configs) — closes the security-review MEDIUM "documentation gap at trust boundary" finding. SoupConfig cross-validators: `_validate_fp8_attention_compat` (requires `quantization_aware='fp8'` BEFORE the MLX gate so the more actionable error fires first — code-review MEDIUM fix); `_validate_nvfp4_compat` (non-MLX + `modality='text'`; Blackwell SM ≥ 12.0 runtime check fires at trainer construction); `_validate_unsloth_bnb_4bit_compat` (requires `backend='unsloth'` + `quantization='4bit'`); `_validate_bnb_4bit_double_quant` (requires `quantization='4bit'` — rejects `none`/`8bit`/Quant-Menu); `_validate_llm_int8_alias` (asserts `quantization='8bit'`, deliberately disjoint from v0.41.0 `load_in_8bit` aliasing); `_validate_quantize_ref_reward` (extended ref-task allowlist `{dpo, ipo, simpo, orpo, bco, kto, preference, grpo, ppo}` per code-review HIGH fix — first-cut omitted grpo + kto + ppo which all have reference policies); `_validate_kv_cache_type_supported` (only `fp8` gated to non-MLX in v0.53.0; q8_0/bf16/f16 pass-through documented at validator site so v0.53.1 contributor sees the gate immediately). `requires_hopper` reads from `_KV_CACHE_METADATA` spec — single source of truth so adding a Hopper-only type means flipping the spec field only (code-review MEDIUM fix). All 7 new bool fields share `_validate_v053_bool_fields` `field_validator(mode='before')` that rejects bool-as-int with explicit `TypeError("v0.53.0 flag must be bool")` and passes `None` through to Pydantic's `default=False` rather than silently coercing it (python-review MEDIUM fix — `fp8_attention: null` in YAML now surfaces as a "valid boolean" ValidationError instead of masquerading as `False`). Known limitations: (1) Every live wiring is deferred to v0.53.1 — `export_advanced_gguf`, `apply_kv_cache_type`, `apply_fp8_attention`, `apply_nvfp4`, `merge_4bit`, `export_torchao` all raise `NotImplementedError` with explicit `v0.53.1` markers. (2) `validate_calibration_data_path` + `validate_quant_config_path` are shape-only this release; CLI dispatch in v0.53.1 MUST add cwd-containment + TOCTOU symlink rejection. (3) `kv_cache_type` MLX permissive policy: only `fp8` is rejected, the other three pass-through; v0.53.1 may narrow further. (4) Hopper SM-capability check is runtime-only — schema accepts `kv_cache_type='fp8'` + `fp8_attention=true` without GPU probe. (5) NVFP4 + Blackwell (SM ≥ 12.0) check is runtime-only. (6) `bnb_4bit_use_double_quant` only gated against `quantization`, not against `quantization_aware` — the latter combination is already rejected by v0.28.0 Quant-Menu + QAT cross-validator. (7) `llm_int8` is an assertion not an aliaser — diverges from v0.41.0 `load_in_8bit` design on purpose. (v0.53.0) - **v0.52.0 — Modality II (TTS + Distillation + BitNet + EBFT-GDPO + MoE quant + reasoning_effort)**: 7 schema-only Parts; live trainer / loss / export wiring deferred to v0.52.1. Every new validator follows the project's established hardening policy: closed allowlist (`SUPPORTED_TTS_FAMILIES`, `CLASSIFIER_TASKS`, `DIVERGENCES`, `BITNET_QUANT_FORMATS`, `BITNET_EXPORT_FORMATS`, `EBFT_VARIANTS`, `GDPO_VARIANTS`, `MOE_EXPERT_QUANT_FORMATS`, `REASONING_EFFORT_LEVELS`, per-family `_FAMILY_EMOTIONS`) wrapped in `frozenset` / `MappingProxyType` so registries cannot be mutated at runtime; `validate_*` helpers reject non-string / bool / empty / null-byte / oversize / unknown inputs with case-insensitive normalisation (matches v0.41.0 `validate_optimizer_name` / v0.50.0 `grpo_variant` / v0.51.0 `hub` policy); float validators (`validate_distill_temperature`, `validate_ebft_temperature`) gate on `math.isfinite` to reject NaN AND `±inf` (matches v0.32.0 `save_lr_finder_report` policy). `field_validator(mode="before")` on `num_labels` (security-review HIGH fix) rejects `bool` before Pydantic's `ge=1` coercion silently treats `True` as `1`. Field validator on `reasoning_effort` routes through the shared `validate_reasoning_effort` helper so the schema and runtime validator agree on what's accepted (security-review MEDIUM fix). SoupConfig cross-validators: `_validate_tts_compat` (requires `task='tts'` + `modality='audio_out'` + non-MLX backend; per-family emotion allowlist via `_FAMILY_EMOTIONS`), `_validate_classifier_compat` (with lazy-import early-return — code-review HIGH fix — so SFT hot path doesn't pay import cost; requires `num_labels` on classifier tasks; rejects classifier-only fields outside the task family with named offenders), `_validate_distill_compat` (requires `teacher_model` when `task='distill'`; rejects distill-only fields outside the task), `_validate_bitnet_compat` (gates to non-MLX + text-modality + task ∈ {sft, pretrain, dpo}), `_validate_ebft_compat` + `_validate_gdpo_compat` (task-family gates), `_validate_moe_expert_quant_compat` (requires `moe_lora=true` to prevent silent no-op), `_validate_reasoning_effort_task_gate` (code-review HIGH fix — rejects `reasoning_effort` + `train_on_eot` outside the SFT-family task set with named offenders; mirrors v0.50.0 GRPO stability task-gate policy). Public `DIVERGENCES` frozenset is derived from `_DIVERGENCE_ALIASES` so adding a new alias updates both the accepted-input set and the error message in lockstep (review fix LOW). `validate_bitnet_export` enforces a closed-allowlist canonical form for `soup export --format `, both of which are CLI-registered with a yellow advisory panel + `Exit(0)` stub (no artifact written until v0.52.1 — the format flag is accepted so existing scripts pinned to v0.52.0 will not break). 6 new YAML recipes appended (5 TTS + Falcon-E BitNet) — every entry is exercised by `tests/test_v0520.py` for `load_config_from_string` round-trip + `_no_null_or_whitespace` model-id check (mirrors v0.51.0 review-fix LOW). Known limitations: (1) Every live trainer / loss / export path is deferred to v0.52.1 — `build_tts_trainer`, `build_classifier_trainer`, `build_distill_trainer`, `build_bitnet_trainer`, `export_bitnet_gguf`, `apply_ebft_loss`, `apply_gdpo_loss`, `apply_moe_expert_quant` all raise `NotImplementedError` with explicit `v0.52.1` markers; schema accepts every new task / quant / variant + the CLI stub for `soup export --format bitnet/tq1_0` prints a deferred-advisory panel and exits 0. (2) `modality='audio_out'` accepted on non-TTS tasks — design choice this release so future audio-output tasks (ASR / V2A) can reuse it; today's runtime trainer dispatch must check `task == 'tts'` to avoid silent routing into the deferred TTS path. (3) Oute emotion allowlist is a tight 6-entry subset (neutral / happy / sad / angry / calm / excited); operators wanting custom emotions will need a v0.52.1 patch to extend `OUTE_EMOTIONS`. (4) `is_bitnet_model` is best-effort heuristic over name prefixes (`bitnet`, `falcon-e`, `1bitllm`, `onebit`); a BitNet checkpoint published under an org without any of those prefixes returns False. This is detection, not gating — the trainer wrapper (v0.52.1) loads the model regardless of the heuristic. (5) `quantization='bitnet_1.58'` gated to task ∈ {sft, pretrain, dpo} — extending to GRPO / PPO / RewardModel requires upstream onebitllms RL kernels not yet shipped. (v0.52.0) diff --git a/pyproject.toml b/pyproject.toml index 6a78ed4..86caf59 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.53.1" +version = "0.53.2" description = "Fine-tune LLMs in one command. No SSH, no config hell." readme = "README.md" license = "Apache-2.0" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index 762e583..ed45c45 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.53.1" +__version__ = "0.53.2" diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index e56bd86..220dce9 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -475,6 +475,19 @@ def train( if cfg.training.quantization_aware: quant_label += " + QAT" + # v0.53.2 review-fix: classifier-family tasks train a sequence-classification + # head, not a causal-LM LoRA — render "head" instead of LoRA r/alpha. + classifier_family = ("classifier", "reranker", "cross_encoder") + if cfg.task in classifier_family: + peft_line = ( + f"Head: [bold]num_labels={cfg.training.num_labels}, " + f"kind={cfg.training.classifier_kind}[/]" + ) + else: + peft_line = ( + f"LoRA: [bold]r={cfg.training.lora.r}, " + f"alpha={cfg.training.lora.alpha}[/]" + ) console.print( Panel( f"Device: [bold]{device_name}[/]\n" @@ -482,7 +495,7 @@ def train( f"Model: [bold]{cfg.base}[/]\n" f"Task: [bold]{cfg.task}[/]\n" f"Backend: [bold]{backend_label}[/]\n" - f"LoRA: [bold]r={cfg.training.lora.r}, alpha={cfg.training.lora.alpha}[/]\n" + f"{peft_line}\n" f"Quant: [bold]{quant_label}[/]", title="Training Setup", ) @@ -736,6 +749,16 @@ def train( from soup_cli.trainer.embedding import EmbeddingTrainerWrapper trainer_wrapper = EmbeddingTrainerWrapper(cfg, **trainer_kwargs) + elif cfg.task == "distill": + # v0.53.2 #133 — knowledge distillation (student + frozen teacher). + from soup_cli.trainer.distill import DistillTrainerWrapper + + trainer_wrapper = DistillTrainerWrapper(cfg, **trainer_kwargs) + elif cfg.task in ("classifier", "reranker", "cross_encoder"): + # v0.53.2 #132 — sequence-classification head. + from soup_cli.trainer.classifier import ClassifierTrainerWrapper + + trainer_wrapper = ClassifierTrainerWrapper(cfg, **trainer_kwargs) else: trainer_wrapper = SFTTrainerWrapper(cfg, **trainer_kwargs) trainer_wrapper.setup(dataset) diff --git a/soup_cli/data/loss_mask.py b/soup_cli/data/loss_mask.py index 4b80cea..265d407 100644 --- a/soup_cli/data/loss_mask.py +++ b/soup_cli/data/loss_mask.py @@ -110,6 +110,8 @@ def build_assistant_only_labels( messages: Sequence[dict], tokenizer: Any, max_length: int = 2048, + *, + include_eot: bool = False, ) -> dict[str, list[int]]: """Build labels where only assistant tokens contribute to loss. @@ -117,6 +119,11 @@ def build_assistant_only_labels( messages: Chat messages list (``{"role": ..., "content": ...}``). tokenizer: HF tokenizer with a ``chat_template`` set. max_length: Truncate to this many tokens. + include_eot: When True (axolotl ``train_on_eot``), extend each + assistant span to include the immediately-following EOS / EOT + token in the unmasked region — so the model learns to predict + the turn terminator. Default False matches HF Trainer's standard + chat-template loss-mask behaviour. (v0.53.2 #137) Returns: ``{"input_ids": [...], "labels": [...], "attention_mask": [...]}`` @@ -125,13 +132,22 @@ def build_assistant_only_labels( Raises: ValueError: empty messages, non-positive max_length, or tokenizer lacking a chat_template. + TypeError: ``include_eot`` not bool. """ + if not isinstance(include_eot, bool): + raise TypeError( + f"include_eot must be bool, got {type(include_eot).__name__}" + ) _check_messages(messages) _validate_max_length(max_length) + eos_token_id = _resolve_eos_token_id(tokenizer) if include_eot else None + preferred = _apply_template_with_mask(tokenizer, messages) if preferred is not None: input_ids, mask = preferred + if include_eot and eos_token_id is not None: + mask = _extend_mask_to_eot(input_ids, mask, eos_token_id) labels = [ tok if flag else IGNORE_INDEX for tok, flag in zip(input_ids, mask) @@ -150,10 +166,62 @@ def build_assistant_only_labels( if msg.get("role") == "assistant": end = min(new_len, len(full_ids)) labels[prev_len:end] = full_ids[prev_len:end] + if include_eot and eos_token_id is not None: + # Extend through the immediately-following EOT/EOS run. + extra = end + while extra < len(full_ids) and full_ids[extra] == eos_token_id: + labels[extra] = full_ids[extra] + extra += 1 prev_len = new_len return _truncate(full_ids, labels, max_length) +def _resolve_eos_token_id(tokenizer: Any) -> Optional[int]: + """Return an int EOS/EOT token id, or None if undetermined. + + Handles tokenizers exposing ``eos_token_id`` as int (most), list[int] + (e.g. Llama 3 with the additional ``<|eot_id|>`` entry — we pick the + first int entry), or anything else (str/None/bool → None). + """ + candidate = getattr(tokenizer, "eos_token_id", None) + if isinstance(candidate, bool): + return None + if isinstance(candidate, int): + return candidate + if isinstance(candidate, list): + for entry in candidate: + if isinstance(entry, int) and not isinstance(entry, bool): + return entry + return None + + +def _extend_mask_to_eot( + input_ids: Sequence[int], mask: Sequence[int], eos_token_id: int +) -> list[int]: + """Mark EOT/EOS tokens immediately following an assistant span as kept. + + Idempotent: a second pass over already-extended output produces the + same result (no extra EOT absorbed downstream of the original span). + """ + result = list(mask) + n = len(input_ids) + i = 0 + while i < n: + if result[i]: + # Walk to the end of this kept span, then absorb trailing EOS. + j = i + while j < n and result[j]: + j += 1 + while j < n and input_ids[j] == eos_token_id: + result[j] = 1 + j += 1 + # j > i guaranteed: the truthy-span walk advanced j at least once. + i = j + else: + i += 1 + return result + + def build_per_message_train_labels( messages: Sequence[dict], tokenizer: Any, diff --git a/soup_cli/data/sft_format.py b/soup_cli/data/sft_format.py index 729bc3a..1e7284e 100644 --- a/soup_cli/data/sft_format.py +++ b/soup_cli/data/sft_format.py @@ -34,8 +34,16 @@ def build_format_row( tokenizer: Any, data_cfg: DataConfig, console: Any | None = None, + training_cfg: Any | None = None, ) -> Callable[[dict], dict]: - """Factory: return the ``format_row`` function appropriate for ``data_cfg``.""" + """Factory: return the ``format_row`` function appropriate for ``data_cfg``. + + v0.53.2 #137: when ``training_cfg.reasoning_effort`` is set the gpt-oss + ``<|reasoning_effort|>...<|/reasoning_effort|>`` control tag is injected + into the system message before formatting. When ``training_cfg.train_on_eot`` + is true the loss mask is extended to include the trailing EOT/EOS token + after each assistant span (axolotl ``train_on_eot``). + """ # v0.36.0 Part C: apply chat-template override BEFORE deciding on path. # Override may turn a templateless tokenizer into a usable one. The # override warning surfaces here (not at sft.py call site) so it fires @@ -48,6 +56,13 @@ def build_format_row( use_train_field = bool(data_cfg.train_on_messages_with_train_field) max_length = int(data_cfg.max_length) + reasoning_effort = ( + getattr(training_cfg, "reasoning_effort", None) if training_cfg else None + ) + include_eot = bool( + getattr(training_cfg, "train_on_eot", False) if training_cfg else False + ) + if (use_responses_only or use_train_field) and not has_template: if console is not None: console.print( @@ -55,21 +70,51 @@ def build_format_row( "has no chat_template — falling back to text path. Pass " "data.chat_template explicitly to enable masking.[/]" ) - return _legacy_text_format_row(tokenizer) + return _wrap_with_reasoning_effort( + _legacy_text_format_row(tokenizer), reasoning_effort + ) if use_train_field: - return _build_per_message_format_row(tokenizer, max_length) - if use_responses_only: - return _build_assistant_only_format_row(tokenizer, max_length) - return _legacy_text_format_row(tokenizer) + inner = _build_per_message_format_row(tokenizer, max_length) + elif use_responses_only: + inner = _build_assistant_only_format_row( + tokenizer, max_length, include_eot=include_eot + ) + else: + inner = _legacy_text_format_row(tokenizer) + return _wrap_with_reasoning_effort(inner, reasoning_effort) + + +def _wrap_with_reasoning_effort( + inner: Callable[[dict], dict], level: Any | None +) -> Callable[[dict], dict]: + """Decorate ``inner`` so each example's messages get a reasoning-effort prefix.""" + if level is None: + return inner + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + def wrapped(example: dict) -> dict: + msgs = example.get("messages") + if isinstance(msgs, list) and msgs: + new_example = { + **example, + "messages": apply_reasoning_effort_prefix(msgs, level), + } + return inner(new_example) + return inner(example) + + return wrapped def _build_assistant_only_format_row( - tokenizer: Any, max_length: int + tokenizer: Any, max_length: int, include_eot: bool = False ) -> Callable[[dict], dict]: def format_row(example: dict) -> dict: return build_assistant_only_labels( - example["messages"], tokenizer, max_length=max_length + example["messages"], + tokenizer, + max_length=max_length, + include_eot=include_eot, ) return format_row diff --git a/soup_cli/trainer/classifier.py b/soup_cli/trainer/classifier.py new file mode 100644 index 0000000..377cccb --- /dev/null +++ b/soup_cli/trainer/classifier.py @@ -0,0 +1,368 @@ +"""Classifier / reranker / cross_encoder trainer (v0.53.2 #132). + +Wraps :class:`transformers.AutoModelForSequenceClassification` for the three +classifier-family tasks declared in v0.52.0 Part B: + +* ``classifier`` — single-input sequence classification. +* ``reranker`` — single-input reranker (typically score head); shares the + classifier head topology. +* ``cross_encoder`` — paired-input scoring (passage / query, NLI, etc.). + +Single-label uses ``CrossEntropyLoss``; multi-label uses ``BCEWithLogitsLoss`` +(automatically when ``training.classifier_kind='multi_label'``). + +Data format expected per row: + +* ``text`` (or ``messages`` joined into a string) — single-input tasks. +* ``text_a`` + ``text_b`` — paired-input ``cross_encoder``. +* ``label`` — int (single-label), list[int] (multi-label), or string in + ``training.label_names``. + +Mirrors the BCO / pretrain wrapper pattern. Lazy imports for heavy deps; +trust_remote_code threaded through the v0.36.0 resolver. +""" + +from __future__ import annotations + +import math +import time +from pathlib import Path +from typing import Any, List, Union + +from rich.console import Console + +from soup_cli.config.schema import SoupConfig + +console = Console() + +# Cap on multi-label list entries — defense against malformed dataset rows +# (security review v0.53.2 H2). Matches v0.52.0 ``_MAX_LABELS=1024``. +_MAX_MULTI_LABEL_ENTRIES: int = 1024 + + +def _row_to_text(row: dict) -> str: + """Extract the single-input text from a row (``text`` or joined messages). + + Raises: + TypeError: non-string ``content`` inside a messages list — silent skip + could poison training data with empty strings (security review + v0.53.2 M3). + ValueError: neither ``text`` nor ``messages`` present. + """ + if "text" in row and isinstance(row["text"], str): + return row["text"] + msgs = row.get("messages") + if isinstance(msgs, list): + parts: list[str] = [] + for msg in msgs: + if not isinstance(msg, dict): + # Non-dict entries are silently skipped — caller's loader + # would normally have produced these; keep loud-fail at the + # row-level (missing text) rather than per-message. + continue + content = msg.get("content", "") + if not isinstance(content, str): + raise TypeError( + "Classifier row messages[i]['content'] must be str, got " + f"{type(content).__name__!r}" + ) + parts.append(content) + return "\n".join(parts) + raise ValueError( + "Classifier row missing 'text' field and no joinable 'messages' list. " + f"Row keys: {sorted(row)!r}" + ) + + +def _row_to_pair(row: dict) -> tuple[str, str]: + """Extract (text_a, text_b) for paired ``cross_encoder`` rows. + + Raises: + TypeError: either field present but not a string (security review M4 — + silent ``str()`` coercion of dicts/lists produced garbage training + text). + ValueError: neither pair of fields present. + """ + if "text_a" in row and "text_b" in row: + a, b = row["text_a"], row["text_b"] + if not isinstance(a, str) or not isinstance(b, str): + raise TypeError( + "cross_encoder rows require 'text_a' and 'text_b' to be str; " + f"got text_a={type(a).__name__}, text_b={type(b).__name__}" + ) + return a, b + if "question" in row and "answer" in row: + q, ans = row["question"], row["answer"] + if not isinstance(q, str) or not isinstance(ans, str): + raise TypeError( + "cross_encoder rows require 'question' and 'answer' to be " + f"str; got question={type(q).__name__}, answer={type(ans).__name__}" + ) + return q, ans + raise ValueError( + "cross_encoder row requires 'text_a' + 'text_b' (or 'question' + " + f"'answer'). Row keys: {sorted(row)!r}" + ) + + +def _normalise_label( + raw: object, + label_names: List[str] | None, + num_labels: int, + multi_label: bool, +) -> Union[int, list[float]]: + """Convert a raw label (int / str / list) to the trainer-expected form. + + Single-label → int in [0, num_labels). Multi-label → list[float] of + length ``num_labels``. + + Raises: + ValueError: invalid index, oversize multi-label list (defense-in-depth + against malformed datasets — security review H2). + TypeError: bool / unsupported scalar type. + """ + if multi_label: + if isinstance(raw, list): + if len(raw) > _MAX_MULTI_LABEL_ENTRIES: + raise ValueError( + f"multi-label list too long: {len(raw)} entries " + f"(max {_MAX_MULTI_LABEL_ENTRIES})" + ) + vec = [0.0] * num_labels + for entry in raw: + idx = _label_index(entry, label_names, num_labels) + vec[idx] = 1.0 + return vec + # Single label silently broadcast to one-hot multi-label. + idx = _label_index(raw, label_names, num_labels) + vec = [0.0] * num_labels + vec[idx] = 1.0 + return vec + return _label_index(raw, label_names, num_labels) + + +def _label_index( + raw: object, label_names: List[str] | None, num_labels: int +) -> int: + if isinstance(raw, bool): + # Project policy (v0.30.0 Candidate / v0.39.0 ReLoRAPolicy / v0.41.0 + # Part B): bool-as-int violations raise TypeError, not ValueError. + raise TypeError(f"label must not be bool, got {raw!r}") + if isinstance(raw, int): + if raw < 0 or raw >= num_labels: + raise ValueError( + f"label index {raw} out of range [0, {num_labels})" + ) + return raw + if isinstance(raw, str): + if label_names is None: + raise ValueError( + f"label is str {raw!r} but training.label_names is unset" + ) + try: + return label_names.index(raw) + except ValueError as exc: + raise ValueError( + f"label {raw!r} not in training.label_names={label_names!r}" + ) from exc + raise TypeError( + f"label must be int / str / list, got {type(raw).__name__}" + ) + + +class ClassifierTrainerWrapper: + """High-level wrapper for classifier / reranker / cross_encoder training.""" + + def __init__( + self, + config: SoupConfig, + device: str = "cuda", + report_to: str = "none", + deepspeed_config: str | None = None, + fsdp_config: dict | None = None, + trust_remote_code: bool = False, + ) -> None: + self.config = config + self.device = device + self.report_to = report_to + self.deepspeed_config = deepspeed_config + self.fsdp_config = fsdp_config + self.trust_remote_code = trust_remote_code + + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + + self.model: Any = None + self.tokenizer: Any = None + self.trainer: Any = None + self._output_dir: str | None = None + + def setup(self, dataset: dict) -> None: + """Load model + tokenizer, tokenise dataset, build HF Trainer.""" + from datasets import Dataset + from transformers import ( + AutoModelForSequenceClassification, + AutoTokenizer, + Trainer, + TrainingArguments, + ) + + cfg = self.config + tcfg = cfg.training + + if tcfg.num_labels is None: + raise ValueError( + f"task={cfg.task!r} requires training.num_labels to be set" + ) + num_labels = int(tcfg.num_labels) + multi_label = (tcfg.classifier_kind == "multi_label") + problem_type = ( + "multi_label_classification" if multi_label else "single_label_classification" + ) + + 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 and self.tokenizer.eos_token is not None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + console.print(f"[dim]Loading classifier model: {cfg.base}[/]") + self.model = AutoModelForSequenceClassification.from_pretrained( + cfg.base, + num_labels=num_labels, + problem_type=problem_type, + trust_remote_code=self._trust_remote_code, + ) + + is_paired = (cfg.task == "cross_encoder") + label_names = ( + list(tcfg.label_names) if tcfg.label_names is not None else None + ) + + def encode(row: dict) -> dict: + if is_paired: + a, b = _row_to_pair(row) + enc = self.tokenizer( + a, b, + truncation=True, + max_length=cfg.data.max_length, + ) + else: + text = _row_to_text(row) + enc = self.tokenizer( + text, + truncation=True, + max_length=cfg.data.max_length, + ) + label = _normalise_label( + row.get("label"), label_names, num_labels, multi_label + ) + enc["labels"] = label + return enc + + raw_train = Dataset.from_list(dataset["train"]) + train_ds = raw_train.map(encode, remove_columns=raw_train.column_names) + eval_ds = None + if "val" in dataset and dataset["val"]: + raw_val = Dataset.from_list(dataset["val"]) + eval_ds = raw_val.map(encode, remove_columns=raw_val.column_names) + + output_dir = Path(cfg.output) + if cfg.experiment_name: + output_dir = output_dir / cfg.experiment_name + output_dir.mkdir(parents=True, exist_ok=True) + + batch_size = tcfg.batch_size if tcfg.batch_size != "auto" else 8 + total_steps = ( + math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps) + * tcfg.epochs + ) + warmup_steps = int(total_steps * tcfg.warmup_ratio) + + args = TrainingArguments( + output_dir=str(output_dir), + num_train_epochs=tcfg.epochs, + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=tcfg.gradient_accumulation_steps, + learning_rate=tcfg.lr, + warmup_steps=warmup_steps, + weight_decay=tcfg.weight_decay, + max_grad_norm=tcfg.max_grad_norm, + optim=tcfg.optimizer, + lr_scheduler_type=tcfg.scheduler, + logging_steps=tcfg.logging_steps, + save_steps=tcfg.save_steps, + save_total_limit=3, + bf16=self.device == "cuda", + report_to=self.report_to, + deepspeed=self.deepspeed_config, + **(self.fsdp_config or {}), + ) + + from transformers import DataCollatorWithPadding + + self.trainer = Trainer( + model=self.model, + args=args, + train_dataset=train_ds, + eval_dataset=eval_ds, + tokenizer=self.tokenizer, + data_collator=DataCollatorWithPadding(tokenizer=self.tokenizer), + ) + self._output_dir = str(output_dir) + + def train( + self, + display: object | None = None, + tracker: object | None = None, + run_id: str = "", + resume_from_checkpoint: str | None = None, + ) -> dict: + if self.trainer is None: + raise RuntimeError( + "ClassifierTrainerWrapper.train() called before setup(). " + "Call setup(dataset) first." + ) + start = time.time() + if display is not None: + from soup_cli.monitoring.callback import SoupTrainerCallback + + self.trainer.add_callback( + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) + ) + self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) + duration = time.time() - start + + self.trainer.save_model(self._output_dir) + self.tokenizer.save_pretrained(self._output_dir) + + logs = self.trainer.state.log_history + train_losses = [entry["loss"] for entry in logs if "loss" in entry] + + hours = int(duration // 3600) + minutes = int((duration % 3600) // 60) + duration_str = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m" + return { + "initial_loss": train_losses[0] if train_losses else 0, + "final_loss": train_losses[-1] if train_losses else 0, + "duration": duration_str, + "duration_secs": duration, + "output_dir": self._output_dir, + "total_steps": self.trainer.state.global_step, + } diff --git a/soup_cli/trainer/distill.py b/soup_cli/trainer/distill.py new file mode 100644 index 0000000..6ff4452 --- /dev/null +++ b/soup_cli/trainer/distill.py @@ -0,0 +1,428 @@ +"""Knowledge-distillation trainer (v0.53.2 #133). + +``task='distill'`` — student model learns from a frozen teacher. + +The training loss is a mix of: + +* the standard SFT cross-entropy on the student logits, AND +* a token-level divergence loss between student and teacher logits, scaled + by ``training.distill_temperature`` (T) following Hinton et al. 2015. + +Four divergence options (mirroring axolotl's distill plugin): + +* ``kl`` (alias of ``forward_kl``) — KL(teacher || student); standard + distillation. +* ``reverse_kl`` — KL(student || teacher); mode-seeking. +* ``js`` — Jensen-Shannon (symmetric). + +The teacher is loaded once, frozen (``requires_grad_(False)``), and +evaluated under ``torch.no_grad()`` to keep VRAM bounded. Student wears +LoRA per the standard PEFT pipeline. +""" + +from __future__ import annotations + +import math +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from rich.console import Console + +from soup_cli.config.schema import SoupConfig + +if TYPE_CHECKING: + import torch as _torch_typ + +console = Console() + +# 50/50 CE / distillation blend — matches Hinton et al. 2015. +# Promote to a schema field (distill_ce_weight) in a follow-up patch. +_CE_WEIGHT: float = 0.5 +_DISTILL_WEIGHT: float = 1.0 - _CE_WEIGHT + + +def _compute_distill_term( + student_logits: "_torch_typ.Tensor", + teacher_logits: "_torch_typ.Tensor", + divergence: str, + temperature: float, +) -> "_torch_typ.Tensor": + """Pure tensor kernel: divergence between student and teacher logits. + + Both logits are ``(batch, seq, vocab)``. Temperature softens the + distributions before the divergence is computed (Hinton). The result is + a scalar mean over the token-level divergences. + + Raises: + TypeError: ``temperature`` not numeric or is bool. + ValueError: ``temperature`` non-finite or non-positive; ``divergence`` + outside the supported set. + """ + import torch + + if isinstance(temperature, bool): + raise TypeError(f"temperature must not be bool, got {temperature!r}") + if not isinstance(temperature, (int, float)): + raise TypeError( + f"temperature must be float, got {type(temperature).__name__}" + ) + if not math.isfinite(float(temperature)) or float(temperature) <= 0: + raise ValueError( + f"temperature must be finite and positive, got {temperature!r}" + ) + temp = float(temperature) + s = student_logits / temp + t = teacher_logits / temp + if divergence == "forward_kl": + # KL(teacher || student). Use kl_div which expects log-probs of the + # student and probs of the teacher. + log_s = torch.log_softmax(s, dim=-1) + p_t = torch.softmax(t, dim=-1) + return torch.nn.functional.kl_div( + log_s, p_t, reduction="batchmean" + ) * (temp * temp) + if divergence == "reverse_kl": + log_t = torch.log_softmax(t, dim=-1) + p_s = torch.softmax(s, dim=-1) + return torch.nn.functional.kl_div( + log_t, p_s, reduction="batchmean" + ) * (temp * temp) + if divergence == "js": + # Jensen-Shannon: 0.5 (KL(p||m) + KL(q||m)), m = 0.5 (p + q). + log_s = torch.log_softmax(s, dim=-1) + log_t = torch.log_softmax(t, dim=-1) + p_s = log_s.exp() + p_t = log_t.exp() + m = 0.5 * (p_s + p_t) + log_m = m.clamp(min=1e-12).log() + kl_pm = torch.nn.functional.kl_div(log_m, p_s, reduction="batchmean") + kl_qm = torch.nn.functional.kl_div(log_m, p_t, reduction="batchmean") + return 0.5 * (kl_pm + kl_qm) * (temp * temp) + raise ValueError(f"Unknown divergence {divergence!r}") + + +class DistillTrainerWrapper: + """High-level wrapper for student/teacher distillation. + + Mirrors :class:`BCOTrainerWrapper` lifecycle (``__init__`` → ``setup`` → + ``train``). Teacher loads once in ``setup``; SFT-shaped dataset is + formatted via the standard ``build_format_row`` factory so the student + sees ``{input_ids, labels, attention_mask}`` rows. The custom + ``compute_loss`` injects the distillation term. + """ + + def __init__( + self, + config: SoupConfig, + device: str = "cuda", + report_to: str = "none", + deepspeed_config: str | None = None, + fsdp_config: dict | None = None, + trust_remote_code: bool = False, + ) -> None: + self.config = config + self.device = device + self.report_to = report_to + self.deepspeed_config = deepspeed_config + self.fsdp_config = fsdp_config + # Raw user-supplied flag — kept so teacher trust_remote_code can be + # resolved separately against its own model id during setup(). + self._raw_trust_remote_code = trust_remote_code + + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code, + resolve_trust_remote_code, + ) + + requires = model_requires_trust_remote_code(config.base) or False + self._trust_remote_code = resolve_trust_remote_code( + config.base, + requested=trust_remote_code, + console=console, + requires_remote_code=requires, + ) + + self.model: Any = None + self.teacher: Any = None + self.tokenizer: Any = None + self.trainer: Any = None + self._output_dir: str | None = None + + def setup(self, dataset: dict) -> None: + """Load student + teacher, build distillation Trainer.""" + from datasets import Dataset + from peft import LoraConfig, TaskType, get_peft_model + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + Trainer, + TrainingArguments, + ) + + from soup_cli.data.sft_format import build_format_row + from soup_cli.utils.distill import validate_divergence + + cfg = self.config + tcfg = cfg.training + + if tcfg.teacher_model is None: + raise ValueError( + "task='distill' requires training.teacher_model to be set" + ) + divergence = validate_divergence(tcfg.distill_divergence or "forward_kl") + temperature = float(tcfg.distill_temperature or 2.0) + + console.print(f"[dim]Loading tokenizer (student/teacher shared): {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 + + console.print(f"[dim]Loading student: {cfg.base}[/]") + dev_map = "cpu" if self.device == "cpu" else "auto" + self.model = AutoModelForCausalLM.from_pretrained( + cfg.base, + trust_remote_code=self._trust_remote_code, + device_map=dev_map, + ) + + # LoRA on the student — bracket with v0.40.6 #67 surgical PEFT patches. + 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, + ) + from soup_cli.utils.peft_wiring import ( + apply_post_lora_patches, + apply_pre_lora_patches, + ) + apply_pre_lora_patches(self.model, cfg.base) + self.model = get_peft_model(self.model, lora_config) + apply_post_lora_patches(self.model) + + # Teacher trust_remote_code resolved INDEPENDENTLY against the teacher + # model id (code-review HIGH-1 fix — student's resolution must not + # auto-trust the teacher). + from soup_cli.utils.trust_remote import ( + model_requires_trust_remote_code as _req, + ) + from soup_cli.utils.trust_remote import ( + resolve_trust_remote_code as _resolve, + ) + + teacher_requires = _req(tcfg.teacher_model) or False + teacher_trc = _resolve( + tcfg.teacher_model, + requested=self._raw_trust_remote_code, + console=console, + requires_remote_code=teacher_requires, + ) + + console.print(f"[dim]Loading teacher (frozen): {tcfg.teacher_model}[/]") + self.teacher = AutoModelForCausalLM.from_pretrained( + tcfg.teacher_model, + trust_remote_code=teacher_trc, + device_map=dev_map, + ) + self.teacher.eval() + for p in self.teacher.parameters(): + p.requires_grad_(False) + + # Cross-tokenizer distillation is not supported — vocab sizes must + # match so KL between logit distributions is well-defined. + teacher_vocab = getattr(self.teacher.config, "vocab_size", None) + student_vocab = getattr(self.model.config, "vocab_size", None) + if ( + teacher_vocab is not None + and student_vocab is not None + and teacher_vocab != student_vocab + ): + raise ValueError( + f"Teacher vocab size ({teacher_vocab}) != student vocab " + f"size ({student_vocab}). Cross-tokenizer distillation is " + "not supported in v0.53.2; use a teacher that shares the " + "student tokenizer family." + ) + + # Dataset prep — reuse the SFT formatter so distill sees + # {input_ids, labels, attention_mask}. v0.53.2 #137: pass training_cfg + # so reasoning_effort + train_on_eot are honored on task='distill'. + format_row = build_format_row( + tokenizer=self.tokenizer, + data_cfg=cfg.data, + console=console, + training_cfg=tcfg, + ) + raw_train = Dataset.from_list(dataset["train"]) + train_ds = raw_train.map( + format_row, remove_columns=raw_train.column_names + ) + eval_ds = None + if "val" in dataset and dataset["val"]: + raw_val = Dataset.from_list(dataset["val"]) + eval_ds = raw_val.map( + format_row, remove_columns=raw_val.column_names + ) + + output_dir = Path(cfg.output) + if cfg.experiment_name: + output_dir = output_dir / cfg.experiment_name + output_dir.mkdir(parents=True, exist_ok=True) + + batch_size = tcfg.batch_size if tcfg.batch_size != "auto" else 4 + total_steps = ( + math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps) + * tcfg.epochs + ) + warmup_steps = int(total_steps * tcfg.warmup_ratio) + + args = TrainingArguments( + output_dir=str(output_dir), + num_train_epochs=tcfg.epochs, + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=tcfg.gradient_accumulation_steps, + learning_rate=tcfg.lr, + warmup_steps=warmup_steps, + weight_decay=tcfg.weight_decay, + max_grad_norm=tcfg.max_grad_norm, + optim=tcfg.optimizer, + lr_scheduler_type=tcfg.scheduler, + logging_steps=tcfg.logging_steps, + save_steps=tcfg.save_steps, + save_total_limit=3, + bf16=self.device == "cuda", + report_to=self.report_to, + remove_unused_columns=False, + deepspeed=self.deepspeed_config, + **(self.fsdp_config or {}), + ) + + teacher_ref = self.teacher + + class _DistillTrainer(Trainer): + def compute_loss( + self, + model, + inputs, + return_outputs: bool = False, + num_items_in_batch=None, + ): + import torch + + labels = inputs.get("labels") + outputs = model(**{k: v for k, v in inputs.items() if k != "labels"}) + student_logits = outputs.logits + + ce_loss = torch.tensor(0.0, device=student_logits.device) + if labels is not None: + shift_logits = student_logits[:, :-1, :].contiguous() + shift_labels = labels[:, 1:].contiguous() + ce_loss = torch.nn.functional.cross_entropy( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1), + ignore_index=-100, + ) + + # Bridge devices: HF Trainer may auto-move the student to + # CUDA while the frozen teacher stays on CPU (or vice versa). + # Move teacher inputs onto the teacher's device, then move + # teacher_logits back onto the student's device for the + # KL kernel. + try: + teacher_device = next(teacher_ref.parameters()).device + except StopIteration: + teacher_device = student_logits.device + teacher_inputs = { + k: (v.to(teacher_device) if hasattr(v, "to") else v) + for k, v in inputs.items() + if k != "labels" + } + with torch.no_grad(): + teacher_out = teacher_ref(**teacher_inputs) + teacher_logits = teacher_out.logits.to(student_logits.device) + + distill_loss = _compute_distill_term( + student_logits, teacher_logits, divergence, temperature + ) + total = _CE_WEIGHT * ce_loss + _DISTILL_WEIGHT * distill_loss + return (total, outputs) if return_outputs else total + + # ``DataCollatorForSeq2Seq`` pads ``input_ids`` and ``attention_mask`` + # via the tokenizer AND pads ``labels`` with ``label_pad_token_id`` + # (-100 = IGNORE_INDEX). ``DataCollatorForLanguageModeling`` does + # NOT pad labels — incorrect for our pre-tokenised loss-masked rows. + from transformers import DataCollatorForSeq2Seq + + self.trainer = _DistillTrainer( + model=self.model, + args=args, + train_dataset=train_ds, + eval_dataset=eval_ds, + tokenizer=self.tokenizer, + data_collator=DataCollatorForSeq2Seq( + tokenizer=self.tokenizer, + label_pad_token_id=-100, + padding=True, + ), + ) + + # v0.40.6 #67 — ReLoRA callback. + from soup_cli.utils.peft_wiring import attach_relora_callback + attach_relora_callback(self.trainer, tcfg) + + self._output_dir = str(output_dir) + + def train( + self, + display: object | None = None, + tracker: object | None = None, + run_id: str = "", + resume_from_checkpoint: str | None = None, + ) -> dict: + if self.trainer is None: + raise RuntimeError( + "DistillTrainerWrapper.train() called before setup(). " + "Call setup(dataset) first." + ) + start = time.time() + if display is not None: + from soup_cli.monitoring.callback import SoupTrainerCallback + + self.trainer.add_callback( + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) + ) + self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) + duration = time.time() - start + + self.trainer.save_model(self._output_dir) + self.tokenizer.save_pretrained(self._output_dir) + + logs = self.trainer.state.log_history + train_losses = [entry["loss"] for entry in logs if "loss" in entry] + + hours = int(duration // 3600) + minutes = int((duration % 3600) // 60) + duration_str = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m" + return { + "initial_loss": train_losses[0] if train_losses else 0, + "final_loss": train_losses[-1] if train_losses else 0, + "duration": duration_str, + "duration_secs": duration, + "output_dir": self._output_dir, + "total_steps": self.trainer.state.global_step, + } diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index a380295..a633c83 100644 --- a/soup_cli/trainer/dpo.py +++ b/soup_cli/trainer/dpo.py @@ -162,6 +162,10 @@ class DPOTrainerWrapper: from soup_cli.utils.peft_wiring import attach_relora_callback attach_relora_callback(self.trainer, tcfg) + # v0.53.2 #135 — GDPO loss hook (no-op if gdpo_variant unset). + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + attach_gdpo_compute_loss(self.trainer, tcfg) + self._output_dir = str(output_dir) def _setup_transformers(self, cfg: SoupConfig, tcfg) -> None: diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index 9ddf646..c0dc754 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -157,6 +157,7 @@ class SFTTrainerWrapper: tokenizer=self.tokenizer, data_cfg=cfg.data, console=console, + training_cfg=tcfg, ) train_ds = Dataset.from_list(dataset["train"]).map( format_row, remove_columns=["messages"] @@ -843,6 +844,10 @@ class SFTTrainerWrapper: from soup_cli.utils.peft_wiring import attach_relora_callback attach_relora_callback(self.trainer, self.config.training) + # v0.53.2 #135 — EBFT compute_loss hook (no-op if ebft_variant unset). + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + attach_ebft_compute_loss(self.trainer, self.config.training) + # Activation offloading (v0.28.0) — wrap train() so saved-tensor hooks # are active only during training (and removed afterwards). from soup_cli.utils.activation_offload import offload_context diff --git a/soup_cli/utils/classifier.py b/soup_cli/utils/classifier.py index ccfebf1..a4c825a 100644 --- a/soup_cli/utils/classifier.py +++ b/soup_cli/utils/classifier.py @@ -159,10 +159,14 @@ def validate_classifier_compat(*, task: str, backend: str, modality: str) -> Non ) -def build_classifier_trainer() -> None: - """Live classifier trainer factory — deferred to v0.52.1.""" - raise NotImplementedError( - "Classifier / reranker / cross_encoder trainer live wiring deferred " - "to v0.52.1. Schema accepts the task but no trainer wrapper is " - "registered yet." - ) +def build_classifier_trainer( + config: object, **kwargs: object +) -> object: + """Live classifier / reranker / cross_encoder trainer factory (v0.53.2 #132). + + Returns a :class:`ClassifierTrainerWrapper`. Lazy import keeps the heavy + transformers/peft surface out of schema-only import paths. + """ + from soup_cli.trainer.classifier import ClassifierTrainerWrapper + + return ClassifierTrainerWrapper(config, **kwargs) # type: ignore[arg-type] diff --git a/soup_cli/utils/distill.py b/soup_cli/utils/distill.py index cf7de89..f8d1b8c 100644 --- a/soup_cli/utils/distill.py +++ b/soup_cli/utils/distill.py @@ -191,10 +191,14 @@ def validate_distill_compat( validate_teacher_model(teacher_model) -def build_distill_trainer() -> None: - """Live distillation trainer factory — deferred to v0.52.1.""" - raise NotImplementedError( - "Distillation trainer (task='distill') live wiring deferred to " - "v0.52.1. Schema accepts the value but no trainer wrapper is " - "registered yet." - ) +def build_distill_trainer( + config: object, **kwargs: object +) -> object: + """Live distillation trainer factory (v0.53.2 #133). + + Returns a :class:`DistillTrainerWrapper`. Lazy import keeps the heavy + transformers/peft surface out of schema-only import paths. + """ + from soup_cli.trainer.distill import DistillTrainerWrapper + + return DistillTrainerWrapper(config, **kwargs) # type: ignore[arg-type] diff --git a/soup_cli/utils/ebft_gdpo.py b/soup_cli/utils/ebft_gdpo.py index 71fdc91..74d681d 100644 --- a/soup_cli/utils/ebft_gdpo.py +++ b/soup_cli/utils/ebft_gdpo.py @@ -34,12 +34,12 @@ _EBFT_METADATA: Mapping[str, EBFTSpec] = MappingProxyType({ "structured": EBFTSpec( name="structured", description="Structured Energy-Based FT (per-token energies)", - live_wired=False, + live_wired=True, # v0.53.2 #135 — kernel + attach hook shipped. ), "strided": EBFTSpec( name="strided", description="Strided Energy-Based FT (block-sampled energies)", - live_wired=False, + live_wired=True, # v0.53.2 #135 ), }) @@ -57,17 +57,17 @@ _GDPO_METADATA: Mapping[str, GDPOSpec] = MappingProxyType({ "standard": GDPOSpec( name="standard", description="Standard GDPO (general preference objective)", - live_wired=False, + live_wired=True, # v0.53.2 #135 — kernel + DPO attach hook shipped. ), "length_normalized": GDPOSpec( name="length_normalized", description="Length-normalized GDPO (SimPO-style normalisation)", - live_wired=False, + live_wired=True, # v0.53.2 #135 ), "margin": GDPOSpec( name="margin", description="Margin-augmented GDPO (DPO + margin term)", - live_wired=False, + live_wired=True, # v0.53.2 #135 ), }) @@ -175,17 +175,298 @@ def validate_gdpo_compat(*, task: str, backend: str) -> None: ) -def apply_ebft_loss() -> None: - """Live EBFT loss kernel — deferred to v0.52.1.""" - raise NotImplementedError( - "EBFT (Energy-Based FT) live loss kernel deferred to v0.52.1. " - "Schema accepts the variant but no loss is wired yet." - ) +def attach_ebft_compute_loss(trainer: object, tcfg: object) -> bool: + """Wrap ``trainer.compute_loss`` so the EBFT term is added to CE (v0.53.2 #135). + + No-op when ``tcfg.ebft_variant`` is None. Otherwise the original + ``compute_loss`` is preserved and called first (for the standard SFT + cross-entropy), and the EBFT kernel is added to the returned loss. + + Idempotent: a sentinel attribute (``_soup_ebft_wrapped``) on the trainer + prevents double-wrapping when ``setup()`` is called twice on the same + trainer instance (security review v0.53.2 H1). + + Returns: + True if the wrap was installed, False otherwise. + """ + variant = getattr(tcfg, "ebft_variant", None) + if variant is None: + return False + if getattr(trainer, "_soup_ebft_wrapped", False): + return False + raw_temp = getattr(tcfg, "ebft_temperature", None) + temperature = float(raw_temp) if raw_temp is not None else 1.0 + canonical_variant = validate_ebft_variant(variant) + original = trainer.compute_loss # type: ignore[attr-defined] + + def wrapped( + model: object, + inputs: dict, + return_outputs: bool = False, + num_items_in_batch: object = None, + ): + result = original( + model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch + ) + ce_loss, outputs = result + labels = inputs.get("labels") + if labels is None: + return (ce_loss, outputs) if return_outputs else ce_loss + ebft_term = apply_ebft_loss( + outputs.logits, + labels, + variant=canonical_variant, + temperature=temperature, + ) + total = ce_loss + ebft_term + return (total, outputs) if return_outputs else total + + trainer.compute_loss = wrapped # type: ignore[attr-defined] + trainer._soup_ebft_wrapped = True # type: ignore[attr-defined] + return True -def apply_gdpo_loss() -> None: - """Live GDPO loss kernel — deferred to v0.52.1.""" - raise NotImplementedError( - "GDPO (Generalized DPO) live loss kernel deferred to v0.52.1. " - "Schema accepts the variant but no loss is wired yet." - ) +def attach_gdpo_compute_loss(trainer: object, tcfg: object) -> bool: + """Wrap TRL's ``DPOTrainer.dpo_loss`` so a GDPO variant is used (v0.53.2 #135). + + No-op when ``tcfg.gdpo_variant`` is None. Replaces the trainer's + ``dpo_loss`` method (the stable TRL hook returning losses, chosen rewards, + rejected rewards) with a thin wrapper that calls :func:`apply_gdpo_loss`. + + Returns: + True if the wrap was installed, False otherwise. + """ + variant = getattr(tcfg, "gdpo_variant", None) + if variant is None: + return False + if getattr(trainer, "_soup_gdpo_wrapped", False): + return False + canonical_variant = validate_gdpo_variant(variant) + raw_beta = getattr(tcfg, "dpo_beta", None) + beta = float(raw_beta) if raw_beta is not None else 0.1 + raw_margin = getattr(tcfg, "dpo_margin", None) + margin = float(raw_margin) if raw_margin is not None else 0.0 + original = getattr(trainer, "dpo_loss", None) + if original is None: + return False + + def wrapped( + policy_chosen_logps, + policy_rejected_logps, + reference_chosen_logps, + reference_rejected_logps, + chosen_lens=None, + rejected_lens=None, + *args: object, + **kwargs: object, + ): + # length_normalized variant pulls lengths from explicit args; TRL's + # callers either pass them positionally (newer TRL with length-norm + # support) or via **kwargs. + lens_c = chosen_lens if chosen_lens is not None else kwargs.get("chosen_lens") + lens_r = ( + rejected_lens + if rejected_lens is not None + else kwargs.get("rejected_lens") + ) + loss = apply_gdpo_loss( + policy_chosen_logps=policy_chosen_logps, + policy_rejected_logps=policy_rejected_logps, + ref_chosen_logps=reference_chosen_logps, + ref_rejected_logps=reference_rejected_logps, + variant=canonical_variant, + beta=beta, + margin=margin, + chosen_lens=lens_c, + rejected_lens=lens_r, + ) + # Recreate TRL's standard return shape: per-sample losses, chosen + # rewards, rejected rewards. We broadcast the mean to a per-sample + # tensor and compute simple rewards = beta * (pi - ref). + per_sample_loss = loss.expand_as(policy_chosen_logps) + chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps).detach() + rejected_rewards = beta * ( + policy_rejected_logps - reference_rejected_logps + ).detach() + return per_sample_loss, chosen_rewards, rejected_rewards + + trainer.dpo_loss = wrapped # type: ignore[attr-defined] + trainer._soup_gdpo_wrapped = True # type: ignore[attr-defined] + return True + + +def apply_ebft_loss( + logits, + labels, + *, + variant: str, + temperature: float, + stride: int = 4, + ignore_index: int = -100, +): + """Energy-Based Fine-Tuning loss kernel (v0.53.2 #135). + + EBFT treats per-token logits as energies (lower = more probable) and + penalises high-energy correct tokens. Two variants: + + * ``structured`` — per-token energy summed over every non-ignored + position, divided by the count. + * ``strided`` — same kernel but only every ``stride``-th position + contributes (faster on long sequences). + + The temperature scales the softmax sharpness: lower temperature + sharpens the implicit distribution, producing harsher gradients. + + Args: + logits: ``(batch, seq, vocab)`` float tensor. + labels: ``(batch, seq)`` long tensor; ``ignore_index`` entries + contribute nothing. + variant: ``"structured"`` or ``"strided"``. + temperature: in ``[1e-4, 100]`` (validated). + stride: positive int used only for the strided variant. + ignore_index: label id treated as padding (default ``-100``). + + Returns: + Scalar tensor (zero-dim) carrying gradient. + + Raises: + ValueError: shape mismatch, unknown variant, invalid temperature, + non-positive stride. + TypeError: ``stride`` not int or is bool. + """ + import torch + + canonical_variant = validate_ebft_variant(variant) + temp = validate_ebft_temperature(temperature) + if isinstance(stride, bool) or not isinstance(stride, int): + raise TypeError(f"stride must be int, got {type(stride).__name__}") + if stride < 1: + raise ValueError(f"stride must be >= 1, got {stride}") + + if logits.dim() != 3: + raise ValueError( + f"logits must be (batch, seq, vocab); got shape {tuple(logits.shape)}" + ) + if labels.dim() != 2: + raise ValueError( + f"labels must be (batch, seq); got shape {tuple(labels.shape)}" + ) + if logits.shape[:2] != labels.shape: + raise ValueError( + f"logits/labels shape mismatch: {tuple(logits.shape[:2])} vs " + f"{tuple(labels.shape)}" + ) + + batch, seq, _ = logits.shape + valid_mask = labels.ne(ignore_index) + if canonical_variant == "strided": + positions = torch.zeros_like(valid_mask) + positions[:, ::stride] = True + valid_mask = valid_mask & positions + + if not valid_mask.any(): + # Preserve grad path by multiplying logits by zero. + return (logits.sum() * 0.0).reshape(()) + + # Per-token energy: negative log-softmax of the correct label, scaled by + # 1 / temperature. Equivalent to a temperature-scaled cross-entropy. + safe_labels = labels.clamp(min=0) + log_probs = torch.log_softmax(logits / temp, dim=-1) + nll = -log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1) + nll = nll * valid_mask.to(nll.dtype) + denom = valid_mask.sum().clamp(min=1).to(nll.dtype) + return nll.sum() / denom + + +def apply_gdpo_loss( + *, + policy_chosen_logps, + policy_rejected_logps, + variant: str, + beta: float, + ref_chosen_logps=None, + ref_rejected_logps=None, + chosen_lens=None, + rejected_lens=None, + margin: float = 0.0, +): + """Generalized DPO loss kernel (v0.53.2 #135). + + Three variants: + + * ``standard`` — closed-form DPO: ``-log σ(β·(Δπ - Δref))``. + Requires reference log-probs. + * ``length_normalized`` — SimPO-style: ``-log σ(β·(π_w/L_w - π_l/L_l))``; + requires ``chosen_lens`` and ``rejected_lens``. + * ``margin`` — DPO with an explicit margin: ``-log σ(β·(Δπ - Δref) - + margin)``. Requires reference log-probs. + + All log-prob tensors are 1-D, batch-shaped (one entry per pair). + + Returns: + Scalar tensor (mean across the batch) carrying gradient. + + Raises: + ValueError: unknown variant, missing reference/lengths, beta bounds, + shape mismatch, non-finite margin. + TypeError: ``beta`` or ``margin`` bool / non-numeric. + """ + import torch + + canonical_variant = validate_gdpo_variant(variant) + + if isinstance(beta, bool): + raise TypeError(f"beta must not be bool, got {beta!r}") + if not isinstance(beta, (int, float)): + raise TypeError(f"beta must be float, got {type(beta).__name__}") + fbeta = float(beta) + if not math.isfinite(fbeta) or fbeta <= 0.0 or fbeta > 100.0: + raise ValueError(f"beta must be in (0, 100], got {beta!r}") + + if isinstance(margin, bool): + raise TypeError(f"margin must not be bool, got {margin!r}") + if not isinstance(margin, (int, float)): + raise TypeError(f"margin must be float, got {type(margin).__name__}") + fmargin = float(margin) + if not math.isfinite(fmargin): + raise ValueError(f"margin must be finite, got {margin!r}") + + if policy_chosen_logps.shape != policy_rejected_logps.shape: + raise ValueError( + f"policy chosen/rejected shape mismatch: " + f"{tuple(policy_chosen_logps.shape)} vs " + f"{tuple(policy_rejected_logps.shape)}" + ) + + if canonical_variant in ("standard", "margin"): + if ref_chosen_logps is None or ref_rejected_logps is None: + raise ValueError( + f"variant {canonical_variant!r} requires reference log-probs " + "(ref_chosen_logps + ref_rejected_logps)" + ) + if ( + ref_chosen_logps.shape != policy_chosen_logps.shape + or ref_rejected_logps.shape != policy_rejected_logps.shape + ): + raise ValueError( + "reference log-probs shape mismatch with policy log-probs" + ) + pi_delta = policy_chosen_logps - policy_rejected_logps + ref_delta = ref_chosen_logps - ref_rejected_logps + logits = fbeta * (pi_delta - ref_delta) + if canonical_variant == "margin": + logits = logits - fmargin + return -torch.nn.functional.logsigmoid(logits).mean() + + # length_normalized + if chosen_lens is None or rejected_lens is None: + raise ValueError( + "variant 'length_normalized' requires chosen_lens and " + "rejected_lens" + ) + chosen_lens_f = torch.clamp(chosen_lens.float(), min=1.0) + rejected_lens_f = torch.clamp(rejected_lens.float(), min=1.0) + chosen_norm = policy_chosen_logps / chosen_lens_f + rejected_norm = policy_rejected_logps / rejected_lens_f + logits = fbeta * (chosen_norm - rejected_norm) + return -torch.nn.functional.logsigmoid(logits).mean() diff --git a/soup_cli/utils/reasoning_effort.py b/soup_cli/utils/reasoning_effort.py index a144398..b1ac68f 100644 --- a/soup_cli/utils/reasoning_effort.py +++ b/soup_cli/utils/reasoning_effort.py @@ -1,16 +1,22 @@ -"""v0.52.0 Part G — gpt-oss reasoning_effort schema helper. +"""gpt-oss reasoning_effort schema helper + live prompt-prefix injector. -Schema-only support for ``training.reasoning_effort: Literal["low","medium","high"]`` -mirroring the unsloth gpt-oss training recipe. Routes through the prompt -prefix at training time; live formatter wiring lands in v0.52.1. +Schema support for ``training.reasoning_effort: Literal["low","medium","high"]`` +mirroring the unsloth gpt-oss training recipe. ``apply_reasoning_effort_prefix`` +injects the ``<|reasoning_effort|>...<|/reasoning_effort|>`` control tag into +the system message at training-time formatter dispatch (v0.53.2 #137). """ from __future__ import annotations +from typing import Mapping + REASONING_EFFORT_LEVELS: frozenset[str] = frozenset({"low", "medium", "high"}) _MAX_REASONING_EFFORT_LEN: int = 16 +_REASONING_OPEN: str = "<|reasoning_effort|>" +_REASONING_CLOSE: str = "<|/reasoning_effort|>" + def validate_reasoning_effort(value: object) -> str: """Validate a reasoning_effort string and return the canonical form.""" @@ -35,3 +41,57 @@ def validate_reasoning_effort(value: object) -> str: f"reasoning_effort {value!r} not supported. Supported: {supported}" ) return canonical + + +def apply_reasoning_effort_prefix( + messages: object, + level: object, +) -> list[dict]: + """Inject the gpt-oss reasoning-effort control tag into the system message. + + Returns a NEW list (input is not mutated; mirrors v0.33.0 #47 immutability + policy). The canonical tag is + ``<|reasoning_effort|>{level}<|/reasoning_effort|>`` and is prepended to + the FIRST system message in ``messages``; if no system message is present + one is inserted at index 0. + + Args: + messages: list of ``{"role": ..., "content": ...}`` dicts. + level: ``"low"`` / ``"medium"`` / ``"high"`` (case-insensitive). + + Raises: + TypeError: ``messages`` not a list, message entry not dict, ``level`` + not a string (bool rejected explicitly). + ValueError: ``messages`` empty, level outside the allowlist, or level + contains null bytes. + """ + canonical_level = validate_reasoning_effort(level) + if not isinstance(messages, list): + raise TypeError( + f"messages must be a list, got {type(messages).__name__}" + ) + if len(messages) == 0: + raise ValueError("messages must not be empty") + + out: list[dict] = [] + injected = False + tag = f"{_REASONING_OPEN}{canonical_level}{_REASONING_CLOSE}" + for msg in messages: + if not isinstance(msg, Mapping): + raise TypeError( + f"each message must be dict-like, got {type(msg).__name__}" + ) + if not injected and msg.get("role") == "system": + copy = dict(msg) + existing = copy.get("content", "") + if not isinstance(existing, str): + existing = str(existing) + copy["content"] = f"{tag}\n{existing}" if existing else tag + out.append(copy) + injected = True + else: + out.append(dict(msg)) + + if not injected: + out.insert(0, {"role": "system", "content": tag}) + return out diff --git a/tests/qa/v053_qa.md b/tests/qa/v053_qa.md index 67ba1e8..5c80250 100644 --- a/tests/qa/v053_qa.md +++ b/tests/qa/v053_qa.md @@ -112,3 +112,82 @@ EOF soup export --model ./merged --format torchao --quant-config ./q.yaml \ --output ./out/torchao ``` + +## v0.53.2 + +### #71 — ONNX export pipeline (`soup export --format onnx`) + +**Status:** PARTIAL PASS — small-model smoke ✅, TinyLlama-1.1B blocked by host +RAM during `onnx.load(load_external_data=True)` post-process. + +**Platform:** Windows 11, Python 3.10, torch 2.5.1+cu121, transformers 4.57.6, +optimum 2.1.0, RTX 3050 4 GB (ONNX export runs on CPU; GPU not used). + +**Smoke 1 — `hf-internal-testing/tiny-random-gpt2` (~5 MB):** + +```python +from optimum.exporters.onnx import main_export +main_export( + "hf-internal-testing/tiny-random-gpt2", + output="./out/tinygpt2_onnx", + task="text-generation", + trust_remote_code=False, +) +``` + +Result: **PASS** — completed in 25.8s, emitted 8 files, `model.onnx` = 0.62 MB, +all tokenizer/config artefacts present. + +**Smoke 2 — `TinyLlama/TinyLlama-1.1B-Chat-v1.0` + LoRA(r=4, +target=`q_proj`,`v_proj`) → merge → ONNX:** + +```python +# load fp16 base, attach tiny LoRA, save adapter, merge_and_unload, save merged +# main_export(merged_dir, task="text-generation-with-past", ...) +``` + +Result: **PARTIAL** — base loaded ✅ (3.2s), LoRA attached ✅, merge ✅ (9.6s), +ONNX trace + serialise ✅. Post-process step +`onnx.load(model.onnx, load_external_data=True)` failed with **MemoryError** +when loading the external-data tensor file back into RAM (TinyLlama-1.1B fp32 +is ~4.4 GB; host had insufficient free RAM to hold the model twice — once on +disk as `model.onnx_data`, once in the `onnx.ModelProto` for post-processing). + +**Verdict:** Pipeline integrity proven by tiny-gpt2. TinyLlama-1.1B full export +is host-RAM-bound (not a bug in `soup_cli` or `optimum`). Re-run on a +≥16 GB free-RAM machine to complete the size-validation pass. + +**Recorded:** 2026-05-13 (Wave 3 of v0.53.2 release). + +### v0.53.2 Live-trainer CPU smokes (Step 6d) + +**Status:** PASS — both new trainer wrappers train end-to-end on CPU. + +**Platform:** Windows 11, Python 3.10, torch 2.5.1+cu121, transformers 4.57.6, +peft 0.x — `hf-internal-testing/tiny-random-gpt2` (safetensors-only base, avoids +the torch < 2.6 `.bin` security gate). + +**`ClassifierTrainerWrapper`:** 4-row sentiment dataset, num_labels=2, +batch_size=2, max_steps=2. setup() built `GPT2ForSequenceClassification` with +`problem_type=single_label_classification`. train() completed 2 steps with +finite loss 0.787 in 0.66s. + +**`DistillTrainerWrapper`:** 4-row chat dataset, teacher=same tiny-gpt2 (frozen, +verified `requires_grad=False`), `distill_divergence=forward_kl`, +`distill_temperature=2.0`, batch_size=2, max_steps=2. setup() loaded student + +teacher, applied PEFT LoRA, built `_DistillTrainer` subclass. train() completed +2 steps with finite loss 3.479 in 1.79s — KL kernel exercised in compute_loss. + +**Bugs surfaced + fixed in Wave 3:** +1. `DataCollatorForLanguageModeling` did not pad pre-tokenised `labels` + (variable-length rows crashed during batching) → switched to + `DataCollatorForSeq2Seq(label_pad_token_id=-100)`. +2. `_DistillTrainer.compute_loss` device-mismatch: HF Trainer auto-moved the + student to CUDA on a CUDA-capable box while the teacher (`device_map="cpu"`) + stayed on CPU; teacher forward raised on cross-device `index_select` → + added `teacher_inputs.to(teacher_device)` + `teacher_logits.to(student_logits.device)` + bridge. Source-level regression guards added in `test_v0532.py` + (`TestDistillSourceLevelGuards`). + +**Recorded:** 2026-05-13. + diff --git a/tests/test_v0520.py b/tests/test_v0520.py index c428085..4920423 100644 --- a/tests/test_v0520.py +++ b/tests/test_v0520.py @@ -282,11 +282,14 @@ class TestClassifierUtils: task="reranker", backend="transformers", modality="vision", ) - def test_build_classifier_trainer_deferred(self): + def test_build_classifier_trainer_lifted_in_v0532(self): + """v0.52.0 shipped as a NotImplementedError stub; v0.53.2 #132 lifts it + to a live factory returning ClassifierTrainerWrapper. The argless call + now raises TypeError (missing ``config``) rather than NotImplementedError.""" from soup_cli.utils.classifier import build_classifier_trainer - with pytest.raises(NotImplementedError, match="v0.52.1"): - build_classifier_trainer() + with pytest.raises(TypeError): + build_classifier_trainer() # type: ignore[call-arg] class TestClassifierSchema: @@ -414,10 +417,13 @@ class TestDistillUtils: task="distill", backend="mlx", teacher_model="t/model", ) - def test_build_distill_trainer_deferred(self): + def test_build_distill_trainer_lifted_in_v0532(self): + """v0.52.0 shipped as a NotImplementedError stub; v0.53.2 #133 lifts it + to a live factory returning DistillTrainerWrapper. The argless call + now raises TypeError (missing ``config``) rather than NotImplementedError.""" from soup_cli.utils.distill import build_distill_trainer - with pytest.raises(NotImplementedError, match="v0.52.1"): + with pytest.raises(TypeError): build_distill_trainer() @@ -661,18 +667,23 @@ class TestEbftGdpoUtils: validate_gdpo_compat(task="sft", backend="transformers") def test_get_ebft_spec(self): + # v0.53.2 #135 lifted EBFT + GDPO live_wired flags from False to True + # (kernel + attach hooks shipped). from soup_cli.utils.ebft_gdpo import get_ebft_spec, get_gdpo_spec - assert get_ebft_spec("structured").live_wired is False - assert get_gdpo_spec("margin").live_wired is False + assert get_ebft_spec("structured").live_wired is True + assert get_gdpo_spec("margin").live_wired is True - def test_apply_ebft_loss_deferred(self): + def test_apply_ebft_loss_lifted_in_v0532(self): + """v0.52.0 shipped both as NotImplementedError stubs; v0.53.2 #135 + lifts them to live tensor kernels. The argless invocation now raises + TypeError (missing required args) rather than NotImplementedError.""" from soup_cli.utils.ebft_gdpo import apply_ebft_loss, apply_gdpo_loss - with pytest.raises(NotImplementedError, match="v0.52.1"): - apply_ebft_loss() - with pytest.raises(NotImplementedError, match="v0.52.1"): - apply_gdpo_loss() + with pytest.raises(TypeError): + apply_ebft_loss() # type: ignore[call-arg] + with pytest.raises(TypeError): + apply_gdpo_loss() # type: ignore[call-arg] class TestEbftGdpoSchema: diff --git a/tests/test_v0532.py b/tests/test_v0532.py new file mode 100644 index 0000000..3dbe3d9 --- /dev/null +++ b/tests/test_v0532.py @@ -0,0 +1,1664 @@ +"""v0.53.2 — Modality II live trainers. + +Covers: +- #137 ``apply_reasoning_effort_prefix`` system-prompt injector + ``train_on_eot`` + parameter on :func:`build_assistant_only_labels` + ``build_format_row`` wiring. +- #135 Live ``apply_ebft_loss`` (structured / strided) and ``apply_gdpo_loss`` + (standard / length_normalized / margin) pure-tensor kernels + + ``attach_ebft_compute_loss`` (SFT) / ``attach_gdpo_compute_loss`` (DPO) hooks. +- #133 ``DistillTrainerWrapper`` + ``build_distill_trainer`` factory + + ``commands/train.py`` ``task='distill'`` routing. +- #132 ``ClassifierTrainerWrapper`` + ``build_classifier_trainer`` factory + + ``commands/train.py`` ``task in {classifier, reranker, cross_encoder}`` routing. + +Deferred to a follow-up patch: +- #71 ONNX export GPU-bound smoke — recorded in ``tests/qa/v053_qa.md``. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +import pytest + +# --------------------------------------------------------------------------- +# #137 — reasoning_effort prompt-prefix injector +# --------------------------------------------------------------------------- + + +class TestReasoningEffortPrefix: + def test_inserts_system_message_when_none(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + messages = [{"role": "user", "content": "hi"}] + out = apply_reasoning_effort_prefix(messages, "high") + assert out[0]["role"] == "system" + assert "<|reasoning_effort|>high<|/reasoning_effort|>" in out[0]["content"] + # User message preserved at index 1 + assert out[1] == {"role": "user", "content": "hi"} + + def test_prepends_to_existing_system_message(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"}, + ] + out = apply_reasoning_effort_prefix(messages, "low") + assert out[0]["role"] == "system" + # Both the tag AND original content are present + assert "<|reasoning_effort|>low<|/reasoning_effort|>" in out[0]["content"] + assert "You are helpful." in out[0]["content"] + # User message unchanged + assert out[1] == {"role": "user", "content": "hi"} + + def test_does_not_mutate_input(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + original = [{"role": "user", "content": "hi"}] + snapshot = [dict(m) for m in original] + _ = apply_reasoning_effort_prefix(original, "medium") + assert original == snapshot + + @pytest.mark.parametrize("level", ["low", "medium", "high"]) + def test_accepts_canonical_levels(self, level: str) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + out = apply_reasoning_effort_prefix([{"role": "user", "content": "x"}], level) + assert f"<|reasoning_effort|>{level}<|/reasoning_effort|>" in out[0]["content"] + + def test_case_insensitive_level(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + out = apply_reasoning_effort_prefix([{"role": "user", "content": "x"}], "HIGH") + # Canonical lower-case is emitted + assert "<|reasoning_effort|>high<|/reasoning_effort|>" in out[0]["content"] + + def test_rejects_unknown_level(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(ValueError, match="not supported"): + apply_reasoning_effort_prefix([{"role": "user", "content": "x"}], "extreme") + + def test_rejects_bool_level(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(TypeError, match="must not be bool"): + apply_reasoning_effort_prefix([{"role": "user", "content": "x"}], True) # type: ignore[arg-type] + + def test_rejects_non_list_messages(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(TypeError, match="messages must be a list"): + apply_reasoning_effort_prefix("hi", "low") # type: ignore[arg-type] + + def test_rejects_empty_messages(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(ValueError, match="empty"): + apply_reasoning_effort_prefix([], "low") + + def test_rejects_non_dict_message(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(TypeError, match="must be dict"): + apply_reasoning_effort_prefix(["not a dict"], "low") # type: ignore[list-item] + + +# --------------------------------------------------------------------------- +# #137 — train_on_eot loss-mask extension +# --------------------------------------------------------------------------- + + +class _FakeTokenizer: + """Minimal tokenizer with chat_template + EOS token for loss-mask tests. + + Renders each message as ``role:content`` joined by a single space, + where ```` is token id 9. Each character becomes a token id (its + Unicode code point) so we can verify mask positions exactly. ``role:`` + prefix is the role's first char + colon (codepoints). + """ + + chat_template = "fake" + eos_token = "" + eos_token_id = 9 + + def apply_chat_template( + self, + messages: Any, + tokenize: bool = True, + add_generation_prompt: bool = False, + return_assistant_tokens_mask: bool = False, + add_special_tokens: bool = True, + **kwargs: Any, + ) -> Any: + ids: list[int] = [] + mask: list[int] = [] + for msg in messages: + role = msg["role"] + content = msg["content"] + is_assistant = role == "assistant" + chunk_ids = [ord(c) for c in f"{role[0]}:{content}"] + ids.extend(chunk_ids) + mask.extend([1 if is_assistant else 0] * len(chunk_ids)) + # EOT token after each message + ids.append(self.eos_token_id) + mask.append(1 if is_assistant else 0) + if not tokenize: + return "".join(chr(i) for i in ids) + if return_assistant_tokens_mask: + return {"input_ids": ids, "assistant_masks": mask} + return ids + + +class _EotBoundaryTokenizer: + """Preferred-path fake whose assistant_mask EXCLUDES the trailing EOT. + + Mirrors a real chat template whose ``{% generation %}`` block wraps only + the assistant content, leaving the EOT token outside the mask — which is + exactly the case ``train_on_eot`` is designed to handle. + """ + + chat_template = "fake" + eos_token = "" + eos_token_id = 9 + + def apply_chat_template( + self, + messages: Any, + tokenize: bool = True, + add_generation_prompt: bool = False, + return_assistant_tokens_mask: bool = False, + return_dict: bool = False, + add_special_tokens: bool = True, + **kwargs: Any, + ) -> Any: + ids: list[int] = [] + mask: list[int] = [] + for msg in messages: + role = msg["role"] + content = msg["content"] + is_assistant = role == "assistant" + chunk_ids = [ord(c) for c in f"{role[0]}:{content}"] + ids.extend(chunk_ids) + mask.extend([1 if is_assistant else 0] * len(chunk_ids)) + # EOT token after each message — ALWAYS unmasked (0) in this fake. + ids.append(self.eos_token_id) + mask.append(0) + if not tokenize: + return "".join(chr(i) for i in ids) + if return_assistant_tokens_mask and return_dict: + return {"input_ids": ids, "assistant_masks": mask} + return ids + + +class TestTrainOnEot: + def test_include_eot_default_false_masks_eot(self) -> None: + """Existing behaviour — EOT *outside* the assistant mask is IGNORE.""" + from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels + + tok = _EotBoundaryTokenizer() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + result = build_assistant_only_labels(messages, tok, max_length=64) + labels = result["labels"] + # Layout from _EotBoundaryTokenizer: u:hia:ok + # mask: 0 0 0 0 0 0 0 1 1 0 + # With include_eot=False the trailing EOT (id=9) at the end is + # IGNORE_INDEX. + assert labels[-1] == IGNORE_INDEX + + def test_include_eot_true_extends_label_to_eos(self) -> None: + from soup_cli.data.loss_mask import ( + IGNORE_INDEX, + build_assistant_only_labels, + ) + + tok = _EotBoundaryTokenizer() + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + without_eot = build_assistant_only_labels(messages, tok, max_length=64) + with_eot = build_assistant_only_labels( + messages, tok, max_length=64, include_eot=True + ) + # With include_eot=True, the EOT (id=9) immediately following an + # assistant span is KEPT (not IGNORE), so the unmasked count is + # strictly higher by exactly the number of assistant turns. + n_kept_no = sum(1 for x in without_eot["labels"] if x != IGNORE_INDEX) + n_kept_yes = sum(1 for x in with_eot["labels"] if x != IGNORE_INDEX) + assert n_kept_yes - n_kept_no == 1 + # Last token (the trailing EOT) is now KEPT. + assert with_eot["labels"][-1] != IGNORE_INDEX + assert with_eot["labels"][-1] == 9 + + def test_include_eot_must_be_bool(self) -> None: + from soup_cli.data.loss_mask import build_assistant_only_labels + + with pytest.raises(TypeError, match="include_eot must be bool"): + build_assistant_only_labels( + [{"role": "user", "content": "x"}, {"role": "assistant", "content": "y"}], + _FakeTokenizer(), + max_length=64, + include_eot="yes", # type: ignore[arg-type] + ) + + +# --------------------------------------------------------------------------- +# #135 — EBFT live loss kernel +# --------------------------------------------------------------------------- + + +def _torch_or_skip(): + try: + import torch # noqa: F401 + + return torch + except Exception: # pragma: no cover - CI without torch + pytest.skip("torch not available") + + +class TestEbftLossLive: + def test_structured_returns_finite_scalar(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + logits = torch.randn(2, 4, 8, requires_grad=True) + labels = torch.tensor([[0, 1, 2, -100], [3, 4, -100, -100]]) + loss = apply_ebft_loss(logits, labels, variant="structured", temperature=1.0) + assert loss.ndim == 0 + assert torch.isfinite(loss) + # Gradient flows + loss.backward() + assert logits.grad is not None + + def test_strided_returns_finite_scalar(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + logits = torch.randn(2, 4, 8) + labels = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]]) + loss = apply_ebft_loss( + logits, labels, variant="strided", temperature=0.5, stride=2 + ) + assert loss.ndim == 0 + assert torch.isfinite(loss) + + def test_temperature_scales_loss(self) -> None: + """Lower temperature sharpens the energy distribution → different loss.""" + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + torch.manual_seed(0) + logits = torch.randn(1, 3, 5) + labels = torch.tensor([[0, 1, 2]]) + hot = apply_ebft_loss(logits, labels, variant="structured", temperature=2.0) + cold = apply_ebft_loss(logits, labels, variant="structured", temperature=0.5) + assert not math.isclose(float(hot), float(cold), rel_tol=1e-6) + + def test_unknown_variant_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(ValueError, match="not supported"): + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="bogus", + temperature=1.0, + ) + + def test_temperature_validated(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(ValueError): + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="structured", + temperature=float("nan"), + ) + + def test_all_ignore_labels_returns_zero(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + logits = torch.randn(2, 3, 4) + labels = torch.full((2, 3), -100, dtype=torch.long) + loss = apply_ebft_loss(logits, labels, variant="structured", temperature=1.0) + assert float(loss) == 0.0 + + def test_shape_mismatch_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(ValueError, match="shape"): + apply_ebft_loss( + torch.randn(2, 4, 8), + torch.tensor([[0, 1, 2]]), # seq=3 vs logits seq=4 + variant="structured", + temperature=1.0, + ) + + +# --------------------------------------------------------------------------- +# #135 — GDPO live loss kernel +# --------------------------------------------------------------------------- + + +class TestGdpoLossLive: + def test_standard_returns_finite_scalar(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + pol_chosen = torch.tensor([-2.0, -1.5, -3.0], requires_grad=True) + pol_rejected = torch.tensor([-3.5, -2.5, -4.0], requires_grad=True) + ref_chosen = torch.tensor([-2.2, -1.8, -3.1]) + ref_rejected = torch.tensor([-3.2, -2.4, -3.9]) + loss = apply_gdpo_loss( + policy_chosen_logps=pol_chosen, + policy_rejected_logps=pol_rejected, + ref_chosen_logps=ref_chosen, + ref_rejected_logps=ref_rejected, + variant="standard", + beta=0.1, + ) + assert loss.ndim == 0 + assert torch.isfinite(loss) + loss.backward() + assert pol_chosen.grad is not None + + def test_length_normalized_uses_lengths(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + pol_chosen = torch.tensor([-10.0, -8.0]) + pol_rejected = torch.tensor([-12.0, -10.0]) + # Length-normalized uses chosen_lens / rejected_lens + loss = apply_gdpo_loss( + policy_chosen_logps=pol_chosen, + policy_rejected_logps=pol_rejected, + variant="length_normalized", + beta=0.5, + chosen_lens=torch.tensor([5.0, 4.0]), + rejected_lens=torch.tensor([6.0, 5.0]), + ) + assert loss.ndim == 0 + assert torch.isfinite(loss) + + def test_margin_includes_margin_term(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + pol_chosen = torch.tensor([-2.0]) + pol_rejected = torch.tensor([-3.0]) + ref_chosen = torch.tensor([-2.0]) + ref_rejected = torch.tensor([-3.0]) + no_margin = apply_gdpo_loss( + policy_chosen_logps=pol_chosen, + policy_rejected_logps=pol_rejected, + ref_chosen_logps=ref_chosen, + ref_rejected_logps=ref_rejected, + variant="margin", + beta=0.1, + margin=0.0, + ) + with_margin = apply_gdpo_loss( + policy_chosen_logps=pol_chosen, + policy_rejected_logps=pol_rejected, + ref_chosen_logps=ref_chosen, + ref_rejected_logps=ref_rejected, + variant="margin", + beta=0.1, + margin=1.0, + ) + # Larger margin → larger loss (harder to satisfy) + assert float(with_margin) > float(no_margin) + + def test_unknown_variant_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(ValueError, match="not supported"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + variant="bogus", + beta=0.1, + ) + + def test_standard_requires_reference(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(ValueError, match="reference"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + variant="standard", + beta=0.1, + ) + + def test_length_normalized_requires_lengths(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(ValueError, match="chosen_lens"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + variant="length_normalized", + beta=0.1, + ) + + def test_beta_validated(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(ValueError, match="beta"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + ref_chosen_logps=t, + ref_rejected_logps=t, + variant="standard", + beta=-0.1, + ) + + def test_beta_rejects_bool(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(TypeError, match="bool"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + ref_chosen_logps=t, + ref_rejected_logps=t, + variant="standard", + beta=True, # type: ignore[arg-type] + ) + + def test_shape_mismatch_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + with pytest.raises(ValueError, match="shape"): + apply_gdpo_loss( + policy_chosen_logps=torch.zeros(2), + policy_rejected_logps=torch.zeros(3), + ref_chosen_logps=torch.zeros(2), + ref_rejected_logps=torch.zeros(2), + variant="standard", + beta=0.1, + ) + + +# --------------------------------------------------------------------------- +# Cross-cutting — stubs lifted +# --------------------------------------------------------------------------- + + +def test_apply_ebft_loss_no_longer_raises_not_implemented() -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + # Should succeed (not raise NotImplementedError) — confirms stub lifted. + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="structured", + temperature=1.0, + ) + + +def test_apply_gdpo_loss_no_longer_raises_not_implemented() -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + ref_chosen_logps=t, + ref_rejected_logps=t, + variant="standard", + beta=0.1, + ) + + +# --------------------------------------------------------------------------- +# #135 — attach hooks (live wiring into SFT + DPO trainers) +# --------------------------------------------------------------------------- + + +class _StubTrainer: + """Minimal trainer surface for hook tests — captures compute_loss calls.""" + + def __init__(self) -> None: + self.call_log: list[tuple] = [] + + def _compute_loss(model, inputs, return_outputs=False, num_items_in_batch=None): + self.call_log.append(("compute_loss", return_outputs)) + torch = _torch_or_skip() + outputs = type("Out", (), {"logits": torch.zeros(1, 2, 3)})() + loss = torch.tensor(0.5) + return (loss, outputs) if return_outputs else loss + + self.compute_loss = _compute_loss + + +class TestAttachEbftComputeLoss: + def test_no_op_when_variant_unset(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + + trainer = _StubTrainer() + tcfg = type("Tcfg", (), {"ebft_variant": None})() + assert attach_ebft_compute_loss(trainer, tcfg) is False + + def test_wraps_when_variant_set(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + + trainer = _StubTrainer() + original = trainer.compute_loss + tcfg = type( + "Tcfg", (), {"ebft_variant": "structured", "ebft_temperature": 1.0} + )() + assert attach_ebft_compute_loss(trainer, tcfg) is True + assert trainer.compute_loss is not original + + # Call the wrapped function — verify it returns a scalar tensor and + # calls the original (loss should be CE + EBFT term). + labels = torch.tensor([[0, 1]]) + loss = trainer.compute_loss(None, {"labels": labels}, return_outputs=False) + assert torch.is_tensor(loss) + # Original was called with return_outputs=True. + assert trainer.call_log[-1] == ("compute_loss", True) + + def test_invalid_variant_rejected(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + + trainer = _StubTrainer() + tcfg = type("Tcfg", (), {"ebft_variant": "bogus"})() + with pytest.raises(ValueError, match="not supported"): + attach_ebft_compute_loss(trainer, tcfg) + + +class _StubDpoTrainer: + """Stand-in for TRL DPOTrainer's dpo_loss surface.""" + + def __init__(self) -> None: + def _dpo_loss(*args, **kwargs): + return None # Original returns 3-tuple; stub returns sentinel. + + self.dpo_loss = _dpo_loss + + +class TestAttachGdpoComputeLoss: + def test_no_op_when_variant_unset(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = _StubDpoTrainer() + tcfg = type("Tcfg", (), {"gdpo_variant": None})() + assert attach_gdpo_compute_loss(trainer, tcfg) is False + + def test_no_op_when_trainer_lacks_dpo_loss(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = object() + tcfg = type("Tcfg", (), {"gdpo_variant": "standard"})() + assert attach_gdpo_compute_loss(trainer, tcfg) is False + + def test_wraps_and_returns_trl_shape(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = _StubDpoTrainer() + tcfg = type( + "Tcfg", + (), + {"gdpo_variant": "standard", "dpo_beta": 0.1, "dpo_margin": 0.0}, + )() + assert attach_gdpo_compute_loss(trainer, tcfg) is True + + pol_c = torch.tensor([-1.0, -2.0]) + pol_r = torch.tensor([-2.0, -3.0]) + ref_c = torch.tensor([-1.1, -2.1]) + ref_r = torch.tensor([-2.1, -3.1]) + losses, chosen_rewards, rejected_rewards = trainer.dpo_loss( + pol_c, pol_r, ref_c, ref_r + ) + assert losses.shape == pol_c.shape + assert chosen_rewards.shape == pol_c.shape + assert rejected_rewards.shape == pol_r.shape + + +# --------------------------------------------------------------------------- +# #133 — DistillTrainerWrapper + distill divergence kernel +# --------------------------------------------------------------------------- + + +class TestDistillDivergenceKernel: + @pytest.mark.parametrize("divergence", ["forward_kl", "reverse_kl", "js"]) + def test_divergence_returns_finite_scalar(self, divergence: str) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + student = torch.randn(2, 4, 8, requires_grad=True) + teacher = torch.randn(2, 4, 8) + out = _compute_distill_term(student, teacher, divergence, temperature=2.0) + assert out.ndim == 0 + assert torch.isfinite(out) + out.backward() + assert student.grad is not None + + def test_unknown_divergence_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + with pytest.raises(ValueError, match="Unknown divergence"): + _compute_distill_term( + torch.zeros(1, 2, 3), torch.zeros(1, 2, 3), "bogus", 1.0 + ) + + def test_identical_logits_zero_kl(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + logits = torch.randn(1, 3, 5) + out = _compute_distill_term(logits, logits.clone(), "forward_kl", 1.0) + assert abs(float(out)) < 1e-5 + + +class TestDistillWrapper: + def test_imports_cleanly(self) -> None: + from soup_cli.trainer.distill import DistillTrainerWrapper # noqa: F401 + + def test_train_before_setup_raises(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.trainer.distill import DistillTrainerWrapper + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: distill + training: + teacher_model: sshleifer/tiny-gpt2 + distill_temperature: 2.0 + distill_divergence: forward_kl + data: + train: ./fake.jsonl + """ + ) + wrapper = DistillTrainerWrapper(cfg, device="cpu") + with pytest.raises(RuntimeError, match="before setup"): + wrapper.train() + + def test_build_distill_trainer_factory_lifted(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.distill import build_distill_trainer + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: distill + training: + teacher_model: sshleifer/tiny-gpt2 + data: + train: ./fake.jsonl + """ + ) + # Lifted from NotImplementedError in v0.53.2 — returns wrapper. + wrapper = build_distill_trainer(cfg, device="cpu") + from soup_cli.trainer.distill import DistillTrainerWrapper + + assert isinstance(wrapper, DistillTrainerWrapper) + + +# --------------------------------------------------------------------------- +# #132 — ClassifierTrainerWrapper + helpers +# --------------------------------------------------------------------------- + + +class TestClassifierWrapperHelpers: + def test_row_to_text_with_text_field(self) -> None: + from soup_cli.trainer.classifier import _row_to_text + + assert _row_to_text({"text": "hello"}) == "hello" + + def test_row_to_text_joins_messages(self) -> None: + from soup_cli.trainer.classifier import _row_to_text + + out = _row_to_text( + { + "messages": [ + {"role": "user", "content": "a"}, + {"role": "assistant", "content": "b"}, + ] + } + ) + assert "a" in out and "b" in out + + def test_row_to_text_missing_field_raises(self) -> None: + from soup_cli.trainer.classifier import _row_to_text + + with pytest.raises(ValueError, match="missing 'text'"): + _row_to_text({"label": 0}) + + def test_row_to_pair_text_ab(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + a, b = _row_to_pair({"text_a": "x", "text_b": "y"}) + assert (a, b) == ("x", "y") + + def test_row_to_pair_question_answer(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + a, b = _row_to_pair({"question": "q", "answer": "a"}) + assert (a, b) == ("q", "a") + + def test_row_to_pair_missing_raises(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + with pytest.raises(ValueError, match="text_a"): + _row_to_pair({"text": "single"}) + + @pytest.mark.parametrize("idx", [0, 1, 2]) + def test_label_index_int_in_range(self, idx: int) -> None: + from soup_cli.trainer.classifier import _label_index + + assert _label_index(idx, None, num_labels=3) == idx + + def test_label_index_int_out_of_range_rejected(self) -> None: + from soup_cli.trainer.classifier import _label_index + + with pytest.raises(ValueError, match="out of range"): + _label_index(5, None, num_labels=3) + + def test_label_index_string_via_label_names(self) -> None: + from soup_cli.trainer.classifier import _label_index + + assert _label_index("pos", ["neg", "pos"], num_labels=2) == 1 + + def test_label_index_string_without_names_rejected(self) -> None: + from soup_cli.trainer.classifier import _label_index + + with pytest.raises(ValueError, match="label_names is unset"): + _label_index("pos", None, num_labels=2) + + def test_label_index_bool_rejected(self) -> None: + from soup_cli.trainer.classifier import _label_index + + # Project policy (v0.30.0 Candidate / v0.39.0 ReLoRAPolicy / v0.41.0 + # Part B): bool-as-int violations raise TypeError. + with pytest.raises(TypeError, match="bool"): + _label_index(True, None, num_labels=2) + + def test_normalise_label_multi_label_from_list(self) -> None: + from soup_cli.trainer.classifier import _normalise_label + + vec = _normalise_label( + [0, 2], label_names=None, num_labels=3, multi_label=True + ) + assert vec == [1.0, 0.0, 1.0] + + +class TestClassifierWrapper: + def test_imports_cleanly(self) -> None: + from soup_cli.trainer.classifier import ClassifierTrainerWrapper # noqa: F401 + + def test_train_before_setup_raises(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.trainer.classifier import ClassifierTrainerWrapper + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: classifier + training: + num_labels: 3 + data: + train: ./fake.jsonl + """ + ) + wrapper = ClassifierTrainerWrapper(cfg, device="cpu") + with pytest.raises(RuntimeError, match="before setup"): + wrapper.train() + + def test_build_classifier_trainer_factory_lifted(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.classifier import build_classifier_trainer + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: classifier + training: + num_labels: 2 + data: + train: ./fake.jsonl + """ + ) + wrapper = build_classifier_trainer(cfg, device="cpu") + from soup_cli.trainer.classifier import ClassifierTrainerWrapper + + assert isinstance(wrapper, ClassifierTrainerWrapper) + + +# --------------------------------------------------------------------------- +# commands/train.py routing — source-level audit +# --------------------------------------------------------------------------- + + +class TestTrainRouting: + """Source-grep audit — looks for the *instantiation site* of each wrapper. + + Bare ``"DistillTrainerWrapper" in src`` would pass on a comment mentioning + the class. ``"DistillTrainerWrapper(cfg, **trainer_kwargs)"`` requires the + actual call expression to be present, which is much harder to satisfy by + accident. + """ + + def test_distill_routed(self) -> None: + from soup_cli.commands import train as train_cmd + + src = __import__("inspect").getsource(train_cmd) + assert 'cfg.task == "distill"' in src + # Require the actual instantiation expression, not just the bare name. + assert "DistillTrainerWrapper(cfg, **trainer_kwargs)" in src + + def test_classifier_family_routed(self) -> None: + from soup_cli.commands import train as train_cmd + + src = __import__("inspect").getsource(train_cmd) + # Tuple membership in the if-branch is the load-bearing pattern. + assert ( + 'cfg.task in ("classifier", "reranker", "cross_encoder")' in src + ) + # Instantiation expression — not just the class name. + assert "ClassifierTrainerWrapper(cfg, **trainer_kwargs)" in src + + +# --------------------------------------------------------------------------- +# #137 — build_format_row wiring of reasoning_effort + train_on_eot +# --------------------------------------------------------------------------- + + +class _Tcfg: + def __init__( + self, + reasoning_effort: Optional[str] = None, + train_on_eot: bool = False, + ) -> None: + self.reasoning_effort = reasoning_effort + self.train_on_eot = train_on_eot + + +def _make_data_cfg( + train_on_responses_only: bool = True, chat_template: Optional[str] = None +): + """Minimal duck-typed DataConfig for sft_format tests.""" + obj = type( + "DataCfg", + (), + { + "train_on_responses_only": train_on_responses_only, + "train_on_messages_with_train_field": False, + "max_length": 64, + "chat_template": chat_template, + }, + )() + return obj + + +class TestFormatRowReasoningEffort: + def test_no_reasoning_effort_passthrough(self) -> None: + from soup_cli.data.sft_format import build_format_row + + format_row = build_format_row( + tokenizer=_EotBoundaryTokenizer(), + data_cfg=_make_data_cfg(), + training_cfg=_Tcfg(reasoning_effort=None), + ) + row = { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + out = format_row(row) + # Should produce {input_ids, labels, attention_mask} + assert set(out) >= {"input_ids", "labels", "attention_mask"} + + def test_reasoning_effort_injected_via_format_row(self) -> None: + from soup_cli.data.sft_format import build_format_row + + captured_messages: list = [] + + class _CapturingTok(_EotBoundaryTokenizer): + def apply_chat_template(self, messages, **kwargs): + captured_messages.append([dict(m) for m in messages]) + return super().apply_chat_template(messages, **kwargs) + + format_row = build_format_row( + tokenizer=_CapturingTok(), + data_cfg=_make_data_cfg(), + training_cfg=_Tcfg(reasoning_effort="high"), + ) + row = { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + _ = format_row(row) + # First captured render should have a system message with the tag. + first = captured_messages[0] + assert first[0]["role"] == "system" + assert "<|reasoning_effort|>high<|/reasoning_effort|>" in first[0]["content"] + + def test_does_not_mutate_original_row(self) -> None: + from soup_cli.data.sft_format import build_format_row + + format_row = build_format_row( + tokenizer=_EotBoundaryTokenizer(), + data_cfg=_make_data_cfg(), + training_cfg=_Tcfg(reasoning_effort="low"), + ) + row = { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + original_msgs = [dict(m) for m in row["messages"]] + _ = format_row(row) + assert row["messages"] == original_msgs + + +class TestFormatRowTrainOnEot: + def test_train_on_eot_extends_loss_mask(self) -> None: + from soup_cli.data.loss_mask import IGNORE_INDEX + from soup_cli.data.sft_format import build_format_row + + tok = _EotBoundaryTokenizer() + row = { + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "ok"}, + ] + } + without = build_format_row( + tokenizer=tok, + data_cfg=_make_data_cfg(), + training_cfg=_Tcfg(train_on_eot=False), + )(row) + # Need a fresh tokenizer instance because the override is module-level + tok2 = _EotBoundaryTokenizer() + with_eot = build_format_row( + tokenizer=tok2, + data_cfg=_make_data_cfg(), + training_cfg=_Tcfg(train_on_eot=True), + )(row) + n_no = sum(1 for x in without["labels"] if x != IGNORE_INDEX) + n_yes = sum(1 for x in with_eot["labels"] if x != IGNORE_INDEX) + assert n_yes - n_no == 1 + + +# --------------------------------------------------------------------------- +# Schema gates — classifier and distill still validate +# --------------------------------------------------------------------------- + + +def test_distill_task_loads_with_teacher() -> None: + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: distill + training: + teacher_model: sshleifer/tiny-gpt2 + distill_temperature: 2.0 + distill_divergence: forward_kl + data: + train: ./fake.jsonl + """ + ) + assert cfg.task == "distill" + assert cfg.training.teacher_model == "sshleifer/tiny-gpt2" + + +def test_classifier_task_requires_num_labels() -> None: + from soup_cli.config.loader import load_config_from_string + + # ``load_config_from_string`` re-raises pydantic validation as ValueError. + with pytest.raises(ValueError, match="num_labels"): + load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: classifier + data: + train: ./fake.jsonl + """ + ) + + +# --------------------------------------------------------------------------- +# Review-fix coverage (v0.53.2 TDD + security review followups) +# --------------------------------------------------------------------------- + + +class TestEbftAttachLabelsNone: + """tdd-review C1 — wrapped compute_loss must tolerate inputs without labels.""" + + def test_labels_missing_returns_ce_only(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + + trainer = _StubTrainer() + tcfg = type( + "Tcfg", + (), + {"ebft_variant": "structured", "ebft_temperature": 1.0}, + )() + attach_ebft_compute_loss(trainer, tcfg) + # No "labels" key — wrapper must fall through to CE without crashing. + loss = trainer.compute_loss(None, {}, return_outputs=False) + assert torch.is_tensor(loss) + + +class TestEbftAttachIdempotent: + def test_double_wrap_is_no_op(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_ebft_compute_loss + + trainer = _StubTrainer() + tcfg = type( + "Tcfg", + (), + {"ebft_variant": "structured", "ebft_temperature": 1.0}, + )() + first = attach_ebft_compute_loss(trainer, tcfg) + second = attach_ebft_compute_loss(trainer, tcfg) + assert first is True + assert second is False + assert getattr(trainer, "_soup_ebft_wrapped", False) is True + + +class TestGdpoAttachIdempotent: + def test_double_wrap_is_no_op(self) -> None: + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = _StubDpoTrainer() + tcfg = type( + "Tcfg", + (), + {"gdpo_variant": "standard", "dpo_beta": 0.1, "dpo_margin": 0.0}, + )() + first = attach_gdpo_compute_loss(trainer, tcfg) + second = attach_gdpo_compute_loss(trainer, tcfg) + assert first is True + assert second is False + assert getattr(trainer, "_soup_gdpo_wrapped", False) is True + + +class TestGdpoAttachLengthNormalizedForwardsLens: + """python-review MEDIUM — length_normalized requires chosen_lens/rejected_lens.""" + + def test_length_normalized_via_kwargs(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = _StubDpoTrainer() + tcfg = type( + "Tcfg", + (), + { + "gdpo_variant": "length_normalized", + "dpo_beta": 0.5, + "dpo_margin": 0.0, + }, + )() + attach_gdpo_compute_loss(trainer, tcfg) + + pol_c = torch.tensor([-10.0, -8.0]) + pol_r = torch.tensor([-12.0, -10.0]) + ref_c = torch.zeros(2) + ref_r = torch.zeros(2) + losses, _, _ = trainer.dpo_loss( + pol_c, + pol_r, + ref_c, + ref_r, + chosen_lens=torch.tensor([5.0, 4.0]), + rejected_lens=torch.tensor([6.0, 5.0]), + ) + assert losses.shape == pol_c.shape + + def test_length_normalized_via_positional(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import attach_gdpo_compute_loss + + trainer = _StubDpoTrainer() + tcfg = type( + "Tcfg", + (), + { + "gdpo_variant": "length_normalized", + "dpo_beta": 0.5, + "dpo_margin": 0.0, + }, + )() + attach_gdpo_compute_loss(trainer, tcfg) + + pol_c = torch.tensor([-10.0, -8.0]) + pol_r = torch.tensor([-12.0, -10.0]) + # Positional chosen_lens / rejected_lens (newer TRL signatures may + # pass these positionally). + losses, _, _ = trainer.dpo_loss( + pol_c, + pol_r, + torch.zeros(2), + torch.zeros(2), + torch.tensor([5.0, 4.0]), + torch.tensor([6.0, 5.0]), + ) + assert losses.shape == pol_c.shape + + +class TestExtendMaskIdempotency: + """tdd-review C2 — _extend_mask_to_eot must be idempotent.""" + + def test_second_pass_no_op(self) -> None: + from soup_cli.data.loss_mask import _extend_mask_to_eot + + ids = [10, 11, 9, 20, 21, 9] + mask = [1, 1, 0, 0, 0, 0] + once = _extend_mask_to_eot(ids, mask, eos_token_id=9) + twice = _extend_mask_to_eot(ids, once, eos_token_id=9) + assert once == twice + + def test_leading_eos_not_marked(self) -> None: + from soup_cli.data.loss_mask import _extend_mask_to_eot + + ids = [9, 10, 11, 9] + mask = [0, 1, 1, 0] + out = _extend_mask_to_eot(ids, mask, eos_token_id=9) + assert out[0] == 0 # leading EOS not absorbed + assert out[3] == 1 # trailing EOS absorbed + + def test_two_assistant_spans_both_get_eos(self) -> None: + from soup_cli.data.loss_mask import _extend_mask_to_eot + + ids = [1, 2, 9, 9, 3, 4, 9] + mask = [0, 1, 0, 0, 0, 1, 0] + out = _extend_mask_to_eot(ids, mask, eos_token_id=9) + assert out[2] == 1 # consecutive EOS absorbed + assert out[3] == 1 + assert out[6] == 1 + + +class TestResolveEosTokenId: + """python-review MEDIUM — handle list/str/None/bool eos_token_id.""" + + def test_int(self) -> None: + from soup_cli.data.loss_mask import _resolve_eos_token_id + + class T: + eos_token_id = 9 + + assert _resolve_eos_token_id(T()) == 9 + + def test_list_picks_first_int(self) -> None: + from soup_cli.data.loss_mask import _resolve_eos_token_id + + class T: + eos_token_id = [128001, 128009] # Llama 3 style + + assert _resolve_eos_token_id(T()) == 128001 + + def test_str_returns_none(self) -> None: + from soup_cli.data.loss_mask import _resolve_eos_token_id + + class T: + eos_token_id = "9" + + assert _resolve_eos_token_id(T()) is None + + def test_none_returns_none(self) -> None: + from soup_cli.data.loss_mask import _resolve_eos_token_id + + class T: + pass + + assert _resolve_eos_token_id(T()) is None + + def test_bool_returns_none(self) -> None: + from soup_cli.data.loss_mask import _resolve_eos_token_id + + class T: + eos_token_id = True + + assert _resolve_eos_token_id(T()) is None + + +class TestComputeDistillTermValidation: + """python-review MEDIUM / tdd-review H1 — temperature validation in kernel.""" + + def test_zero_temperature_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + with pytest.raises(ValueError, match="positive"): + _compute_distill_term( + torch.zeros(1, 2, 3), + torch.zeros(1, 2, 3), + "forward_kl", + 0.0, + ) + + def test_negative_temperature_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + with pytest.raises(ValueError, match="positive"): + _compute_distill_term( + torch.zeros(1, 2, 3), + torch.zeros(1, 2, 3), + "forward_kl", + -1.0, + ) + + def test_nan_temperature_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + with pytest.raises(ValueError, match="finite"): + _compute_distill_term( + torch.zeros(1, 2, 3), + torch.zeros(1, 2, 3), + "forward_kl", + float("nan"), + ) + + def test_bool_temperature_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.trainer.distill import _compute_distill_term + + with pytest.raises(TypeError, match="bool"): + _compute_distill_term( + torch.zeros(1, 2, 3), + torch.zeros(1, 2, 3), + "forward_kl", + True, # type: ignore[arg-type] + ) + + +class TestReasoningEffortNullByte: + """tdd-review H3 — null-byte in level rejected.""" + + def test_null_byte_rejected(self) -> None: + from soup_cli.utils.reasoning_effort import apply_reasoning_effort_prefix + + with pytest.raises(ValueError, match="null"): + apply_reasoning_effort_prefix( + [{"role": "user", "content": "x"}], "lo\x00w" + ) + + +class TestEbftStrideValidation: + """tdd-review M2 — stride bool / non-positive rejected.""" + + def test_stride_bool_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(TypeError, match="stride"): + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="strided", + temperature=1.0, + stride=True, # type: ignore[arg-type] + ) + + def test_stride_zero_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(ValueError, match="stride"): + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="strided", + temperature=1.0, + stride=0, + ) + + def test_stride_negative_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_ebft_loss + + with pytest.raises(ValueError, match="stride"): + apply_ebft_loss( + torch.randn(1, 2, 3), + torch.tensor([[0, 1]]), + variant="strided", + temperature=1.0, + stride=-1, + ) + + +class TestGdpoMarginValidation: + """tdd-review M3 — margin bool / NaN rejected.""" + + def test_margin_bool_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(TypeError, match="margin"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + ref_chosen_logps=t, + ref_rejected_logps=t, + variant="margin", + beta=0.1, + margin=True, # type: ignore[arg-type] + ) + + def test_margin_nan_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + t = torch.zeros(2) + with pytest.raises(ValueError, match="finite"): + apply_gdpo_loss( + policy_chosen_logps=t, + policy_rejected_logps=t, + ref_chosen_logps=t, + ref_rejected_logps=t, + variant="margin", + beta=0.1, + margin=float("nan"), + ) + + +class TestGdpoRefShapeMismatch: + """tdd-review M4 — ref logps shape mismatch vs policy.""" + + def test_ref_shape_mismatch_rejected(self) -> None: + torch = _torch_or_skip() + from soup_cli.utils.ebft_gdpo import apply_gdpo_loss + + with pytest.raises(ValueError, match="shape"): + apply_gdpo_loss( + policy_chosen_logps=torch.zeros(3), + policy_rejected_logps=torch.zeros(3), + ref_chosen_logps=torch.zeros(2), + ref_rejected_logps=torch.zeros(3), + variant="standard", + beta=0.1, + ) + + +class TestRowToTextRejectsNonStrContent: + """security-review M3 — non-str content raises rather than silent skip.""" + + def test_non_str_content_raises(self) -> None: + from soup_cli.trainer.classifier import _row_to_text + + with pytest.raises(TypeError, match="must be str"): + _row_to_text( + { + "messages": [ + {"role": "user", "content": ["multi", "modal"]}, + ] + } + ) + + def test_non_dict_message_silently_skipped(self) -> None: + from soup_cli.trainer.classifier import _row_to_text + + out = _row_to_text( + {"messages": ["not-a-dict", {"role": "user", "content": "real"}]} + ) + assert "real" in out + + +class TestRowToPairRejectsNonStr: + """security-review M4 — unchecked str() coercion fix.""" + + def test_non_str_text_a_rejected(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + with pytest.raises(TypeError, match="text_a"): + _row_to_pair({"text_a": {"d": "ict"}, "text_b": "y"}) + + def test_non_str_text_b_rejected(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + with pytest.raises(TypeError, match="text_b"): + _row_to_pair({"text_a": "x", "text_b": [1, 2]}) + + def test_non_str_question_rejected(self) -> None: + from soup_cli.trainer.classifier import _row_to_pair + + with pytest.raises(TypeError, match="question"): + _row_to_pair({"question": 42, "answer": "ok"}) + + +class TestLabelIndexFurtherCoverage: + """tdd-review M5 / M6 / M9.""" + + def test_single_label_list_raises(self) -> None: + from soup_cli.trainer.classifier import _normalise_label + + with pytest.raises(TypeError): + _normalise_label( + [0, 1], label_names=None, num_labels=3, multi_label=False + ) + + def test_string_not_in_label_names(self) -> None: + from soup_cli.trainer.classifier import _label_index + + with pytest.raises(ValueError, match="not in training.label_names"): + _label_index("unknown", ["pos", "neg"], num_labels=2) + + @pytest.mark.parametrize("bad", [None, 1.5, object()]) + def test_unsupported_type_rejected(self, bad: object) -> None: + from soup_cli.trainer.classifier import _label_index + + with pytest.raises(TypeError, match="must be int"): + _label_index(bad, None, num_labels=3) + + +class TestMultiLabelListCap: + """security-review H2 — uncapped multi-label list DoS defense.""" + + def test_oversize_list_rejected(self) -> None: + from soup_cli.trainer.classifier import _normalise_label + + # _MAX_MULTI_LABEL_ENTRIES is 1024. + big = [0] * 2000 + with pytest.raises(ValueError, match="too long"): + _normalise_label( + big, label_names=None, num_labels=3, multi_label=True + ) + + +class TestFactoryUnknownKwarg: + """tdd-review L1 — factories must reject unknown kwargs loudly.""" + + def test_build_distill_trainer_unknown_kwarg(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.distill import build_distill_trainer + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: distill + training: + teacher_model: sshleifer/tiny-gpt2 + data: + train: ./fake.jsonl + """ + ) + with pytest.raises(TypeError): + build_distill_trainer(cfg, device="cpu", nonexistent=True) + + def test_build_classifier_trainer_unknown_kwarg(self) -> None: + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils.classifier import build_classifier_trainer + + cfg = load_config_from_string( + """ + base: sshleifer/tiny-gpt2 + task: classifier + training: + num_labels: 2 + data: + train: ./fake.jsonl + """ + ) + with pytest.raises(TypeError): + build_classifier_trainer(cfg, device="cpu", nonexistent=True) + + +class TestDistillSourceLevelGuards: + """Regression guards for v0.53.2 wave 3 smoke-surfaced bugs.""" + + def test_uses_seq2seq_collator(self) -> None: + """``DataCollatorForLanguageModeling`` does NOT pad labels — would crash + on variable-length pre-tokenised rows. v0.53.2 wave 3 switched to + ``DataCollatorForSeq2Seq``. Check the import statement specifically + (the deprecated class name still appears in an explanatory comment).""" + import inspect + + from soup_cli.trainer import distill as distill_mod + + src = inspect.getsource(distill_mod) + # The actual import line must reference the Seq2Seq collator. + assert "from transformers import DataCollatorForSeq2Seq" in src + # No import of the wrong collator. + assert "from transformers import DataCollatorForLanguageModeling" not in src + + def test_compute_loss_bridges_teacher_device(self) -> None: + """Bug surfaced during CPU smoke on a CUDA-capable box: HF Trainer + auto-moved the student to CUDA while the teacher stayed where loaded. + compute_loss must bridge teacher inputs onto the teacher's device and + teacher_logits back onto the student's device.""" + import inspect + + from soup_cli.trainer import distill as distill_mod + + src = inspect.getsource(distill_mod) + assert "teacher_device = next(teacher_ref.parameters()).device" in src + assert ".to(student_logits.device)" in src + + +# --------------------------------------------------------------------------- +# Step 6f — failure-mode smoke: every cross-validator triggered by bad YAML +# --------------------------------------------------------------------------- + + +class TestFailureModeSmoke: + """One YAML per cross-validator path the v0.53.2 surfaces add or trigger. + + Goal: prove the validators name the actual problem so users grep + successfully, not "ValidationError" generic. + """ + + def test_classifier_missing_num_labels(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="num_labels"): + load_config_from_string( + "base: x\ntask: classifier\ndata:\n train: ./f.jsonl\n" + ) + + def test_classifier_label_names_length_mismatch(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="num_labels"): + load_config_from_string( + "base: x\ntask: classifier\ntraining:\n" + " num_labels: 3\n label_names: [a, b]\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_classifier_fields_outside_classifier_task(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="num_labels|classifier"): + load_config_from_string( + "base: x\ntask: sft\ntraining:\n num_labels: 3\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_distill_missing_teacher(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="teacher_model"): + load_config_from_string( + "base: x\ntask: distill\ndata:\n train: ./f.jsonl\n" + ) + + def test_distill_fields_outside_distill_task(self) -> None: + from soup_cli.config.loader import load_config_from_string + + # teacher_model set on a non-distill task → rejected with named field. + with pytest.raises(ValueError, match="teacher_model|distill"): + load_config_from_string( + "base: x\ntask: sft\ntraining:\n teacher_model: y\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_reasoning_effort_on_non_sft_family_task(self) -> None: + from soup_cli.config.loader import load_config_from_string + + # reasoning_effort is SFT-family-only; reject on grpo. + with pytest.raises(ValueError, match="reasoning_effort"): + load_config_from_string( + "base: x\ntask: grpo\ntraining:\n reasoning_effort: high\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_train_on_eot_on_non_sft_family_task(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="train_on_eot"): + load_config_from_string( + "base: x\ntask: grpo\ntraining:\n train_on_eot: true\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_ebft_temperature_without_variant(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="ebft_variant"): + load_config_from_string( + "base: x\ntask: sft\ntraining:\n ebft_temperature: 1.0\n" + "data:\n train: ./f.jsonl\n" + ) + + def test_gdpo_variant_on_sft_rejected(self) -> None: + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="gdpo_variant"): + load_config_from_string( + "base: x\ntask: sft\ntraining:\n gdpo_variant: standard\n" + "data:\n train: ./f.jsonl\n" + )