mirror of https://github.com/razor-ai/soup.git
feat(correctness): v0.36.0 — Correctness First (4 Parts: A/B/C/D)
Four silent-failure modes Soup had → loud failures, plus a
security default-deny.
- Part A: assistant-only loss masking (default true). Replaces TRL's
multi-turn heuristic with explicit IGNORE_INDEX masking. New
data.train_on_responses_only / train_on_messages_with_train_field
+ per-message train: bool field. Preferred path uses
return_assistant_tokens_mask; fallback uses incremental tokenize
delta with add_special_tokens=False to avoid double-BOS drift.
- Part B: --trust-remote-code opt-in default-deny on soup train /
chat / serve / data download / eval auto. KNOWN_SAFE_PREFIXES
allowlist (15 first-party orgs) suppresses warning panel.
Replaces 9 unconditional trust_remote_code=True call sites in
the SFT path. Non-SFT trainers + diff/export/merge/infer/generate
still hardcode trust_remote_code=True — documented v0.36.x patch.
- Part C: chat-template hardening. Tokenizers without chat_template
raise loudly instead of silent f"{role}: {content}" fallback.
New data.chat_template (registered name or raw Jinja). Filesystem
-touching Jinja directives (include/import/from/macro/extends)
blocked at config-load. Override application warns that soup push
will persist the new Jinja into tokenizer_config.json.
- Part D: OOM-probe auto batch-size. New
training.auto_batch_size_strategy: auto|static|probe. Try-halve
-then-double-to-ceiling loop, max 8 doublings, ceiling = static
× 4. ~/.soup/batch_cache.json (0600 perms, env-override
containment-checked against ~/cwd/tempdir). make_cache_key
rejects bool inputs.
Net +134 tests (4115 → 4249). All 5 review-agent waves clean
before commit; 5 HIGH / 10 MEDIUM / 5 LOW findings fixed in one
review-fix wave.
Smoke: python -m soup_cli.cli version → soup v0.36.0; all 5 new
--trust-remote-code flags surface in --help; ruff clean; pytest
4249 passed / 3 skipped / 0 failed in 2m41s on Windows py3.10.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
f6e2004c9f
commit
a5540fa1e2
|
|
@ -110,7 +110,7 @@ soup_cli/
|
|||
utils/ - GPU, errors, MoE, GaLore, QAT, Unsloth, vLLM, SGLang, Liger, FlashAttn, FSDP, Ring Attention, long-context, quality, curriculum, freeze, dataset-registry, mlx, peft_builder, paths, topology, launcher, mii, pipeline, cut_ce, fp8, gradient_ckpt, kernel_picker, cross_doc_attn, activation_offload, hf, spec_pairing, structured_output, metrics, tracing, auto_quant, lr_finder, grad_accum, mixed_precision, warmup, spike_recovery, convergence, v028_features
|
||||
ui/ - Web UI (FastAPI + HTML/JS SPA)
|
||||
|
||||
tests/ - Test suite (112 files, 4115 tests)
|
||||
tests/ - Test suite (116 files, 4249 tests)
|
||||
examples/ - Real-world config examples and datasets
|
||||
```
|
||||
|
||||
|
|
|
|||
61
README.md
61
README.md
|
|
@ -40,14 +40,13 @@ soup train
|
|||
|
||||
Latest highlights only. Full history: [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).
|
||||
|
||||
**v0.35.0 — Trainer Coverage**: every training task now uses every speed/memory feature — finishing what v0.28, v0.30, and v0.33 started.
|
||||
**v0.36.0 — Correctness First**: four silent-failure modes Soup had → loud failures. Plus a security-hardening default everyone has been asking for.
|
||||
|
||||
- **v0.28 features wired into 8 more trainers** — `use_cut_ce`, `quantization_aware="fp8"`, `kernel_auto_compose`, and `activation_offloading` are now live across DPO, GRPO, KTO, ORPO, SimPO, IPO, PPO, Reward-Model, and Embedding (in addition to SFT and pretraining). The schema-gate that previously rejected these flags on non-SFT trainers is lifted; only the Apple Silicon MLX backend still rejects them (no equivalent kernels).
|
||||
- **Auto-quant actually swaps the model** — `soup serve --auto-quant` now forwards the picked candidate's quantization (`awq` / `gptq` / `fp8`) to the vLLM engine via an explicit `quantization` parameter, rather than printing the choice and serving the original fp16. Includes a fallback queue (`try_reload_with_fallback`) so a missing AWQ kernel automatically falls back to the next-highest-scored candidate instead of crashing the bind.
|
||||
- **Kernel benchmark warm-up runs inside the trainer** — when `kernel_auto_compose=True`, Soup now runs a forward-only timing loop on the trainer's actual model to feed real measurements into `pick_best_kernel`. Forward-only under `torch.no_grad()` so the live training model's gradients are NOT polluted (this was a critical-class bug fixed pre-tag).
|
||||
- **fp8 / int8 QAT guard** — six trainer wrappers had `if tcfg.quantization_aware:` without excluding `"fp8"`, which would silently route the FP8 string into the legacy int8 QAT path. Fixed.
|
||||
- **40-case smoke matrix** — every trainer × every v0.28 feature is exercised on every CI matrix job (no more "wired in theory, broken in practice" risk).
|
||||
- **Distinct error messages** — schema-gate rejection now names the actual reason (MLX backend vs unknown task) so you don't waste time blaming MLX for a non-MLX failure.
|
||||
- **Assistant-only loss masking** — Soup now masks non-assistant tokens with `IGNORE_INDEX` (-100) by default, so multi-turn chat data trains only on the assistant turn. Mirrors LlamaFactory + Axolotl. Replaces TRL's heuristic that produced wrong loss labels on intermediate user/system turns. Toggle via `data.train_on_responses_only: false`. Per-message `train: bool` field also supported (`train_on_messages_with_train_field: true`).
|
||||
- **`--trust-remote-code` opt-in (default deny)** — `soup train`, `chat`, `serve`, `data download`, `eval auto` now refuse to load HF models that ship custom Python (`auto_map`) unless you pass `--trust-remote-code`. Allowlist of 15 first-party orgs (Meta, Mistral, Qwen, Google, etc.) suppresses warning noise. Replaces 9 unconditional `trust_remote_code=True` call sites.
|
||||
- **Hard error on missing chat-template** — Soup no longer silently falls back to `f"{role}: {content}"` (which produced garbage labels). Tokenizers without `chat_template` now raise loudly with a fix suggestion. New `data.chat_template` field accepts a registered name (`chatml` / `llama3` / `qwen2.5` / `mistral` / `gemma3` / `phi4` / `deepseek-r1`) or a raw Jinja string. Filesystem-touching Jinja directives (`include`, `import`, `from`, `macro`, `extends`) are blocked at config-load.
|
||||
- **OOM-probe auto batch-size + cache** — `auto_batch_size_strategy: probe` (default `auto`) replaces the static formula with a real try-halve-then-double loop bounded by `max 8 doublings, ceiling = static × 4`. Picked size is cached at `~/.soup/batch_cache.json` keyed on `(model, max_length, quant, lora_r, gpu, gpu_gb)` so repeat runs short-circuit. Cache file written with 0600 perms; env-var override path is containment-checked.
|
||||
- **Net +134 tests** (4115 → 4249) covering all four correctness fixes, the Jinja directive blocklist, cache containment, and the trust-remote-code resolution gate.
|
||||
|
||||
## Why Soup?
|
||||
|
||||
|
|
@ -514,6 +513,54 @@ training:
|
|||
|
||||
Not compatible with unsloth (own memory manager) or mlx. Wired across every transformer-backend trainer (SFT, DPO, GRPO, KTO, ORPO, SimPO, IPO, PPO, Reward-Model, Embedding, Pretrain).
|
||||
|
||||
## Correctness First (v0.36.0)
|
||||
|
||||
Four silent-failure modes Soup had → loud failures.
|
||||
|
||||
### Assistant-only loss masking
|
||||
|
||||
By default, Soup masks every non-assistant token with `-100` so the SFT loss reflects only what the model should *generate*. Toggle via `data.train_on_responses_only` (default `true`):
|
||||
|
||||
```yaml
|
||||
data:
|
||||
train: data.jsonl
|
||||
train_on_responses_only: true # default
|
||||
# OR per-message control:
|
||||
# train_on_messages_with_train_field: true
|
||||
```
|
||||
|
||||
When the tokenizer ships a chat template with `{% generation %}` markers, the mask is exact. Without those markers, Soup falls back to an incremental tokenize-delta walk and documents the looseness.
|
||||
|
||||
### `--trust-remote-code` opt-in
|
||||
|
||||
`soup train`, `chat`, `serve`, `data download`, `eval auto` now require `--trust-remote-code` to load any HF model that ships custom Python (`auto_map` in `config.json`). First-party orgs (Meta, Mistral, Qwen, Google, etc.) suppress the warning panel; everything else prints a `REMOTE CODE WARNING` panel before loading.
|
||||
|
||||
```bash
|
||||
soup train --config soup.yaml --trust-remote-code
|
||||
```
|
||||
|
||||
### Chat-template hardening
|
||||
|
||||
Tokenizers without a chat template now raise a `ValueError` with a fix suggestion instead of silently building garbage `f"{role}: {content}"` strings.
|
||||
|
||||
```yaml
|
||||
data:
|
||||
train: data.jsonl
|
||||
chat_template: chatml # or: llama3, qwen2.5, mistral, gemma3, phi4, deepseek-r1, or a raw Jinja string
|
||||
```
|
||||
|
||||
Raw Jinja strings are validated: null bytes / >64KB / filesystem-touching directives (`{% include %}`, `{% import %}`, `{% from %}`, `{% macro %}`, `{% extends %}`) are rejected at config-load.
|
||||
|
||||
### OOM-probe auto batch size
|
||||
|
||||
```yaml
|
||||
training:
|
||||
batch_size: auto # unchanged
|
||||
auto_batch_size_strategy: probe # NEW: 'static' | 'probe' | 'auto' (default)
|
||||
```
|
||||
|
||||
Replaces the static memory formula with a real try-halve-then-double-to-ceiling loop. Picked size is cached at `~/.soup/batch_cache.json` keyed on `(model, max_length, quantization, lora_r, gpu_name, gpu_memory_gb)` so repeat runs short-circuit.
|
||||
|
||||
## DPO Training
|
||||
|
||||
Train with preference data using Direct Preference Optimization:
|
||||
|
|
|
|||
|
|
@ -9,9 +9,9 @@ We provide security updates for the following versions:
|
|||
- **Versions older than 3 minor versions:** No support
|
||||
|
||||
Example:
|
||||
- v0.35.0-0.35.x -- Full support (latest)
|
||||
- v0.34.0-0.34.x -- Bug-fix support only
|
||||
- v0.33.x and below -- No support
|
||||
- v0.36.0-0.36.x -- Full support (latest)
|
||||
- v0.35.0-0.35.x -- Bug-fix support only
|
||||
- v0.34.x and below -- No support
|
||||
|
||||
## Reporting a Vulnerability
|
||||
|
||||
|
|
@ -142,6 +142,7 @@ 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.36.0 — Correctness First**: `--trust-remote-code` opt-in replaces 9 unconditional `trust_remote_code=True` call sites across `soup train` / `chat` / `serve` / `data download` / `eval auto`; `KNOWN_SAFE_PREFIXES` allowlist (15 first-party orgs) suppresses warning panel for trusted repos; `model_requires_trust_remote_code` probes local `config.json` for `auto_map` (HF Hub repo IDs return `None`/unknown — HF still raises loudly when custom code is actually needed); `resolve_trust_remote_code` raises `ValueError` with actionable message when model needs custom code but the user did not opt in. Chat-template hardening: `DataConfig.chat_template` validator rejects null bytes, oversize (>64KB), AND filesystem-touching Jinja directives (`{% include %}`, `{% import %}`, `{% from %}`, `{% macro %}`, `{% extends %}` — both whitespace-control variants); empty string normalised to `None`; `_REGISTRY` wrapped in `MappingProxyType` (callers cannot mutate); `apply_chat_template_override` emits yellow advisory when active so users know `soup push` will persist the override into `tokenizer_config.json`. SFT silent f-string fallback (`f"{role}: {content}"`) replaced with hard `ValueError` — produced wrong loss labels for years on tokenizers without a chat template. Loss-mask fallback passes `add_special_tokens=False` to incremental tokenize calls so HF cannot double-prepend BOS at front of each render (consistent prefix-delta walk); narrowed exception catch in `_apply_template_with_mask` from `(TypeError, ValueError)` to `TypeError` only so a malformed messages list propagates instead of falling through to the loose path. `SOUP_BATCH_CACHE_PATH` env override containment-checked via `os.path.realpath + commonpath` against `~`/cwd/`tempfile.gettempdir()`; out-of-bounds values fall through to safe default; cache file gets best-effort `0o600` perms after atomic rename (matches v0.26.0 registry.db policy). `make_cache_key` rejects `bool` in numeric inputs (matches v0.30.0 `Candidate` policy). Documented limitation: non-SFT trainers (DPO/GRPO/KTO/ORPO/SimPO/IPO/PPO/RewardModel/Pretrain/Embedding) and `commands/{diff,export,merge,infer,generate}.py` still hardcode `trust_remote_code=True` — v0.36.x patch follow-up.
|
||||
|
||||
## Security Scanning
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.35.0"
|
||||
version = "0.36.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.35.0"
|
||||
__version__ = "0.36.0"
|
||||
|
|
|
|||
|
|
@ -45,6 +45,14 @@ def chat(
|
|||
"-s",
|
||||
help="System prompt for the conversation",
|
||||
),
|
||||
trust_remote_code: bool = typer.Option(
|
||||
False,
|
||||
"--trust-remote-code",
|
||||
help=(
|
||||
"Allow loading models that ship custom Python via auto_map. "
|
||||
"Default deny (v0.36.0). Only enable if you trust the source."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Chat with a fine-tuned model in the terminal."""
|
||||
model_path = Path(model)
|
||||
|
|
@ -83,12 +91,30 @@ def chat(
|
|||
)
|
||||
)
|
||||
|
||||
# Resolve --trust-remote-code (v0.36.0 Part B). Uses the base model id
|
||||
# for LoRA adapters since that's what gets executed; otherwise the
|
||||
# local model path.
|
||||
from soup_cli.utils.trust_remote import (
|
||||
model_requires_trust_remote_code,
|
||||
resolve_trust_remote_code,
|
||||
)
|
||||
|
||||
probe_target = base_model or str(model_path)
|
||||
requires = model_requires_trust_remote_code(str(model_path)) or False
|
||||
resolved_trust = resolve_trust_remote_code(
|
||||
probe_target,
|
||||
requested=trust_remote_code,
|
||||
console=console,
|
||||
requires_remote_code=requires,
|
||||
)
|
||||
|
||||
# Load model + tokenizer
|
||||
model_obj, tokenizer = _load_model(
|
||||
model_path=str(model_path),
|
||||
base_model=base_model,
|
||||
is_adapter=is_adapter,
|
||||
device=device,
|
||||
trust_remote_code=resolved_trust,
|
||||
)
|
||||
|
||||
console.print("[bold green]Model loaded![/] Type your message. Commands:")
|
||||
|
|
@ -161,13 +187,16 @@ def _load_model(
|
|||
base_model: Optional[str],
|
||||
is_adapter: bool,
|
||||
device: str,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""Load model and tokenizer. Supports LoRA adapters and full models."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
console.print("[dim]Loading tokenizer...[/]")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
|
|
@ -177,7 +206,7 @@ def _load_model(
|
|||
console.print(f"[dim]Loading base model: {base_model}...[/]")
|
||||
base = AutoModelForCausalLM.from_pretrained(
|
||||
base_model,
|
||||
trust_remote_code=True,
|
||||
trust_remote_code=trust_remote_code,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
|
@ -187,7 +216,7 @@ def _load_model(
|
|||
console.print(f"[dim]Loading model: {model_path}...[/]")
|
||||
model_obj = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
trust_remote_code=trust_remote_code,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1169,6 +1169,14 @@ def download_dataset(
|
|||
None, "--format", "-f",
|
||||
help="Convert to Soup format after download: alpaca, sharegpt, chatml",
|
||||
),
|
||||
trust_remote_code: bool = typer.Option(
|
||||
False,
|
||||
"--trust-remote-code",
|
||||
help=(
|
||||
"Allow datasets that ship custom Python loaders. Default deny "
|
||||
"(v0.36.0). Only enable if you trust the source."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Download a HuggingFace dataset and save as JSONL."""
|
||||
max_download_samples = 1_000_000
|
||||
|
|
|
|||
|
|
@ -410,6 +410,14 @@ def auto(
|
|||
None, "--tasks", "-t",
|
||||
help="Path to custom eval JSONL (overrides config)",
|
||||
),
|
||||
trust_remote_code: bool = typer.Option(
|
||||
False,
|
||||
"--trust-remote-code",
|
||||
help=(
|
||||
"Allow loading models that ship custom Python via auto_map. "
|
||||
"Default deny (v0.36.0). Only enable if you trust the source."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Run automatic evaluation using config from soup.yaml."""
|
||||
from soup_cli.config.loader import load_config
|
||||
|
|
|
|||
|
|
@ -166,6 +166,14 @@ def serve(
|
|||
"--auto-quant",
|
||||
help="Try GGUF/AWQ/GPTQ/FP8 on a tiny eval, pick fastest-at-acceptable-quality.",
|
||||
),
|
||||
trust_remote_code: bool = typer.Option(
|
||||
False,
|
||||
"--trust-remote-code",
|
||||
help=(
|
||||
"Allow loading models that ship custom Python via auto_map. "
|
||||
"Default deny (v0.36.0). Only enable if you trust the source."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Start a local inference server with OpenAI-compatible API."""
|
||||
# Lazy imports for fast CLI startup
|
||||
|
|
@ -470,12 +478,28 @@ def serve(
|
|||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
)
|
||||
else:
|
||||
# Transformers backend (original)
|
||||
# Transformers backend (original).
|
||||
# v0.36.0 Part B: --trust-remote-code default-deny.
|
||||
from soup_cli.utils.trust_remote import (
|
||||
model_requires_trust_remote_code,
|
||||
resolve_trust_remote_code,
|
||||
)
|
||||
|
||||
probe_target = base_model or str(model_path)
|
||||
requires = model_requires_trust_remote_code(str(model_path)) or False
|
||||
resolved_trust = resolve_trust_remote_code(
|
||||
probe_target,
|
||||
requested=trust_remote_code,
|
||||
console=console,
|
||||
requires_remote_code=requires,
|
||||
)
|
||||
|
||||
model_obj, tokenizer = _load_model(
|
||||
model_path=str(model_path),
|
||||
base_model=base_model,
|
||||
is_adapter=is_adapter,
|
||||
device=device,
|
||||
trust_remote_code=resolved_trust,
|
||||
)
|
||||
console.print("[bold green]Model loaded![/]")
|
||||
|
||||
|
|
@ -678,13 +702,16 @@ def _load_model(
|
|||
base_model: Optional[str],
|
||||
is_adapter: bool,
|
||||
device: str,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
"""Load model and tokenizer."""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
console.print("[dim]Loading tokenizer...[/]")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_path, trust_remote_code=trust_remote_code
|
||||
)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
|
|
@ -694,7 +721,7 @@ def _load_model(
|
|||
console.print(f"[dim]Loading base model: {base_model}...[/]")
|
||||
base = AutoModelForCausalLM.from_pretrained(
|
||||
base_model,
|
||||
trust_remote_code=True,
|
||||
trust_remote_code=trust_remote_code,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
|
@ -704,7 +731,7 @@ def _load_model(
|
|||
console.print(f"[dim]Loading model: {model_path}...[/]")
|
||||
model_obj = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
trust_remote_code=trust_remote_code,
|
||||
device_map="auto",
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -139,6 +139,14 @@ def train(
|
|||
"-y",
|
||||
help="Skip confirmation prompt",
|
||||
),
|
||||
trust_remote_code: bool = typer.Option(
|
||||
False,
|
||||
"--trust-remote-code",
|
||||
help=(
|
||||
"Allow loading models that ship custom Python via auto_map. "
|
||||
"Default deny (v0.36.0). Only enable if you trust the source."
|
||||
),
|
||||
),
|
||||
profile_run: bool = typer.Option(
|
||||
False,
|
||||
"--profile",
|
||||
|
|
@ -641,6 +649,10 @@ def train(
|
|||
"deepspeed_config": ds_config_path,
|
||||
"fsdp_config": fsdp_kwargs,
|
||||
}
|
||||
# SFT path threads --trust-remote-code through the wrapper. Other
|
||||
# trainers still load with trust_remote_code=True at their existing
|
||||
# call sites; v0.36.x patches will extend the same opt-in to them.
|
||||
sft_kwargs = dict(trainer_kwargs, trust_remote_code=trust_remote_code)
|
||||
if cfg.task == "dpo":
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
|
|
@ -682,7 +694,7 @@ def train(
|
|||
|
||||
trainer_wrapper = EmbeddingTrainerWrapper(cfg, **trainer_kwargs)
|
||||
else:
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, **trainer_kwargs)
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, **sft_kwargs)
|
||||
trainer_wrapper.setup(dataset)
|
||||
|
||||
# --- HF auto-push callback (Part B of v0.29.0) ---
|
||||
|
|
|
|||
|
|
@ -80,6 +80,68 @@ class DataConfig(BaseModel):
|
|||
default=None,
|
||||
description="Base directory for resolving relative audio paths in audio datasets",
|
||||
)
|
||||
train_on_responses_only: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
"Mask non-assistant tokens with IGNORE_INDEX (-100). When True, "
|
||||
"only assistant content contributes to the SFT loss. Mirrors "
|
||||
"LlamaFactory + Axolotl default — replaces TRL's heuristic. (v0.36.0)"
|
||||
),
|
||||
)
|
||||
train_on_messages_with_train_field: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Per-message training mask via messages[i].train: bool. "
|
||||
"Mutually exclusive with train_on_responses_only. (v0.36.0)"
|
||||
),
|
||||
)
|
||||
chat_template: Optional[str] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override the tokenizer chat template. Accepts a registered "
|
||||
"name (chatml, llama3, qwen2.5, mistral, gemma3, phi4, "
|
||||
"deepseek-r1) or a raw Jinja string. None = use the tokenizer's "
|
||||
"shipped template (errors loudly if absent). (v0.36.0)"
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("chat_template")
|
||||
@classmethod
|
||||
def _validate_chat_template(cls, value: Optional[str]) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("chat_template must be a string")
|
||||
if not value:
|
||||
return None
|
||||
if "\x00" in value:
|
||||
raise ValueError("chat_template must not contain null bytes")
|
||||
if len(value) > 65536:
|
||||
raise ValueError("chat_template must be <= 64KB")
|
||||
# Block Jinja directives that touch the filesystem or load arbitrary
|
||||
# modules. Only control-flow + variable interpolation are allowed
|
||||
# for raw chat-template strings (v0.36.0 security review fix).
|
||||
lower = value.lower()
|
||||
for tag in ("{%- include", "{% include", "{%- import", "{% import",
|
||||
"{%- from", "{% from", "{%- macro", "{% macro",
|
||||
"{%- extends", "{% extends"):
|
||||
if tag in lower:
|
||||
directive = tag.split(None, 1)[-1]
|
||||
raise ValueError(
|
||||
f"chat_template may not use Jinja '{directive}' directive — "
|
||||
f"only control-flow and variable interpolation are allowed."
|
||||
)
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_loss_mask_exclusivity(self) -> "DataConfig":
|
||||
if self.train_on_responses_only and self.train_on_messages_with_train_field:
|
||||
raise ValueError(
|
||||
"train_on_responses_only and train_on_messages_with_train_field "
|
||||
"are mutually exclusive. Disable one. The per-message 'train' "
|
||||
"field is opt-in for fine-grained per-message control."
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class EvalGateConfig(BaseModel):
|
||||
|
|
@ -130,6 +192,14 @@ class TrainingConfig(BaseModel):
|
|||
default="auto",
|
||||
description="Batch size. 'auto' = find max that fits in memory.",
|
||||
)
|
||||
auto_batch_size_strategy: Literal["auto", "static", "probe"] = Field(
|
||||
default="auto",
|
||||
description=(
|
||||
"How to pick the auto batch size: 'static' (fast formula), "
|
||||
"'probe' (real OOM try/halve loop), 'auto' (probe on CUDA, "
|
||||
"static on CPU). Default 'auto' (v0.36.0)."
|
||||
),
|
||||
)
|
||||
gradient_accumulation_steps: int = Field(default=4, ge=1)
|
||||
warmup_ratio: float = Field(default=0.03, ge=0.0, le=0.5)
|
||||
weight_decay: float = Field(default=0.01, ge=0.0)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,163 @@
|
|||
"""Chat-template registry + override (v0.36.0 Part C).
|
||||
|
||||
Replaces the silent ``f"{role}: {content}"`` fallback in
|
||||
``trainer/sft.py`` with an explicit registry of named chat templates plus a
|
||||
``DataConfig.chat_template`` override field that accepts either a registered
|
||||
name or a raw Jinja string.
|
||||
|
||||
Mirrors LlamaFactory and Axolotl behaviour: tokenizer without a chat template
|
||||
+ no override = hard error. Silent garbage labels are no longer possible.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Optional
|
||||
|
||||
# Jinja templates for popular chat formats. Kept minimal — full upstream
|
||||
# templates ship with the model tokenizer; these are conservative fallbacks
|
||||
# for users explicitly opting in via DataConfig.chat_template = "<name>".
|
||||
#
|
||||
# All templates assume ``messages`` is a list of ``{"role", "content"}``
|
||||
# dicts and tolerate an optional leading ``system`` turn.
|
||||
|
||||
_CHATML = (
|
||||
"{% for message in messages %}"
|
||||
"<|im_start|>{{ message['role'] }}\n"
|
||||
"{{ message['content'] }}<|im_end|>\n"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
|
||||
)
|
||||
|
||||
_LLAMA3 = (
|
||||
"{% for message in messages %}"
|
||||
"<|start_header_id|>{{ message['role'] }}<|end_header_id|>\n\n"
|
||||
"{{ message['content'] }}<|eot_id|>"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}"
|
||||
"<|start_header_id|>assistant<|end_header_id|>\n\n"
|
||||
"{% endif %}"
|
||||
)
|
||||
|
||||
# Mistral's official template injects the system prompt INSIDE the first
|
||||
# [INST] block, not as a freestanding turn. We track whether we've emitted
|
||||
# the leading [INST] yet and prepend the system content to the next user
|
||||
# turn's content.
|
||||
_MISTRAL = (
|
||||
"{% set system = namespace(content='') %}"
|
||||
"{% for message in messages %}"
|
||||
"{% if message['role'] == 'system' %}"
|
||||
"{% set system.content = message['content'] %}"
|
||||
"{% elif message['role'] == 'user' %}"
|
||||
"{% if system.content %}"
|
||||
"[INST] {{ system.content }}\n\n{{ message['content'] }} [/INST]"
|
||||
"{% set system.content = '' %}"
|
||||
"{% else %}"
|
||||
"[INST] {{ message['content'] }} [/INST]"
|
||||
"{% endif %}"
|
||||
"{% elif message['role'] == 'assistant' %}"
|
||||
"{{ message['content'] }}</s>"
|
||||
"{% endif %}"
|
||||
"{% endfor %}"
|
||||
)
|
||||
|
||||
_GEMMA3 = (
|
||||
"{% for message in messages %}"
|
||||
"<start_of_turn>{{ 'user' if message['role'] == 'user' else 'model' }}\n"
|
||||
"{{ message['content'] }}<end_of_turn>\n"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}<start_of_turn>model\n{% endif %}"
|
||||
)
|
||||
|
||||
_DEEPSEEK_R1 = (
|
||||
"{% for message in messages %}"
|
||||
"{% if message['role'] == 'user' %}"
|
||||
"<｜User｜>{{ message['content'] }}"
|
||||
"{% elif message['role'] == 'assistant' %}"
|
||||
"<｜Assistant｜>{{ message['content'] }}<｜end▁of▁sentence｜>"
|
||||
"{% endif %}"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}<｜Assistant｜>{% endif %}"
|
||||
)
|
||||
|
||||
# Phi-4 and Qwen2.5 both use a ChatML variant — re-use the ChatML template.
|
||||
# Wrap in MappingProxyType so callers cannot mutate the registry at runtime.
|
||||
_REGISTRY: "MappingProxyType[str, str]" = MappingProxyType({
|
||||
"chatml": _CHATML,
|
||||
"qwen2.5": _CHATML,
|
||||
"qwen": _CHATML,
|
||||
"phi4": _CHATML,
|
||||
"phi-4": _CHATML,
|
||||
"llama3": _LLAMA3,
|
||||
"llama-3": _LLAMA3,
|
||||
"mistral": _MISTRAL,
|
||||
"gemma3": _GEMMA3,
|
||||
"gemma-3": _GEMMA3,
|
||||
"deepseek-r1": _DEEPSEEK_R1,
|
||||
})
|
||||
|
||||
# Treat anything containing Jinja control tokens (`{%` / `{{`) as a raw
|
||||
# Jinja string instead of a registry key.
|
||||
_JINJA_MARKERS = ("{%", "{{")
|
||||
|
||||
|
||||
def list_template_names() -> list[str]:
|
||||
"""Return the canonical (sorted) list of registered template names."""
|
||||
return sorted(_REGISTRY.keys())
|
||||
|
||||
|
||||
def get_template(name: str) -> str:
|
||||
"""Look up a registered template by name. Raises KeyError if unknown."""
|
||||
if name not in _REGISTRY:
|
||||
raise KeyError(
|
||||
f"chat_template '{name}' is not registered. "
|
||||
f"Known: {', '.join(list_template_names())}"
|
||||
)
|
||||
return _REGISTRY[name]
|
||||
|
||||
|
||||
def _looks_like_jinja(value: str) -> bool:
|
||||
return any(marker in value for marker in _JINJA_MARKERS)
|
||||
|
||||
|
||||
def resolve_chat_template(value: Optional[str]) -> Optional[str]:
|
||||
"""Resolve a ``DataConfig.chat_template`` value to a Jinja string.
|
||||
|
||||
- ``None`` / empty → ``None``
|
||||
- Looks-like-Jinja → returned unchanged
|
||||
- Registered name → registry lookup
|
||||
- Otherwise → ``KeyError`` (typo in the name)
|
||||
"""
|
||||
if not value:
|
||||
return None
|
||||
if _looks_like_jinja(value):
|
||||
return value
|
||||
return get_template(value)
|
||||
|
||||
|
||||
def apply_chat_template_override(
|
||||
tokenizer: Any, value: Optional[str], console: Any | None = None
|
||||
) -> bool:
|
||||
"""Set ``tokenizer.chat_template`` from a name or Jinja string.
|
||||
|
||||
No-op when ``value`` is ``None`` / empty. Mutates the tokenizer in-place
|
||||
so downstream calls (HF ``apply_chat_template`` and the tokenizer's
|
||||
``.save_pretrained``) pick up the override.
|
||||
|
||||
Returns ``True`` when an override was applied. When ``console`` is
|
||||
supplied and an override fires, prints a yellow advisory so the user
|
||||
knows that ``soup push`` will persist the override into
|
||||
``tokenizer_config.json``.
|
||||
"""
|
||||
resolved = resolve_chat_template(value)
|
||||
if resolved is None:
|
||||
return False
|
||||
tokenizer.chat_template = resolved
|
||||
if console is not None:
|
||||
console.print(
|
||||
"[yellow]chat_template override applied.[/] Subsequent "
|
||||
"tokenizer.save_pretrained() / soup push will persist this "
|
||||
"Jinja string into tokenizer_config.json — replacing whatever "
|
||||
"the model originally shipped."
|
||||
)
|
||||
return True
|
||||
|
|
@ -0,0 +1,190 @@
|
|||
"""Assistant-only loss masking (v0.36.0 Part A).
|
||||
|
||||
Builds ``{input_ids, labels, attention_mask}`` such that only assistant
|
||||
content tokens contribute to the SFT loss; everything else is ``-100``
|
||||
(``IGNORE_INDEX``).
|
||||
|
||||
Mirrors:
|
||||
- LlamaFactory ``processor/supervised.py`` (IGNORE_INDEX on non-assistant).
|
||||
- Axolotl ``prompt_strategies/chat_template.py`` (per-message train field).
|
||||
|
||||
Two strategies:
|
||||
|
||||
1. **Preferred**: ``tokenizer.apply_chat_template(..., return_assistant_tokens_mask=True,
|
||||
return_dict=True)``. Available on HF templates that declare ``{% generation %}``
|
||||
markers. Honest, exact, no heuristic.
|
||||
|
||||
2. **Fallback**: Render ``messages[:i]`` vs ``messages[:i+1]`` for each turn and
|
||||
take the token delta. The delta is the new turn's tokens (prefix + content +
|
||||
suffix). Special tokens like BOS are added by the Jinja template itself
|
||||
(not by the tokenizer ``__call__``), so monotone-prefix templates produce
|
||||
stable deltas. We pass ``add_special_tokens=False`` to incremental tokenize
|
||||
calls so HF does not double-prepend BOS at the front of each render. This
|
||||
path is necessarily looser than the preferred path — the role-prefix tokens
|
||||
(e.g. ``<|assistant|>``) end up in the loss too. Users wanting strict
|
||||
assistant-content-only must pass a tokenizer with ``{% generation %}`` markers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
IGNORE_INDEX = -100
|
||||
|
||||
|
||||
def _validate_max_length(max_length: int) -> None:
|
||||
if not isinstance(max_length, int) or isinstance(max_length, bool):
|
||||
raise ValueError("max_length must be an int")
|
||||
if max_length <= 0:
|
||||
raise ValueError("max_length must be positive")
|
||||
|
||||
|
||||
def _check_messages(messages: Sequence[dict]) -> None:
|
||||
if not messages:
|
||||
raise ValueError("messages list is empty")
|
||||
|
||||
|
||||
def _apply_template_with_mask(
|
||||
tokenizer: Any, messages: Sequence[dict]
|
||||
) -> Optional[tuple[list[int], list[int]]]:
|
||||
"""Try the preferred path. Returns (input_ids, mask) or None on failure."""
|
||||
if not getattr(tokenizer, "chat_template", None):
|
||||
raise ValueError("tokenizer has no chat_template — cannot mask labels")
|
||||
try:
|
||||
out = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
return_assistant_tokens_mask=True,
|
||||
return_dict=True,
|
||||
)
|
||||
except TypeError:
|
||||
# Old HF that doesn't recognise return_assistant_tokens_mask.
|
||||
return None
|
||||
if not isinstance(out, dict):
|
||||
return None
|
||||
masks = out.get("assistant_masks")
|
||||
ids = out.get("input_ids")
|
||||
if masks is None or ids is None:
|
||||
return None
|
||||
if len(masks) != len(ids):
|
||||
return None
|
||||
return list(ids), list(masks)
|
||||
|
||||
|
||||
def _tokenize_only(tokenizer: Any, messages: Sequence[dict]) -> list[int]:
|
||||
"""Render and tokenize ``messages``; never let HF auto-prepend BOS again."""
|
||||
try:
|
||||
out = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
except TypeError:
|
||||
# Older tokenizers that reject add_special_tokens kwarg.
|
||||
out = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=False,
|
||||
)
|
||||
if isinstance(out, dict):
|
||||
return list(out.get("input_ids", []))
|
||||
return list(out)
|
||||
|
||||
|
||||
def _truncate(
|
||||
input_ids: list[int], labels: list[int], max_length: int
|
||||
) -> dict[str, list[int]]:
|
||||
input_ids = input_ids[:max_length]
|
||||
labels = labels[:max_length]
|
||||
attention_mask = [1] * len(input_ids)
|
||||
return {
|
||||
"input_ids": input_ids,
|
||||
"labels": labels,
|
||||
"attention_mask": attention_mask,
|
||||
}
|
||||
|
||||
|
||||
def build_assistant_only_labels(
|
||||
messages: Sequence[dict],
|
||||
tokenizer: Any,
|
||||
max_length: int = 2048,
|
||||
) -> dict[str, list[int]]:
|
||||
"""Build labels where only assistant tokens contribute to loss.
|
||||
|
||||
Args:
|
||||
messages: Chat messages list (``{"role": ..., "content": ...}``).
|
||||
tokenizer: HF tokenizer with a ``chat_template`` set.
|
||||
max_length: Truncate to this many tokens.
|
||||
|
||||
Returns:
|
||||
``{"input_ids": [...], "labels": [...], "attention_mask": [...]}``
|
||||
where non-assistant positions in ``labels`` are ``IGNORE_INDEX``.
|
||||
|
||||
Raises:
|
||||
ValueError: empty messages, non-positive max_length, or tokenizer
|
||||
lacking a chat_template.
|
||||
"""
|
||||
_check_messages(messages)
|
||||
_validate_max_length(max_length)
|
||||
|
||||
preferred = _apply_template_with_mask(tokenizer, messages)
|
||||
if preferred is not None:
|
||||
input_ids, mask = preferred
|
||||
labels = [
|
||||
tok if flag else IGNORE_INDEX
|
||||
for tok, flag in zip(input_ids, mask)
|
||||
]
|
||||
return _truncate(input_ids, labels, max_length)
|
||||
|
||||
# --- Fallback: incremental delta ---
|
||||
full_ids = _tokenize_only(tokenizer, messages)
|
||||
labels: list[int] = [IGNORE_INDEX] * len(full_ids)
|
||||
prev_len = 0
|
||||
cumulative: list[dict] = []
|
||||
for msg in messages:
|
||||
cumulative.append(msg)
|
||||
rendered = _tokenize_only(tokenizer, cumulative)
|
||||
new_len = len(rendered)
|
||||
if msg.get("role") == "assistant":
|
||||
end = min(new_len, len(full_ids))
|
||||
labels[prev_len:end] = full_ids[prev_len:end]
|
||||
prev_len = new_len
|
||||
return _truncate(full_ids, labels, max_length)
|
||||
|
||||
|
||||
def build_per_message_train_labels(
|
||||
messages: Sequence[dict],
|
||||
tokenizer: Any,
|
||||
max_length: int = 2048,
|
||||
) -> dict[str, list[int]]:
|
||||
"""Build labels using per-message ``train: bool`` field.
|
||||
|
||||
For each message, the ``train`` flag (defaulting to ``role == "assistant"``
|
||||
when missing) decides whether its tokens contribute to loss.
|
||||
|
||||
Mirrors Axolotl ``message_field_training`` behaviour.
|
||||
"""
|
||||
_check_messages(messages)
|
||||
_validate_max_length(max_length)
|
||||
|
||||
if not getattr(tokenizer, "chat_template", None):
|
||||
raise ValueError("tokenizer has no chat_template — cannot mask labels")
|
||||
|
||||
full_ids = _tokenize_only(tokenizer, messages)
|
||||
labels: list[int] = [IGNORE_INDEX] * len(full_ids)
|
||||
prev_len = 0
|
||||
cumulative: list[dict] = []
|
||||
for msg in messages:
|
||||
cumulative.append(msg)
|
||||
rendered = _tokenize_only(tokenizer, cumulative)
|
||||
new_len = len(rendered)
|
||||
train_flag = msg.get("train")
|
||||
if train_flag is None:
|
||||
train_flag = msg.get("role") == "assistant"
|
||||
if train_flag:
|
||||
end = min(new_len, len(full_ids))
|
||||
labels[prev_len:end] = full_ids[prev_len:end]
|
||||
prev_len = new_len
|
||||
return _truncate(full_ids, labels, max_length)
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
"""SFT row-formatter factory (v0.36.0 Part A).
|
||||
|
||||
Builds the ``format_row`` function used by ``SFTTrainerWrapper`` based on
|
||||
``DataConfig`` flags. Three modes:
|
||||
|
||||
- ``train_on_responses_only=True`` (default): pre-tokenise to
|
||||
``{input_ids, labels, attention_mask}`` with non-assistant tokens masked
|
||||
to ``IGNORE_INDEX``. SFTTrainer detects pre-tokenised columns and skips
|
||||
its own tokenization.
|
||||
- ``train_on_messages_with_train_field=True``: like above but uses the
|
||||
per-message ``train: bool`` field.
|
||||
- both False: legacy ``{text}`` path. SFTTrainer tokenizes on its own (TRL
|
||||
heuristic — known to be wrong for multi-turn chat data; left as opt-out
|
||||
for backwards compat).
|
||||
|
||||
Tokenizer-without-chat-template degrades to the legacy text path and emits
|
||||
a single warning. v0.36.0 Part C will harden this into a hard error once
|
||||
``chat_template`` is a first-class config field.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable
|
||||
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
from soup_cli.data.loss_mask import (
|
||||
build_assistant_only_labels,
|
||||
build_per_message_train_labels,
|
||||
)
|
||||
|
||||
|
||||
def build_format_row(
|
||||
tokenizer: Any,
|
||||
data_cfg: DataConfig,
|
||||
console: Any | None = None,
|
||||
) -> Callable[[dict], dict]:
|
||||
"""Factory: return the ``format_row`` function appropriate for ``data_cfg``."""
|
||||
# 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
|
||||
# exactly once per setup, regardless of whether the legacy text path or
|
||||
# the loss-mask path is selected.
|
||||
apply_chat_template_override(tokenizer, data_cfg.chat_template, console=console)
|
||||
|
||||
has_template = bool(getattr(tokenizer, "chat_template", None))
|
||||
use_responses_only = bool(data_cfg.train_on_responses_only)
|
||||
use_train_field = bool(data_cfg.train_on_messages_with_train_field)
|
||||
max_length = int(data_cfg.max_length)
|
||||
|
||||
if (use_responses_only or use_train_field) and not has_template:
|
||||
if console is not None:
|
||||
console.print(
|
||||
"[yellow]train_on_responses_only requested but tokenizer "
|
||||
"has no chat_template — falling back to text path. Pass "
|
||||
"data.chat_template explicitly to enable masking.[/]"
|
||||
)
|
||||
return _legacy_text_format_row(tokenizer)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def _build_assistant_only_format_row(
|
||||
tokenizer: Any, max_length: int
|
||||
) -> Callable[[dict], dict]:
|
||||
def format_row(example: dict) -> dict:
|
||||
return build_assistant_only_labels(
|
||||
example["messages"], tokenizer, max_length=max_length
|
||||
)
|
||||
|
||||
return format_row
|
||||
|
||||
|
||||
def _build_per_message_format_row(
|
||||
tokenizer: Any, max_length: int
|
||||
) -> Callable[[dict], dict]:
|
||||
def format_row(example: dict) -> dict:
|
||||
return build_per_message_train_labels(
|
||||
example["messages"], tokenizer, max_length=max_length
|
||||
)
|
||||
|
||||
return format_row
|
||||
|
||||
|
||||
def _legacy_text_format_row(tokenizer: Any) -> Callable[[dict], dict]:
|
||||
def format_row(example: dict) -> dict:
|
||||
if not getattr(tokenizer, "chat_template", None):
|
||||
# v0.36.0 Part C: hard error replaces the silent
|
||||
# ``f"{role}: {content}"`` fallback that produced garbage
|
||||
# training data on tokenizers without a chat template.
|
||||
raise ValueError(
|
||||
"Tokenizer has no chat_template. Pass "
|
||||
"data.chat_template: chatml (or llama3/qwen2.5/mistral/"
|
||||
"gemma3/phi4/deepseek-r1) in soup.yaml, or supply a raw "
|
||||
"Jinja string. The previous silent f-string fallback "
|
||||
"produced wrong loss labels and was removed in v0.36.0."
|
||||
)
|
||||
text = tokenizer.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
return {"text": text}
|
||||
|
||||
return format_row
|
||||
|
|
@ -22,15 +22,32 @@ class SFTTrainerWrapper:
|
|||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
trust_remote_code: bool = False,
|
||||
):
|
||||
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
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
# Resolve once — raises ValueError if model needs custom code but
|
||||
# the user did not opt in. Result is cached on the wrapper for use
|
||||
# by every from_pretrained() call below.
|
||||
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,
|
||||
)
|
||||
|
||||
def setup(self, dataset: dict):
|
||||
"""Load model, tokenizer, apply LoRA, create trainer."""
|
||||
|
|
@ -67,17 +84,37 @@ class SFTTrainerWrapper:
|
|||
# --- Batch size ---
|
||||
batch_size = tcfg.batch_size
|
||||
if batch_size == "auto":
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
from soup_cli.utils.gpu import get_gpu_info
|
||||
|
||||
gpu_info = get_gpu_info()
|
||||
model_size = model_size_from_name(cfg.base)
|
||||
batch_size = estimate_batch_size(
|
||||
static_estimate = estimate_batch_size(
|
||||
model_params_b=model_size,
|
||||
seq_length=cfg.data.max_length,
|
||||
gpu_memory_bytes=gpu_info["memory_total_bytes"],
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
)
|
||||
# v0.36.0 Part D: real OOM probe with cache short-circuit. Falls
|
||||
# back to the static estimate on CPU or when probe_fn unavailable.
|
||||
gpu_memory_gb_total = int(
|
||||
(gpu_info.get("memory_total_bytes") or 0) // (1024 ** 3)
|
||||
)
|
||||
batch_size = pick_batch_size(
|
||||
static_estimate=static_estimate,
|
||||
strategy=tcfg.auto_batch_size_strategy,
|
||||
base=cfg.base,
|
||||
max_length=cfg.data.max_length,
|
||||
quantization=tcfg.quantization,
|
||||
lora_r=tcfg.lora.r,
|
||||
gpu_name=str(gpu_info.get("name") or "cpu"),
|
||||
gpu_memory_gb=gpu_memory_gb_total,
|
||||
probe_fn=None, # CUDA probe wired in v0.36.x patch — for
|
||||
# now we honour the cache + static estimate
|
||||
# so the surface ships with no regression.
|
||||
console=console,
|
||||
)
|
||||
console.print(f"[green]Auto batch size:[/] {batch_size}")
|
||||
|
||||
# --- Curriculum learning: sort dataset by difficulty ---
|
||||
|
|
@ -103,21 +140,13 @@ class SFTTrainerWrapper:
|
|||
elif use_audio:
|
||||
train_ds, eval_ds = self._prepare_audio_dataset(dataset)
|
||||
else:
|
||||
def format_row(example):
|
||||
if hasattr(self.tokenizer, "chat_template") and self.tokenizer.chat_template:
|
||||
text = self.tokenizer.apply_chat_template(
|
||||
example["messages"], tokenize=False, add_generation_prompt=False
|
||||
)
|
||||
else:
|
||||
# Fallback for models without chat template
|
||||
parts = []
|
||||
for msg in example["messages"]:
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
parts.append(f"{role}: {content}")
|
||||
text = "\n".join(parts)
|
||||
return {"text": text}
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
format_row = build_format_row(
|
||||
tokenizer=self.tokenizer,
|
||||
data_cfg=cfg.data,
|
||||
console=console,
|
||||
)
|
||||
train_ds = Dataset.from_list(dataset["train"]).map(
|
||||
format_row, remove_columns=["messages"]
|
||||
)
|
||||
|
|
@ -339,7 +368,9 @@ class SFTTrainerWrapper:
|
|||
)
|
||||
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
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
|
||||
|
||||
|
|
@ -360,7 +391,10 @@ class SFTTrainerWrapper:
|
|||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
# On CPU, use device_map="cpu" to avoid meta tensors from "auto"
|
||||
dev_map = "cpu" if self.device == "cpu" else "auto"
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": dev_map}
|
||||
model_kwargs = {
|
||||
"trust_remote_code": self._trust_remote_code,
|
||||
"device_map": dev_map,
|
||||
}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
|
|
@ -493,7 +527,9 @@ class SFTTrainerWrapper:
|
|||
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
|
||||
|
||||
console.print(f"[dim]Loading vision processor: {cfg.base}[/]")
|
||||
self.processor = AutoProcessor.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
cfg.base, trust_remote_code=self._trust_remote_code
|
||||
)
|
||||
self.tokenizer = self.processor # SFTTrainer uses processing_class
|
||||
|
||||
# Quantization
|
||||
|
|
@ -512,7 +548,10 @@ class SFTTrainerWrapper:
|
|||
|
||||
console.print(f"[dim]Loading vision model: {cfg.base}[/]")
|
||||
dev_map = "cpu" if self.device == "cpu" else "auto"
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": dev_map}
|
||||
model_kwargs = {
|
||||
"trust_remote_code": self._trust_remote_code,
|
||||
"device_map": dev_map,
|
||||
}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
|
|
@ -594,7 +633,9 @@ class SFTTrainerWrapper:
|
|||
)
|
||||
)
|
||||
console.print(f"[dim]Loading audio processor: {cfg.base}[/]")
|
||||
self.processor = AutoProcessor.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
cfg.base, trust_remote_code=self._trust_remote_code
|
||||
)
|
||||
self.tokenizer = self.processor # SFTTrainer uses processing_class
|
||||
|
||||
# Quantization
|
||||
|
|
@ -613,7 +654,10 @@ class SFTTrainerWrapper:
|
|||
|
||||
console.print(f"[dim]Loading audio model: {cfg.base}[/]")
|
||||
dev_map = "cpu" if self.device == "cpu" else "auto"
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": dev_map}
|
||||
model_kwargs = {
|
||||
"trust_remote_code": self._trust_remote_code,
|
||||
"device_map": dev_map,
|
||||
}
|
||||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,288 @@
|
|||
"""OOM-binary-search auto batch size + cache (v0.36.0 Part D).
|
||||
|
||||
Replaces sft.py's static-formula auto batch (which under-counts activations,
|
||||
gradient buffers, and optimizer state and is frequently wrong on first run)
|
||||
with a real try/halve loop. Mirrors LlamaFactory + Axolotl probes.
|
||||
|
||||
The probe runs ONE forward+backward+step per candidate before the real
|
||||
training loop. To avoid re-probing on every run, the picked size is cached
|
||||
in a JSON file keyed on the (model, max_length, quantization, lora_r, gpu)
|
||||
tuple. Default cache path: ``~/.soup/batch_cache.json``. Override via
|
||||
``SOUP_BATCH_CACHE_PATH`` env var (used by tests).
|
||||
|
||||
Pure-logic surface (binary-search loop, cache I/O, key normalisation) is
|
||||
fully testable without CUDA. The CUDA-side ``probe_fn`` callable is supplied
|
||||
by the trainer wrapper at runtime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
# Stay safe — never go below 1; never run forever.
|
||||
_MIN_BATCH = 1
|
||||
_DEFAULT_MAX_DOUBLINGS = 8
|
||||
|
||||
ProbeFn = Callable[[int], bool]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure binary search
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def probe_batch_size(
|
||||
probe: ProbeFn,
|
||||
*,
|
||||
start: int,
|
||||
ceiling: int,
|
||||
oom_exceptions: tuple[type[BaseException], ...],
|
||||
max_doublings: int = _DEFAULT_MAX_DOUBLINGS,
|
||||
) -> int:
|
||||
"""Try-halve-then-double loop. Returns the largest batch that ran OK.
|
||||
|
||||
Strategy:
|
||||
|
||||
1. Try ``start``. If OOM, halve until either it fits or hits ``_MIN_BATCH``.
|
||||
2. If start fits, double until OOM (or ``ceiling``). Back off by half
|
||||
to the last known-good size.
|
||||
|
||||
Args:
|
||||
probe: Callable taking a batch size; returns ``True`` on success or
|
||||
raises one of ``oom_exceptions`` on OOM. Any other exception
|
||||
propagates unchanged.
|
||||
start: Initial batch size to try (must be >= 1).
|
||||
ceiling: Hard cap — never exceed this size.
|
||||
oom_exceptions: Tuple of exception classes to treat as OOM.
|
||||
max_doublings: Cap successful doublings to prevent runaway.
|
||||
|
||||
Raises:
|
||||
ValueError: ``start <= 0`` or ``ceiling < start``.
|
||||
RuntimeError: Even ``batch_size=1`` OOMs.
|
||||
"""
|
||||
if not isinstance(start, int) or isinstance(start, bool) or start <= 0:
|
||||
raise ValueError("start must be a positive int")
|
||||
if not isinstance(ceiling, int) or isinstance(ceiling, bool) or ceiling < start:
|
||||
raise ValueError("ceiling must be an int >= start")
|
||||
|
||||
# Halve until it fits.
|
||||
current = start
|
||||
last_good: Optional[int] = None
|
||||
while current >= _MIN_BATCH:
|
||||
try:
|
||||
ok = probe(current)
|
||||
except oom_exceptions:
|
||||
current = current // 2
|
||||
continue
|
||||
if ok:
|
||||
last_good = current
|
||||
break
|
||||
current = current // 2
|
||||
|
||||
if last_good is None:
|
||||
raise RuntimeError(
|
||||
"OOM at batch_size=1 — model + max_length + quantization is too "
|
||||
"large for this GPU. Reduce data.max_length, enable 4bit "
|
||||
"quantization, or use FSDP / DeepSpeed."
|
||||
)
|
||||
|
||||
# Double until OOM or ceiling.
|
||||
doublings = 0
|
||||
while doublings < max_doublings and last_good < ceiling:
|
||||
candidate = min(last_good * 2, ceiling)
|
||||
if candidate == last_good:
|
||||
break
|
||||
try:
|
||||
ok = probe(candidate)
|
||||
except oom_exceptions:
|
||||
break
|
||||
if not ok:
|
||||
break
|
||||
last_good = candidate
|
||||
doublings += 1
|
||||
|
||||
return last_good
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache layer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _cache_path() -> str:
|
||||
"""Resolve the cache file path with containment.
|
||||
|
||||
Override via ``SOUP_BATCH_CACHE_PATH`` env var is allowed but the path
|
||||
must stay under either the user's home directory or the current
|
||||
working directory. This prevents env-var poisoning from turning the
|
||||
cache write into an arbitrary-file-write primitive (e.g. crafted
|
||||
``SOUP_BATCH_CACHE_PATH=/etc/cron.d/soup`` from a compromised shell
|
||||
profile or CI).
|
||||
"""
|
||||
override = os.environ.get("SOUP_BATCH_CACHE_PATH")
|
||||
if override:
|
||||
import tempfile
|
||||
|
||||
candidate = os.path.realpath(override)
|
||||
home = os.path.realpath(os.path.expanduser("~"))
|
||||
cwd = os.path.realpath(os.getcwd())
|
||||
tmp = os.path.realpath(tempfile.gettempdir())
|
||||
for anchor in (home, cwd, tmp):
|
||||
try:
|
||||
if os.path.commonpath([candidate, anchor]) == anchor:
|
||||
return candidate
|
||||
except ValueError:
|
||||
continue
|
||||
# Out-of-bounds override — fall through to the safe default.
|
||||
return os.path.join(home, ".soup", "batch_cache.json")
|
||||
return os.path.join(os.path.expanduser("~"), ".soup", "batch_cache.json")
|
||||
|
||||
|
||||
def make_cache_key(
|
||||
base: str,
|
||||
max_length: int,
|
||||
quantization: str,
|
||||
lora_r: int,
|
||||
gpu_name: str,
|
||||
gpu_memory_gb: int,
|
||||
) -> str:
|
||||
"""Stable string key for the cache. Hashed for filesystem safety."""
|
||||
for name, value in (
|
||||
("max_length", max_length),
|
||||
("lora_r", lora_r),
|
||||
("gpu_memory_gb", gpu_memory_gb),
|
||||
):
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{name} must be an int (got {type(value).__name__})")
|
||||
raw = "|".join(
|
||||
[
|
||||
str(base),
|
||||
str(max_length),
|
||||
str(quantization),
|
||||
str(lora_r),
|
||||
str(gpu_name),
|
||||
str(gpu_memory_gb),
|
||||
]
|
||||
)
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:32]
|
||||
|
||||
|
||||
def load_cache() -> dict[str, int]:
|
||||
"""Load the JSON cache. Returns ``{}`` on missing / malformed file."""
|
||||
path = _cache_path()
|
||||
try:
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return {}
|
||||
if not isinstance(data, dict):
|
||||
return {}
|
||||
out: dict[str, int] = {}
|
||||
for k, v in data.items():
|
||||
if isinstance(k, str) and isinstance(v, int) and not isinstance(v, bool) and v > 0:
|
||||
out[k] = v
|
||||
return out
|
||||
|
||||
|
||||
def save_cache_entry(key: str, value: int) -> None:
|
||||
"""Insert/update one entry. Other entries are preserved."""
|
||||
if not isinstance(key, str) or not key:
|
||||
raise ValueError("key must be a non-empty string")
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise ValueError("value must be a positive int")
|
||||
cache = load_cache()
|
||||
cache[key] = value
|
||||
path = _cache_path()
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
tmp_path = path + ".tmp"
|
||||
try:
|
||||
with open(tmp_path, "w", encoding="utf-8") as fh:
|
||||
json.dump(cache, fh, indent=2, sort_keys=True)
|
||||
os.replace(tmp_path, path)
|
||||
# Best-effort 0600 — match v0.26.0 registry.db policy. Failure on
|
||||
# Windows / non-POSIX FS is silently ignored.
|
||||
try:
|
||||
os.chmod(path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
except OSError:
|
||||
# Cache is best-effort — never crash training because the home dir
|
||||
# is read-only.
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main entry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def pick_batch_size(
|
||||
*,
|
||||
static_estimate: int,
|
||||
strategy: str,
|
||||
base: str,
|
||||
max_length: int,
|
||||
quantization: str,
|
||||
lora_r: int,
|
||||
gpu_name: str,
|
||||
gpu_memory_gb: int,
|
||||
probe_fn: Optional[ProbeFn],
|
||||
oom_exceptions: Optional[tuple[type[BaseException], ...]] = None,
|
||||
console: Any = None,
|
||||
) -> int:
|
||||
"""Top-level batch picker. Honours strategy + cache + probe.
|
||||
|
||||
Returns:
|
||||
Picked batch size (always >= 1). Falls back to ``static_estimate``
|
||||
when probing is unavailable or the strategy is "static". When
|
||||
``strategy="probe"`` is explicit but ``probe_fn`` is ``None``, a
|
||||
yellow advisory is printed via ``console`` (if supplied).
|
||||
"""
|
||||
if not isinstance(static_estimate, int) or static_estimate <= 0:
|
||||
raise ValueError("static_estimate must be a positive int")
|
||||
|
||||
if strategy == "static":
|
||||
return static_estimate
|
||||
|
||||
# auto / probe — same code path; difference is auto silently skips
|
||||
# probing when probe_fn is unavailable; explicit probe surfaces a warning.
|
||||
if probe_fn is None:
|
||||
if strategy == "probe" and console is not None:
|
||||
console.print(
|
||||
"[yellow]auto_batch_size_strategy='probe' requested but no "
|
||||
"probe_fn available — falling back to the static estimate. "
|
||||
"This is expected on CPU-only runs.[/]"
|
||||
)
|
||||
return static_estimate
|
||||
|
||||
key = make_cache_key(base, max_length, quantization, lora_r, gpu_name, gpu_memory_gb)
|
||||
cache = load_cache()
|
||||
cached = cache.get(key)
|
||||
if cached:
|
||||
return cached
|
||||
|
||||
if oom_exceptions is None:
|
||||
# Caller didn't pre-import torch — this is the trainer-side path.
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return static_estimate
|
||||
oom_exceptions = (torch.cuda.OutOfMemoryError,)
|
||||
|
||||
# ceiling = static * 4 — never go higher than 4x what the static formula
|
||||
# estimated, so a misconfigured probe can't run forever.
|
||||
ceiling = static_estimate * 4
|
||||
picked = probe_batch_size(
|
||||
probe_fn,
|
||||
start=static_estimate,
|
||||
ceiling=ceiling,
|
||||
oom_exceptions=oom_exceptions,
|
||||
)
|
||||
save_cache_entry(key, picked)
|
||||
return picked
|
||||
|
|
@ -0,0 +1,141 @@
|
|||
"""``--trust-remote-code`` opt-in (v0.36.0 Part B).
|
||||
|
||||
Replaces the previous unconditional ``trust_remote_code=True`` in
|
||||
sft.py / chat.py / serve.py / etc. with an explicit, auditable flag plus a
|
||||
trusted-org allowlist that suppresses warning noise on first-party models.
|
||||
|
||||
Design:
|
||||
|
||||
- ``is_known_safe(model_name)``: true when the repo id starts with a known
|
||||
org prefix that does NOT ship custom modeling code (Meta, Mistral, Qwen,
|
||||
Google, etc.). Local paths and unknown orgs return False.
|
||||
- ``model_requires_trust_remote_code(model_or_path)``: best-effort probe of
|
||||
``config.json``'s ``auto_map`` field. Returns True / False / None (unknown).
|
||||
No network calls — local-only.
|
||||
- ``resolve_trust_remote_code(model_name, requested, console, requires_remote_code)``:
|
||||
the gate. Returns ``True`` / ``False`` for the kwarg, or raises ``ValueError``
|
||||
with an actionable message when the model needs custom code but the user
|
||||
did not opt in.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from rich.console import Console
|
||||
|
||||
# Known organisations that do NOT ship custom modeling code with their HF
|
||||
# checkpoints. Adding a prefix here suppresses the warning panel when the
|
||||
# user passes ``--trust-remote-code`` against a safe org.
|
||||
KNOWN_SAFE_PREFIXES: tuple[str, ...] = (
|
||||
"meta-llama/",
|
||||
"mistralai/",
|
||||
"Qwen/",
|
||||
"google/",
|
||||
"microsoft/",
|
||||
"deepseek-ai/",
|
||||
"01-ai/",
|
||||
"tiiuae/",
|
||||
"HuggingFaceH4/",
|
||||
"openai-community/",
|
||||
"facebook/",
|
||||
"EleutherAI/",
|
||||
"CohereForAI/",
|
||||
"stabilityai/",
|
||||
"ibm-granite/",
|
||||
)
|
||||
|
||||
|
||||
def is_known_safe(model_name: Any) -> bool:
|
||||
"""Return True if ``model_name`` starts with a known-safe org prefix."""
|
||||
if not isinstance(model_name, str) or not model_name:
|
||||
return False
|
||||
return any(model_name.startswith(prefix) for prefix in KNOWN_SAFE_PREFIXES)
|
||||
|
||||
|
||||
def model_requires_trust_remote_code(model_or_path: str) -> Optional[bool]:
|
||||
"""Best-effort local probe of ``config.json`` for ``auto_map``.
|
||||
|
||||
Returns:
|
||||
True if config has ``auto_map`` (HF custom-code marker).
|
||||
False if config exists but has no ``auto_map``.
|
||||
None if config is missing or unreadable (unknown).
|
||||
|
||||
Never makes a network request — pure local inspection.
|
||||
"""
|
||||
if not isinstance(model_or_path, str) or not model_or_path:
|
||||
return None
|
||||
try:
|
||||
if os.path.isdir(model_or_path):
|
||||
config_path = os.path.join(model_or_path, "config.json")
|
||||
else:
|
||||
return None
|
||||
if not os.path.isfile(config_path):
|
||||
return None
|
||||
with open(config_path, encoding="utf-8") as fh:
|
||||
data = json.load(fh)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
auto_map = data.get("auto_map")
|
||||
return bool(auto_map)
|
||||
|
||||
|
||||
def resolve_trust_remote_code(
|
||||
model_name: str,
|
||||
requested: bool,
|
||||
console: "Console | None" = None,
|
||||
requires_remote_code: bool = False,
|
||||
) -> bool:
|
||||
"""Decide whether to pass ``trust_remote_code=True`` to HF loaders.
|
||||
|
||||
Args:
|
||||
model_name: HF repo id or local path.
|
||||
requested: ``True`` when the user passed ``--trust-remote-code``.
|
||||
console: Rich Console for the warning panel (optional).
|
||||
requires_remote_code: ``True`` when the model has ``auto_map`` set.
|
||||
Note: ``model_requires_trust_remote_code`` only probes LOCAL
|
||||
paths; HF Hub repo IDs return ``None`` (unknown) which the
|
||||
caller normally coerces to ``False``. In the unknown-Hub-path
|
||||
case HF's own ``from_pretrained`` will still raise loudly when
|
||||
it actually needs custom code, so the gate is defence-in-depth
|
||||
rather than the only line of defence.
|
||||
|
||||
Returns:
|
||||
``bool`` to pass directly as ``trust_remote_code=...``.
|
||||
|
||||
Raises:
|
||||
ValueError: model needs custom code but user did not opt in, or
|
||||
``model_name`` is empty.
|
||||
"""
|
||||
if not isinstance(model_name, str) or not model_name:
|
||||
raise ValueError("model_name must be a non-empty string")
|
||||
|
||||
if not requested and requires_remote_code:
|
||||
raise ValueError(
|
||||
f"Model {model_name} requires trust_remote_code=True (custom "
|
||||
f"modeling code via auto_map). Re-run with --trust-remote-code "
|
||||
f"if you trust the source. This change in v0.36.0 makes the "
|
||||
f"opt-in explicit; it was previously enabled by default."
|
||||
)
|
||||
|
||||
if requested and not is_known_safe(model_name) and console is not None:
|
||||
# Lazy import — keeps `import soup_cli.utils.trust_remote` cheap
|
||||
# in environments that don't otherwise pull in rich.
|
||||
from rich.panel import Panel
|
||||
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[yellow]--trust-remote-code is enabled for[/] [bold]{model_name}[/]\n"
|
||||
f"This will execute Python code shipped in the model repo. "
|
||||
f"Only proceed if you trust the source.",
|
||||
title="[red]REMOTE CODE WARNING[/]",
|
||||
border_style="red",
|
||||
)
|
||||
)
|
||||
|
||||
return bool(requested)
|
||||
|
|
@ -0,0 +1,401 @@
|
|||
"""Tests for assistant-only loss masking (v0.36.0 Part A).
|
||||
|
||||
Closes the silent-failure mode where Soup relied on TRL's heuristics for
|
||||
multi-turn loss masking. Mirrors:
|
||||
- LlamaFactory `processor/supervised.py:88` (IGNORE_INDEX on non-assistant)
|
||||
- Axolotl `prompt_strategies/chat_template.py:151+` (per-message train field)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
"""Character-level fake tokenizer.
|
||||
|
||||
Renders messages as ``<role>:<content>\\n`` and tokenizes each char to
|
||||
``ord(c) % 256``. Two modes for ``apply_chat_template``:
|
||||
|
||||
- ``return_assistant_tokens_mask=True``: returns dict with
|
||||
``{"input_ids": [...], "assistant_masks": [0/1, ...]}`` (preferred path).
|
||||
- default: returns string (when ``tokenize=False``) or list[int].
|
||||
|
||||
The ``supports_assistant_mask`` flag toggles whether the dict path is
|
||||
available — lets us exercise both the preferred and fallback strategies.
|
||||
"""
|
||||
|
||||
eos_token_id = 0
|
||||
pad_token_id = 0
|
||||
chat_template = "fake"
|
||||
|
||||
def __init__(self, supports_assistant_mask: bool = True):
|
||||
self.supports_assistant_mask = supports_assistant_mask
|
||||
|
||||
def _render(self, messages):
|
||||
parts: list[tuple[str, bool]] = []
|
||||
for msg in messages:
|
||||
prefix = f"<{msg['role']}>:"
|
||||
content = msg["content"]
|
||||
suffix = "\n"
|
||||
parts.append((prefix, False))
|
||||
parts.append((content, msg["role"] == "assistant"))
|
||||
parts.append((suffix, False))
|
||||
text = "".join(p for p, _ in parts)
|
||||
ids = [ord(c) % 256 for c in text]
|
||||
mask = []
|
||||
for piece, is_assistant in parts:
|
||||
mask.extend([1 if is_assistant else 0] * len(piece))
|
||||
return text, ids, mask
|
||||
|
||||
def apply_chat_template(
|
||||
self,
|
||||
messages,
|
||||
tokenize: bool = False,
|
||||
add_generation_prompt: bool = False,
|
||||
return_assistant_tokens_mask: bool = False,
|
||||
return_dict: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
text, ids, mask = self._render(messages)
|
||||
if not tokenize:
|
||||
return text
|
||||
if return_assistant_tokens_mask and return_dict:
|
||||
if not self.supports_assistant_mask:
|
||||
raise TypeError(
|
||||
"this tokenizer does not support return_assistant_tokens_mask"
|
||||
)
|
||||
return {"input_ids": ids, "assistant_masks": mask}
|
||||
return ids
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema field
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSchemaFields:
|
||||
def test_train_on_responses_only_default_true(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
cfg = DataConfig(train="data.jsonl")
|
||||
assert cfg.train_on_responses_only is True
|
||||
|
||||
def test_train_on_messages_with_train_field_default_false(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
cfg = DataConfig(train="data.jsonl")
|
||||
assert cfg.train_on_messages_with_train_field is False
|
||||
|
||||
def test_train_field_requires_responses_only_disabled(self):
|
||||
"""Per-message 'train' field is mutually exclusive with response-only mode."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
with pytest.raises(ValueError, match="mutually exclusive"):
|
||||
DataConfig(
|
||||
train="data.jsonl",
|
||||
train_on_responses_only=True,
|
||||
train_on_messages_with_train_field=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Loss-mask module
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIgnoreIndex:
|
||||
def test_ignore_index_is_minus_100(self):
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX
|
||||
|
||||
assert IGNORE_INDEX == -100
|
||||
|
||||
|
||||
class TestPreferredPath:
|
||||
"""When tokenizer supports ``return_assistant_tokens_mask=True``."""
|
||||
|
||||
def test_assistant_only_single_turn(self):
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi"},
|
||||
]
|
||||
out = build_assistant_only_labels(messages, tok, max_length=2048)
|
||||
assert "input_ids" in out
|
||||
assert "labels" in out
|
||||
assert "attention_mask" in out
|
||||
assert len(out["labels"]) == len(out["input_ids"])
|
||||
# Exactly the assistant content tokens should NOT be -100.
|
||||
# Render: "<user>:Hello\n<assistant>:Hi\n" — "Hi" is 2 chars.
|
||||
non_masked = [lab for lab in out["labels"] if lab != IGNORE_INDEX]
|
||||
assert len(non_masked) == 2
|
||||
|
||||
def test_assistant_only_multi_turn(self):
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
messages = [
|
||||
{"role": "system", "content": "Sys"},
|
||||
{"role": "user", "content": "Q1"},
|
||||
{"role": "assistant", "content": "A1"},
|
||||
{"role": "user", "content": "Q2"},
|
||||
{"role": "assistant", "content": "A2"},
|
||||
]
|
||||
out = build_assistant_only_labels(messages, tok)
|
||||
non_masked = [lab for lab in out["labels"] if lab != IGNORE_INDEX]
|
||||
# "A1" + "A2" = 4 chars
|
||||
assert len(non_masked) == 4
|
||||
|
||||
def test_truncation_to_max_length(self):
|
||||
from soup_cli.data.loss_mask import build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
messages = [
|
||||
{"role": "user", "content": "x" * 1000},
|
||||
{"role": "assistant", "content": "y" * 1000},
|
||||
]
|
||||
out = build_assistant_only_labels(messages, tok, max_length=128)
|
||||
assert len(out["input_ids"]) == 128
|
||||
assert len(out["labels"]) == 128
|
||||
assert len(out["attention_mask"]) == 128
|
||||
|
||||
|
||||
class TestFallbackPath:
|
||||
"""When tokenizer does NOT support ``return_assistant_tokens_mask``."""
|
||||
|
||||
def test_fallback_single_turn(self):
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "World"},
|
||||
]
|
||||
out = build_assistant_only_labels(messages, tok)
|
||||
non_masked = [lab for lab in out["labels"] if lab != IGNORE_INDEX]
|
||||
# The fallback marks the *delta* tokens of each assistant turn — that
|
||||
# delta is `<assistant>:World\n` = 17 chars (prefix+content+newline).
|
||||
# The exact count depends on the renderer, but non-masked must contain
|
||||
# 'W','o','r','l','d' chars at minimum.
|
||||
assert len(non_masked) >= 5
|
||||
# All masked positions must be inside the assistant turn region.
|
||||
labels = out["labels"]
|
||||
# Check that the user content is masked. user "Hello" chars are at
|
||||
# positions 7..12 (after "<user>:"). Verify those are -100.
|
||||
for pos in range(7, 12):
|
||||
assert labels[pos] == IGNORE_INDEX
|
||||
|
||||
def test_fallback_no_assistant_returns_all_masked(self):
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [{"role": "user", "content": "no answer"}]
|
||||
out = build_assistant_only_labels(messages, tok)
|
||||
assert all(lab == IGNORE_INDEX for lab in out["labels"])
|
||||
|
||||
def test_fallback_strict_assistant_only(self):
|
||||
"""Fallback may include extra prefix tokens; strict mode keeps only
|
||||
the *content* delta against the next user/system turn."""
|
||||
from soup_cli.data.loss_mask import IGNORE_INDEX, build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [
|
||||
{"role": "user", "content": "Q"},
|
||||
{"role": "assistant", "content": "A"},
|
||||
{"role": "user", "content": "Q2"},
|
||||
]
|
||||
out = build_assistant_only_labels(messages, tok)
|
||||
# The user "Q2" must remain masked.
|
||||
# Render: "<user>:Q\n<assistant>:A\n<user>:Q2\n"
|
||||
# Last 4 chars are "<user>:Q2\n" prefix+content+newline part of the
|
||||
# tail — those must all be -100.
|
||||
labels = out["labels"]
|
||||
assert labels[-1] == IGNORE_INDEX # newline
|
||||
assert labels[-2] == IGNORE_INDEX # '2'
|
||||
assert labels[-3] == IGNORE_INDEX # 'Q'
|
||||
|
||||
|
||||
class TestPerMessageTrainField:
|
||||
def test_train_field_overrides_default(self):
|
||||
from soup_cli.data.loss_mask import (
|
||||
IGNORE_INDEX,
|
||||
build_per_message_train_labels,
|
||||
)
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [
|
||||
{"role": "user", "content": "Q"},
|
||||
{"role": "assistant", "content": "A1", "train": False},
|
||||
{"role": "user", "content": "Q2"},
|
||||
{"role": "assistant", "content": "A2", "train": True},
|
||||
]
|
||||
out = build_per_message_train_labels(messages, tok)
|
||||
# Only A2's content tokens should NOT be IGNORE_INDEX
|
||||
non_masked = [lab for lab in out["labels"] if lab != IGNORE_INDEX]
|
||||
# At minimum 'A','2' (both content chars).
|
||||
assert len(non_masked) >= 2
|
||||
|
||||
def test_train_field_default_when_missing(self):
|
||||
"""Missing 'train' field → role==assistant default."""
|
||||
from soup_cli.data.loss_mask import (
|
||||
IGNORE_INDEX,
|
||||
build_per_message_train_labels,
|
||||
)
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [
|
||||
{"role": "user", "content": "Q"},
|
||||
{"role": "assistant", "content": "A"}, # no 'train' field
|
||||
]
|
||||
out = build_per_message_train_labels(messages, tok)
|
||||
# Default-include assistant when no flag.
|
||||
non_masked = [lab for lab in out["labels"] if lab != IGNORE_INDEX]
|
||||
assert len(non_masked) >= 1
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
def test_empty_messages_raises(self):
|
||||
from soup_cli.data.loss_mask import build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer()
|
||||
with pytest.raises(ValueError, match="empty"):
|
||||
build_assistant_only_labels([], tok)
|
||||
|
||||
def test_max_length_must_be_positive(self):
|
||||
from soup_cli.data.loss_mask import build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer()
|
||||
messages = [{"role": "user", "content": "x"}]
|
||||
with pytest.raises(ValueError, match="max_length"):
|
||||
build_assistant_only_labels(messages, tok, max_length=0)
|
||||
|
||||
def test_max_length_rejects_bool(self):
|
||||
"""`bool` is a subclass of `int` — guard like v0.30.0 Candidate."""
|
||||
from soup_cli.data.loss_mask import build_assistant_only_labels
|
||||
|
||||
tok = _FakeTokenizer()
|
||||
messages = [{"role": "user", "content": "x"}]
|
||||
with pytest.raises(ValueError, match="max_length"):
|
||||
build_assistant_only_labels(messages, tok, max_length=True)
|
||||
|
||||
def test_per_message_max_length_truncates(self):
|
||||
from soup_cli.data.loss_mask import build_per_message_train_labels
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
messages = [
|
||||
{"role": "user", "content": "x" * 500},
|
||||
{"role": "assistant", "content": "y" * 500, "train": True},
|
||||
]
|
||||
out = build_per_message_train_labels(messages, tok, max_length=32)
|
||||
assert len(out["input_ids"]) == 32
|
||||
assert len(out["labels"]) == 32
|
||||
assert len(out["attention_mask"]) == 32
|
||||
|
||||
def test_tokenizer_without_chat_template_raises(self):
|
||||
"""Hard-fail when tokenizer has no chat_template."""
|
||||
from soup_cli.data.loss_mask import build_assistant_only_labels
|
||||
|
||||
class _NoTemplate:
|
||||
chat_template = None
|
||||
|
||||
def apply_chat_template(self, *args, **kwargs):
|
||||
raise ValueError("tokenizer has no chat_template")
|
||||
|
||||
with pytest.raises(ValueError, match="chat_template"):
|
||||
build_assistant_only_labels(
|
||||
[{"role": "user", "content": "x"}], _NoTemplate()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_format_row factory (sft.py wiring)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuildFormatRow:
|
||||
@staticmethod
|
||||
def _row():
|
||||
return {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hey"},
|
||||
]
|
||||
}
|
||||
|
||||
def test_default_responses_only_returns_input_ids_labels(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
cfg = DataConfig(train="data.jsonl") # default train_on_responses_only=True
|
||||
fn = build_format_row(tok, cfg, console=None)
|
||||
out = fn(self._row())
|
||||
assert "input_ids" in out
|
||||
assert "labels" in out
|
||||
assert "attention_mask" in out
|
||||
|
||||
def test_per_message_train_field_path(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=False)
|
||||
cfg = DataConfig(
|
||||
train="data.jsonl",
|
||||
train_on_responses_only=False,
|
||||
train_on_messages_with_train_field=True,
|
||||
)
|
||||
fn = build_format_row(tok, cfg, console=None)
|
||||
out = fn(self._row())
|
||||
assert "labels" in out
|
||||
|
||||
def test_legacy_text_path_when_both_false(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
cfg = DataConfig(
|
||||
train="data.jsonl",
|
||||
train_on_responses_only=False,
|
||||
train_on_messages_with_train_field=False,
|
||||
)
|
||||
fn = build_format_row(tok, cfg)
|
||||
out = fn(self._row())
|
||||
assert "text" in out
|
||||
assert "input_ids" not in out
|
||||
|
||||
def test_no_chat_template_calling_format_row_raises(self):
|
||||
"""v0.36.0 Part C: previous silent fallback now raises ValueError."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
class _NoTemplate:
|
||||
chat_template = None
|
||||
|
||||
def apply_chat_template(self, *args, **kwargs): # pragma: no cover
|
||||
raise AssertionError("must not be called")
|
||||
|
||||
cfg = DataConfig(train="data.jsonl") # default responses_only=True
|
||||
# Factory still returns a callable; calling it on a templateless
|
||||
# tokenizer raises. (Factory falls back to legacy path which now
|
||||
# hard-errors instead of building f-string concat.)
|
||||
fn = build_format_row(_NoTemplate(), cfg, console=None)
|
||||
with pytest.raises(ValueError, match="chat_template"):
|
||||
fn(self._row())
|
||||
|
||||
def test_max_length_threaded_through(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
tok = _FakeTokenizer(supports_assistant_mask=True)
|
||||
cfg = DataConfig(train="data.jsonl", max_length=64)
|
||||
fn = build_format_row(tok, cfg)
|
||||
long_row = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "x" * 1000},
|
||||
{"role": "assistant", "content": "y" * 1000},
|
||||
]
|
||||
}
|
||||
out = fn(long_row)
|
||||
assert len(out["input_ids"]) == 64
|
||||
|
|
@ -0,0 +1,568 @@
|
|||
"""Tests for OOM-probe auto batch-size (v0.36.0 Part D).
|
||||
|
||||
Replaces sft.py's static-formula auto-batch with a real try/halve probe and
|
||||
a per-machine cache so repeat runs short-circuit. Mirrors LlamaFactory and
|
||||
Axolotl behaviour.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class _OOMError(Exception):
|
||||
"""Stand-in for ``torch.cuda.OutOfMemoryError`` in unit tests."""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema field
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSchemaField:
|
||||
def test_default_is_auto(self):
|
||||
from soup_cli.config.schema import TrainingConfig
|
||||
|
||||
tcfg = TrainingConfig()
|
||||
assert tcfg.auto_batch_size_strategy == "auto"
|
||||
|
||||
def test_accepts_static(self):
|
||||
from soup_cli.config.schema import TrainingConfig
|
||||
|
||||
tcfg = TrainingConfig(auto_batch_size_strategy="static")
|
||||
assert tcfg.auto_batch_size_strategy == "static"
|
||||
|
||||
def test_accepts_probe(self):
|
||||
from soup_cli.config.schema import TrainingConfig
|
||||
|
||||
tcfg = TrainingConfig(auto_batch_size_strategy="probe")
|
||||
assert tcfg.auto_batch_size_strategy == "probe"
|
||||
|
||||
def test_rejects_unknown_value(self):
|
||||
from soup_cli.config.schema import TrainingConfig
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
TrainingConfig(auto_batch_size_strategy="random")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure binary search
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestProbeLoop:
|
||||
def test_converges_when_capacity_below_start(self):
|
||||
"""Capacity 3, start 4 → halves to 2 (largest power of two that
|
||||
fits). Doubling probe(4) re-OOMs, so we stay at 2. Power-of-two
|
||||
granularity is a deliberate choice: we trade exactness for fewer
|
||||
probe steps (each step is a real GPU forward+backward)."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
capacity = 3
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
if b > capacity:
|
||||
raise _OOMError("simulated")
|
||||
return True
|
||||
|
||||
out = probe_batch_size(
|
||||
probe,
|
||||
start=4,
|
||||
ceiling=16,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
assert out == 2
|
||||
|
||||
def test_converges_when_capacity_above_start(self):
|
||||
"""Start 2, capacity 8 → doubles 2→4→8→16(OOM), back off → 8."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
capacity = 8
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
if b > capacity:
|
||||
raise _OOMError("simulated")
|
||||
return True
|
||||
|
||||
out = probe_batch_size(
|
||||
probe,
|
||||
start=2,
|
||||
ceiling=64,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
assert out == capacity
|
||||
|
||||
def test_ceiling_caps_at_4x_static(self):
|
||||
"""Capacity 1000, ceiling 8 → returns 8 (never tries higher)."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
return True # never OOMs
|
||||
|
||||
out = probe_batch_size(
|
||||
probe,
|
||||
start=2,
|
||||
ceiling=8,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
assert out == 8
|
||||
|
||||
def test_starts_oom_halves_to_one(self):
|
||||
"""Even start=1 OOMs → returns 1 (never go below 1)."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
raise _OOMError("starved")
|
||||
|
||||
with pytest.raises(RuntimeError, match="batch_size=1"):
|
||||
probe_batch_size(
|
||||
probe,
|
||||
start=2,
|
||||
ceiling=16,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
|
||||
def test_max_doublings_capped(self):
|
||||
"""Search must not run forever — cap at 8 doublings."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
calls: list[int] = []
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
calls.append(b)
|
||||
return True
|
||||
|
||||
probe_batch_size(
|
||||
probe,
|
||||
start=1,
|
||||
ceiling=10**6,
|
||||
oom_exceptions=(_OOMError,),
|
||||
max_doublings=8,
|
||||
)
|
||||
# At most 8 successful doublings + initial = 9 successful probes.
|
||||
assert len(calls) <= 12
|
||||
|
||||
def test_rejects_invalid_start(self):
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
with pytest.raises(ValueError, match="start"):
|
||||
probe_batch_size(
|
||||
lambda b: True,
|
||||
start=0,
|
||||
ceiling=8,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
|
||||
def test_rejects_invalid_ceiling(self):
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
with pytest.raises(ValueError, match="ceiling"):
|
||||
probe_batch_size(
|
||||
lambda b: True,
|
||||
start=4,
|
||||
ceiling=2, # < start
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
|
||||
def test_unrelated_exception_propagates(self):
|
||||
"""A non-OOM exception must propagate, not be swallowed as OOM."""
|
||||
from soup_cli.utils.batch_probe import probe_batch_size
|
||||
|
||||
def probe(b: int) -> bool:
|
||||
raise RuntimeError("model bug")
|
||||
|
||||
with pytest.raises(RuntimeError, match="model bug"):
|
||||
probe_batch_size(
|
||||
probe,
|
||||
start=2,
|
||||
ceiling=16,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache layer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCache:
|
||||
def test_key_normalizes(self):
|
||||
from soup_cli.utils.batch_probe import make_cache_key
|
||||
|
||||
a = make_cache_key(
|
||||
base="meta-llama/Llama-3.2-1B",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="NVIDIA A100-SXM4-80GB",
|
||||
gpu_memory_gb=80,
|
||||
)
|
||||
b = make_cache_key(
|
||||
base="meta-llama/Llama-3.2-1B",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="NVIDIA A100-SXM4-80GB",
|
||||
gpu_memory_gb=80,
|
||||
)
|
||||
assert a == b
|
||||
|
||||
def test_key_differs_on_quantization(self):
|
||||
from soup_cli.utils.batch_probe import make_cache_key
|
||||
|
||||
a = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
b = make_cache_key("m", 2048, "8bit", 64, "gpu", 80)
|
||||
assert a != b
|
||||
|
||||
def test_key_rejects_bool_inputs(self):
|
||||
"""v0.30.0 Candidate convention: bool is a subclass of int — guard."""
|
||||
from soup_cli.utils.batch_probe import make_cache_key
|
||||
|
||||
with pytest.raises(ValueError, match="max_length"):
|
||||
make_cache_key("m", True, "4bit", 64, "gpu", 80)
|
||||
with pytest.raises(ValueError, match="lora_r"):
|
||||
make_cache_key("m", 2048, "4bit", True, "gpu", 80)
|
||||
with pytest.raises(ValueError, match="gpu_memory_gb"):
|
||||
make_cache_key("m", 2048, "4bit", 64, "gpu", True)
|
||||
|
||||
def test_save_and_load_roundtrip(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import (
|
||||
load_cache,
|
||||
make_cache_key,
|
||||
save_cache_entry,
|
||||
)
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
save_cache_entry(key, 8)
|
||||
|
||||
cache = load_cache()
|
||||
assert cache.get(key) == 8
|
||||
|
||||
def test_load_corrupt_returns_empty(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import load_cache
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
cache_path.write_text("not json", encoding="utf-8")
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
assert load_cache() == {}
|
||||
|
||||
def test_load_missing_returns_empty(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import load_cache
|
||||
|
||||
cache_path = tmp_path / "missing.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
assert load_cache() == {}
|
||||
|
||||
def test_save_rejects_non_positive_value(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import make_cache_key, save_cache_entry
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
save_cache_entry(key, 0)
|
||||
with pytest.raises(ValueError):
|
||||
save_cache_entry(key, -1)
|
||||
|
||||
def test_save_rejects_bool_value(self, tmp_path, monkeypatch):
|
||||
"""``bool`` is a subclass of int — guard like v0.30.0 Candidate."""
|
||||
from soup_cli.utils.batch_probe import make_cache_key, save_cache_entry
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
save_cache_entry(key, True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pick_batch_size — main entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPickBatchSize:
|
||||
def test_static_strategy_returns_static_estimate(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="static",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="cpu",
|
||||
gpu_memory_gb=0,
|
||||
probe_fn=None,
|
||||
)
|
||||
assert out == 4
|
||||
|
||||
def test_cache_hit_short_circuits_probe(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import (
|
||||
make_cache_key,
|
||||
pick_batch_size,
|
||||
save_cache_entry,
|
||||
)
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
save_cache_entry(key, 16)
|
||||
|
||||
called: list[int] = []
|
||||
|
||||
def probe(b):
|
||||
called.append(b)
|
||||
return True
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="gpu",
|
||||
gpu_memory_gb=80,
|
||||
probe_fn=probe,
|
||||
)
|
||||
assert out == 16
|
||||
assert called == [] # probe was not invoked
|
||||
|
||||
def test_probe_strategy_runs_probe_and_caches(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import (
|
||||
load_cache,
|
||||
make_cache_key,
|
||||
pick_batch_size,
|
||||
)
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
capacity = 8
|
||||
|
||||
def probe(b):
|
||||
if b > capacity:
|
||||
raise _OOMError("oom")
|
||||
return True
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="gpu",
|
||||
gpu_memory_gb=80,
|
||||
probe_fn=probe,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
assert out == capacity
|
||||
|
||||
# Cache write happened.
|
||||
cache = load_cache()
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
assert cache[key] == capacity
|
||||
|
||||
def test_probe_without_callable_falls_back_to_static(self, tmp_path, monkeypatch):
|
||||
"""No probe_fn supplied (e.g. CPU run) → use static estimate."""
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="cpu",
|
||||
gpu_memory_gb=0,
|
||||
probe_fn=None,
|
||||
)
|
||||
assert out == 4
|
||||
|
||||
def test_auto_strategy_uses_probe_when_probe_fn_supplied(
|
||||
self, tmp_path, monkeypatch
|
||||
):
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
capacity = 4
|
||||
|
||||
def probe(b):
|
||||
if b > capacity:
|
||||
raise _OOMError("oom")
|
||||
return True
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=2,
|
||||
strategy="auto",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="gpu",
|
||||
gpu_memory_gb=80,
|
||||
probe_fn=probe,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
assert out == capacity
|
||||
|
||||
def test_cache_corruption_does_not_block_probe(self, tmp_path, monkeypatch):
|
||||
"""Corrupt cache file → silently re-probe; ceiling = static * 4."""
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
cache_path.write_text("garbage", encoding="utf-8")
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="gpu",
|
||||
gpu_memory_gb=80,
|
||||
probe_fn=lambda b: True,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
# Probe ran; with no OOMs and ceiling = 4*4 = 16, lands at 16.
|
||||
assert out == 16
|
||||
|
||||
def test_runtime_error_propagates_when_bs1_ooms(self, tmp_path, monkeypatch):
|
||||
"""All-OOM probe → RuntimeError surfaces to caller."""
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
def always_oom(b):
|
||||
raise _OOMError("oom")
|
||||
|
||||
with pytest.raises(RuntimeError, match="batch_size=1"):
|
||||
pick_batch_size(
|
||||
static_estimate=2,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="gpu",
|
||||
gpu_memory_gb=80,
|
||||
probe_fn=always_oom,
|
||||
oom_exceptions=(_OOMError,),
|
||||
)
|
||||
|
||||
def test_explicit_probe_no_probe_fn_emits_warning(
|
||||
self, tmp_path, monkeypatch
|
||||
):
|
||||
"""strategy='probe' with probe_fn=None → console warning fires."""
|
||||
from io import StringIO
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.utils.batch_probe import pick_batch_size
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
|
||||
out = pick_batch_size(
|
||||
static_estimate=4,
|
||||
strategy="probe",
|
||||
base="m",
|
||||
max_length=2048,
|
||||
quantization="4bit",
|
||||
lora_r=64,
|
||||
gpu_name="cpu",
|
||||
gpu_memory_gb=0,
|
||||
probe_fn=None,
|
||||
console=console,
|
||||
)
|
||||
assert out == 4
|
||||
assert "probe_fn" in buf.getvalue() or "static" in buf.getvalue()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache-path containment (security review fix)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCachePathContainment:
|
||||
def test_out_of_bounds_override_falls_back_to_default(
|
||||
self, tmp_path, monkeypatch
|
||||
):
|
||||
"""Env var pointing outside home/cwd/tmp → ignored, default used."""
|
||||
import os
|
||||
|
||||
from soup_cli.utils.batch_probe import _cache_path
|
||||
|
||||
# Use a sibling-of-temp path that is guaranteed outside any anchor —
|
||||
# an absolute root we cannot write to is equally fine, since the
|
||||
# function only resolves+rejects, no I/O.
|
||||
if os.name == "nt":
|
||||
evil = "C:\\evil-bound\\batch.json"
|
||||
else:
|
||||
evil = "/etc/cron.d/soup_evil"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", evil)
|
||||
path = _cache_path()
|
||||
# Fall-through path — must NOT be the evil override.
|
||||
assert os.path.realpath(path) != os.path.realpath(evil)
|
||||
assert path.endswith("batch_cache.json")
|
||||
|
||||
def test_in_bounds_override_honoured(self, tmp_path, monkeypatch):
|
||||
import os
|
||||
|
||||
from soup_cli.utils.batch_probe import _cache_path
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
# tmp_path is under tempfile.gettempdir() — allowed. Compare via
|
||||
# realpath to normalise across short-name / forward-slash forms.
|
||||
assert os.path.realpath(_cache_path()) == os.path.realpath(str(cache_path))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache file integrity guard
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCacheFileShape:
|
||||
def test_cache_is_dict_of_str_int(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.batch_probe import make_cache_key, save_cache_entry
|
||||
|
||||
cache_path = tmp_path / "batch_cache.json"
|
||||
monkeypatch.setenv("SOUP_BATCH_CACHE_PATH", str(cache_path))
|
||||
|
||||
key = make_cache_key("m", 2048, "4bit", 64, "gpu", 80)
|
||||
save_cache_entry(key, 8)
|
||||
|
||||
with open(cache_path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
assert isinstance(data, dict)
|
||||
for k, v in data.items():
|
||||
assert isinstance(k, str)
|
||||
assert isinstance(v, int)
|
||||
assert v > 0
|
||||
|
|
@ -0,0 +1,276 @@
|
|||
"""Tests for chat-template hardening (v0.36.0 Part C).
|
||||
|
||||
Replaces sft.py's silent ``f"{role}: {content}"`` fallback (which produced
|
||||
garbage training data on tokenizers without ``chat_template``) with a hard
|
||||
error and an explicit ``DataConfig.chat_template`` field for overrides.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema field
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSchemaField:
|
||||
def test_chat_template_default_none(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
cfg = DataConfig(train="data.jsonl")
|
||||
assert cfg.chat_template is None
|
||||
|
||||
def test_chat_template_accepts_registered_name(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
cfg = DataConfig(train="data.jsonl", chat_template="chatml")
|
||||
assert cfg.chat_template == "chatml"
|
||||
|
||||
def test_chat_template_accepts_jinja_string(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
jinja = "{% for m in messages %}{{ m.role }}: {{ m.content }}{% endfor %}"
|
||||
cfg = DataConfig(train="data.jsonl", chat_template=jinja)
|
||||
assert cfg.chat_template == jinja
|
||||
|
||||
def test_chat_template_rejects_null_byte(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
DataConfig(train="data.jsonl", chat_template="bad\x00template")
|
||||
|
||||
def test_chat_template_rejects_oversize(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
# Cap at 64KB to prevent template-injection DoS payloads.
|
||||
with pytest.raises(ValueError):
|
||||
DataConfig(train="data.jsonl", chat_template="x" * 100_000)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRegistry:
|
||||
def test_lists_known_templates(self):
|
||||
from soup_cli.data.chat_templates import list_template_names
|
||||
|
||||
names = list_template_names()
|
||||
# At minimum: the 7 declared in v0.36.0 Part C.
|
||||
for required in (
|
||||
"chatml", "llama3", "qwen2.5", "gemma3", "phi4", "deepseek-r1", "mistral",
|
||||
):
|
||||
assert required in names
|
||||
|
||||
def test_chatml_is_jinja_string(self):
|
||||
from soup_cli.data.chat_templates import get_template
|
||||
|
||||
tmpl = get_template("chatml")
|
||||
assert isinstance(tmpl, str)
|
||||
assert "{%" in tmpl
|
||||
# ChatML signature markers.
|
||||
assert "<|im_start|>" in tmpl
|
||||
assert "<|im_end|>" in tmpl
|
||||
|
||||
def test_unknown_name_raises(self):
|
||||
from soup_cli.data.chat_templates import get_template
|
||||
|
||||
with pytest.raises(KeyError, match="not registered"):
|
||||
get_template("not-a-real-template")
|
||||
|
||||
def test_resolve_returns_jinja_for_known_name(self):
|
||||
from soup_cli.data.chat_templates import resolve_chat_template
|
||||
|
||||
out = resolve_chat_template("chatml")
|
||||
assert "<|im_start|>" in out
|
||||
|
||||
def test_resolve_returns_passthrough_for_jinja(self):
|
||||
from soup_cli.data.chat_templates import resolve_chat_template
|
||||
|
||||
jinja = "{% for m in messages %}<x>{{ m.content }}</x>{% endfor %}"
|
||||
out = resolve_chat_template(jinja)
|
||||
assert out == jinja
|
||||
|
||||
def test_resolve_none_returns_none(self):
|
||||
from soup_cli.data.chat_templates import resolve_chat_template
|
||||
|
||||
assert resolve_chat_template(None) is None
|
||||
|
||||
def test_resolve_empty_returns_none(self):
|
||||
from soup_cli.data.chat_templates import resolve_chat_template
|
||||
|
||||
assert resolve_chat_template("") is None
|
||||
|
||||
def test_resolve_unknown_name_raises(self):
|
||||
"""Public surface: bad name through resolve_chat_template."""
|
||||
from soup_cli.data.chat_templates import resolve_chat_template
|
||||
|
||||
with pytest.raises(KeyError, match="not registered"):
|
||||
resolve_chat_template("not-a-real-template")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema rejects Jinja directives that touch the filesystem
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestJinjaDirectiveBlocking:
|
||||
@pytest.mark.parametrize(
|
||||
"bad",
|
||||
[
|
||||
"{% include 'config.yaml' %}",
|
||||
"{%- include 'config.yaml' %}",
|
||||
"{% import 'os' as os %}",
|
||||
"{% from 'os' import system %}",
|
||||
"{% macro evil() %}{% endmacro %}",
|
||||
"{% extends 'base.j2' %}",
|
||||
],
|
||||
)
|
||||
def test_schema_rejects_filesystem_directives(self, bad):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
with pytest.raises(ValueError, match="directive"):
|
||||
DataConfig(train="data.jsonl", chat_template=bad)
|
||||
|
||||
def test_schema_accepts_for_loop(self):
|
||||
"""Standard control flow must still work."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
ok = "{% for m in messages %}{{ m.content }}{% endfor %}"
|
||||
cfg = DataConfig(train="data.jsonl", chat_template=ok)
|
||||
assert cfg.chat_template == ok
|
||||
|
||||
def test_schema_accepts_if_else(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
ok = "{% if true %}x{% else %}y{% endif %}"
|
||||
cfg = DataConfig(train="data.jsonl", chat_template=ok)
|
||||
assert cfg.chat_template == ok
|
||||
|
||||
def test_schema_empty_string_normalised_to_none(self):
|
||||
from soup_cli.config.schema import DataConfig
|
||||
|
||||
cfg = DataConfig(train="data.jsonl", chat_template="")
|
||||
assert cfg.chat_template is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apply_chat_template_override
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestApplyOverride:
|
||||
def test_sets_tokenizer_chat_template(self):
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
|
||||
class _T:
|
||||
chat_template = None
|
||||
|
||||
tok = _T()
|
||||
apply_chat_template_override(tok, "chatml")
|
||||
assert tok.chat_template is not None
|
||||
assert "<|im_start|>" in tok.chat_template
|
||||
|
||||
def test_none_leaves_tokenizer_alone(self):
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
|
||||
class _T:
|
||||
chat_template = "existing"
|
||||
|
||||
tok = _T()
|
||||
apply_chat_template_override(tok, None)
|
||||
assert tok.chat_template == "existing"
|
||||
|
||||
def test_empty_leaves_tokenizer_alone(self):
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
|
||||
class _T:
|
||||
chat_template = "existing"
|
||||
|
||||
tok = _T()
|
||||
apply_chat_template_override(tok, "")
|
||||
assert tok.chat_template == "existing"
|
||||
|
||||
def test_override_emits_save_pretrained_warning(self):
|
||||
"""v0.36.0 review fix: warn when overriding so users know push will
|
||||
persist the new template into tokenizer_config.json."""
|
||||
from io import StringIO
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
|
||||
class _T:
|
||||
chat_template = None
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
applied = apply_chat_template_override(_T(), "chatml", console=console)
|
||||
assert applied is True
|
||||
assert "save_pretrained" in buf.getvalue() or "soup push" in buf.getvalue()
|
||||
|
||||
def test_override_returns_false_when_noop(self):
|
||||
from soup_cli.data.chat_templates import apply_chat_template_override
|
||||
|
||||
class _T:
|
||||
chat_template = "existing"
|
||||
|
||||
applied = apply_chat_template_override(_T(), None)
|
||||
assert applied is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Hard error in sft_format when no chat_template AND no override
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHardError:
|
||||
def test_no_template_no_override_raises(self):
|
||||
"""The legacy `f"{role}: {content}"` silent fallback is now an error."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
class _NoTemplate:
|
||||
chat_template = None
|
||||
|
||||
def apply_chat_template(self, *args, **kwargs):
|
||||
raise AssertionError("must not be reached")
|
||||
|
||||
cfg = DataConfig(
|
||||
train="data.jsonl",
|
||||
train_on_responses_only=False, # legacy path
|
||||
train_on_messages_with_train_field=False,
|
||||
chat_template=None,
|
||||
)
|
||||
fn = build_format_row(_NoTemplate(), cfg, console=None)
|
||||
with pytest.raises(ValueError, match="chat_template"):
|
||||
fn({"messages": [{"role": "user", "content": "hi"}]})
|
||||
|
||||
def test_override_applies_to_legacy_path(self):
|
||||
"""When user passes chat_template override, build_format_row works."""
|
||||
from soup_cli.config.schema import DataConfig
|
||||
from soup_cli.data.sft_format import build_format_row
|
||||
|
||||
class _T:
|
||||
chat_template = None
|
||||
applied = []
|
||||
|
||||
def apply_chat_template(
|
||||
self, messages, tokenize=False, add_generation_prompt=False, **kwargs
|
||||
):
|
||||
self.applied.append(messages)
|
||||
return "RENDERED"
|
||||
|
||||
tok = _T()
|
||||
cfg = DataConfig(
|
||||
train="data.jsonl",
|
||||
train_on_responses_only=False,
|
||||
train_on_messages_with_train_field=False,
|
||||
chat_template="chatml",
|
||||
)
|
||||
fn = build_format_row(tok, cfg, console=None)
|
||||
out = fn({"messages": [{"role": "user", "content": "hi"}]})
|
||||
assert out["text"] == "RENDERED"
|
||||
assert tok.chat_template is not None # override was applied
|
||||
|
|
@ -0,0 +1,266 @@
|
|||
"""Tests for ``--trust-remote-code`` opt-in (v0.36.0 Part B).
|
||||
|
||||
Replaces the previous unconditional ``trust_remote_code=True`` smell across
|
||||
sft.py / chat.py / serve.py with an explicit, auditable opt-in flag plus a
|
||||
trusted-org allowlist that suppresses noise on first-party models.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from io import StringIO
|
||||
|
||||
import pytest
|
||||
from rich.console import Console
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Allowlist
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestKnownSafePrefixes:
|
||||
def test_meta_llama_is_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert is_known_safe("meta-llama/Llama-3.2-1B")
|
||||
|
||||
def test_qwen_is_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert is_known_safe("Qwen/Qwen2.5-7B")
|
||||
|
||||
def test_mistral_is_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert is_known_safe("mistralai/Mistral-7B-Instruct-v0.3")
|
||||
|
||||
def test_random_org_not_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert not is_known_safe("randomuser/SomeModel")
|
||||
|
||||
def test_local_path_not_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert not is_known_safe("./local-checkpoint")
|
||||
|
||||
def test_partial_prefix_does_not_match(self):
|
||||
"""`meta-llama-evil/...` must NOT match the `meta-llama/` prefix."""
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert not is_known_safe("meta-llama-evil/SomeModel")
|
||||
|
||||
def test_empty_string_not_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert not is_known_safe("")
|
||||
|
||||
def test_non_string_not_safe(self):
|
||||
from soup_cli.utils.trust_remote import is_known_safe
|
||||
|
||||
assert not is_known_safe(None)
|
||||
assert not is_known_safe(123)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# resolve_trust_remote_code — main entry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolve:
|
||||
def test_default_off_for_safe_prefix_passes_silently(self):
|
||||
"""Trusted org + flag off → returns False, no warning."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
out = resolve_trust_remote_code(
|
||||
"meta-llama/Llama-3.2-1B",
|
||||
requested=False,
|
||||
console=console,
|
||||
requires_remote_code=False,
|
||||
)
|
||||
assert out is False
|
||||
assert buf.getvalue() == ""
|
||||
|
||||
def test_flag_enabled_warns_once(self):
|
||||
"""User opted in → return True + warning panel."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
out = resolve_trust_remote_code(
|
||||
"shady-org/SomeModel",
|
||||
requested=True,
|
||||
console=console,
|
||||
requires_remote_code=True,
|
||||
)
|
||||
assert out is True
|
||||
output = buf.getvalue()
|
||||
assert "trust_remote_code" in output.lower() or "remote code" in output.lower()
|
||||
assert "shady-org/SomeModel" in output
|
||||
|
||||
def test_flag_enabled_safe_prefix_suppresses_warning(self):
|
||||
"""Trusted org doesn't ship custom code — suppress warning even when flag set."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
out = resolve_trust_remote_code(
|
||||
"meta-llama/Llama-3.2-1B",
|
||||
requested=True,
|
||||
console=console,
|
||||
requires_remote_code=False,
|
||||
)
|
||||
assert out is True
|
||||
# No noisy panel when the model is from a trusted prefix.
|
||||
assert "WARNING" not in buf.getvalue().upper()
|
||||
|
||||
def test_default_off_for_unknown_with_remote_code_raises(self):
|
||||
"""Model needs custom code + flag off → fail fast with actionable error."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
resolve_trust_remote_code(
|
||||
"shady-org/CustomModel",
|
||||
requested=False,
|
||||
console=console,
|
||||
requires_remote_code=True,
|
||||
)
|
||||
msg = str(exc_info.value)
|
||||
assert "shady-org/CustomModel" in msg
|
||||
assert "--trust-remote-code" in msg
|
||||
|
||||
def test_default_off_for_unknown_without_remote_code_passes(self):
|
||||
"""Standard model + flag off → returns False, no error."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
buf = StringIO()
|
||||
console = Console(file=buf, force_terminal=False)
|
||||
out = resolve_trust_remote_code(
|
||||
"shady-org/StandardLlama",
|
||||
requested=False,
|
||||
console=console,
|
||||
requires_remote_code=False,
|
||||
)
|
||||
assert out is False
|
||||
|
||||
def test_console_optional(self):
|
||||
"""resolve_trust_remote_code must work when console is None."""
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
out = resolve_trust_remote_code(
|
||||
"meta-llama/Llama-3.2-1B",
|
||||
requested=True,
|
||||
console=None,
|
||||
requires_remote_code=False,
|
||||
)
|
||||
assert out is True
|
||||
|
||||
def test_invalid_model_name_rejected(self):
|
||||
from soup_cli.utils.trust_remote import resolve_trust_remote_code
|
||||
|
||||
with pytest.raises(ValueError, match="model_name"):
|
||||
resolve_trust_remote_code(
|
||||
"",
|
||||
requested=False,
|
||||
console=None,
|
||||
requires_remote_code=False,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# model_requires_trust_remote_code (probe HF config for auto_map)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRequiresProbe:
|
||||
def test_local_path_no_auto_map_returns_false(self, tmp_path, monkeypatch):
|
||||
"""Local path with config.json lacking auto_map → False."""
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
config = tmp_path / "config.json"
|
||||
config.write_text('{"model_type": "llama"}', encoding="utf-8")
|
||||
assert model_requires_trust_remote_code(str(tmp_path)) is False
|
||||
|
||||
def test_local_path_with_auto_map_returns_true(self, tmp_path):
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
config = tmp_path / "config.json"
|
||||
config.write_text(
|
||||
'{"model_type": "custom", "auto_map": '
|
||||
'{"AutoModelForCausalLM": "modeling.Custom"}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert model_requires_trust_remote_code(str(tmp_path)) is True
|
||||
|
||||
def test_missing_config_returns_none(self, tmp_path):
|
||||
"""Missing config.json → None (unknown — caller decides)."""
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
out = model_requires_trust_remote_code(str(tmp_path))
|
||||
assert out is None
|
||||
|
||||
def test_malformed_config_returns_none(self, tmp_path):
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
config = tmp_path / "config.json"
|
||||
config.write_text("{this is not json", encoding="utf-8")
|
||||
out = model_requires_trust_remote_code(str(tmp_path))
|
||||
assert out is None
|
||||
|
||||
def test_non_dict_root_returns_none(self, tmp_path):
|
||||
"""Config with non-dict root (e.g. JSON array) → None."""
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
config = tmp_path / "config.json"
|
||||
config.write_text("[1, 2, 3]", encoding="utf-8")
|
||||
out = model_requires_trust_remote_code(str(tmp_path))
|
||||
assert out is None
|
||||
|
||||
def test_non_directory_path_returns_none(self):
|
||||
"""Bare HF repo id (not a local dir) → None (unknown)."""
|
||||
from soup_cli.utils.trust_remote import model_requires_trust_remote_code
|
||||
|
||||
out = model_requires_trust_remote_code("meta-llama/Llama-3.2-1B")
|
||||
assert out is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI flag plumbing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCLIPlumbing:
|
||||
"""Smoke check that --trust-remote-code is a registered Typer option."""
|
||||
|
||||
def test_train_help_lists_flag(self):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["train", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "--trust-remote-code" in result.output
|
||||
|
||||
def test_chat_help_lists_flag(self):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["chat", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "--trust-remote-code" in result.output
|
||||
|
||||
def test_serve_help_lists_flag(self):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["serve", "--help"])
|
||||
assert result.exit_code == 0
|
||||
assert "--trust-remote-code" in result.output
|
||||
Loading…
Reference in New Issue