From a5540fa1e2f7695828a50beda96556f313664441 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Thu, 30 Apr 2026 11:52:51 +0500 Subject: [PATCH] =?UTF-8?q?feat(correctness):=20v0.36.0=20=E2=80=94=20Corr?= =?UTF-8?q?ectness=20First=20(4=20Parts:=20A/B/C/D)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- CONTRIBUTING.md | 2 +- README.md | 61 +++- SECURITY.md | 7 +- pyproject.toml | 2 +- soup_cli/__init__.py | 2 +- soup_cli/commands/chat.py | 35 +- soup_cli/commands/data.py | 8 + soup_cli/commands/eval.py | 8 + soup_cli/commands/serve.py | 35 +- soup_cli/commands/train.py | 14 +- soup_cli/config/schema.py | 70 ++++ soup_cli/data/chat_templates.py | 163 +++++++++ soup_cli/data/loss_mask.py | 190 +++++++++++ soup_cli/data/sft_format.py | 107 ++++++ soup_cli/trainer/sft.py | 86 +++-- soup_cli/utils/batch_probe.py | 288 ++++++++++++++++ soup_cli/utils/trust_remote.py | 141 ++++++++ tests/test_assistant_mask.py | 401 ++++++++++++++++++++++ tests/test_batch_probe.py | 568 ++++++++++++++++++++++++++++++++ tests/test_chat_template.py | 276 ++++++++++++++++ tests/test_trust_remote_code.py | 266 +++++++++++++++ 21 files changed, 2688 insertions(+), 42 deletions(-) create mode 100644 soup_cli/data/chat_templates.py create mode 100644 soup_cli/data/loss_mask.py create mode 100644 soup_cli/data/sft_format.py create mode 100644 soup_cli/utils/batch_probe.py create mode 100644 soup_cli/utils/trust_remote.py create mode 100644 tests/test_assistant_mask.py create mode 100644 tests/test_batch_probe.py create mode 100644 tests/test_chat_template.py create mode 100644 tests/test_trust_remote_code.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9e6c274..a3dc7ed 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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 ``` diff --git a/README.md b/README.md index ed2e6ef..ff83a71 100644 --- a/README.md +++ b/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: diff --git a/SECURITY.md b/SECURITY.md index 9ceea80..f370f23 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 54c5dfd..3102e66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index 5497c4b..eac931d 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.35.0" +__version__ = "0.36.0" diff --git a/soup_cli/commands/chat.py b/soup_cli/commands/chat.py index 9097e24..9b29cb2 100644 --- a/soup_cli/commands/chat.py +++ b/soup_cli/commands/chat.py @@ -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, ) diff --git a/soup_cli/commands/data.py b/soup_cli/commands/data.py index 08a1104..71d4b13 100644 --- a/soup_cli/commands/data.py +++ b/soup_cli/commands/data.py @@ -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 diff --git a/soup_cli/commands/eval.py b/soup_cli/commands/eval.py index d306ff3..6fc772a 100644 --- a/soup_cli/commands/eval.py +++ b/soup_cli/commands/eval.py @@ -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 diff --git a/soup_cli/commands/serve.py b/soup_cli/commands/serve.py index d3fa42d..ea389f2 100644 --- a/soup_cli/commands/serve.py +++ b/soup_cli/commands/serve.py @@ -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, ) diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index 9a5503c..e932e1a 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -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) --- diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 13d5337..ac6fddb 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -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) diff --git a/soup_cli/data/chat_templates.py b/soup_cli/data/chat_templates.py new file mode 100644 index 0000000..ec1874e --- /dev/null +++ b/soup_cli/data/chat_templates.py @@ -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 = "". +# +# 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'] }}" + "{% endif %}" + "{% endfor %}" +) + +_GEMMA3 = ( + "{% for message in messages %}" + "{{ 'user' if message['role'] == 'user' else 'model' }}\n" + "{{ message['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}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 diff --git a/soup_cli/data/loss_mask.py b/soup_cli/data/loss_mask.py new file mode 100644 index 0000000..4b80cea --- /dev/null +++ b/soup_cli/data/loss_mask.py @@ -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) diff --git a/soup_cli/data/sft_format.py b/soup_cli/data/sft_format.py new file mode 100644 index 0000000..729bc3a --- /dev/null +++ b/soup_cli/data/sft_format.py @@ -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 diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index 81a02e2..2b92a40 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -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 diff --git a/soup_cli/utils/batch_probe.py b/soup_cli/utils/batch_probe.py new file mode 100644 index 0000000..fba50d3 --- /dev/null +++ b/soup_cli/utils/batch_probe.py @@ -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 diff --git a/soup_cli/utils/trust_remote.py b/soup_cli/utils/trust_remote.py new file mode 100644 index 0000000..310faa2 --- /dev/null +++ b/soup_cli/utils/trust_remote.py @@ -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) diff --git a/tests/test_assistant_mask.py b/tests/test_assistant_mask.py new file mode 100644 index 0000000..8e0b13f --- /dev/null +++ b/tests/test_assistant_mask.py @@ -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 ``:\\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: ":Hello\n: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 `: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 ":"). 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: ":Q\n:A\n:Q2\n" + # Last 4 chars are ":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 diff --git a/tests/test_batch_probe.py b/tests/test_batch_probe.py new file mode 100644 index 0000000..7ecea5f --- /dev/null +++ b/tests/test_batch_probe.py @@ -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 diff --git a/tests/test_chat_template.py b/tests/test_chat_template.py new file mode 100644 index 0000000..3c9addb --- /dev/null +++ b/tests/test_chat_template.py @@ -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 %}{{ m.content }}{% 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 diff --git a/tests/test_trust_remote_code.py b/tests/test_trust_remote_code.py new file mode 100644 index 0000000..92295f3 --- /dev/null +++ b/tests/test_trust_remote_code.py @@ -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