RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(perf): MiniLLM/MoLE KV-cache + deploy-measure live factories + live-codec TTS (v0.71.22) 

#263 MiniLLM on-policy KV-cache (PEFT-unwrap probe activates the cache for LoRA
students; per-step single-token forward), #262 serve --mole per-adapter KV cache
(fresh per generate, no cross-request leak, byte-identical to no-cache), #143
deploy-autopilot live generator factories (baseline scored once + up-front
candidate validation; injected seams retained), #265-partial live-codec TTS
(soundfile.info pre-probe + O_NOFOLLOW; SNAC Orpheus encode validated).

Review: 1 HIGH + 5 MEDIUM + ~10 LOW fixed across 2 review waves + verification +
step-6 live smoke (Windows + RTX 3050). Tests 14084 -> 14184 (+100 in
tests/test_v07122.py; 293 files). Full suite 14067 passed / 117 skipped, exit 0.
This commit is contained in:
Alpamys 2026-06-10 21:30:01 +05:00
parent 9fd356b8ae
commit ccd5c80e4d
13 changed files with 2385 additions and 104 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
CHANGELOG.md
View File

@ -12,6 +12,61 @@ reproducing 70+ versions of notes.
## [Unreleased]
## [0.71.22] - 2026-06-10
### Added
- **Perf & measure polish** — a 4-issue patch tightening four live paths from
the recent BETA lifts. Pure code, validated on Windows + RTX 3050.
- **MiniLLM on-policy KV-cache (closes #263).** The on-policy distillation
rollout (`soup train` with `training.minillm_on_policy: true`) now threads
`past_key_values` so each step forwards only the new token instead of
re-feeding the whole prefix — resolving the O(L²) per-step cost from
v0.71.18. A LoRA student (the common distill case) activates the cache
too: the new `_supports_kv_cache` probe unwraps the PEFT model via
`get_base_model()` before deciding. The teacher is always cached; the
student cache respects the retained autograd graph and degrades gracefully
if a model returns no cache mid-loop.
- **`soup serve --mole` KV-cache (closes #262).** Each of the N task adapters
in a served MoLE now keeps its own KV cache in lockstep, created fresh per
`generate()` call (never stored on the instance, so there is no
cross-request leak). Top-k zero-weight adapters are still skipped, and the
output is byte-identical to the no-cache path on a real MoLE.
- **Deploy-autopilot live measure factories (closes #143).** `soup deploy
autopilot --measure` ships a first-party transformers loader factory (lazy
import, per-candidate quant config via the Quant Menu loader; `before` =
base, `after` = quantised) replacing the inject-only test hooks. The
baseline is now scored **once** and the whole candidate list is
**pre-validated up front**, so a typo in `--measure-candidates` raises
before any model load instead of burning N live loads or doubling peak
VRAM.
- **Live-codec TTS via SNAC, partial (#265-partial).** The live-codec
encode path (`data.format='audio'`) is validated for **Orpheus**:
`load_audio_mono` now probes `soundfile.info` (duration + byte cap)
*before* `soundfile.read` (no multi-GB decode into RAM) and reads through
an `O_NOFOLLOW` file descriptor; a real SNAC-backed encode of a 24 kHz wav
produced 42 Orpheus codec tokens.
### Fixed
- MiniLLM on-policy KV-cache was silently disabled for LoRA students (the
PEFT wrapper hid the base model's `past_key_values` support) — now probed
via `get_base_model()`.
- Deploy-measure no longer re-scores the baseline once per candidate or burns
live model loads on a bad candidate (per-candidate validation moved up front).
- `load_audio_mono` capped audio duration only *after* decoding into RAM —
the cap is now checked from `soundfile.info` before reading.
### Known limitations
- KV-cache correctness is validated (cache == no-cache equality on real tiny
artifacts) but large-model throughput gains were not measured on the 4 GB
dev box.
- **#265 stays open** — the live-codec `data.format='audio'` SNAC encode path
is validated for Orpheus only; the other four TTS families keep their
per-family codec dependency gate.
- The deploy-measure first-party factory's real quantized (bitsandbytes 4-bit)
load is CUDA + bitsandbytes-gated; on Windows / no-bnb the injected test
seams are the validated path.
- The MoLE serve KV-cache assumes single-sequence (`B == 1`) decode.
## [0.71.21] - 2026-06-10
### Added

4
CONTRIBUTING.md
View File

@ -120,7 +120,7 @@ src/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 (292 files, 14084 tests)
tests/ - Test suite (293 files, 14184 tests)
examples/ - Real-world config examples and datasets
```
@ -150,7 +150,7 @@ pytest tests/test_data.py::test_detect_alpaca_format -v
pytest tests/ --cov=soup_cli --cov-report=html
```
### Test Files (292 files)
### Test Files (293 files)
> A representative sample of the suite below. The full table lives in
> [`.claude/CLAUDE.md`](.claude/CLAUDE.md); run `pytest tests/ -v` for the complete list.

31
README.md
View File

@ -49,23 +49,22 @@ infrastructure instead of improving models. Soup fixes that.
## What's New
**v0.71.21 — Precision & rollout lift (BETA, hardware-gated).** Five long-deferred stubs go live:
**v0.71.22 — Perf & measure polish.** Four live paths from the recent BETA lifts get tighter:
- **Multi-turn agent rollouts for GRPO** — `training.rollout_backend: openenv` +
`training.rollout_func: my_module:my_fn` runs a live multi-turn rollout at the start of GRPO
training; the rows your function returns replace the prompt dataset. ART / RULER / NeMo-Gym
adapters ship behind honest dependency gates.
- **FP8 attention + NVFP4 training** — `training.fp8_attention: true` converts the attention
projections to torchao float8 (Hopper-gated) and `training.nvfp4: true` applies torchao NVFP4
quantization (Blackwell-gated). Unsupported hardware degrades to a clear advisory, never a crash.
- **vLLM sleep mode for RL** — `training.vllm_sleep_mode: true` puts the vLLM engine on standby
between GRPO rollouts (vLLM ≥ 0.7), freeing VRAM for the training step.
- **Apple-adapter conversion is live** — `soup apple-adapter <dir> --direction hf-to-mlx | mlx-to-hf`
converts PEFT LoRA safetensors ↔ mlx-lm adapters with a numerically-equal round trip
(rank / scale / dropout / num_layers carried).
- **Llama-4 expert delinearization is live** — `soup delinearize-llama4 <src> --target <out>`
reshapes fused `[E*din, dout]` expert weights to `[E, din, dout]` shard-by-shard and copies the
JSON sidecars so the target stays loadable.
- **MiniLLM on-policy distillation is fast** — the on-policy rollout
(`training.minillm_on_policy: true`) now threads a KV cache so each step forwards only the new
token instead of re-feeding the whole prefix, resolving the earlier O(L²) cost. A LoRA student
activates the cache too (the PEFT wrapper is unwrapped before the cache check).
- **`soup serve --mole` uses a KV cache** — each task adapter in a served MoLE keeps its own cache
in lockstep, created fresh per request (no cross-request leak). Output is byte-identical to the
previous no-cache path.
- **`soup deploy autopilot --measure` is real and frugal** — a first-party transformers loader
factory loads each candidate live (per-candidate quant config), scores the baseline **once**, and
pre-validates the whole `--measure-candidates` list up front, so a typo fails before any model
load instead of burning N loads or doubling peak VRAM.
- **Live-codec TTS via SNAC (Orpheus)** — encoding raw audio at train time (`data.format: audio`)
is validated for Orpheus: audio is duration- and byte-capped from `soundfile.info` **before**
decoding into RAM, and read through a symlink-safe file descriptor.
Full history: [CHANGELOG.md](CHANGELOG.md) &middot; [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).

30
pyproject.toml
View File

@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "soup-cli"
version = "0.71.21"
version = "0.71.22"
description = "Fine-tune and post-train LLMs in one command. No SSH, no config hell."
readme = "README.md"
license = "Apache-2.0"
@ -22,7 +22,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
# Core install (`pip install soup-cli`) is intentionally light: the CLI, config
# system, and data tools — no PyTorch. v0.71.0 split the heavy training stack
# system, and data tools — no PyTorch. v0.71.0 split the heavy training stack
# (torch / transformers / peft / trl / datasets / bitsandbytes / accelerate)
# into the `[train]` extra below.
dependencies = [
@ -35,7 +35,7 @@ dependencies = [
]
[project.optional-dependencies]
# v0.71.0 — heavy training stack. `pip install 'soup-cli[train]'` to fine-tune.
# v0.71.0 — heavy training stack. `pip install 'soup-cli[train]'` to fine-tune.
# These were core dependencies through v0.70.0; pins are unchanged.
train = [
"torch>=2.0.0",
@ -46,7 +46,7 @@ train = [
"bitsandbytes>=0.41.0",
"accelerate>=0.25.0",
]
# v0.71.0 — convenience meta-extra pulling the main optional stacks.
# v0.71.0 — convenience meta-extra pulling the main optional stacks.
all = ["soup-cli[train,serve,ui,data]"]
eval = ["lm-eval>=0.4.0"]
data = ["datasketch>=1.6.0"]
@ -74,31 +74,31 @@ sglang = ["sglang>=0.2.0", "fastapi>=0.104.0", "uvicorn>=0.24.0"]
mlx = ["mlx>=0.20.0", "mlx-lm>=0.20.0"]
cce = ["cut-cross-entropy>=24.10.0"]
tui = ["textual>=0.50.0"]
# v0.53.8 #89 — bundle MLflow / SwanLab / Trackio for `--tracker` users.
# v0.53.8 #89 — bundle MLflow / SwanLab / Trackio for `--tracker` users.
trackers = ["mlflow>=2.0.0", "swanlab>=0.3.0", "trackio>=0.0.1"]
# v0.53.8 #85 — fsspec backends for remote dataset loading (s3 / gs / az / oci).
# v0.53.8 #85 — fsspec backends for remote dataset loading (s3 / gs / az / oci).
remote = ["fsspec>=2024.1.0", "s3fs>=2024.1.0", "gcsfs>=2024.1.0", "adlfs>=2024.1.0"]
# v0.53.10 #150 — bundle scikit-optimize so `soup data mix --optimize` runs the
# v0.53.10 #150 — bundle scikit-optimize so `soup data mix --optimize` runs the
# Bayesian-style loop instead of falling back to the v0.48.0 Dirichlet sampler.
mix = ["scikit-optimize>=0.9.0"]
# v0.53.10 #113 — production-grade data quality: langdetect (language) +
# v0.53.10 #113 — production-grade data quality: langdetect (language) +
# presidio-analyzer (PII). Llama-Guard-3-1B is documented as a manual recipe
# (license + ~600 MB weight blob too large to bundle by default).
data-pro = ["langdetect>=1.0.9", "presidio-analyzer>=2.2.0"]
# v0.71.2 #179/#185 — ed25519 detached signing for `soup attest` / `soup
# v0.71.2 #179/#185 — ed25519 detached signing for `soup attest` / `soup
# adapters sign`. Pure-offline; Sigstore keyless (OIDC + Fulcio/Rekor network)
# stays infra-blocked and is NOT bundled here.
sign = ["cryptography>=41.0.0"]
# v0.71.3 #181 — reportlab PDF rendering for `soup train --annex-xi *.pdf`.
# v0.71.3 #181 — reportlab PDF rendering for `soup train --annex-xi *.pdf`.
pdf = ["reportlab>=4.0.0"]
# v0.71.3 #180 — codecarbon offline energy/CO2 measurement for
# v0.71.3 #180 — codecarbon offline energy/CO2 measurement for
# `soup train --track-energy` (offline; no IP-geolocation network call).
carbon = ["codecarbon>=2.0.0"]
# v0.71.13 #225/#227 — prompt-program / tool-schema compilers
# v0.71.13 #225/#227 — prompt-program / tool-schema compilers
# (`soup compile` / `soup compile-tools`). Lazy-imported with a friendly
# ImportError; not installed by default (heavy dependency trees).
compile = ["dspy-ai>=2.5.0", "textgrad>=0.1.0", "gepa>=0.0.1"]
# v0.71.18 #16 — serverless cloud GPU training (`soup train --cloud modal`).
# v0.71.18 #16 — serverless cloud GPU training (`soup train --cloud modal`).
# Lazy-imported; only needed for `--cloud-submit` (plan-only render needs no
# dependency). Modal auth is via `modal setup`.
modal = ["modal>=0.60.0"]
@ -113,7 +113,7 @@ Issues = "https://github.com/MakazhanAlpamys/Soup/issues"
[tool.hatch.build.targets.wheel]
packages = ["src/soup_cli"]
# v0.53.8 #93 — include bundled fixture JSONLs as package data so
# v0.53.8 #93 — include bundled fixture JSONLs as package data so
# `soup data demo` works in zipapp / namespace-package installs.
# Hatchling's ``packages = ["src/soup_cli"]`` already recurses into the
# package directory, so we use the artifacts directive (NOT
@ -146,7 +146,7 @@ follow_imports = "silent"
testpaths = ["tests"]
markers = [
"smoke: slow smoke tests that download models and run training (run with: pytest -m smoke)",
"unit: fast isolated tests — no subprocess, network, filesystem, or real model load",
"unit: fast isolated tests — no subprocess, network, filesystem, or real model load",
"integration: tests that touch real subprocess, SQLite, filesystem, or HTTP",
]
addopts = "-m 'not smoke' --cov=soup_cli --cov-fail-under=77 --cov-report=term-missing:skip-covered"

4
src/soup_cli/__init__.py
View File

@ -1,3 +1,3 @@
"""Soup CLI — Fine-tune and post-train LLMs in one command."""
"""Soup CLI — Fine-tune and post-train LLMs in one command."""
__version__ = "0.71.21"
__version__ = "0.71.22"

44
src/soup_cli/commands/deploy.py
View File

@ -4,7 +4,7 @@ from __future__ import annotations
import os
from pathlib import Path
from typing import TYPE_CHECKING, Callable, List, Optional
from typing import TYPE_CHECKING, List, Optional
import typer
@ -792,25 +792,21 @@ def _run_deploy_autopilot_measure(
f"against {Path(tasks_file).name}...[/]"
)
def _placeholder_before(prompt: str) -> str:
# The full v0.46.1 live measurement plumbs in real
# transformers / vllm generators. v0.53.1 ships the orchestrator
# surface; callers / smoke runs can monkeypatch this in.
return ""
def _placeholder_after_factory(candidate: str) -> Callable[[str], str]:
def _gen(prompt: str) -> str:
return ""
return _gen
# Pull injected generators if the caller registered them via env (escape
# hatch for tests + advanced operator workflows)
# Injected generators (test seam / advanced operator workflows) still win;
# otherwise v0.71.22 #143 first-party transformers factories load the real
# base (before) + per-candidate Quant-Menu-quantized model (after), lazily
# on first prompt so a cache hit never loads a model.
from soup_cli.utils import deploy_measure as _dm
before_gen = getattr(_dm, "_DEPLOY_MEASURE_BEFORE_GEN", None) or _placeholder_before
after_factory = (
getattr(_dm, "_DEPLOY_MEASURE_AFTER_FACTORY", None)
or _placeholder_after_factory
)
injected_before = getattr(_dm, "_DEPLOY_MEASURE_BEFORE_GEN", None)
injected_after = getattr(_dm, "_DEPLOY_MEASURE_AFTER_FACTORY", None)
if injected_before is None and injected_after is None:
console.print(
"[dim]Using live transformers generators (greedy decode; models "
"load lazily per candidate)...[/]"
)
before_gen = injected_before or _dm.build_before_generator(base)
after_factory = injected_after or _dm.build_after_generator_factory(base)
try:
results, cache_hit = run_measure(
@ -824,10 +820,18 @@ def _run_deploy_autopilot_measure(
except (TypeError, ValueError, FileNotFoundError) as exc:
console.print(f"[red]Measure failed:[/] {escape(str(exc))}")
raise typer.Exit(1) from exc
except (RuntimeError, ImportError, OSError) as exc:
# Model-load failures from the live factories (missing quant kernel,
# OOM, network) — friendly exit, no traceback dump.
console.print(f"[red]Live measure failed:[/] {escape(str(exc))}")
raise typer.Exit(1) from exc
console.print(render_measure_table(results))
if cache_hit:
console.print("[dim](cache hit — re-run with --no-cache to refresh)[/]")
console.print(
"[dim](cache hit — delete ~/.soup/deploy_autopilot_cache.json or "
"change --tasks to refresh)[/]"
)
best = pick_best(results)
if best is not None:
console.print(

33
src/soup_cli/trainer/tts.py
View File

@ -64,20 +64,27 @@ class TTSTrainerWrapper(SFTTrainerWrapper):
data_format = getattr(cfg.data, "format", None)
if data_format in _LIVE_CODEC_FORMATS:
# Live-codec mode: encoding raw audio needs the family's codec.
# Live-codec mode (v0.71.22 #265-partial): encode raw audio into
# codec-token strings at setup time, then fall through to the
# validated pre-encoded CE path. The per-family codec dep gate
# fires first (friendly pip hint); families without a validated
# encoder raise inside encode_tts_dataset (tracked in #265 —
# Orpheus/SNAC is the v0.71.22 live family).
self._require_tts_codec(family)
# If the codec IS importable (never validated on the maintainer's
# box — no codecs/large models installed), the encode-then-CE path
# would run here. It is intentionally surfaced as an explicit
# not-yet-validated error so we never silently ship an unrun path.
pkg = tts_codec_package(family)
raise RuntimeError(
f"TTS live-codec mode (data.format='audio') for family "
f"'{family}' requires encoding audio with the {pkg!r} codec at "
"train time. This path is hardware/dependency-gated and not "
"yet validated. Use the pre-encoded chat workflow (encode "
"audio to codec tokens offline, then train with "
"data.format=chat) for a runnable TTS fine-tune."
from soup_cli.utils.tts_codec import encode_tts_dataset
console.print(
f"[green]TTS live-codec:[/] encoding audio with the "
f"{tts_codec_package(family)!r} codec (family={family})"
)
# device is set by the SFTTrainerWrapper.__init__; getattr keeps a
# bare object.__new__ test fixture (no device) working — default
# None falls through to CPU encoding in encode_tts_dataset.
dataset = encode_tts_dataset(
dataset,
family,
device=getattr(self, "device", None),
console=console,
)
# Pre-encoded chat mode: plain SFT cross-entropy over the codec tokens.

273
src/soup_cli/utils/deploy_measure.py
View File

@ -16,6 +16,7 @@ from __future__ import annotations
import hashlib
import json
import os
import re
import stat
import tempfile
from dataclasses import asdict, dataclass
@ -31,12 +32,21 @@ DEFAULT_MAJOR_THRESHOLD: float = 0.05
_MAX_CACHE_BYTES: int = 1 * 1024 * 1024 # 1 MB cap on cache file
_MAX_CANDIDATES: int = 32
_MAX_BASE_LEN: int = 512 # mirrors v0.40.5 reward_model / v0.62.0 --base policy
# v0.71.22 #143 — closed candidate allowlist for the first-party generator
# factories. Mirrors the TrainingConfig.quantization Quant Menu surface
# ("none" = the unquantized baseline; HQQ uses the hqq:Nbit shape).
MEASURABLE_QUANT_CANDIDATES: frozenset[str] = frozenset(
{"none", "4bit", "8bit", "gptq", "awq", "aqlm", "eetq", "mxfp4", "fp8"}
)
_HQQ_CANDIDATE_RE = re.compile(r"^hqq:[12348]bit$")
# Test + advanced-operator escape hatch: when set on this module, the deploy
# CLI uses these callables instead of the v0.46.1 model-loading generators.
# Documented as v0.53.1 deferral — see the v0.53.1 Known Limitations entry
# in plan.md. NOT a public API; the live transformers / vLLM generators land
# alongside v0.53.2.
# CLI uses these callables INSTEAD of the first-party transformers factories
# below (v0.71.22 #143 lifted the v0.53.1 placeholder deferral — the live
# loaders are ``build_before_generator`` / ``build_after_generator_factory``).
# NOT a public API; kept as the test seam.
_DEPLOY_MEASURE_BEFORE_GEN: Optional[Callable[[str], str]] = None
_DEPLOY_MEASURE_AFTER_FACTORY: Optional[
Callable[[str], Callable[[str], str]]
@ -227,12 +237,21 @@ def measure_candidate(
*,
candidate: str,
tasks_file: str,
before_gen: Callable[[str], str],
before_gen: Optional[Callable[[str], str]] = None,
after_gen: Callable[[str], str],
before_score: Optional[float] = None,
) -> MeasureResult:
"""Score one quant candidate against the baseline.
Returns a :class:`MeasureResult` with classify_delta verdict.
The baseline is deterministic across candidates, so the orchestrator
(:func:`run_measure`) computes it ONCE and threads the float in via
``before_score`` — this avoids re-running full greedy generation over
every task once per candidate (N× redundant GPU time) and keeps the
baseline model out of VRAM while a quantized candidate is loaded.
Direct / back-compat callers may instead pass ``before_gen`` and the
baseline is scored here. Exactly one of the two must be supplied.
"""
if isinstance(candidate, bool) or not isinstance(candidate, str):
raise TypeError("candidate must be a non-bool str")
@ -241,7 +260,18 @@ def measure_candidate(
if "\x00" in candidate:
raise ValueError("candidate must not contain null bytes")
before = _score_tasks(tasks_file, before_gen)
if before_score is not None:
if isinstance(before_score, bool) or not isinstance(
before_score, (int, float)
):
raise TypeError("before_score must be a non-bool number")
before = float(before_score)
elif before_gen is not None:
before = _score_tasks(tasks_file, before_gen)
else:
raise ValueError(
"measure_candidate requires either before_gen or before_score"
)
after = _score_tasks(tasks_file, after_gen)
delta = after - before
if delta >= 0:
@ -276,6 +306,221 @@ def pick_best(results: Sequence[MeasureResult]) -> Optional[MeasureResult]:
return max(results, key=lambda r: r.delta)
# --- v0.71.22 #143 — first-party transformers generator factories -----------
def validate_measure_candidate(candidate: object) -> str:
"""Validate + canonicalise a quant candidate for the live measure loop.
Accepts the closed :data:`MEASURABLE_QUANT_CANDIDATES` allowlist plus the
``hqq:{1,2,3,4,8}bit`` shape. Case-insensitive; returns the lowercase
canonical form. Bool rejected before the str check (project policy).
"""
if isinstance(candidate, bool):
raise TypeError("candidate must not be bool")
if not isinstance(candidate, str):
raise TypeError(f"candidate must be str, got {type(candidate).__name__}")
if not candidate:
raise ValueError("candidate must be non-empty")
if "\x00" in candidate:
raise ValueError("candidate must not contain null bytes")
canonical = candidate.lower()
if canonical in MEASURABLE_QUANT_CANDIDATES or _HQQ_CANDIDATE_RE.match(
canonical
):
return canonical
supported = ", ".join(sorted(MEASURABLE_QUANT_CANDIDATES))
# Truncate the echoed candidate (mirrors longlora._truncate_for_message
# policy) so a pathologically long value can't bloat the error message.
shown = candidate if len(candidate) <= 64 else candidate[:64] + "…"
raise ValueError(
f"candidate {shown!r} is not a measurable quant. Supported: "
f"{supported}, hqq:Nbit (N in 1/2/3/4/8)"
)
def _check_base_id(base: object) -> str:
"""Shape-validate a base model id / path for the factory builders."""
# NOTE: this raises ValueError (not TypeError) for bool/non-str wrong-type
# input — DELIBERATELY inconsistent with validate_measure_candidate (which
# raises TypeError for the same class). The ValueError behaviour is pinned
# by tests (test_v07122.test_builders_validate_max_new_tokens); changing it
# would break them, so we keep it and document the seam here.
if isinstance(base, bool) or not isinstance(base, str):
raise ValueError("base must be a non-bool string")
if not base:
raise ValueError("base must be non-empty")
if "\x00" in base:
raise ValueError("base must not contain null bytes")
if len(base) > _MAX_BASE_LEN:
raise ValueError(f"base exceeds {_MAX_BASE_LEN} chars")
return base
def _check_max_new_tokens(value: object) -> int:
# NOTE: raises ValueError (not TypeError) for bool/non-int wrong-type input
# — deliberately inconsistent with validate_measure_candidate's TypeError;
# the ValueError behaviour is pinned by tests (see _check_base_id note).
if isinstance(value, bool) or not isinstance(value, int):
raise ValueError("max_new_tokens must be a non-bool int")
if value < 1:
raise ValueError(f"max_new_tokens must be >= 1, got {value}")
return value
def _free_accelerator_memory() -> None:
"""Best-effort GC + CUDA cache flush between candidate model loads."""
import gc
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception: # noqa: BLE001 — torch missing / CUDA error: best-effort
pass
def _import_transformers():
"""Lazy transformers import — module-level seam for tests."""
import transformers
return transformers
def _load_measure_model(
base: str,
*,
quantization: str,
device: Optional[str] = None,
trust_remote_code: bool = False,
):
"""Load ``base`` with the candidate quant via the Quant Menu loader.
Returns a ``(model, tokenizer, device)`` triple shaped for
:func:`soup_cli.utils.live_eval.make_generator`'s ``loaded=`` kwarg.
``quantization='none'`` loads the plain baseline. Heavy imports are lazy.
"""
from soup_cli.config.schema import TrainingConfig
from soup_cli.utils.live_eval import resolve_device
from soup_cli.utils.quant_menu import build_quantization_config_for_loader
transformers = _import_transformers()
dev = resolve_device(device)
tcfg = TrainingConfig(quantization=quantization)
quant_config = build_quantization_config_for_loader(
tcfg=tcfg, base=base, console=None
)
tokenizer = transformers.AutoTokenizer.from_pretrained(
base, trust_remote_code=trust_remote_code
)
if getattr(tokenizer, "pad_token", None) is None:
tokenizer.pad_token = tokenizer.eos_token
if quant_config is not None:
# Quantized loads place weights themselves (bnb/gptq/awq need
# device_map on CUDA; a post-hoc .to() would break bnb layouts).
model = transformers.AutoModelForCausalLM.from_pretrained(
base,
quantization_config=quant_config,
device_map="auto" if dev == "cuda" else None,
trust_remote_code=trust_remote_code,
)
else:
model = transformers.AutoModelForCausalLM.from_pretrained(
base, trust_remote_code=trust_remote_code
).to(dev)
model.eval()
return model, tokenizer, dev
def _lazy_generator(
base: str,
*,
quantization: str,
device: Optional[str],
max_new_tokens: int,
trust_remote_code: bool,
) -> Callable[[str], str]:
"""Closure that loads the model on FIRST call (cache hits never load)."""
holder: dict[str, Callable[[str], str]] = {}
def _gen(prompt: str) -> str:
if "fn" not in holder:
_free_accelerator_memory()
loaded = _load_measure_model(
base,
quantization=quantization,
device=device,
trust_remote_code=trust_remote_code,
)
from soup_cli.utils import live_eval
holder["fn"] = live_eval.make_generator(
base, loaded=loaded, max_new_tokens=max_new_tokens
)
return holder["fn"](prompt)
return _gen
def build_before_generator(
base: str,
*,
device: Optional[str] = None,
max_new_tokens: int = 64,
trust_remote_code: bool = False,
) -> Callable[[str], str]:
"""First-party baseline generator for ``deploy autopilot --measure``.
Lazy-loads the UNQUANTIZED ``base`` (greedy decode via
:func:`soup_cli.utils.live_eval.make_generator`) on the first prompt —
a cache-hit measure run never touches the model. (v0.71.22 #143)
"""
_check_base_id(base)
_check_max_new_tokens(max_new_tokens)
return _lazy_generator(
base,
quantization="none",
device=device,
max_new_tokens=max_new_tokens,
trust_remote_code=trust_remote_code,
)
def build_after_generator_factory(
base: str,
*,
device: Optional[str] = None,
max_new_tokens: int = 64,
trust_remote_code: bool = False,
) -> Callable[[str], Callable[[str], str]]:
"""First-party per-candidate quantized generator factory (#143).
``factory(candidate)`` validates the candidate eagerly (fail-fast at the
measure loop boundary) and returns a generator that lazy-loads ``base``
with that candidate's quant config from the Quant Menu loader on first
use. CUDA memory from the previous candidate is freed before each load.
"""
_check_base_id(base)
_check_max_new_tokens(max_new_tokens)
def factory(candidate: str) -> Callable[[str], str]:
canonical = validate_measure_candidate(candidate)
return _lazy_generator(
base,
quantization=canonical,
device=device,
max_new_tokens=max_new_tokens,
trust_remote_code=trust_remote_code,
)
# Marker so run_measure can pre-validate the whole candidate list up front
# for the first-party path (M2) without touching injected test seams.
factory._soup_first_party_measure_factory = True # type: ignore[attr-defined]
return factory
def run_measure(
*,
profile_name: str,
@ -303,6 +548,15 @@ def run_measure(
f"too many candidates ({len(candidates)}; cap {_MAX_CANDIDATES})"
)
# M2 — pre-validate the ENTIRE candidate list up front so a typo'd Nth
# candidate fails BEFORE any (expensive) model load, rather than burning a
# live load + eval on every preceding candidate. Only enforced when the
# first-party factory is in play; an injected test seam factory is left
# untouched (it may legitimately use non-Quant-Menu candidate names).
if getattr(after_gen_factory, "_soup_first_party_measure_factory", False):
for candidate in candidates:
validate_measure_candidate(candidate)
tasks_sha = sha_of_file(tasks_file)
key = compute_cache_key(
base_sha=base_sha, profile_name=profile_name, tasks_sha=tasks_sha,
@ -320,13 +574,16 @@ def run_measure(
except (TypeError, ValueError):
pass
# Cache miss — run the eval loop
# Cache miss — compute the deterministic baseline ONCE (M3), then run each
# candidate against it. This keeps the baseline model out of VRAM while a
# quantized candidate is loaded and avoids N× redundant baseline scoring.
before_score = _score_tasks(tasks_file, before_gen)
results: list[MeasureResult] = []
for candidate in candidates:
after_gen = after_gen_factory(candidate)
result = measure_candidate(
candidate=candidate, tasks_file=tasks_file,
before_gen=before_gen, after_gen=after_gen,
after_gen=after_gen, before_score=before_score,
)
results.append(result)

119
src/soup_cli/utils/minillm.py
View File

@ -235,6 +235,39 @@ def _multinomial_sample(probs):
return torch.multinomial(probs, num_samples=1)
def _supports_kv_cache(model: object) -> bool:
"""Capability probe: does ``model.forward`` explicitly take a KV cache?
Requires BOTH ``past_key_values`` and ``use_cache`` to be declared as
explicit parameters (every HF ``*ForCausalLM`` does). A bare ``**kwargs``
does NOT count — a model that merely swallows the kwarg would silently
ignore the cache and the rollout would feed it a single token per step
against an empty context. PEFT wrappers
(``PeftModel.forward(*args, **kwargs)``) — the common live distill case —
are probed through ``get_base_model()`` so a LoRA student still gets the
cached path (mirrors :func:`mole_routing._model_supports_cache`). The real
forwards already go through the wrapper; only the probe is unwrapped. Never
raises.
"""
import inspect
target = model
get_base = getattr(model, "get_base_model", None)
if callable(get_base):
try:
target = get_base()
except Exception: # noqa: BLE001 — probe must never raise
target = model
forward = getattr(target, "forward", None)
if forward is None:
return False
try:
params = inspect.signature(forward).parameters
except (TypeError, ValueError):
return False
return "past_key_values" in params and "use_cache" in params
def minillm_on_policy_rollout(
student_model,
teacher_model,
@ -245,6 +278,7 @@ def minillm_on_policy_rollout(
max_new_tokens: int,
temperature: float = 1.0,
sample_fn=None,
use_cache: bool = True,
):
"""True on-policy MiniLLM teacher-mixed rollout (v0.71.18 #257).
@ -263,6 +297,28 @@ def minillm_on_policy_rollout(
reverse-KL is averaged over the rollout length so longer rollouts don't
dominate the gradient.
KV-cache (v0.71.22 #263): when ``use_cache`` is True (default) and BOTH
models declare ``past_key_values``/``use_cache`` on their forward
(:func:`_supports_kv_cache`), each step forwards only the new token
against the cached context — O(L) compute instead of the O(L²) full
re-forward. Forward logits are identical either way (modulo float
noise), so the loss trajectory matches the no-cache path. Gradients are
mathematically equivalent too — both paths are full
backprop-through-rollout — so only the retained-graph *shape* differs:
- **Teacher** cache lives under ``torch.no_grad()`` — trivially safe.
- **Student** cache is grad-carrying: cached KVs from earlier steps stay
in the autograd graph, so the cached path keeps ONE shared graph
spanning the whole rollout, whereas the no-cache path builds L separate
per-step graphs (each its own fresh full-prefix re-forward). Both
backprop through every step; the cached shared-graph form has retained
activations of ~one full-sequence forward — at or below the no-cache
peak.
Models whose forward does not declare the cache params (or that ignore
``use_cache`` and return no ``past_key_values``) transparently keep the
legacy full re-forward path.
Returns ``(loss, num_steps)`` where ``loss`` is a scalar differentiable
w.r.t. the student parameters and scaled by ``temperature**2`` (the
Hinton convention shared with the offline term).
@ -279,6 +335,8 @@ def minillm_on_policy_rollout(
raise ValueError(f"max_new_tokens must be >= 1, got {max_new_tokens}")
if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
raise TypeError("temperature must be a non-bool number")
if not isinstance(use_cache, bool):
raise TypeError("use_cache must be bool")
t = float(temperature)
if not math.isfinite(t) or t <= 0.0:
raise ValueError("temperature must be finite and positive")
@ -290,13 +348,51 @@ def minillm_on_policy_rollout(
cur_mask = torch.ones_like(input_ids)
else:
cur_mask = attention_mask
cache_ok = (
use_cache
and _supports_kv_cache(student_model)
and _supports_kv_cache(teacher_model)
)
s_past = None
t_past = None
step_ids = cur_ids # full prompt on the first step; the delta afterwards
kl_terms = []
for _ in range(int(max_new_tokens)):
s_out = student_model(input_ids=cur_ids, attention_mask=cur_mask)
if cache_ok:
s_out = student_model(
input_ids=step_ids,
attention_mask=cur_mask,
past_key_values=s_past,
use_cache=True,
)
else:
s_out = student_model(input_ids=cur_ids, attention_mask=cur_mask)
s_logits = s_out.logits[:, -1, :] / t # [B, Vs] (grad ON)
with torch.no_grad():
te_out = teacher_model(input_ids=cur_ids, attention_mask=cur_mask)
if cache_ok:
te_out = teacher_model(
input_ids=step_ids,
attention_mask=cur_mask,
past_key_values=t_past,
use_cache=True,
)
else:
te_out = teacher_model(
input_ids=cur_ids, attention_mask=cur_mask
)
te_logits = te_out.logits[:, -1, :] / t # [B, Vt] (grad OFF)
if cache_ok:
new_s_past = getattr(s_out, "past_key_values", None)
new_t_past = getattr(te_out, "past_key_values", None)
if new_s_past is None or new_t_past is None:
# Model ignored use_cache — degrade to full re-forwards for
# the remaining steps (cur_ids is maintained either way).
cache_ok = False
s_past = None
t_past = None
else:
s_past = new_s_past
t_past = new_t_past
# MiniLLM assumes a shared vocab; clamp to the common min so a
# mismatched teacher never index-errors (cross-tokenizer = ULD path).
common = min(s_logits.shape[-1], te_logits.shape[-1])
@ -315,6 +411,7 @@ def minillm_on_policy_rollout(
next_tok = next_tok.to(cur_ids.dtype)
cur_ids = torch.cat([cur_ids, next_tok], dim=1)
cur_mask = torch.cat([cur_mask, torch.ones_like(next_tok)], dim=1)
step_ids = next_tok if cache_ok else cur_ids
stacked = torch.stack(kl_terms, dim=1) # [B, L]
if config.length_normalize:
@ -348,11 +445,14 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
lazily from ``pretrain_anchor_path``, scaled by
``pretrain_anchor_weight`` (prevents the student drifting away from
coherent language).
- :meth:`on_policy_term` — true on-policy teacher-mixed autoregressive
rollout (sample → teacher-score) via
:func:`minillm_on_policy_rollout` (shipped v0.71.18 #257; KV-cache
accelerated v0.71.22 #263).
Honest scope: the teacher-mix is the offline distribution-blend
analog of MiniLLM's on-policy teacher-mixed *sampling*; a full
autoregressive rollout loop (sample → teacher-score) is a larger
follow-up documented as a known limitation.
The offline :meth:`distill_term` teacher-mix and the live on-policy
rollout are both available; pick per the ``training.minillm_on_policy``
flag.
"""
def __init__(
@ -394,6 +494,7 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
attention_mask=None,
*,
sample_fn=None,
use_cache: bool = True,
):
"""True on-policy teacher-mixed rollout reverse-KL (v0.71.18 #257).
@ -401,6 +502,11 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
rollout from the teacher/student mixture and returns the
length-normalised reverse-KL along that path. Returns the scalar
loss only (the step count is internal to the rollout).
``use_cache`` (v0.71.22 #263) threads through to
:func:`minillm_on_policy_rollout` — cache-capable models forward only
the new token per step (see the rollout docstring for the gradient
trade-off).
"""
loss, _ = minillm_on_policy_rollout(
student_model,
@ -411,6 +517,7 @@ class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
max_new_tokens=self.config.rollout_length,
temperature=self.temperature,
sample_fn=sample_fn,
use_cache=use_cache,
)
return loss

165
src/soup_cli/utils/mole_routing.py
View File

@ -33,7 +33,7 @@ import math
import os
import stat
import tempfile
from dataclasses import dataclass
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any, Optional, Tuple
@ -61,42 +61,42 @@ _MAX_ADAPTER_PATH_LEN = 4096
# ---------------------------------------------------------------------------
def _check_int(value: object, field: str, lo: int, hi: int) -> int:
def _check_int(value: object, field_name: str, lo: int, hi: int) -> int:
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{field} must be int")
raise TypeError(f"{field_name} must be int")
if value < lo:
raise ValueError(f"{field} {value} below floor {lo}")
raise ValueError(f"{field_name} {value} below floor {lo}")
if value > hi:
raise ValueError(f"{field} {value} above cap {hi}")
raise ValueError(f"{field_name} {value} above cap {hi}")
return value
def _check_finite_positive(
value: object, field: str, lo: float, hi: float
value: object, field_name: str, lo: float, hi: float
) -> float:
if isinstance(value, bool):
raise TypeError(f"{field} must not be bool")
raise TypeError(f"{field_name} must not be bool")
if not isinstance(value, (int, float)):
raise TypeError(f"{field} must be numeric")
raise TypeError(f"{field_name} must be numeric")
val = float(value)
if not math.isfinite(val):
raise ValueError(f"{field} must be finite")
raise ValueError(f"{field_name} must be finite")
if val < lo:
raise ValueError(f"{field} {val} below floor {lo}")
raise ValueError(f"{field_name} {val} below floor {lo}")
if val > hi:
raise ValueError(f"{field} {val} above cap {hi}")
raise ValueError(f"{field_name} {val} above cap {hi}")
return val
def _check_str_field(value: object, field: str) -> str:
def _check_str_field(value: object, field_name: str) -> str:
if isinstance(value, bool):
raise TypeError(f"{field} must not be bool")
raise TypeError(f"{field_name} must not be bool")
if not isinstance(value, str):
raise TypeError(f"{field} must be str")
raise TypeError(f"{field_name} must be str")
if not value:
raise ValueError(f"{field} must be non-empty")
raise ValueError(f"{field_name} must be non-empty")
if "\x00" in value:
raise ValueError(f"{field} must not contain null bytes")
raise ValueError(f"{field_name} must not contain null bytes")
return value
@ -441,6 +441,56 @@ def load_mole_manifest(directory: str) -> MoleServeManifest:
return _manifest_from_dict(data)
def _model_supports_cache(model: object) -> bool:
"""Capability probe: can this model take a per-call KV cache? (#262)
Requires BOTH ``past_key_values`` and ``use_cache`` as EXPLICIT params on
the forward signature — a bare ``**kwargs`` does not count (a model that
merely swallows the kwarg would silently ignore the cache). PEFT wrappers
(``PeftModel.forward(*args, **kwargs)``) are probed through
``get_base_model()`` so real HF causal LMs get the cached path. Never
raises.
"""
import inspect
target = model
get_base = getattr(model, "get_base_model", None)
if callable(get_base):
try:
target = get_base()
except Exception: # noqa: BLE001 — probe must never raise
target = model
forward = getattr(target, "forward", None)
if forward is None:
return False
try:
params = inspect.signature(forward).parameters
except (TypeError, ValueError):
return False
return "past_key_values" in params and "use_cache" in params
@dataclass
class _MoleKvState:
"""Per-``generate()`` KV-cache state — one cache per adapter + the base.
``caches[i]`` / ``lens[i]`` track adapter ``i``'s cache object and how
many tokens of the current sequence it has seen. An adapter skipped by
top-k masking simply falls behind; when it becomes active again it is fed
the missed tokens in one catch-up forward (caches stay in lockstep with
the sequence without paying for inactive adapters every step).
``enabled`` flips off when the model ignores ``use_cache`` (no
``past_key_values`` on the output) — the call degrades to the legacy
full-re-forward path.
"""
enabled: bool = True
base_cache: Any = None
base_len: int = 0
caches: dict = field(default_factory=dict)
lens: dict = field(default_factory=dict)
class LoadedMole:
"""Runtime serve-time MoLE: base + N frozen task LoRAs + the trained gate.
@ -450,10 +500,14 @@ class LoadedMole:
blended by the per-token gate weights. At decode we only need the LAST
token's blend per step.
Generation recomputes the full sequence each step (no shared KV cache —
each task adapter would need its own cache, so recompute keeps the blend
correct). This is fine for short demo generations on a tiny model; large
serving deployments should expect linear-per-step cost.
KV-cache (v0.71.22 #262): cache-capable models (probe:
:func:`_model_supports_cache`) keep one KV cache per adapter plus one for
the base/router forward, all scoped to a single ``generate()`` call. Each
step then feeds only the new token (or, for an adapter that was skipped
by top-k masking, the tokens it missed — a catch-up delta) instead of
re-forwarding the whole sequence through every adapter. Models without
explicit cache support keep the legacy full-re-forward path, which is
fine for short demo generations on a tiny model.
"""
def __init__(
@ -507,6 +561,68 @@ class LoadedMole:
model.set_adapter(self.adapter_names[0])
return blended
def _blended_last_logits_cached(self, input_ids, attention_mask, state):
"""KV-cached per-token blend (#262) — returns ``[B, V]``.
Feeds each forward only the tokens its cache has not seen yet
(``input_ids[:, cached_len:]``) with the FULL attention mask, so the
base/router and every active adapter pay one-token cost per step
instead of re-processing the whole sequence. Falls back to the legacy
path (and disables the state) when the model ignores ``use_cache``.
"""
import torch
model = self.model
seq_len = int(input_ids.shape[1])
base_delta = input_ids[:, state.base_len:]
with torch.no_grad(), model.disable_adapter():
base_out = model(
input_ids=base_delta,
attention_mask=attention_mask,
past_key_values=state.base_cache,
use_cache=True,
output_hidden_states=True,
)
new_base_cache = getattr(base_out, "past_key_values", None)
if new_base_cache is None:
state.enabled = False
return self._blended_last_logits(input_ids, attention_mask)
state.base_cache = new_base_cache
state.base_len = seq_len
# hidden_states cover only the fed tokens; -1 is the newest token.
router_hidden = base_out.hidden_states[-1][:, -1, :] # [B, H]
weights = self.gate(router_hidden.to(self.gate.gate.weight.dtype)) # [B, N]
col_max = weights.abs().amax(dim=tuple(range(weights.dim() - 1))) # [N]
active = [i for i in range(len(self.adapter_names)) if float(col_max[i]) > 0.0]
if not active: # degenerate (softmax should always keep >=1) — run all
active = list(range(len(self.adapter_names)))
blended = None
for i in active:
model.set_adapter(self.adapter_names[i])
# Catch-up delta: an adapter skipped by top-k on earlier steps is
# behind the sequence — feed it everything it missed at once.
delta_i = input_ids[:, state.lens.get(i, 0):]
with torch.no_grad():
out_i = model(
input_ids=delta_i,
attention_mask=attention_mask,
past_key_values=state.caches.get(i),
use_cache=True,
)
cache_i = getattr(out_i, "past_key_values", None)
if cache_i is None:
state.enabled = False
model.set_adapter(self.adapter_names[0])
return self._blended_last_logits(input_ids, attention_mask)
state.caches[i] = cache_i
state.lens[i] = seq_len
logits_i = out_i.logits[:, -1, :].detach() # [B, V]
w_i = weights[..., i : i + 1].to(logits_i.dtype)
term = w_i * logits_i
blended = term if blended is None else blended + term
model.set_adapter(self.adapter_names[0])
return blended
def generate(
self,
input_ids,
@ -522,6 +638,9 @@ class LoadedMole:
Single-sequence (``B == 1``): the EOS stop inspects row 0, so a batched
caller would stop the whole batch on the first row's EOS. The serve
handler always sends one sequence; ``B > 1`` is rejected.
Cache state (#262) is local to this call — two concurrent or
sequential ``generate()`` calls never share KV caches.
"""
import torch
@ -538,8 +657,12 @@ class LoadedMole:
do_sample = isinstance(temperature, (int, float)) and not isinstance(
temperature, bool
) and float(temperature) > 0.0
state = _MoleKvState() if _model_supports_cache(self.model) else None
for _ in range(int(max_new_tokens)):
logits = self._blended_last_logits(seq, attn) # [B, V]
if state is not None and state.enabled:
logits = self._blended_last_logits_cached(seq, attn, state)
else:
logits = self._blended_last_logits(seq, attn) # [B, V]
if do_sample:
next_id = _sample_next(logits, float(temperature), float(top_p))
else:

373
src/soup_cli/utils/tts_codec.py Normal file
View File

@ -0,0 +1,373 @@
"""v0.71.22 #265-partial — live-codec TTS audio encoding (Orpheus via SNAC).
Lifts the v0.71.20 ``data.format='audio'`` hardware gate into a real
encode-at-train-time path. Rows shaped ``{"audio": <path>, "messages": [...]}``
(the v0.17.0 audio format — paths are resolved + containment-checked by the
data loader) are encoded into discrete codec-token strings that become the
assistant turn, after which training proceeds through the validated
pre-encoded SFT cross-entropy path.
The generic pipeline (``encode_tts_dataset`` → ``encode_tts_row`` →
per-family encoder) is family-agnostic; v0.71.22 ships ONE validated encoder:
* **Orpheus / SNAC** (``pip install snac``) — 24 kHz SNAC produces 3
codebooks at a 1/2/4 frame ratio; Orpheus interleaves them 7 tokens per
frame with per-slot codebook offsets and renders each code as
``<custom_token_N>`` where ``N = code + 10 + slot*4096`` (the official
Orpheus id layout: ``id = 128256 + N = 128266 + code + slot*4096`` on the
Orpheus tokenizer, whose ``<custom_token_i>`` maps to ``128256 + i``).
The other four families (sesame_csm / llasa / spark / oute) keep their
per-family codec dep gate; their encoders are tracked in #265 and raise a
friendly ``RuntimeError`` pointing at the offline pre-encode workflow.
Security / robustness:
- Heavy imports (numpy / soundfile / torch / snac) are lazy — module import
stays light for the CLI hot path.
- Audio paths: null-byte rejection, symlink rejection (``os.lstat``,
defence-in-depth — the loader already containment-checks under
``data.audio_dir``), duration cap (``_MAX_AUDIO_SECONDS``).
- Pure interleave kernel validates the 1/2/4 codebook ratio, per-code bounds
``[0, 4096)``, and rejects bool codes (project bool-as-int policy).
- Rows are deep-copied — the caller's dataset is never mutated (v0.33.0 #47
immutability policy).
"""
from __future__ import annotations
import os
import stat
from typing import Any, Callable, Optional
from soup_cli.utils.tts import tts_codec_package, validate_tts_family
# SNAC checkpoint for the Orpheus family (24 kHz, 3 codebooks).
SNAC_MODEL_ID = "hubertsiuzdak/snac_24khz"
SNAC_SAMPLE_RATE = 24_000
# Orpheus token layout: <custom_token_N> with N = code + OFFSET + slot*4096.
# Slots 0..9 are Orpheus control tokens; audio codes start at 10.
ORPHEUS_CODE_OFFSET = 10
ORPHEUS_CODEBOOK_SIZE = 4096
_ORPHEUS_FRAME_SLOTS = 7
# Families whose live-codec encoder is implemented AND validated.
LIVE_CODEC_FAMILIES: frozenset[str] = frozenset({"orpheus"})
_MAX_AUDIO_SECONDS = 600.0 # 10 min cap — defends against runaway encodes
# Byte cap (~600 s of 24 kHz stereo float32 + headroom) — rejected before any
# read so a runaway file never materialises into RAM.
_MAX_AUDIO_BYTES = 512 * 1024 * 1024
# Process-level SNAC model cache (one load per device). NOTE: single-threaded
# assumption — the live-codec trainer encode path runs on one thread, so this
# unbounded, unguarded cache is intentional (no eviction / no lock needed).
_SNAC_CACHE: dict[str, Any] = {}
# ---------------------------------------------------------------------------
# Pure interleave kernel (no torch — unit-testable on plain lists)
# ---------------------------------------------------------------------------
def _check_codes(codes, name: str, expected_len: int) -> list[int]:
out = []
for code in codes:
if isinstance(code, bool) or not isinstance(code, int):
raise TypeError(f"{name} codes must be non-bool ints")
if not (0 <= code < ORPHEUS_CODEBOOK_SIZE):
raise ValueError(
f"{name} code {code} out of range [0, {ORPHEUS_CODEBOOK_SIZE})"
)
out.append(code)
if len(out) != expected_len:
raise ValueError(
f"{name} codebook length {len(out)} != expected {expected_len} "
"(SNAC codebooks must have the 1/2/4 frame ratio)"
)
return out
def interleave_orpheus_codes(coarse, medium, fine) -> list[int]:
"""Interleave the 3 SNAC codebooks into Orpheus token indices.
``coarse``/``medium``/``fine`` are the per-frame code sequences at the
SNAC 1/2/4 ratio (lengths T / 2T / 4T). Returns 7 token indices per
frame in the official Orpheus slot order, each offset by
``ORPHEUS_CODE_OFFSET + slot * ORPHEUS_CODEBOOK_SIZE`` so the rendered
``<custom_token_N>`` strings map onto the Orpheus tokenizer's audio ids.
"""
coarse = list(coarse)
if not coarse:
raise ValueError("coarse codebook must be non-empty")
frames = len(coarse)
coarse = _check_codes(coarse, "coarse", frames)
medium = _check_codes(list(medium), "medium", 2 * frames)
fine = _check_codes(list(fine), "fine", 4 * frames)
size = ORPHEUS_CODEBOOK_SIZE
base = ORPHEUS_CODE_OFFSET
out: list[int] = []
for i in range(frames):
out.append(coarse[i] + base)
out.append(medium[2 * i] + base + size)
out.append(fine[4 * i] + base + 2 * size)
out.append(fine[4 * i + 1] + base + 3 * size)
out.append(medium[2 * i + 1] + base + 4 * size)
out.append(fine[4 * i + 2] + base + 5 * size)
out.append(fine[4 * i + 3] + base + 6 * size)
return out
def orpheus_tokens_to_string(indices) -> str:
"""Render Orpheus token indices as a ``<custom_token_N>`` string.
Indices must be non-bool ints (mirrors :func:`_check_codes` — project
bool-as-int policy; no float/bool coercion).
"""
indices = list(indices)
if not indices:
raise ValueError("token index list must be non-empty")
for n in indices:
if isinstance(n, bool) or not isinstance(n, int):
raise TypeError("token indices must be non-bool ints")
return "".join(f"<custom_token_{n}>" for n in indices)
# ---------------------------------------------------------------------------
# Audio loading (lazy soundfile / numpy)
# ---------------------------------------------------------------------------
def load_audio_mono(path: str, *, target_sr: int = SNAC_SAMPLE_RATE):
"""Load an audio file as mono float32 at ``target_sr``.
Lazy-imports soundfile (friendly ImportError naming the install).
Off-rate audio is resampled with basic linear interpolation — adequate
for codec encoding of speech; operators wanting studio-grade resampling
should resample offline. Returns a 1-D ``np.float32`` array.
"""
if not isinstance(path, str) or not path:
raise ValueError("audio path must be a non-empty string")
if "\x00" in path:
raise ValueError("audio path must not contain null bytes")
if not os.path.exists(path):
raise FileNotFoundError(
f"audio file not found: {os.path.basename(path)!r}"
)
# Symlink rejection — defence-in-depth; the loader already containment-
# checks rows under data.audio_dir (v0.17.0 policy).
if stat.S_ISLNK(os.lstat(path).st_mode):
raise ValueError("audio path must not be a symlink")
# Byte-size cap before any read — defends against a runaway file.
if os.path.getsize(path) > _MAX_AUDIO_BYTES:
raise ValueError(
f"audio file exceeds the {_MAX_AUDIO_BYTES} byte live-codec cap; "
"split the clip or pre-encode offline"
)
try:
import soundfile
except ImportError as exc: # pragma: no cover - dep present in CI
raise ImportError(
"Live-codec TTS needs the 'soundfile' package to read audio. "
"Install it with `pip install soundfile` (or `pip install "
"'soup-cli[audio]'`)."
) from exc
import numpy as np
# Probe metadata FIRST — reject an over-long clip before materialising the
# whole file into RAM (the post-read check below stays as defence-in-depth).
info = soundfile.info(path)
if info.samplerate and info.frames / float(info.samplerate) > _MAX_AUDIO_SECONDS:
raise ValueError(
f"audio is {info.frames / float(info.samplerate):.1f} seconds — "
f"exceeds the {_MAX_AUDIO_SECONDS:.0f} seconds live-codec cap; "
"split the clip or pre-encode offline"
)
# Open via O_NOFOLLOW and hand the fd to soundfile — closes the TOCTOU
# window between the lstat symlink check above and the read (project
# v0.65/v0.67 O_NOFOLLOW reader policy). soundfile.read accepts an open
# file object.
fd = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
try:
with os.fdopen(fd, "rb") as handle:
data, sr = soundfile.read(handle, dtype="float32", always_2d=True)
except OSError as exc:
# O_NOFOLLOW raises ELOOP if the path became a symlink after the lstat.
raise ValueError(f"audio path could not be read: {exc}") from exc
duration = data.shape[0] / float(sr)
if duration > _MAX_AUDIO_SECONDS:
raise ValueError(
f"audio is {duration:.1f} seconds — exceeds the "
f"{_MAX_AUDIO_SECONDS:.0f} seconds live-codec cap; split the "
"clip or pre-encode offline"
)
mono = data.mean(axis=1)
if int(sr) != int(target_sr):
n_out = max(1, int(round(len(mono) * target_sr / float(sr))))
x_old = np.linspace(0.0, 1.0, num=len(mono), endpoint=False)
x_new = np.linspace(0.0, 1.0, num=n_out, endpoint=False)
mono = np.interp(x_new, x_old, mono)
return mono.astype(np.float32)
# ---------------------------------------------------------------------------
# Orpheus / SNAC encoder
# ---------------------------------------------------------------------------
def _get_snac_model(device: Optional[str] = None):
"""Load (once per device) the SNAC codec model. Lazy snac import."""
dev = device or "cpu"
cached = _SNAC_CACHE.get(dev)
if cached is not None:
return cached
try:
from snac import SNAC
except ImportError as exc:
raise ImportError(
"TTS family 'orpheus' live-codec training requires the 'snac' "
"package (audio codec). Install it with `pip install snac`, or "
"pre-encode your audio to codec tokens offline and train with "
"data.format=chat."
) from exc
model = SNAC.from_pretrained(SNAC_MODEL_ID).eval().to(dev)
_SNAC_CACHE[dev] = model
return model
def encode_audio_orpheus(
path: str,
*,
snac_model: Optional[Any] = None,
device: Optional[str] = None,
) -> str:
"""Encode one audio file into the Orpheus ``<custom_token_N>`` string.
``snac_model`` is injectable (tests / pre-loaded models); when ``None``
the 24 kHz SNAC checkpoint is loaded once per process per device.
"""
audio = load_audio_mono(path, target_sr=SNAC_SAMPLE_RATE)
import torch
model = snac_model if snac_model is not None else _get_snac_model(device)
wav = torch.from_numpy(audio).reshape(1, 1, -1)
if device:
wav = wav.to(device)
with torch.no_grad():
codes = model.encode(wav)
coarse = [int(c) for c in codes[0][0].tolist()]
medium = [int(c) for c in codes[1][0].tolist()]
fine = [int(c) for c in codes[2][0].tolist()]
return orpheus_tokens_to_string(
interleave_orpheus_codes(coarse, medium, fine)
)
def tts_encoder_for_family(
family: str, *, device: Optional[str] = None
) -> Callable[[str], str]:
"""Return the live audio→codec-string encoder for ``family``.
Only the families in :data:`LIVE_CODEC_FAMILIES` have a validated
encoder (v0.71.22 ships Orpheus/SNAC). The remaining families raise a
friendly ``RuntimeError`` naming the offline workflow — their encoders
are tracked in #265 (the per-family codec dep gate fires earlier, at
trainer setup).
"""
canonical = validate_tts_family(family)
if canonical == "orpheus":
def _encode(path: str) -> str:
return encode_audio_orpheus(path, device=device)
return _encode
pkg = tts_codec_package(canonical)
raise RuntimeError(
f"Live-codec encoding for TTS family '{canonical}' is not yet "
f"implemented (v0.71.22 ships the Orpheus/SNAC encoder; the "
f"'{canonical}' encoder via {pkg!r} is tracked in #265). Pre-encode "
"your audio to codec tokens offline and train with data.format=chat."
)
# ---------------------------------------------------------------------------
# Dataset mapping
# ---------------------------------------------------------------------------
def encode_tts_row(row: dict, encoder: Callable[[str], str]) -> dict:
"""Encode one ``{"audio", "messages"}`` row into a chat row.
The codec-token string becomes the FINAL assistant turn: an existing
trailing assistant message has its content replaced (the audio is the
training target), otherwise an assistant turn is appended. Returns a NEW
deep-copied row without the ``audio`` key — the caller's row is never
mutated.
"""
import copy
if not isinstance(row, dict):
raise TypeError(f"row must be a dict, got {type(row).__name__}")
audio = row.get("audio")
if not isinstance(audio, str) or not audio:
raise ValueError("live-codec TTS row must have a non-empty 'audio' path")
messages = row.get("messages")
if not isinstance(messages, (list, tuple)):
raise ValueError("live-codec TTS row must have a 'messages' list")
codec_string = encoder(audio)
new_messages = [copy.deepcopy(m) for m in messages]
if new_messages and isinstance(new_messages[-1], dict) and (
new_messages[-1].get("role") == "assistant"
):
new_messages[-1]["content"] = codec_string
else:
new_messages.append({"role": "assistant", "content": codec_string})
out = {
key: copy.deepcopy(value)
for key, value in row.items()
if key not in ("audio", "messages")
}
out["messages"] = new_messages
return out
def encode_tts_dataset(
dataset: dict,
family: str,
*,
encoder: Optional[Callable[[str], str]] = None,
device: Optional[str] = None,
console: Optional[Any] = None,
) -> dict:
"""Encode every ``train`` / ``val`` row's audio into codec-token chat rows.
Returns a NEW dataset dict (input never mutated). ``encoder`` is
injectable for tests; by default it is resolved per-family via
:func:`tts_encoder_for_family` (which raises for not-yet-implemented
families — #265).
"""
if not isinstance(dataset, dict):
raise TypeError(
f"dataset must be a dict, got {type(dataset).__name__}"
)
canonical = validate_tts_family(family)
encode = encoder if encoder is not None else tts_encoder_for_family(
canonical, device=device
)
new_dataset = dict(dataset)
total = 0
for split in ("train", "val"):
rows = new_dataset.get(split)
if not isinstance(rows, (list, tuple)):
continue
new_dataset[split] = [encode_tts_row(row, encode) for row in rows]
total += len(rows)
if console is not None:
console.print(
f"[green]TTS live-codec:[/] encoded {total} row(s) "
f"(family={canonical})"
)
return new_dataset

15
tests/test_v07120.py
View File

@ -200,7 +200,20 @@ class TestTtsTrainerSetup:
with pytest.raises(RuntimeError, match="sparktts"):
w.setup({"train": []})
def test_live_codec_mode_orpheus_names_snac(self):
def test_live_codec_mode_orpheus_names_snac(self, monkeypatch):
# v0.71.22 #265 lifted the unconditional not-yet-validated gate —
# orpheus is live when snac IS installed, so force the missing-dep
# branch to keep exercising the friendly pip hint.
import importlib.util as ilu
real_find_spec = ilu.find_spec
monkeypatch.setattr(
ilu,
"find_spec",
lambda name, *a, **k: None
if name == "snac"
else real_find_spec(name, *a, **k),
)
w = self._wrapper(data_format="audio", family="orpheus")
with pytest.raises(RuntimeError, match="snac"):
w.setup({"train": []})

1343
tests/test_v07122.py Normal file
View File

File diff suppressed because it is too large Load Diff