diff --git a/.gitignore b/.gitignore index 246f51c..9b639a1 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,9 @@ report.xml .env *.key +# Internal docs (local only — not for repo) +docs/ + # Internal plan + Claude Code local dev instructions (not for repo) .claude/plan.md .claude/CLAUDE.md diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3319caa..ad91e57 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -107,11 +107,11 @@ soup_cli/ cans/ - Shareable .can artifact format + run/publish orchestrator (v0.26.0 + v0.33.0) data/traces/ - Trace-to-Preference harvester (v0.26.0) data/collators.py - CrossDocCollator for sample packing (v0.33.0) - 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, multipack_sampler, multipack, neat_packing, jinja_analyzer, quant_menu, relora, peft_patches - templates/ - 16 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0) + 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, multipack_sampler, multipack, neat_packing, jinja_analyzer, quant_menu, relora, peft_patches, dpo_variants + templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (132 files, 4538 tests) +tests/ - Test suite (136 files, 4656 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index 2813d8e..9b16f0b 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.39.0 — LoRA Quality**: PEFT surface improvements that LlamaFactory and Axolotl maintain. Faster init, periodic adapter refresh, per-module rank, surgical architecture patches, and a cleaner template registry. +**v0.40.0 — Preference Variety**: BCO trainer + a unified preference-loss surface so DPO / SimPO / ORPO / IPO / BCO live behind one config knob. Adds two opt-in DPO controls (β-schedule + ref-model regen) and a forward-looking multi-objective preference-loss surface. -- **PiSSA init** — set `training.lora.init_strategy: pissa` for SVD-initialized LoRA pairs. Faster early convergence vs random init at the cost of one extra SVD pass on the first epoch. `init_strategy: olora` is also accepted (equivalent to legacy `use_olora=True`, which auto-aligns for back-compat). -- **ReLoRA callback** — set `training.relora_steps: 500` to magnitude-prune LoRA adapter weights every 500 steps and clear optimizer state. Useful for very long training runs where the LoRA capacity saturates. Bounds: `relora_warmup_ratio` [0,1], `relora_prune_ratio` (0,1), `relora_reset_optimizer: true`. Wired into the SFT trainer; multi-trainer expansion deferred to v0.39.1. -- **Per-pattern LoRA rank** — `training.lora.rank_pattern: {q_proj: 8, v_proj: 16}` and `alpha_pattern` give different ranks per target module pattern. Useful in MoE configs where expert FFNs need lower rank than attention. Caps: 256 keys × value 1024. -- **Surgical PEFT patches** — Gemma 4 `ClippableLinear` is auto-swapped to plain `nn.Linear` so PEFT's matcher recognises it; fused-MoE 3-D expert weights have `lora_dropout` zeroed to dodge `ParamWrapper` crashes. Both are gated by architecture detection (regex word-boundary on the model name) and never run on unrelated models. -- **Template registry** — the 16 built-in templates now live as `soup_cli/templates/*.yaml` with a `manifest.json` index. `soup init --template ` reads the YAML; the inline copies in `schema.py` stay as a back-compat fallback (deprecation pointing at v0.41.0+). -- **Net +164 tests** (4374 → 4538) across PiSSA mutual-exclusion, ReLoRA frozen-policy + real-optimizer-state reset, rank_pattern bounds + null-byte rejection, regex word-boundary Gemma 4 detection, and template-registry path-traversal + manifest-tampering containment. +- **BCO Trainer** — set `task: bco` for Binary Classifier Optimization. Same input format as DPO (`prompt + chosen + rejected`); rows are internally split to TRL's BCO unpaired schema. New `training.bco_beta` (default 0.1, gt=0). Template: `soup init --template bco`. +- **Unified preference dispatcher** — set `task: preference` + `training.preference_loss: dpo|simpo|orpo|ipo|bco` to pick the loss without renaming your task. Legacy `task: dpo`, `task: simpo`, etc remain first-class — the new surface is additive, not a breaking collapse. Useful for hyperparameter sweeps over the loss type itself. +- **KL-controlled DPO variants** — anneal β over training with `training.dpo_beta_schedule: linear|cosine|exponential` + `training.dpo_beta_end`. Periodically refresh the frozen reference model with the current student via `training.dpo_ref_regen_epochs: 2`. Both gated to DPO-family tasks (`dpo`, `ipo`, or `preference` with `preference_loss in {dpo, ipo}`); transformers backend only. +- **Multi-objective preference loss** — define `training.preference_loss_weights: {dpo: 0.7, bco: 0.3}` to blend losses. 2–5 entries, weights must sum to 1. Schema-level surface ships now; live runtime weighted-loss combination deferred to v0.40.1 (`PreferenceTrainerWrapper.setup` raises `NotImplementedError` with a friendly message until then — same stub-then-live pattern as v0.27.0 MII / v0.37.0 multipack / v0.38.0 quant menu / v0.39.0 ReLoRA). +- **Net +118 tests** (4538 → 4656) across BCO trainer + dispatcher + β schedule math + ref-model regen TOCTOU + multi-objective schema bounds. ## Why Soup? @@ -86,6 +85,7 @@ soup init --template kto # KTO unpaired preference alignment soup init --template orpo # ORPO (no reference model needed) soup init --template simpo # SimPO length-normalized preference soup init --template ipo # IPO regularized preference +soup init --template bco # BCO binary classifier preference (v0.40.0) soup init --template rlhf # full RLHF pipeline (SFT→RM→PPO) soup init --template pretrain # continued pre-training on raw text soup init --template moe # MoE fine-tuning (ScatterMoE LoRA) @@ -697,6 +697,75 @@ training: quantization: 4bit ``` +## Preference Variety — BCO + Unified Dispatcher + KL Variants + +Five preference losses live behind one config knob. Pick a loss without +renaming your task, anneal β over training, and periodically refresh the +frozen reference. + +### BCO (Binary Classifier Optimization) + +Same input format as DPO; rows are split internally to TRL's BCO +unpaired schema (`{prompt, completion, label}`). + +```yaml +task: bco +data: + train: ./data/preferences.jsonl + format: dpo +training: + bco_beta: 0.1 +``` + +### Unified preference dispatcher + +Use `task: preference` + `training.preference_loss` to swap losses +without touching `task`. Hyperparameter sweeps over the loss type +itself become trivial. + +```yaml +task: preference +data: + train: ./data/preferences.jsonl + format: dpo +training: + preference_loss: dpo # or simpo, orpo, ipo, bco +``` + +Legacy `task: dpo` / `task: simpo` / etc. remain first-class — the +unified surface is additive. + +### KL-controlled DPO variants + +Anneal β over training, periodically refresh the reference model: + +```yaml +task: dpo # or task: preference + preference_loss: dpo, or task: ipo +training: + dpo_beta: 0.1 + dpo_beta_schedule: linear # linear | cosine | exponential + dpo_beta_end: 0.01 + dpo_ref_regen_epochs: 2 # copy student → ref model every 2 epochs +``` + +Both controls are gated to DPO-family tasks (`dpo`, `ipo`, or +`preference` with `preference_loss in {dpo, ipo}`); transformers +backend only. + +### Multi-objective preference loss (schema-only in v0.40.0) + +```yaml +task: preference +training: + preference_loss_weights: {dpo: 0.7, bco: 0.3} +``` + +Schema validates 2–5 entries summing to 1. Live runtime weighted-loss +combination is wired in v0.40.1; v0.40.0 fails fast with an actionable +`NotImplementedError` if you actually try to train (same stub-then-live +pattern as v0.27.0 MII / v0.37.0 multipack / v0.38.0 quant menu / +v0.39.0 ReLoRA). + ## GRPO Training (Reasoning) Train reasoning models with Group Relative Policy Optimization (DeepSeek-R1 style): diff --git a/SECURITY.md b/SECURITY.md index c997d1c..ef1c021 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -9,7 +9,8 @@ We provide security updates for the following versions: - **Versions older than 3 minor versions:** No support Example: -- v0.39.0-0.39.x -- Full support (latest) +- v0.40.0-0.40.x -- Full support (latest) +- v0.39.0-0.39.x -- Bug-fix support only - v0.38.0-0.38.x -- Bug-fix support only - v0.37.x and below -- No support @@ -142,6 +143,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.40.0 — Preference Variety**: New `task='bco'` (Binary Classifier Optimization) and `task='preference'` (unified dispatcher). New schema fields: `bco_beta` (gt=0), `preference_loss: Literal[dpo,simpo,orpo,ipo,bco]|None`, `preference_loss_weights: Optional[Dict[str,float]]`, `dpo_beta_schedule: Literal[linear,cosine,exponential]|None`, `dpo_beta_end: float, gt=0|None`, `dpo_ref_regen_epochs: int [1,1000]|None`. Cross-validators: `_validate_preference_dispatcher` rejects setting either `preference_loss` or `preference_loss_weights` outside `task='preference'` (closes ordering-dependency between Part B/D validators); `_validate_dpo_variants_supported_tasks` gates β-schedule + ref-regen to DPO-family tasks (`dpo`, `ipo`, or `preference` + `preference_loss in {dpo, ipo}`); rejected on mlx backend with distinct error message (matches v0.34.0 distinct-reason policy); `_validate_preference_loss_weights` enforces 2–5 entries (single-entry rejected with actionable message pointing at scalar `preference_loss`), key allowlist `{dpo, simpo, orpo, ipo, bco}`, explicit null-byte rejection on keys (matches v0.39.0 rank_pattern policy), per-value bounds `(0, 1]`, weights must sum to 1.0 (±1e-6), mutually exclusive with scalar `preference_loss`, rejected on mlx backend. `compute_beta_at_step` rejects `bool` on `step` and `total_steps` (project bool-as-int policy from v0.30.0). `BetaScheduleCallback` resolves `total_steps` lazily in `on_train_begin` so the schedule sees the real `state.max_steps` populated by HF Trainer (closes a first-cut silent-no-op bug where total_steps=0 emitted beta_end for every step). `RefModelRegenCallback._regenerate` uses `strict=True` on `load_state_dict` and logs at WARNING on mismatch (closes a first-cut silent partial-copy hazard where strict=False could produce a hybrid old-base + new-LoRA reference); epoch 0 regen suppressed (avoids copying untrained student); trainer `.beta` assignment swallow narrowed to `AttributeError` only. `PreferenceTrainerWrapper._make_inner_cfg` uses `model_copy` (not `model_dump`+`model_validate`) so re-validation never sees an inconsistent intermediate state and the caller's `cfg` is never mutated (mirrors v0.33.0 #47 immutability policy). `_split_dpo_rows_to_bco` skipped-row count emitted at DEBUG so production silent-degradation is inspectable (mirrors v0.33.0 #47 CrossDocCollator policy). Multi-objective live runtime weighted-loss combination is deferred to v0.40.1: `PreferenceTrainerWrapper.setup` raises `NotImplementedError` with a friendly message naming the deferred-version follow-up (mirrors v0.27.0 MII / v0.37.0 multipack / v0.38.0 quant menu / v0.39.0 ReLoRA stub-then-live pattern). Known limitation: `BCOTrainerWrapper._setup_transformers` still hardcodes `trust_remote_code=True` (v0.36.0 #63 known-gap family carry-over across non-SFT trainers). - **v0.39.0 — LoRA Quality**: `LoraConfig.init_strategy: Literal["random","pissa","olora"]` rejects unknown strategies; PiSSA + DoRA / VeRA combinations rejected at config-load. `model_validator(mode="before")` aligns `use_olora=True` → `init_strategy="olora"` via dict-copy (no caller mutation; matches v0.33.0 #47 immutability policy). `rank_pattern`/`alpha_pattern: Optional[Dict[str, int]]` capped at 256 keys × value (0, 1024], rejects `bool` (subclass of `int` — matches v0.30.0 `Candidate` policy), null bytes in keys, empty keys; cross-validator rejects with `use_vera=True`. `ReLoRAPolicy` is `@dataclass(frozen=True)` (post-construction mutation raises `FrozenInstanceError`); bounds: `steps ∈ [1, 1e7]`, `warmup_ratio ∈ [0, 1]`, `prune_ratio ∈ (0, 1)` (strict — prevents zero-everything footgun). `magnitude_prune_tensor` strict `0 < prune_ratio < 1` rejection, non-Tensor input raises `TypeError`, empty / single-element tensor short-circuits (avoids `kthvalue(_, 0)` runtime crash). `_validate_relora_supported_tasks` cross-validator rejects `relora_steps` with `task != "sft"` and `backend=mlx` with distinct error messages (matches v0.34.0 distinct-reason policy); multi-trainer expansion deferred to v0.39.1. `is_gemma4_model` uses a word-boundary regex (`(?:^|[^a-z0-9])gemma-?4(?:[^a-z0-9]|$)`) so `"ungemma4ed"` / `"my-gemma4ish"` no longer over-match; null-byte rejection on `model_name`. `apply_gemma4_clippable_patch` weight-copy fallback logs at DEBUG instead of silent random-init; the patch is gated by `is_gemma4_model(cfg.base)` in `sft.py` before invocation so non-Gemma4 trainings never traverse the module tree. `apply_surgical_patches` rejects empty / null-byte `model_name` with `ValueError`. `templates/load_template` containment: filename re-validated via `_validate_name` (rejects `..`/`/`/`\\`/null/empty); `os.path.realpath + os.path.commonpath` containment check on the resolved path against `_templates_dir()` so a tampered `manifest.json` cannot read files outside the package directory (mirrors v0.26.0 registry policy). Tampered-manifest `ValueError` from `_validate_name` caught and falls back to inline (no propagating exception). 256 KB file-size cap. Inline `TEMPLATES` carries an explicit deprecation comment pointing at the canonical YAML registry; `tests/test_templates_yaml.py` asserts byte-equality of all 16 inline ↔ YAML pairs to prevent silent drift. Planned removal: v0.41.0+. - **v0.38.0 — Quant Menu**: `TrainingConfig.quantization` Literal extended with `gptq` / `awq` / `hqq:1bit`..`hqq:8bit` (no `hqq:7bit` — HQQ doesn't support it) / `aqlm` / `eetq` / `mxfp4` / `fp8`; Pydantic rejects every other string at config-load. `validate_gptq_checkpoint` and `validate_awq_checkpoint` probe local paths for `quantize_config.json` / `quant_config.json`; HF repo IDs fall through; null-byte rejection + non-string `TypeError` on the ref. `_validate_prequantized_no_qat` rejects every pre-quantized format combined with `quantization_aware` (int8 QAT or `'fp8'`) — pre-quantized weights carry their own scale and QAT/FP8 prepare would silently corrupt them (mirrors LlamaFactory `quantization.py:117/199/211`). `_validate_bnb_quant_storage_only_with_4bit` rejects `bnb_4bit_quant_storage` on every non-BNB-4bit format (silent no-op otherwise); allowed dtypes: `Literal["uint8", "float16", "bfloat16", "float32"]`. `_validate_quant_menu_supported_tasks` restricts the new formats to `task='sft'` on `backend='transformers'` in v0.38.0 with distinct MLX-backend vs unsupported-task error messages (matches v0.34.0 distinct-reason policy). `check_quant_distributed_compat` hard-fails HQQ/EETQ/AQLM × {FSDP, ZeRO-3} (sourced from LlamaFactory `quantization.py:199/211` plus AQLM dequant constraints); warning-tier (not error) for BNB-4bit + FSDP without `bnb_4bit_quant_storage` so users see the silent perf cliff; unknown `quantization` raises `ValueError` (no silent pass) and the check is wired into `commands/train.py` startup. `parse_hqq_bits` rejects unsupported bit-rates and malformed `hqq:` strings before any kernel build. - **v0.37.0 — Multipack**: `validate_multipack_architecture` raises `ValueError` on unknown arch (loud-fail vs Axolotl's silent-miss footgun); 18-arch frozen allowlist (Llama 3.x / Qwen 2/3 / Mistral / Gemma 2/3 / Phi 3/4 / DeepSeek V2/V3 / Mixtral / Falcon / StableLM / SmolLM2). FFD packer caps at `_MAX_FFD_ITEMS=1_000_000` (algorithm is O(N²) worst-case — defence against adversarial dataset DoS); `bool` rejection on every numeric input (`max_len`, per-element `lengths`, `batch_max_len`, `batch_size`, `seed`, `max_seq_length`) matches v0.30.0+ `Candidate` policy; generator-input materialisation prevents silent empty-bin output when validation exhausts the iterator. `MultipackBatchSampler` rejects empty `lengths`, non-positive `batch_max_len`/`batch_size`, items larger than `batch_max_len`. `build_multipack_sampler_for_lengths` rejects `tcfg.batch_size="auto"` with actionable message (must be resolved upstream). `_validate_multipack_packing_exclusive` cross-validator on `TrainingConfig` rejects both `multipack` and `packing` set; `_validate_multipack_supported_tasks` on `SoupConfig` restricts multipack to `sft`/`pretrain` on `transformers` backend with distinct error messages for MLX backend vs unsupported task (matches v0.34.0 distinct-reason policy). `build_4d_attention_mask` caps allocations at `_MAX_MASK_ELEMENTS=2³¹` cells (~8GB float32) — defence against `max_length=1M` × `batch_size=8` OOM; `tag_sub_sequences` capped at `_MAX_BOUNDARY_SEGMENTS=1_000_000`. Mask builder rejects non-floating dtypes (was silent `np.finfo` ValueError), non-2D `seq_pos_ids`, negative segment IDs; padding (id=0) tokens are fully masked, including diagonal, so softmax is well-defined. `select_packing_strategy` rejects non-bool `flash_attn_available`. `JinjaTemplateAnalyzer` parses chat templates via `Environment.parse` only — never renders, so a crafted soup.yaml cannot trigger SSRF / filesystem reads; 128KB template cap, null-byte rejection, `TemplateSyntaxError` re-raised as `ValueError`. `DEFAULT_MESSAGE_FIELDS` is a `frozenset` (runtime-immutable); `JinjaTemplateAnalyzer.message_fields` returns a defensive copy. diff --git a/pyproject.toml b/pyproject.toml index 6cadefe..3f4456b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.39.0" +version = "0.40.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 117a548..8f7b2c5 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.39.0" +__version__ = "0.40.0" diff --git a/soup_cli/commands/sweep.py b/soup_cli/commands/sweep.py index bd2e1cf..604ae5a 100644 --- a/soup_cli/commands/sweep.py +++ b/soup_cli/commands/sweep.py @@ -299,6 +299,7 @@ def _set_nested_param(config_dict: dict, key: str, value) -> dict: "simpo_gamma": "training.simpo_gamma", "cpo_alpha": "training.cpo_alpha", "ipo_tau": "training.ipo_tau", + "bco_beta": "training.bco_beta", "loraplus_lr_ratio": "training.loraplus_lr_ratio", "use_dora": "training.lora.use_dora", "use_galore": "training.use_galore", @@ -392,6 +393,14 @@ def _run_single(base_cfg, params: dict, run_name: str, config_path: Path) -> dic from soup_cli.trainer.ipo import IPOTrainerWrapper trainer_wrapper = IPOTrainerWrapper(cfg, device=device) + elif cfg.task == "bco": + from soup_cli.trainer.bco import BCOTrainerWrapper + + trainer_wrapper = BCOTrainerWrapper(cfg, device=device) + elif cfg.task == "preference": + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + trainer_wrapper = PreferenceTrainerWrapper(cfg, device=device) elif cfg.task == "reward_model": from soup_cli.trainer.reward_model import RewardModelTrainerWrapper diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index 7d5cb32..d83fe62 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -700,6 +700,14 @@ def train( from soup_cli.trainer.ipo import IPOTrainerWrapper trainer_wrapper = IPOTrainerWrapper(cfg, **trainer_kwargs) + elif cfg.task == "bco": + from soup_cli.trainer.bco import BCOTrainerWrapper + + trainer_wrapper = BCOTrainerWrapper(cfg, **trainer_kwargs) + elif cfg.task == "preference": + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + trainer_wrapper = PreferenceTrainerWrapper(cfg, **trainer_kwargs) elif cfg.task == "reward_model": from soup_cli.trainer.reward_model import RewardModelTrainerWrapper diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index f972d3a..3379c1a 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -404,6 +404,56 @@ class TrainingConfig(BaseModel): ipo_tau: float = Field( default=0.1, gt=0, description="IPO tau — regularization strength" ) + # BCO-specific (Binary Classifier Optimization, v0.40.0 Part A) + bco_beta: float = Field( + default=0.1, gt=0, description="BCO beta — KL penalty coefficient" + ) + # Unified preference loss dispatcher (v0.40.0 Part B). + # Set when task='preference'. Legacy task strings ('dpo', 'simpo', ...) + # remain first-class and are unaffected. + preference_loss: Optional[Literal["dpo", "simpo", "orpo", "ipo", "bco"]] = Field( + default=None, + description=( + "Preference loss for task='preference'. One of: dpo, simpo, orpo, " + "ipo, bco. Mutually exclusive with task in {dpo, simpo, orpo, ipo, bco}." + ), + ) + # KL-controlled DPO variants (v0.40.0 Part C). + dpo_beta_schedule: Optional[Literal["linear", "cosine", "exponential"]] = Field( + default=None, + description=( + "Anneal DPO β over training. None = constant β (default). Requires " + "dpo_beta_end. DPO-family tasks only (dpo, ipo, preference+dpo)." + ), + ) + dpo_beta_end: Optional[float] = Field( + default=None, + gt=0, + description=( + "Target β at the end of training when dpo_beta_schedule is set. " + "Must be > 0. The starting β is dpo_beta." + ), + ) + dpo_ref_regen_epochs: Optional[int] = Field( + default=None, + ge=1, + le=1000, + description=( + "Replace the frozen ref model with the current student every N " + "epochs. None = never regen (default). DPO-family tasks only." + ), + ) + # Multi-objective preference loss (v0.40.0 Part D). + preference_loss_weights: Optional[dict[str, float]] = Field( + default=None, + description=( + "Weighted blend of preference losses, e.g. {'dpo': 0.7, 'bco': 0.3}. " + "Each weight ∈ (0, 1]; weights must sum to 1.0 (±1e-6). All keys " + "must be members of {dpo, simpo, orpo, ipo, bco}. Requires " + "task='preference'; mutually exclusive with preference_loss (the " + "scalar form). Capped at 5 components (the supported set)." + ), + ) # GRPO-specific grpo_beta: float = Field( default=0.1, gt=0, description="GRPO beta — KL penalty coefficient" @@ -900,7 +950,7 @@ class SoupConfig(BaseModel): base: str = Field(..., description="Base model name or path (HF model ID)") task: Literal[ "sft", "dpo", "grpo", "ppo", "reward_model", "kto", "orpo", "simpo", "ipo", - "pretrain", "embedding", + "bco", "preference", "pretrain", "embedding", ] = Field( default="sft", description="Training task type" ) @@ -1056,6 +1106,154 @@ class SoupConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_preference_dispatcher(self) -> "SoupConfig": + """v0.40.0 Part B — task='preference' requires preference_loss OR + preference_loss_weights (Part D). Setting either field outside + task='preference' is rejected to keep the two config surfaces disjoint. + """ + loss = self.training.preference_loss + weights = self.training.preference_loss_weights + if self.task == "preference": + if loss is None and weights is None: + raise ValueError( + "task='preference' requires either training.preference_loss " + "(in {dpo, simpo, orpo, ipo, bco}) or " + "training.preference_loss_weights (multi-objective dict)." + ) + return self + if loss is not None: + raise ValueError( + f"training.preference_loss={loss!r} is only meaningful for " + f"task='preference'; got task={self.task!r}. Either set " + "task='preference' or remove preference_loss." + ) + if weights is not None: + raise ValueError( + "training.preference_loss_weights is only meaningful for " + f"task='preference'; got task={self.task!r}. Either set " + "task='preference' or remove preference_loss_weights." + ) + return self + + @model_validator(mode="after") + def _validate_dpo_variants_supported_tasks(self) -> "SoupConfig": + """v0.40.0 Part C — β-schedule + ref-model regen are DPO-family only. + + Allowed: task in {dpo, ipo} OR (task='preference' AND + preference_loss in {dpo, ipo}). Rejected on mlx backend. + """ + tcfg = self.training + sched = tcfg.dpo_beta_schedule + end = tcfg.dpo_beta_end + regen = tcfg.dpo_ref_regen_epochs + if sched is None and end is None and regen is None: + return self + # End/schedule mutual requirement. + if sched is not None and end is None: + raise ValueError( + "dpo_beta_schedule requires dpo_beta_end (the target β at " + "end of training). Set dpo_beta_end or remove dpo_beta_schedule." + ) + if end is not None and sched is None: + raise ValueError( + "dpo_beta_end requires dpo_beta_schedule. Set " + "dpo_beta_schedule in {linear, cosine, exponential} or " + "remove dpo_beta_end." + ) + # Backend gate. + if self.backend == "mlx": + raise ValueError( + "DPO variants (dpo_beta_schedule / dpo_ref_regen_epochs) are " + "not supported on the mlx backend in v0.40.0 (TRL trainer " + "internals required). Use backend='transformers'." + ) + # Task gate — DPO family only. + family_ok = self.task in ("dpo", "ipo") or ( + self.task == "preference" + and tcfg.preference_loss in ("dpo", "ipo") + ) + if not family_ok: + raise ValueError( + f"DPO variants (dpo_beta_schedule / dpo_ref_regen_epochs) " + f"require task in {{dpo, ipo}} or task='preference' with " + f"preference_loss in {{dpo, ipo}}; got task={self.task!r}, " + f"preference_loss={tcfg.preference_loss!r}." + ) + return self + + @model_validator(mode="after") + def _validate_preference_loss_weights(self) -> "SoupConfig": + """v0.40.0 Part D — multi-objective preference_loss_weights gate. + + Allowed: task='preference' only. Mutually exclusive with the scalar + preference_loss. Validates value bounds + sum-to-1 + key allowlist. + """ + tcfg = self.training + weights = tcfg.preference_loss_weights + if weights is None: + return self + if not isinstance(weights, dict): + raise ValueError( + "preference_loss_weights must be a dict, e.g. " + "{'dpo': 0.7, 'bco': 0.3}." + ) + if self.task != "preference": + raise ValueError( + "preference_loss_weights requires task='preference'; got " + f"task={self.task!r}." + ) + if tcfg.preference_loss is not None: + raise ValueError( + "preference_loss_weights and (scalar) preference_loss are " + "mutually exclusive — pick one." + ) + if self.backend == "mlx": + raise ValueError( + "preference_loss_weights is not supported on the mlx backend " + "in v0.40.0. Use backend='transformers'." + ) + if not (2 <= len(weights) <= 5): + raise ValueError( + f"preference_loss_weights must have between 2 and 5 entries " + "(single-entry blends are equivalent to the scalar " + "preference_loss field; use that instead); got " + f"{len(weights)}." + ) + allowed = {"dpo", "simpo", "orpo", "ipo", "bco"} + for key in weights: + if not isinstance(key, str): + raise ValueError( + f"preference_loss_weights keys must be strings; " + f"got {type(key).__name__}." + ) + if "\x00" in key: + raise ValueError( + "preference_loss_weights keys cannot contain null bytes." + ) + unknown = set(weights.keys()) - allowed + if unknown: + raise ValueError( + f"preference_loss_weights keys must be in {sorted(allowed)}; " + f"unknown: {sorted(unknown)}." + ) + for key, value in weights.items(): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"preference_loss_weights[{key!r}] must be a number; " + f"got {type(value).__name__}." + ) + if not (0 < float(value) <= 1): + raise ValueError( + f"preference_loss_weights[{key!r}]={value!r} must be in (0, 1]." + ) + total = sum(float(v) for v in weights.values()) + if abs(total - 1.0) > 1e-6: + raise ValueError( + f"preference_loss_weights must sum to 1.0 (±1e-6); got {total!r}." + ) + return self + @model_validator(mode="after") def _validate_mlx_task_support(self) -> "SoupConfig": """MLX backend only supports sft, dpo, and grpo tasks (v0.25.0). @@ -1275,6 +1473,36 @@ training: quantization: 4bit orpo_beta: 0.1 +output: ./output +""", + "bco": """# Soup template: BCO (Binary Classifier Optimization) +# Preference alignment via binary classification of chosen vs rejected. +# +# Data format (JSONL) — same as DPO: +# {"prompt": "What is 2+2?", "chosen": "4", "rejected": "Fish"} + +base: meta-llama/Llama-3.1-8B-Instruct +task: bco +# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]' + +data: + train: ./data/preference_train.jsonl + format: dpo + val_split: 0.1 + max_length: 2048 + +training: + epochs: 3 + lr: 1e-5 + batch_size: auto + gradient_accumulation_steps: 4 + lora: + r: 64 + alpha: 16 + target_modules: auto + quantization: 4bit + bco_beta: 0.1 + output: ./output """, "simpo": """# Soup template: SimPO (Simple Preference Optimization) diff --git a/soup_cli/templates/bco.yaml b/soup_cli/templates/bco.yaml new file mode 100644 index 0000000..a166953 --- /dev/null +++ b/soup_cli/templates/bco.yaml @@ -0,0 +1,29 @@ +# Soup template: BCO (Binary Classifier Optimization) +# Preference alignment via binary classification of chosen vs rejected. +# +# Data format (JSONL) — same as DPO: +# {"prompt": "What is 2+2?", "chosen": "4", "rejected": "Fish"} + +base: meta-llama/Llama-3.1-8B-Instruct +task: bco +# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]' + +data: + train: ./data/preference_train.jsonl + format: dpo + val_split: 0.1 + max_length: 2048 + +training: + epochs: 3 + lr: 1e-5 + batch_size: auto + gradient_accumulation_steps: 4 + lora: + r: 64 + alpha: 16 + target_modules: auto + quantization: 4bit + bco_beta: 0.1 + +output: ./output diff --git a/soup_cli/templates/manifest.json b/soup_cli/templates/manifest.json index 0ef34f3..288904e 100644 --- a/soup_cli/templates/manifest.json +++ b/soup_cli/templates/manifest.json @@ -1,6 +1,7 @@ { "templates": { "audio": "audio.yaml", + "bco": "bco.yaml", "chat": "chat.yaml", "code": "code.yaml", "embedding": "embedding.yaml", diff --git a/soup_cli/trainer/bco.py b/soup_cli/trainer/bco.py new file mode 100644 index 0000000..e40cbfe --- /dev/null +++ b/soup_cli/trainer/bco.py @@ -0,0 +1,327 @@ +"""BCO (Binary Classifier Optimization) trainer — wraps trl.BCOTrainer. + +v0.40.0 Part A. Mirrors the ORPO / SimPO / IPO wrapper pattern. + +Data format (same as DPO): {"prompt": ..., "chosen": ..., "rejected": ...}. +Each row is internally split into two rows for TRL's BCOTrainer +(``{"prompt", "completion", "label"}``): one chosen-as-completion with +``label=True`` and one rejected-as-completion with ``label=False``. +""" + +from __future__ import annotations + +import math +import time +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +from rich.console import Console + +from soup_cli.config.schema import SoupConfig +from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name + +if TYPE_CHECKING: + from soup_cli.config.schema import TrainingConfig # noqa: F401 + +console = Console() + + +def _split_dpo_rows_to_bco(rows: list[dict]) -> list[dict]: + """Convert DPO-shaped rows to TRL BCO unpaired format. + + Each input row produces two output rows: one with ``label=True`` + (chosen) and one with ``label=False`` (rejected). Rows missing any + required field are skipped; the count is emitted at DEBUG so production + silent-degradation is inspectable (mirrors v0.33.0 #47 CrossDocCollator + DEBUG-on-fallback policy). + """ + out: list[dict] = [] + skipped = 0 + for row in rows: + prompt = row.get("prompt") + chosen = row.get("chosen") + rejected = row.get("rejected") + if prompt is None or chosen is None or rejected is None: + skipped += 1 + continue + out.append({"prompt": prompt, "completion": chosen, "label": True}) + out.append({"prompt": prompt, "completion": rejected, "label": False}) + if skipped: + import logging + + logging.getLogger(__name__).debug( + "BCO: skipped %d row(s) missing prompt/chosen/rejected.", skipped, + ) + return out + + +class BCOTrainerWrapper: + """High-level wrapper for BCO training from SoupConfig. + + BCO uses a binary classifier objective (chosen vs rejected). Soup + accepts the DPO data format and adapts it to TRL's BCOTrainer + unpaired ``{prompt, completion, label}`` schema. + """ + + def __init__( + self, + config: SoupConfig, + device: str = "cuda", + report_to: str = "none", + deepspeed_config: Optional[str] = None, + fsdp_config: Optional[dict] = None, + ): + self.config = config + self.device = device + self.report_to = report_to + self.deepspeed_config = deepspeed_config + self.fsdp_config = fsdp_config + self.model = None + self.tokenizer = None + self.trainer = None + self._output_dir: Optional[str] = None + + def setup(self, dataset: dict) -> None: + """Load model, tokenizer, apply LoRA, create BCO trainer.""" + from datasets import Dataset + from trl import BCOConfig, BCOTrainer + + from soup_cli.trainer.sft import _enable_hf_transfer_progress + + _enable_hf_transfer_progress() + + cfg = self.config + tcfg = cfg.training + use_unsloth = cfg.backend == "unsloth" + + if use_unsloth: + self._setup_unsloth(cfg, tcfg) + else: + self._setup_transformers(cfg, tcfg) + + trainable, total = self.model.get_nb_trainable_parameters() + pct = 100 * trainable / total + console.print( + f"[green]LoRA applied:[/] {trainable:,} trainable" + f" / {total:,} total ({pct:.2f}%)" + ) + + # --- Batch size --- + batch_size = tcfg.batch_size + if batch_size == "auto": + 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( + 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, + ) + # BCO sees ~2x rows per sample (chosen + rejected) → halve. + batch_size = max(1, batch_size // 2) + console.print(f"[green]Auto batch size (BCO):[/] {batch_size}") + + # --- Dataset (split DPO rows into BCO unpaired) --- + train_rows = _split_dpo_rows_to_bco(dataset["train"]) + if not train_rows: + raise ValueError( + "BCO training: dataset['train'] produced no usable rows. " + "Each row must contain 'prompt', 'chosen', and 'rejected'." + ) + train_ds = Dataset.from_list(train_rows) + eval_ds = None + if "val" in dataset and dataset["val"]: + val_rows = _split_dpo_rows_to_bco(dataset["val"]) + if val_rows: + eval_ds = Dataset.from_list(val_rows) + + # --- Output dir --- + output_dir = Path(cfg.output) + if cfg.experiment_name: + output_dir = output_dir / cfg.experiment_name + output_dir.mkdir(parents=True, exist_ok=True) + + # --- Calculate warmup steps from ratio --- + total_steps = ( + math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps) + * tcfg.epochs + ) + warmup_steps = int(total_steps * tcfg.warmup_ratio) + + # --- BCO config --- + bco_config = BCOConfig( + output_dir=str(output_dir), + num_train_epochs=tcfg.epochs, + per_device_train_batch_size=batch_size, + gradient_accumulation_steps=tcfg.gradient_accumulation_steps, + learning_rate=tcfg.lr, + warmup_steps=warmup_steps, + weight_decay=tcfg.weight_decay, + max_grad_norm=tcfg.max_grad_norm, + optim=tcfg.optimizer, + lr_scheduler_type=tcfg.scheduler, + logging_steps=tcfg.logging_steps, + save_steps=tcfg.save_steps, + save_total_limit=3, + bf16=self.device == "cuda", + report_to=self.report_to, + remove_unused_columns=False, + deepspeed=self.deepspeed_config, + **(self.fsdp_config or {}), + beta=tcfg.bco_beta, + max_length=cfg.data.max_length, + max_prompt_length=cfg.data.max_length // 2, + **( + {"neftune_noise_alpha": tcfg.neftune_alpha} + if tcfg.neftune_alpha is not None + else {} + ), + ) + + # --- Trainer --- + self.trainer = BCOTrainer( + model=self.model, + args=bco_config, + train_dataset=train_ds, + eval_dataset=eval_ds, + processing_class=self.tokenizer, + ) + + self._output_dir = str(output_dir) + + def _setup_transformers(self, cfg: SoupConfig, tcfg: "TrainingConfig") -> None: + """Load model via standard transformers + peft pipeline.""" + from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training + from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig + + console.print(f"[dim]Loading tokenizer: {cfg.base}[/]") + self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True) + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + bnb_config = None + if tcfg.quantization == "4bit": + from soup_cli.utils.gpu import get_compute_dtype + + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=get_compute_dtype(), + bnb_4bit_use_double_quant=True, + ) + elif tcfg.quantization == "8bit": + bnb_config = BitsAndBytesConfig(load_in_8bit=True) + + console.print(f"[dim]Loading model: {cfg.base}[/]") + dev_map = "cpu" if self.device == "cpu" else "auto" + model_kwargs = {"trust_remote_code": True, "device_map": dev_map} + if bnb_config: + model_kwargs["quantization_config"] = bnb_config + + self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs) + + if tcfg.quantization in ("4bit", "8bit"): + self.model = prepare_model_for_kbit_training(self.model) + + target_modules = tcfg.lora.target_modules + if target_modules == "auto": + target_modules = None + + lora_config = LoraConfig( + r=tcfg.lora.r, + lora_alpha=tcfg.lora.alpha, + lora_dropout=tcfg.lora.dropout, + target_modules=target_modules, + task_type=TaskType.CAUSAL_LM, + bias="none", + use_dora=tcfg.lora.use_dora, + use_rslora=tcfg.lora.use_rslora, + ) + self.model = get_peft_model(self.model, lora_config) + + # QAT — int8 only; "fp8" handled by apply_v028_speed_memory below. + if tcfg.quantization_aware and tcfg.quantization_aware != "fp8": + from soup_cli.utils.qat import prepare_model_for_qat + + self.model = prepare_model_for_qat(self.model) + + # v0.35.0 #60 — multi-trainer wiring of v0.28.0 speed/memory features. + from soup_cli.utils.v028_features import apply_v028_speed_memory + apply_v028_speed_memory( + model=self.model, tcfg=tcfg, base_model=cfg.base, + console=console, device=self.device, backend=cfg.backend, + ) + + def _setup_unsloth(self, cfg: SoupConfig, tcfg: "TrainingConfig") -> None: + """Load model via unsloth FastLanguageModel.""" + from soup_cli.utils.unsloth import load_model_and_tokenizer + + console.print(f"[dim]Loading model via [bold]unsloth[/]: {cfg.base}[/]") + self.model, self.tokenizer = load_model_and_tokenizer( + model_name=cfg.base, + max_seq_length=cfg.data.max_length, + quantization=tcfg.quantization, + lora_r=tcfg.lora.r, + lora_alpha=tcfg.lora.alpha, + lora_dropout=tcfg.lora.dropout, + target_modules=tcfg.lora.target_modules, + ) + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + def train( + self, + display: Optional[object] = None, + tracker: Optional[object] = None, + run_id: str = "", + resume_from_checkpoint: Optional[str] = None, + ) -> dict: + """Run BCO training and return results summary.""" + if self.trainer is None: + raise RuntimeError( + "BCOTrainerWrapper.train() called before setup(). " + "Call setup(dataset) first." + ) + start = time.time() + + if display: + from soup_cli.monitoring.callback import SoupTrainerCallback + + self.trainer.add_callback( + SoupTrainerCallback( + display, tracker=tracker, run_id=run_id, + loss_watchdog=self.config.training.loss_watchdog, + loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, + loss_watchdog_patience=self.config.training.loss_watchdog_patience, + ) + ) + + from soup_cli.utils.v028_features import activation_offloading_context + + with activation_offloading_context( + self.config.training, self._output_dir, + ): + self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) + duration = time.time() - start + + self.trainer.save_model(self._output_dir) + self.tokenizer.save_pretrained(self._output_dir) + + logs = self.trainer.state.log_history + train_losses = [entry["loss"] for entry in logs if "loss" in entry] + + hours = int(duration // 3600) + minutes = int((duration % 3600) // 60) + duration_str = f"{hours}h {minutes}m" if hours > 0 else f"{minutes}m" + + return { + "initial_loss": train_losses[0] if train_losses else 0, + "final_loss": train_losses[-1] if train_losses else 0, + "duration": duration_str, + "duration_secs": duration, + "output_dir": self._output_dir, + "total_steps": self.trainer.state.global_step, + } diff --git a/soup_cli/trainer/dpo.py b/soup_cli/trainer/dpo.py index 6556d91..b316468 100644 --- a/soup_cli/trainer/dpo.py +++ b/soup_cli/trainer/dpo.py @@ -247,6 +247,23 @@ class DPOTrainerWrapper: ) ) + # v0.40.0 Part C — DPO variants (β-schedule + ref-model regen). + # total_steps=0 is the lazy-resolve sentinel; BetaScheduleCallback's + # on_train_begin reads state.max_steps once HF Trainer has populated it. + from soup_cli.utils.dpo_variants import build_dpo_variant_callbacks + + tcfg = self.config.training + variant_cbs = build_dpo_variant_callbacks( + beta_start=tcfg.dpo_beta, + beta_end=tcfg.dpo_beta_end, + schedule=tcfg.dpo_beta_schedule, + total_steps=0, + ref_regen_epochs=tcfg.dpo_ref_regen_epochs, + ) + for cb in variant_cbs: + cb.attach(self.trainer) + self.trainer.add_callback(cb) + from soup_cli.utils.v028_features import activation_offloading_context with activation_offloading_context( diff --git a/soup_cli/trainer/ipo.py b/soup_cli/trainer/ipo.py index 3591f35..919f2bb 100644 --- a/soup_cli/trainer/ipo.py +++ b/soup_cli/trainer/ipo.py @@ -253,6 +253,23 @@ class IPOTrainerWrapper: ) ) + # v0.40.0 Part C — DPO variants (β-schedule + ref-model regen). + # total_steps=0 is the lazy-resolve sentinel; BetaScheduleCallback's + # on_train_begin reads state.max_steps once HF Trainer has populated it. + from soup_cli.utils.dpo_variants import build_dpo_variant_callbacks + + tcfg = self.config.training + variant_cbs = build_dpo_variant_callbacks( + beta_start=tcfg.dpo_beta, + beta_end=tcfg.dpo_beta_end, + schedule=tcfg.dpo_beta_schedule, + total_steps=0, + ref_regen_epochs=tcfg.dpo_ref_regen_epochs, + ) + for cb in variant_cbs: + cb.attach(self.trainer) + self.trainer.add_callback(cb) + from soup_cli.utils.v028_features import activation_offloading_context with activation_offloading_context( diff --git a/soup_cli/trainer/preference.py b/soup_cli/trainer/preference.py new file mode 100644 index 0000000..11e0871 --- /dev/null +++ b/soup_cli/trainer/preference.py @@ -0,0 +1,154 @@ +"""Unified preference loss dispatcher (v0.40.0 Part B). + +Provides a single entry-point for preference-style training that routes to +the per-loss wrappers (DPO / SimPO / ORPO / IPO / BCO). The legacy +``task: dpo``, ``task: simpo``, ... config forms remain first-class — +this module is purely additive. + +Usage: + + task: preference + training: + preference_loss: dpo # or simpo, orpo, ipo, bco + +The dispatcher constructs a temporary :class:`SoupConfig` view with the +matching legacy task so the underlying TRL trainer is happy, then forwards +``setup`` / ``train`` to the right wrapper. +""" + +from __future__ import annotations + +from typing import Optional + +from soup_cli.config.schema import SoupConfig + +_SUPPORTED_LOSSES: frozenset[str] = frozenset({"dpo", "simpo", "orpo", "ipo", "bco"}) + + +def is_multi_objective_preference(cfg: SoupConfig) -> bool: + """True when the config asks for a weighted blend of preference losses. + + v0.40.0 Part D — schema-level surface only; live runtime weighted + combination is deferred to v0.40.1 (mirrors the project's + stub-then-live pattern from v0.27.0 MII / v0.37.0 multipack / + v0.38.0 quant menu / v0.39.0 ReLoRA). + """ + weights = cfg.training.preference_loss_weights + return weights is not None and len(weights) >= 1 + + +def get_loss_weights(cfg: SoupConfig) -> Optional[dict]: + """Return a defensive copy of ``training.preference_loss_weights``.""" + weights = cfg.training.preference_loss_weights + if weights is None: + return None + return dict(weights) + + +def resolve_preference_loss(cfg: SoupConfig) -> Optional[str]: + """Return the preference loss name for a config, or ``None`` if N/A. + + - ``task='preference'`` → ``training.preference_loss``. + - Legacy ``task in {dpo, simpo, orpo, ipo, bco}`` → that name (identity). + - Anything else → ``None``. + """ + if cfg.task == "preference": + return cfg.training.preference_loss + if cfg.task in _SUPPORTED_LOSSES: + return cfg.task + return None + + +def _make_inner_cfg(cfg: SoupConfig, loss: str) -> SoupConfig: + """Build a SoupConfig view with task= so the inner wrapper can run. + + Uses ``model_copy`` to round-trip without re-running validators on an + intermediate inconsistent state. The caller's cfg is never mutated. + """ + inner_training = cfg.training.model_copy( + update={"preference_loss": None, "preference_loss_weights": None} + ) + return cfg.model_copy(update={"task": loss, "training": inner_training}) + + +class PreferenceTrainerWrapper: + """Dispatcher wrapper for ``task: preference`` configs. + + Forwards to DPO / SimPO / ORPO / IPO / BCO wrappers based on + ``training.preference_loss``. + """ + + def __init__( + self, + config: SoupConfig, + device: str = "cuda", + report_to: str = "none", + deepspeed_config: Optional[str] = None, + fsdp_config: Optional[dict] = None, + ): + self.config = config + self.device = device + self.report_to = report_to + self.deepspeed_config = deepspeed_config + self.fsdp_config = fsdp_config + self._inner = None + + def _build_inner(self): + loss = self.config.training.preference_loss + if loss not in _SUPPORTED_LOSSES: + raise ValueError( + f"Unknown preference_loss={loss!r}. " + f"Supported: {sorted(_SUPPORTED_LOSSES)}" + ) + inner_cfg = _make_inner_cfg(self.config, loss) + kwargs = { + "device": self.device, + "report_to": self.report_to, + "deepspeed_config": self.deepspeed_config, + "fsdp_config": self.fsdp_config, + } + if loss == "dpo": + from soup_cli.trainer.dpo import DPOTrainerWrapper + return DPOTrainerWrapper(inner_cfg, **kwargs) + if loss == "simpo": + from soup_cli.trainer.simpo import SimPOTrainerWrapper + return SimPOTrainerWrapper(inner_cfg, **kwargs) + if loss == "orpo": + from soup_cli.trainer.orpo import ORPOTrainerWrapper + return ORPOTrainerWrapper(inner_cfg, **kwargs) + if loss == "ipo": + from soup_cli.trainer.ipo import IPOTrainerWrapper + return IPOTrainerWrapper(inner_cfg, **kwargs) + # BCO — last branch by allowlist exhaustion. + from soup_cli.trainer.bco import BCOTrainerWrapper + return BCOTrainerWrapper(inner_cfg, **kwargs) + + def setup(self, dataset: dict) -> None: + # v0.40.0 Part D — schema-level multi-objective shipped; live + # weighted-loss combination deferred to v0.40.1 (TRL preference + # trainers do not expose a clean compute_loss override hook; + # subclassing each one is tracked separately). + if is_multi_objective_preference(self.config): + raise NotImplementedError( + "preference_loss_weights (multi-objective preference loss) " + "is config-level only in v0.40.0. Live runtime weighted " + "combination is deferred to v0.40.1 (subclassing TRL " + "preference trainers to override compute_loss). For now, " + "use the scalar 'preference_loss' field instead." + ) + if self._inner is None: + self._inner = self._build_inner() + self._inner.setup(dataset) + + def train(self, **kwargs) -> dict: + if self._inner is None: + raise RuntimeError( + "PreferenceTrainerWrapper.train() called before setup(). " + "Call setup(dataset) first." + ) + return self._inner.train(**kwargs) + + @property + def trainer(self): + """Expose the underlying HF Trainer for callbacks (HF push, eval-gate).""" + return getattr(self._inner, "trainer", None) diff --git a/soup_cli/utils/dpo_variants.py b/soup_cli/utils/dpo_variants.py new file mode 100644 index 0000000..3a37366 --- /dev/null +++ b/soup_cli/utils/dpo_variants.py @@ -0,0 +1,250 @@ +"""KL-controlled DPO variants (v0.40.0 Part C). + + +Two opt-in controls for DPO-family preference training: + +1. **β schedule** — anneal the DPO ``beta`` coefficient over training. + Three shapes: linear, cosine (1/2 (1 + cos(pi t)) ramp), exponential + (geometric decay between ``beta_start`` and ``beta_end``). + +2. **Reference-model regeneration** — every ``every_n_epochs``, replace + the frozen ref-model weights with a deep-copy of the current student. + Useful for self-improving loops where the policy quickly outpaces + the original reference. + +Both helpers are duck-typed callbacks (no ``transformers`` import at +module scope) so they cost nothing on a torch-less interpreter and stay +unit-testable on CI without GPUs. +""" + +from __future__ import annotations + +import math +from typing import Optional + +# Allowed schedule shapes, kept as a frozenset for runtime immutability. +SUPPORTED_SCHEDULES: frozenset[str] = frozenset({"linear", "cosine", "exponential"}) + + +def _validate_finite_positive(name: str, value: float) -> None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{name} must be a number; got {type(value).__name__}") + fvalue = float(value) + if math.isnan(fvalue) or math.isinf(fvalue): + raise ValueError(f"{name}={value!r} must be finite") + if fvalue <= 0: + raise ValueError(f"{name}={value!r} must be > 0") + + +def compute_beta_at_step( + beta_start: float, + beta_end: float, + step: int, + total_steps: int, + schedule: str, +) -> float: + """Return β at ``step`` for the chosen schedule. + + Endpoint contract (matches HF lr_scheduler convention): + step ≤ 0 → ``beta_start`` + step ≥ total_steps → ``beta_end`` + total_steps == 0 → ``beta_end`` (degenerate case; nothing to anneal) + """ + _validate_finite_positive("beta_start", beta_start) + _validate_finite_positive("beta_end", beta_end) + if isinstance(total_steps, bool) or not isinstance(total_steps, int): + raise ValueError( + f"total_steps must be int; got {type(total_steps).__name__}" + ) + if isinstance(step, bool) or not isinstance(step, int): + # Defence-in-depth: bool is a subclass of int. Project policy + # (v0.30.0 Candidate) rejects bool for int fields. + raise ValueError(f"step must be int; got {type(step).__name__}") + if total_steps < 0: + raise ValueError(f"total_steps={total_steps!r} must be >= 0") + if schedule not in SUPPORTED_SCHEDULES: + raise ValueError( + f"Unknown schedule={schedule!r}; supported: {sorted(SUPPORTED_SCHEDULES)}" + ) + if total_steps == 0: + return float(beta_end) + if step <= 0: + return float(beta_start) + if step >= total_steps: + return float(beta_end) + progress = step / total_steps # 0..1 + if schedule == "linear": + return float(beta_start + (beta_end - beta_start) * progress) + if schedule == "cosine": + # Cosine ramp: 1/2 (1 + cos(pi * progress)) goes 1→0 over [0,1]. + weight = 0.5 * (1.0 + math.cos(math.pi * progress)) + return float(beta_end + (beta_start - beta_end) * weight) + # exponential — geometric interpolation; both endpoints are > 0 by guard. + log_start = math.log(beta_start) + log_end = math.log(beta_end) + return float(math.exp(log_start + (log_end - log_start) * progress)) + + +class BetaScheduleCallback: + """Duck-typed ``TrainerCallback``: updates ``trainer.beta`` per step. + + The HF ``TrainerCallback`` signature is matched without importing + ``transformers`` at module scope so this object is constructible in + a torch-less environment for unit testing. + + Use ``attach(trainer)`` once, then add the callback to the trainer. + Without ``attach``, ``on_step_begin`` is a no-op (defence-in-depth). + """ + + def __init__( + self, + beta_start: float, + beta_end: float, + total_steps: int, + schedule: str, + ) -> None: + # Validate immediately so callers don't see a deferred crash mid-train. + _validate_finite_positive("beta_start", beta_start) + _validate_finite_positive("beta_end", beta_end) + if schedule not in SUPPORTED_SCHEDULES: + raise ValueError( + f"Unknown schedule={schedule!r}; supported: " + f"{sorted(SUPPORTED_SCHEDULES)}" + ) + self.beta_start = float(beta_start) + self.beta_end = float(beta_end) + self.total_steps = int(total_steps) + self.schedule = schedule + self._trainer = None + + def attach(self, trainer) -> None: + self._trainer = trainer + + def on_train_begin(self, args, state, control, **kwargs) -> None: + # HF Trainer populates state.max_steps inside _inner_training_loop, + # which runs before on_train_begin. Resolve total_steps lazily from + # the live trainer state so the wrappers don't need to compute it + # ahead of time (the wrappers' pre-train computation is racy and + # was returning 0 in v0.40.0 first cut). + if self.total_steps <= 0: + live_max = int(getattr(state, "max_steps", 0) or 0) + if live_max > 0: + self.total_steps = live_max + + def on_step_begin(self, args, state, control, **kwargs) -> None: + if self._trainer is None: + return + if self.total_steps <= 0: + # Couldn't resolve a meaningful horizon; skip the schedule + # rather than emit beta_end for every step. + return + new_beta = compute_beta_at_step( + beta_start=self.beta_start, + beta_end=self.beta_end, + step=int(getattr(state, "global_step", 0)), + total_steps=self.total_steps, + schedule=self.schedule, + ) + # TRL DPOTrainer exposes .beta directly. Narrow the swallow to + # AttributeError only — a TypeError here would indicate a real bug + # in compute_beta_at_step that should surface in tests, not be hidden. + try: + self._trainer.beta = new_beta + except AttributeError: + return + + +class RefModelRegenCallback: + """Duck-typed callback: deep-copy student weights into ref_model on epoch. + + On every Nth epoch (1-indexed; epoch 0 is skipped to avoid copying + untrained weights), copies the current ``trainer.model`` state_dict + into ``trainer.ref_model``. Falls back to a no-op when ``ref_model`` + is missing (some preference trainers, e.g. ORPO, are reference-free). + """ + + def __init__(self, every_n_epochs: int) -> None: + if isinstance(every_n_epochs, bool) or not isinstance(every_n_epochs, int): + raise TypeError( + f"every_n_epochs must be int; got {type(every_n_epochs).__name__}" + ) + if every_n_epochs < 1: + raise ValueError( + f"every_n_epochs={every_n_epochs!r} must be >= 1" + ) + self.every_n_epochs = every_n_epochs + self._trainer = None + self.regen_count = 0 + + def attach(self, trainer) -> None: + self._trainer = trainer + + def on_epoch_end(self, args, state, control, **kwargs) -> None: + if self._trainer is None: + return + epoch = float(getattr(state, "epoch", 0.0)) + # Skip epoch 0 entirely — copying untrained student is a footgun. + # round() handles HF Trainer's float epoch counters (e.g. 1.999...). + epoch_int = int(round(epoch)) + if epoch_int < 1: + return + if epoch_int % self.every_n_epochs != 0: + return + self._regenerate() + + def _regenerate(self) -> None: + trainer = self._trainer + if trainer is None: + return + ref_model = getattr(trainer, "ref_model", None) + student = getattr(trainer, "model", None) + if ref_model is None or student is None: + return + try: + state_dict = student.state_dict() + except AttributeError: + return + # strict=True so a key/shape mismatch (e.g. PEFT-wrapped student vs + # bare ref) surfaces loudly instead of silently producing a + # half-copied reference. The except below logs at WARNING and the + # callback continues — this is an optimisation, not a safety gate, + # but operators should know when it failed. + try: + ref_model.load_state_dict(state_dict, strict=True) + self.regen_count += 1 + except (RuntimeError, TypeError) as exc: + import logging + + logging.getLogger(__name__).warning( + "RefModelRegenCallback: load_state_dict failed (%s); ref " + "model not updated this epoch.", type(exc).__name__, + ) + return + + +def build_dpo_variant_callbacks( + *, + beta_start: float, + beta_end: Optional[float], + schedule: Optional[str], + total_steps: int, + ref_regen_epochs: Optional[int], +) -> list: + """Build the DPO-variant callback list for a given config slice. + + Returns an empty list when no variants are enabled. Callers should + ``attach(trainer)`` and ``trainer.add_callback(cb)`` for each entry. + """ + callbacks: list = [] + if schedule is not None and beta_end is not None: + callbacks.append( + BetaScheduleCallback( + beta_start=beta_start, + beta_end=beta_end, + total_steps=total_steps, + schedule=schedule, + ) + ) + if ref_regen_epochs is not None: + callbacks.append(RefModelRegenCallback(every_n_epochs=ref_regen_epochs)) + return callbacks diff --git a/soup_cli/utils/v028_features.py b/soup_cli/utils/v028_features.py index 6a78876..e272220 100644 --- a/soup_cli/utils/v028_features.py +++ b/soup_cli/utils/v028_features.py @@ -184,6 +184,8 @@ def supports_v028_features(task: str) -> bool: "orpo", "simpo", "ipo", + "bco", + "preference", "ppo", "reward_model", "embedding", diff --git a/tests/test_bco.py b/tests/test_bco.py new file mode 100644 index 0000000..82aa1cc --- /dev/null +++ b/tests/test_bco.py @@ -0,0 +1,296 @@ +"""Tests for BCO (Binary Classifier Optimization) — v0.40.0 Part A. + +Mirrors ORPO/SimPO/IPO test layout: schema, data format gate, template, +trainer wrapper init + routing (train + sweep), edge cases. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock +from unittest.mock import patch as mock_patch + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import TEMPLATES, SoupConfig + +# ─── Schema Tests ─────────────────────────────────────────────────────────── + + +class TestBCOConfig: + """Test BCO task config validation.""" + + def test_bco_task_accepted(self): + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + assert cfg.task == "bco" + + def test_bco_beta_default(self): + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + assert cfg.training.bco_beta == 0.1 + + def test_bco_beta_custom(self): + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"bco_beta": 0.05}, + ) + assert cfg.training.bco_beta == pytest.approx(0.05) + + def test_bco_beta_must_be_positive(self): + with pytest.raises(ValidationError, match="bco_beta"): + SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"bco_beta": 0}, + ) + + def test_bco_beta_negative_rejected(self): + with pytest.raises(ValidationError, match="bco_beta"): + SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"bco_beta": -0.1}, + ) + + def test_bco_full_config(self): + cfg = SoupConfig( + base="meta-llama/Llama-3.1-8B-Instruct", + task="bco", + data={"train": "./data.jsonl", "format": "dpo", "max_length": 2048}, + training={ + "epochs": 3, + "lr": 1e-5, + "bco_beta": 0.2, + "lora": {"r": 64, "alpha": 16}, + "quantization": "4bit", + }, + ) + assert cfg.task == "bco" + assert cfg.training.bco_beta == pytest.approx(0.2) + + +# ─── Data Format Tests ────────────────────────────────────────────────────── + + +class TestBCODataFormat: + """Test that BCO uses the DPO data format (prompt+chosen+rejected).""" + + def test_dpo_format_works_for_bco(self): + from soup_cli.data.formats import detect_format + + data = [{"prompt": "Q", "chosen": "A", "rejected": "B"}] + assert detect_format(data) == "dpo" + + +# ─── Split helper Tests ───────────────────────────────────────────────────── + + +class TestSplitDpoRowsToBco: + def test_two_rows_become_four_with_correct_labels(self): + from soup_cli.trainer.bco import _split_dpo_rows_to_bco + + rows = [ + {"prompt": "p1", "chosen": "c1", "rejected": "r1"}, + {"prompt": "p2", "chosen": "c2", "rejected": "r2"}, + ] + out = _split_dpo_rows_to_bco(rows) + assert len(out) == 4 + assert out[0] == {"prompt": "p1", "completion": "c1", "label": True} + assert out[1] == {"prompt": "p1", "completion": "r1", "label": False} + assert out[2] == {"prompt": "p2", "completion": "c2", "label": True} + assert out[3] == {"prompt": "p2", "completion": "r2", "label": False} + + def test_empty_input_returns_empty_list(self): + from soup_cli.trainer.bco import _split_dpo_rows_to_bco + + assert _split_dpo_rows_to_bco([]) == [] + + @pytest.mark.parametrize( + "row", + [ + {"chosen": "c", "rejected": "r"}, + {"prompt": "p", "rejected": "r"}, + {"prompt": "p", "chosen": "c"}, + {}, + {"unrelated": "field"}, + ], + ) + def test_missing_required_field_skipped(self, row): + from soup_cli.trainer.bco import _split_dpo_rows_to_bco + + assert _split_dpo_rows_to_bco([row]) == [] + + def test_extra_keys_ignored(self): + from soup_cli.trainer.bco import _split_dpo_rows_to_bco + + out = _split_dpo_rows_to_bco( + [{"prompt": "p", "chosen": "c", "rejected": "r", "extra": "x"}] + ) + assert len(out) == 2 + assert all("extra" not in row for row in out) + + def test_skipped_rows_logged_at_debug(self, caplog): + import logging + + from soup_cli.trainer.bco import _split_dpo_rows_to_bco + + caplog.set_level(logging.DEBUG, logger="soup_cli.trainer.bco") + _split_dpo_rows_to_bco([{"prompt": "p", "chosen": "c"}]) + assert any("skipped" in r.message.lower() for r in caplog.records) + + +# ─── Template Tests ───────────────────────────────────────────────────────── + + +class TestBCOTemplate: + def test_bco_template_exists(self): + assert "bco" in TEMPLATES + + def test_bco_template_valid_yaml(self): + import yaml + + config = yaml.safe_load(TEMPLATES["bco"]) + assert config["task"] == "bco" + assert config["training"]["bco_beta"] == 0.1 + assert config["data"]["format"] == "dpo" + + def test_bco_template_valid_config(self): + import yaml + + raw = yaml.safe_load(TEMPLATES["bco"]) + cfg = SoupConfig(**raw) + assert cfg.task == "bco" + assert cfg.training.bco_beta == 0.1 + + +# ─── Trainer Wrapper Tests ────────────────────────────────────────────────── + + +class TestBCOTrainerWrapper: + def test_bco_import_exists(self): + from soup_cli.trainer.bco import BCOTrainerWrapper + + assert BCOTrainerWrapper is not None + + def test_bco_wrapper_init(self): + from soup_cli.trainer.bco import BCOTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + wrapper = BCOTrainerWrapper(cfg, device="cpu") + assert wrapper.config.task == "bco" + assert wrapper.device == "cpu" + assert wrapper.model is None + assert wrapper.trainer is None + + def test_bco_wrapper_init_with_options(self): + from soup_cli.trainer.bco import BCOTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + wrapper = BCOTrainerWrapper( + cfg, device="cuda", report_to="wandb", deepspeed_config="ds.json", + ) + assert wrapper.report_to == "wandb" + assert wrapper.deepspeed_config == "ds.json" + + def test_bco_train_before_setup_raises(self): + from soup_cli.trainer.bco import BCOTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + wrapper = BCOTrainerWrapper(cfg, device="cpu") + with pytest.raises(RuntimeError, match="setup"): + wrapper.train() + + +# ─── Routing Tests ────────────────────────────────────────────────────────── + + +class TestBCOTrainRouting: + """Test that train + sweep route to BCO trainer.""" + + def test_sweep_routes_to_bco_trainer(self): + from soup_cli.commands.sweep import _run_single + + cfg = SoupConfig( + base="some-model", + task="bco", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + + fake_dataset = { + "train": [{"prompt": "Q?", "chosen": "A", "rejected": "B"}], + } + fake_result = { + "initial_loss": 1.0, + "final_loss": 0.5, + "total_steps": 10, + "duration_secs": 60.0, + "output_dir": "./output", + "duration": "1m", + } + fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0} + + with mock_patch( + "soup_cli.data.loader.load_dataset", return_value=fake_dataset, + ), mock_patch( + "soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU"), + ), mock_patch( + "soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info, + ), mock_patch( + "soup_cli.experiment.tracker.ExperimentTracker", + ) as mock_tracker_cls, mock_patch( + "soup_cli.monitoring.display.TrainingDisplay", + ), mock_patch( + "soup_cli.trainer.bco.BCOTrainerWrapper.setup", + ), mock_patch( + "soup_cli.trainer.bco.BCOTrainerWrapper.train", + return_value=fake_result, + ) as mock_train: + mock_tracker = MagicMock() + mock_tracker.start_run.return_value = "run-bco-1" + mock_tracker_cls.return_value = mock_tracker + result = _run_single(cfg, {}, "bco_run_1", None) + + mock_train.assert_called_once() + assert result["run_id"] == "run-bco-1" + + +# ─── Sweep Shortcut Tests ─────────────────────────────────────────────────── + + +class TestBCOSweepParams: + def test_bco_beta_shortcut(self): + from soup_cli.commands.sweep import _set_nested_param + + config = {"training": {"bco_beta": 0.1}} + _set_nested_param(config, "bco_beta", 0.05) + assert config["training"]["bco_beta"] == 0.05 + + def test_bco_beta_shortcut_creates_nested_key(self): + from soup_cli.commands.sweep import _set_nested_param + + config = {} + _set_nested_param(config, "bco_beta", 0.2) + assert config["training"]["bco_beta"] == pytest.approx(0.2) diff --git a/tests/test_dpo_variants.py b/tests/test_dpo_variants.py new file mode 100644 index 0000000..2ad1340 --- /dev/null +++ b/tests/test_dpo_variants.py @@ -0,0 +1,439 @@ +"""Tests for v0.40.0 Part C — KL-controlled DPO variants. + +Adds two opt-in DPO controls: + * ``dpo_beta_schedule``: anneal β over training (linear / cosine / exponential). + * ``dpo_ref_regen_epochs``: replace the frozen ref model with the current + student every N epochs. + +Both are SFT-trainer-style additive flags; the existing constant-β, +constant-ref-model path is unchanged when both flags are unset. +""" + +from __future__ import annotations + +import math +from unittest.mock import MagicMock + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import SoupConfig + +# ─── Schema bounds ────────────────────────────────────────────────────────── + + +class TestDPOVariantsConfig: + def _base(self, **training): + return SoupConfig( + base="some-model", + task="dpo", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"dpo_beta": 0.1, **training}, + ) + + @pytest.mark.parametrize("sched", ["linear", "cosine", "exponential"]) + def test_beta_schedule_accepted(self, sched): + cfg = self._base(dpo_beta_schedule=sched, dpo_beta_end=0.01) + assert cfg.training.dpo_beta_schedule == sched + assert cfg.training.dpo_beta_end == pytest.approx(0.01) + + def test_beta_schedule_unknown_rejected(self): + with pytest.raises(ValidationError, match="dpo_beta_schedule"): + self._base(dpo_beta_schedule="random", dpo_beta_end=0.01) + + def test_beta_schedule_requires_end(self): + with pytest.raises(ValidationError, match="dpo_beta_end"): + self._base(dpo_beta_schedule="linear") + + def test_beta_end_must_be_positive(self): + with pytest.raises(ValidationError, match="dpo_beta_end"): + self._base(dpo_beta_schedule="linear", dpo_beta_end=0) + + def test_beta_end_alone_rejected(self): + """Setting end without schedule is meaningless.""" + with pytest.raises(ValidationError, match="dpo_beta_schedule"): + self._base(dpo_beta_end=0.01) + + def test_ref_regen_epochs_positive(self): + cfg = self._base(dpo_ref_regen_epochs=2) + assert cfg.training.dpo_ref_regen_epochs == 2 + + def test_ref_regen_epochs_zero_rejected(self): + with pytest.raises(ValidationError, match="dpo_ref_regen_epochs"): + self._base(dpo_ref_regen_epochs=0) + + def test_ref_regen_epochs_negative_rejected(self): + with pytest.raises(ValidationError, match="dpo_ref_regen_epochs"): + self._base(dpo_ref_regen_epochs=-1) + + def test_ref_regen_epochs_too_large_rejected(self): + """Bound at 1000 — runaway values are almost certainly typos.""" + with pytest.raises(ValidationError, match="dpo_ref_regen_epochs"): + self._base(dpo_ref_regen_epochs=10_000) + + def test_dpo_variants_only_for_dpo_family(self): + """β-schedule + ref-regen require a DPO-family trainer.""" + with pytest.raises(ValidationError, match="dpo|ipo|preference"): + SoupConfig( + base="some-model", + task="sft", + data={"train": "./data.jsonl"}, + training={ + "dpo_beta_schedule": "linear", + "dpo_beta_end": 0.01, + }, + ) + + def test_ref_regen_only_for_dpo_family(self): + with pytest.raises(ValidationError, match="dpo|ipo|preference"): + SoupConfig( + base="some-model", + task="orpo", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"dpo_ref_regen_epochs": 2}, + ) + + def test_dpo_variants_allowed_on_ipo(self): + cfg = SoupConfig( + base="some-model", + task="ipo", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"dpo_beta_schedule": "cosine", "dpo_beta_end": 0.01}, + ) + assert cfg.training.dpo_beta_schedule == "cosine" + + def test_dpo_variants_allowed_on_preference_dpo(self): + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={ + "preference_loss": "dpo", + "dpo_beta_schedule": "linear", + "dpo_beta_end": 0.05, + }, + ) + assert cfg.training.dpo_beta_end == pytest.approx(0.05) + + def test_dpo_variants_rejected_on_preference_orpo(self): + with pytest.raises(ValidationError, match="dpo|ipo"): + SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={ + "preference_loss": "orpo", + "dpo_beta_schedule": "linear", + "dpo_beta_end": 0.05, + }, + ) + + def test_dpo_variants_rejected_on_mlx(self): + with pytest.raises(ValidationError, match="mlx"): + SoupConfig( + base="some-model", + task="dpo", + backend="mlx", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"dpo_beta_schedule": "linear", "dpo_beta_end": 0.01}, + ) + + +# ─── β schedule math ──────────────────────────────────────────────────────── + + +class TestBetaSchedule: + def test_linear_endpoints(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=100, schedule="linear", + ) == pytest.approx(0.1) + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=100, total_steps=100, schedule="linear", + ) == pytest.approx(0.01) + + def test_linear_midpoint(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + # Schema requires beta_end > 0; mid of (0.1, 0.02) = 0.06. + mid = compute_beta_at_step( + beta_start=0.1, beta_end=0.02, step=50, total_steps=100, schedule="linear", + ) + assert mid == pytest.approx(0.06) + + def test_cosine_endpoints(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=100, schedule="cosine", + ) == pytest.approx(0.1) + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=100, total_steps=100, schedule="cosine", + ) == pytest.approx(0.01) + + def test_cosine_midpoint(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + # Cosine: at midpoint we expect (start+end)/2. + mid = compute_beta_at_step( + beta_start=0.2, beta_end=0.02, step=50, total_steps=100, schedule="cosine", + ) + assert mid == pytest.approx(0.11, abs=1e-6) + + def test_exponential_endpoints(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=100, + schedule="exponential", + ) == pytest.approx(0.1) + end = compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=100, total_steps=100, + schedule="exponential", + ) + assert end == pytest.approx(0.01, rel=1e-4) + + def test_clamps_at_total_steps(self): + """Step beyond total_steps clamps to beta_end (not extrapolation).""" + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=200, total_steps=100, schedule="linear", + ) == pytest.approx(0.01) + + def test_negative_step_clamps_to_start(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=-5, total_steps=100, schedule="linear", + ) == pytest.approx(0.1) + + def test_invalid_schedule_raises(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + with pytest.raises(ValueError, match="schedule"): + compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=100, + schedule="garbage", + ) + + def test_zero_total_steps_returns_end(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + assert compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=0, schedule="linear", + ) == pytest.approx(0.01) + + def test_negative_total_steps_rejected(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + with pytest.raises(ValueError, match="total_steps"): + compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=0, total_steps=-1, + schedule="linear", + ) + + def test_non_finite_betas_rejected(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + for bad in (float("nan"), float("inf"), -1.0): + with pytest.raises(ValueError, match="finite|> 0"): + compute_beta_at_step( + beta_start=bad, beta_end=0.01, step=0, total_steps=100, + schedule="linear", + ) + + def test_step_bool_rejected(self): + from soup_cli.utils.dpo_variants import compute_beta_at_step + + with pytest.raises(ValueError, match="step"): + compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=True, total_steps=100, + schedule="linear", + ) + + +# ─── BetaScheduleCallback ────────────────────────────────────────────────── + + +class TestBetaScheduleCallback: + def test_on_step_begin_writes_beta_on_trainer(self): + from soup_cli.utils.dpo_variants import BetaScheduleCallback + + cb = BetaScheduleCallback( + beta_start=0.1, beta_end=0.01, total_steps=100, schedule="linear", + ) + trainer = MagicMock() + trainer.beta = 0.1 + state = MagicMock(global_step=50) + cb.on_step_begin(args=None, state=state, control=None, model=None) + cb.attach(trainer) + cb.on_step_begin(args=None, state=state, control=None, model=None) + # After attachment, β is updated. + assert trainer.beta == pytest.approx(0.055, abs=1e-6) + + def test_callback_no_trainer_attached_is_noop(self): + """Without a trainer, callback updates nothing.""" + from soup_cli.utils.dpo_variants import BetaScheduleCallback + + cb = BetaScheduleCallback( + beta_start=0.1, beta_end=0.01, total_steps=100, schedule="linear", + ) + state = MagicMock(global_step=50) + # Should not raise. + cb.on_step_begin(args=None, state=state, control=None, model=None) + + +# ─── RefModelRegenCallback ───────────────────────────────────────────────── + + +class TestRefModelRegenCallback: + def test_fires_on_target_epoch(self): + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + cb = RefModelRegenCallback(every_n_epochs=2) + trainer = MagicMock() + cb.attach(trainer) + state = MagicMock(epoch=2.0) + cb.on_epoch_end(args=None, state=state, control=None, model=None) + assert cb.regen_count == 1 + + def test_does_not_fire_off_target(self): + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + cb = RefModelRegenCallback(every_n_epochs=2) + trainer = MagicMock() + cb.attach(trainer) + state = MagicMock(epoch=1.0) + cb.on_epoch_end(args=None, state=state, control=None, model=None) + assert cb.regen_count == 0 + + def test_skip_at_epoch_zero(self): + """Regen at epoch 0 would copy untrained student → undesirable.""" + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + cb = RefModelRegenCallback(every_n_epochs=1) + trainer = MagicMock() + cb.attach(trainer) + state = MagicMock(epoch=0.0) + cb.on_epoch_end(args=None, state=state, control=None, model=None) + assert cb.regen_count == 0 + + def test_invalid_period_int_below_one_rejected(self): + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + for bad in (0, -1): + with pytest.raises(ValueError, match="every_n_epochs"): + RefModelRegenCallback(every_n_epochs=bad) + + def test_invalid_period_non_int_rejected(self): + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + for bad in (0.5, "2", True): + with pytest.raises(TypeError, match="every_n_epochs"): + RefModelRegenCallback(every_n_epochs=bad) + + def test_no_trainer_attached_is_noop(self): + from soup_cli.utils.dpo_variants import RefModelRegenCallback + + cb = RefModelRegenCallback(every_n_epochs=2) + state = MagicMock(epoch=2.0) + cb.on_epoch_end(args=None, state=state, control=None, model=None) + assert cb.regen_count == 0 + + +# ─── build_dpo_variant_callbacks ──────────────────────────────────────────── + + +class TestBuildDPOVariantCallbacks: + def test_returns_empty_list_when_no_variants(self): + from soup_cli.utils.dpo_variants import build_dpo_variant_callbacks + + cbs = build_dpo_variant_callbacks( + beta_start=0.1, beta_end=None, schedule=None, + total_steps=0, ref_regen_epochs=None, + ) + assert cbs == [] + + def test_returns_beta_only(self): + from soup_cli.utils.dpo_variants import ( + BetaScheduleCallback, + build_dpo_variant_callbacks, + ) + + cbs = build_dpo_variant_callbacks( + beta_start=0.1, beta_end=0.01, schedule="linear", + total_steps=0, ref_regen_epochs=None, + ) + assert len(cbs) == 1 + assert isinstance(cbs[0], BetaScheduleCallback) + + def test_returns_regen_only(self): + from soup_cli.utils.dpo_variants import ( + RefModelRegenCallback, + build_dpo_variant_callbacks, + ) + + cbs = build_dpo_variant_callbacks( + beta_start=0.1, beta_end=None, schedule=None, + total_steps=0, ref_regen_epochs=2, + ) + assert len(cbs) == 1 + assert isinstance(cbs[0], RefModelRegenCallback) + + def test_returns_both_callbacks(self): + from soup_cli.utils.dpo_variants import ( + BetaScheduleCallback, + RefModelRegenCallback, + build_dpo_variant_callbacks, + ) + + cbs = build_dpo_variant_callbacks( + beta_start=0.1, beta_end=0.01, schedule="linear", + total_steps=0, ref_regen_epochs=2, + ) + assert len(cbs) == 2 + kinds = {type(cb) for cb in cbs} + assert kinds == {BetaScheduleCallback, RefModelRegenCallback} + + +class TestBetaScheduleLazyTotalSteps: + def test_on_train_begin_resolves_total_steps_from_state(self): + from soup_cli.utils.dpo_variants import BetaScheduleCallback + + cb = BetaScheduleCallback( + beta_start=0.1, beta_end=0.01, total_steps=0, schedule="linear", + ) + # Initially 0 — sentinel. + assert cb.total_steps == 0 + state = MagicMock(max_steps=200) + cb.on_train_begin(args=None, state=state, control=None) + assert cb.total_steps == 200 + + def test_on_step_begin_skips_when_total_steps_unresolvable(self): + """No max_steps available → don't fall through to compute_beta_at_step.""" + from soup_cli.utils.dpo_variants import BetaScheduleCallback + + cb = BetaScheduleCallback( + beta_start=0.1, beta_end=0.01, total_steps=0, schedule="linear", + ) + trainer = MagicMock() + trainer.beta = 0.1 + cb.attach(trainer) + state = MagicMock(global_step=10) + cb.on_step_begin(args=None, state=state, control=None) + # beta unchanged because total_steps still 0. + assert trainer.beta == 0.1 + + def test_math_module_used_in_cosine(self): + """Sanity: math module imported at top is exercised by cosine schedule.""" + from soup_cli.utils.dpo_variants import compute_beta_at_step + + # Cosine at progress=0 should equal beta_start exactly (cos(0)=1). + result = compute_beta_at_step( + beta_start=0.1, beta_end=0.01, step=1, total_steps=10**9, + schedule="cosine", + ) + # Very near beta_start since progress ~ 1e-9. + assert math.isclose(result, 0.1, rel_tol=1e-6) diff --git a/tests/test_performance.py b/tests/test_performance.py index 4a99b3c..63a12bc 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -589,8 +589,8 @@ class TestLongContextTemplate: assert "use_flash_attn" in TEMPLATES["longcontext"] def test_longcontext_template_count(self): - """Should now have 16 templates (15 + tool-calling added in v0.25.0).""" - assert len(TEMPLATES) == 16 + """Should now have 17 templates (16 + bco added in v0.40.0).""" + assert len(TEMPLATES) == 17 # ─── Trainer fsdp_config Parameter Tests ───────────────────────────────── diff --git a/tests/test_preference_dispatcher.py b/tests/test_preference_dispatcher.py new file mode 100644 index 0000000..ca98db1 --- /dev/null +++ b/tests/test_preference_dispatcher.py @@ -0,0 +1,273 @@ +"""Tests for v0.40.0 Part B — unified preference loss dispatcher. + +Adds ``task: preference`` + ``training.preference_loss`` Literal +(dpo/simpo/orpo/ipo/bco). Existing per-task forms (``task: dpo`` etc) keep +working unchanged — the unified surface is *additive*, not a breaking +collapse. Backward-compat ``resolve_preference_loss`` maps legacy task +strings to their preference_loss equivalents for callers that want a +single dispatch entry-point. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock +from unittest.mock import patch as mock_patch + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import SoupConfig + +# ─── Schema Tests ─────────────────────────────────────────────────────────── + + +class TestPreferenceTaskField: + def test_preference_task_accepted(self): + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + assert cfg.task == "preference" + assert cfg.training.preference_loss == "dpo" + + @pytest.mark.parametrize("loss", ["dpo", "simpo", "orpo", "ipo", "bco"]) + def test_preference_loss_each_value(self, loss): + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": loss}, + ) + assert cfg.training.preference_loss == loss + + def test_preference_loss_unknown_rejected(self): + with pytest.raises(ValidationError, match="preference_loss"): + SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "garbage"}, + ) + + def test_preference_task_requires_preference_loss(self): + """task=preference without preference_loss must error.""" + with pytest.raises(ValidationError, match="preference_loss"): + SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + + def test_preference_loss_default_none_for_other_tasks(self): + """Non-preference tasks default preference_loss=None.""" + cfg = SoupConfig( + base="some-model", + task="dpo", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + assert cfg.training.preference_loss is None + + def test_preference_loss_set_outside_preference_task_rejected(self): + """preference_loss is meaningful only when task=preference.""" + with pytest.raises(ValidationError, match="preference_loss"): + SoupConfig( + base="some-model", + task="dpo", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + + +# ─── Resolver ────────────────────────────────────────────────────────────── + + +class TestResolvePreferenceLoss: + def test_resolve_preference_task_returns_loss(self): + from soup_cli.trainer.preference import resolve_preference_loss + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "simpo"}, + ) + assert resolve_preference_loss(cfg) == "simpo" + + @pytest.mark.parametrize( + "task,expected", + [ + ("dpo", "dpo"), + ("simpo", "simpo"), + ("orpo", "orpo"), + ("ipo", "ipo"), + ("bco", "bco"), + ], + ) + def test_legacy_task_maps_to_loss(self, task, expected): + from soup_cli.trainer.preference import resolve_preference_loss + + cfg = SoupConfig( + base="some-model", + task=task, + data={"train": "./data.jsonl", "format": "dpo"}, + ) + assert resolve_preference_loss(cfg) == expected + + def test_resolve_non_preference_task_returns_none(self): + from soup_cli.trainer.preference import resolve_preference_loss + + cfg = SoupConfig( + base="some-model", + task="sft", + data={"train": "./data.jsonl"}, + ) + assert resolve_preference_loss(cfg) is None + + +# ─── Wrapper Tests ───────────────────────────────────────────────────────── + + +class TestPreferenceTrainerWrapper: + def test_import_exists(self): + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + assert PreferenceTrainerWrapper is not None + + @pytest.mark.parametrize( + "loss,wrapper_path", + [ + ("dpo", "soup_cli.trainer.dpo.DPOTrainerWrapper"), + ("simpo", "soup_cli.trainer.simpo.SimPOTrainerWrapper"), + ("orpo", "soup_cli.trainer.orpo.ORPOTrainerWrapper"), + ("ipo", "soup_cli.trainer.ipo.IPOTrainerWrapper"), + ("bco", "soup_cli.trainer.bco.BCOTrainerWrapper"), + ], + ) + def test_dispatcher_routes_to_correct_wrapper(self, loss, wrapper_path): + """PreferenceTrainerWrapper.setup() must delegate to the right wrapper. + + Asserts the inner cfg sent to the wrapper has task=loss and that + preference_loss has been cleared (defends _make_inner_cfg's contract). + """ + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": loss}, + ) + with mock_patch(wrapper_path) as mock_inner: + mock_instance = MagicMock() + mock_instance.train.return_value = { + "initial_loss": 1.0, "final_loss": 0.5, + "duration": "1m", "duration_secs": 60.0, + "output_dir": "./out", "total_steps": 10, + } + mock_inner.return_value = mock_instance + + wrapper = PreferenceTrainerWrapper(cfg, device="cpu") + wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]}) + wrapper.train() + + mock_inner.assert_called_once() + inner_cfg = mock_inner.call_args[0][0] + assert inner_cfg.task == loss + assert inner_cfg.training.preference_loss is None + mock_instance.setup.assert_called_once() + mock_instance.train.assert_called_once() + + def test_setup_does_not_mutate_caller_cfg(self): + """_make_inner_cfg must return a copy; cfg.task unchanged after setup.""" + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + with mock_patch("soup_cli.trainer.dpo.DPOTrainerWrapper") as mock_inner: + mock_inner.return_value = MagicMock() + wrapper = PreferenceTrainerWrapper(cfg, device="cpu") + wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]}) + # Caller's cfg untouched. + assert cfg.task == "preference" + assert cfg.training.preference_loss == "dpo" + + def test_train_before_setup_raises(self): + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + wrapper = PreferenceTrainerWrapper(cfg, device="cpu") + with pytest.raises(RuntimeError, match="setup"): + wrapper.train() + + def test_unknown_loss_raises_at_setup(self): + """Defence-in-depth: schema gate prevents this, but if it slipped, raise.""" + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + # Mutate after construction to bypass schema validation. + cfg.training.preference_loss = "garbage" + wrapper = PreferenceTrainerWrapper(cfg, device="cpu") + with pytest.raises(ValueError, match="preference_loss"): + wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]}) + + +# ─── Train Routing ───────────────────────────────────────────────────────── + + +class TestPreferenceTrainRouting: + def test_sweep_routes_to_preference_wrapper(self): + from soup_cli.commands.sweep import _run_single + + cfg = SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss": "dpo"}, + ) + fake_dataset = {"train": [{"prompt": "Q", "chosen": "A", "rejected": "B"}]} + fake_result = { + "initial_loss": 1.0, "final_loss": 0.5, + "duration": "1m", "duration_secs": 60.0, + "output_dir": "./out", "total_steps": 10, + } + fake_gpu_info = {"memory_total": "0 MB", "memory_total_bytes": 0} + + with mock_patch( + "soup_cli.data.loader.load_dataset", return_value=fake_dataset, + ), mock_patch( + "soup_cli.utils.gpu.detect_device", return_value=("cpu", "CPU"), + ), mock_patch( + "soup_cli.utils.gpu.get_gpu_info", return_value=fake_gpu_info, + ), mock_patch( + "soup_cli.experiment.tracker.ExperimentTracker", + ) as mock_tracker_cls, mock_patch( + "soup_cli.monitoring.display.TrainingDisplay", + ), mock_patch( + "soup_cli.trainer.preference.PreferenceTrainerWrapper.setup", + ), mock_patch( + "soup_cli.trainer.preference.PreferenceTrainerWrapper.train", + return_value=fake_result, + ) as mock_train: + mock_tracker = MagicMock() + mock_tracker.start_run.return_value = "run-pref-1" + mock_tracker_cls.return_value = mock_tracker + result = _run_single(cfg, {}, "pref_run_1", None) + + mock_train.assert_called_once() + assert result["run_id"] == "run-pref-1" diff --git a/tests/test_preference_multi.py b/tests/test_preference_multi.py new file mode 100644 index 0000000..68ca7ca --- /dev/null +++ b/tests/test_preference_multi.py @@ -0,0 +1,179 @@ +"""Tests for v0.40.0 Part D — multi-objective preference loss. + +Adds ``training.preference_loss_weights: dict[str, float]`` for blending +preference losses (e.g. ``{"dpo": 0.7, "bco": 0.3}``). The schema-level +surface ships in v0.40.0; the live runtime weighted combination is +deferred to v0.40.1 (mirrors the project's stub-then-live pattern from +v0.27.0 MII / v0.37.0 multipack / v0.38.0 quant menu / v0.39.0 ReLoRA). +""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from soup_cli.config.schema import SoupConfig + +# ─── Schema bounds ────────────────────────────────────────────────────────── + + +def _base(**training): + return SoupConfig( + base="some-model", + task="preference", + data={"train": "./data.jsonl", "format": "dpo"}, + training=training, + ) + + +class TestPreferenceLossWeightsConfig: + def test_two_loss_blend_accepted(self): + cfg = _base(preference_loss_weights={"dpo": 0.7, "bco": 0.3}) + assert cfg.training.preference_loss_weights == {"dpo": 0.7, "bco": 0.3} + + def test_three_loss_blend_accepted(self): + cfg = _base( + preference_loss_weights={"dpo": 0.5, "bco": 0.3, "simpo": 0.2}, + ) + assert sum(cfg.training.preference_loss_weights.values()) == pytest.approx(1.0) + + def test_single_entry_rejected_use_scalar_form_instead(self): + """Single-entry blends are equivalent to scalar preference_loss; reject.""" + with pytest.raises(ValidationError, match="2 and 5"): + _base(preference_loss_weights={"dpo": 1.0}) + + def test_more_than_five_entries_rejected(self): + with pytest.raises(ValidationError, match="2 and 5"): + _base( + preference_loss_weights={ + "dpo": 0.2, "bco": 0.2, "simpo": 0.2, + "orpo": 0.2, "ipo": 0.1, "extra": 0.1, + }, + ) + + def test_unknown_key_rejected(self): + with pytest.raises(ValidationError, match="unknown"): + _base(preference_loss_weights={"dpo": 0.5, "garbage": 0.5}) + + def test_null_byte_in_key_rejected(self): + with pytest.raises(ValidationError, match="null byte"): + _base(preference_loss_weights={"dpo": 0.5, "bco\x00": 0.5}) + + def test_empty_dict_rejected(self): + with pytest.raises(ValidationError, match="2 and 5"): + _base(preference_loss_weights={}) + + def test_weights_must_sum_to_one(self): + with pytest.raises(ValidationError, match="sum"): + _base(preference_loss_weights={"dpo": 0.5, "bco": 0.4}) + + def test_weight_zero_rejected(self): + with pytest.raises(ValidationError, match=r"\(0, 1\]"): + _base(preference_loss_weights={"dpo": 1.0, "bco": 0.0}) + + def test_weight_negative_rejected(self): + with pytest.raises(ValidationError, match=r"\(0, 1\]"): + _base(preference_loss_weights={"dpo": 1.5, "bco": -0.5}) + + def test_weight_above_one_rejected(self): + # Two values both > 1 — the per-value bound (0,1] fires before the + # sum gate. + with pytest.raises(ValidationError, match=r"\(0, 1\]"): + _base(preference_loss_weights={"dpo": 1.1, "bco": 1.1}) + + def test_weight_bool_coerced_to_zero_then_rejected(self): + """Pydantic coerces True/False → 1.0/0.0 before model_validator sees + them, so a False weight is rejected by the (0, 1] gate (not the bool + guard). Either rejection path is acceptable; this test pins the + observed behaviour rather than the rejection mechanism.""" + with pytest.raises(ValidationError, match=r"\(0, 1\]|must be a number"): + _base(preference_loss_weights={"dpo": True, "bco": False}) + + def test_requires_preference_task(self): + with pytest.raises(ValidationError, match="preference"): + SoupConfig( + base="some-model", + task="dpo", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss_weights": {"dpo": 1.0}}, + ) + + def test_mutually_exclusive_with_scalar_loss(self): + with pytest.raises(ValidationError, match="mutually exclusive"): + _base( + preference_loss="dpo", + preference_loss_weights={"dpo": 0.7, "bco": 0.3}, + ) + + def test_rejected_on_mlx(self): + with pytest.raises(ValidationError, match="mlx"): + SoupConfig( + base="some-model", + task="preference", + backend="mlx", + data={"train": "./data.jsonl", "format": "dpo"}, + training={"preference_loss_weights": {"dpo": 0.7, "bco": 0.3}}, + ) + + +# ─── Helper API ───────────────────────────────────────────────────────────── + + +class TestMultiObjectiveHelpers: + def test_is_multi_objective_true(self): + from soup_cli.trainer.preference import is_multi_objective_preference + + cfg = _base(preference_loss_weights={"dpo": 0.7, "bco": 0.3}) + assert is_multi_objective_preference(cfg) is True + + def test_is_multi_objective_false_scalar(self): + from soup_cli.trainer.preference import is_multi_objective_preference + + cfg = _base(preference_loss="dpo") + assert is_multi_objective_preference(cfg) is False + + def test_is_multi_objective_false_legacy(self): + from soup_cli.trainer.preference import is_multi_objective_preference + + cfg = SoupConfig( + base="some-model", + task="dpo", + data={"train": "./data.jsonl", "format": "dpo"}, + ) + assert is_multi_objective_preference(cfg) is False + + def test_get_loss_weights_returns_copy(self): + """Defensive copy so caller mutation cannot affect cfg.""" + from soup_cli.trainer.preference import get_loss_weights + + cfg = _base(preference_loss_weights={"dpo": 0.7, "bco": 0.3}) + weights = get_loss_weights(cfg) + assert weights == {"dpo": 0.7, "bco": 0.3} + weights["dpo"] = 999.0 + # Re-fetch — original cfg unchanged. + assert get_loss_weights(cfg) == {"dpo": 0.7, "bco": 0.3} + + def test_get_loss_weights_none_when_not_set(self): + from soup_cli.trainer.preference import get_loss_weights + + cfg = _base(preference_loss="dpo") + assert get_loss_weights(cfg) is None + + +# ─── Live wiring stub-then-live ───────────────────────────────────────────── + + +class TestMultiObjectiveDeferred: + def test_setup_raises_with_actionable_message(self): + """v0.40.0 ships schema only; live runtime wiring deferred to v0.40.1. + + ``setup`` must raise a ``NotImplementedError`` that names the + deferred-version follow-up so users know whether to wait or + switch to the scalar form. + """ + from soup_cli.trainer.preference import PreferenceTrainerWrapper + + cfg = _base(preference_loss_weights={"dpo": 0.7, "bco": 0.3}) + wrapper = PreferenceTrainerWrapper(cfg, device="cpu") + with pytest.raises(NotImplementedError, match="v0.40.1"): + wrapper.setup({"train": [{"prompt": "p", "chosen": "c", "rejected": "r"}]})