From e406688f2db4ce5389f12a59aee06e31d5442c91 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Mon, 27 Apr 2026 18:33:39 +0500 Subject: [PATCH] feat(training): stability auto-tuning live (v0.33.0 Part E) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #56, #57, #58, #59. #56 Live --find-lr in-process LR-sweep: - New utils/lr_finder.run_lr_sweep(model, dataloader, schedule, optimizer_factory, device): per-step LR mutation + forward + backward, records loss until exhaustion or NaN/Inf divergence. - commands/train.py wires it via _live_lr_sweep_from_config (loads model + tokenizer + first N rows of cfg.data.train), with synthetic-curve fallback so users without GPU/torch still get a parseable report. #57 Loss-spike recovery hint: - SoupTrainerCallback gains spike_recovery / spike_recovery_max_attempts / spike_recovery_lr_decay; on watchdog fire writes output_dir/spike_recovery.json with previous_lr, recommended_lr (per SpikeRecoveryStrategy.compute_new_lr), should_recover, attempts. A wrapper / re-launch can resume with the decayed LR. Live optimizer rewind is intentionally NOT done — HF Trainer has no safe public API for mid-loop optimizer-state mutation; the JSON hint is the contract. #58 auto_mixed_precision push to TrainingArguments: - New SFTTrainerWrapper._resolve_mixed_precision: when tcfg.auto_mixed_precision is True, queries torch.cuda compute capability and calls pick_mixed_precision(base, cc) to set bf16=/fp16= flags. CPU short-circuits to (False, False). When the flag is False, legacy default preserved (bf16=cuda). #59 Grad-accum advisory (Phase 1): - SoupTrainerCallback gains grad_accum_auto_tune / grad_accum_pressure_threshold / grad_accum_total_vram_gb / grad_accum_current_steps / grad_accum_current_batch. - on_log samples torch.cuda.max_memory_allocated each step; if GradAccumMonitor.should_adjust crosses the threshold once, prints (batch, accum) -> (new_batch, new_accum) advisory and short-circuits (one-shot). Phase 2 (live DataLoader rebuild) needs a small TRL upstream PR — tracked as a known limitation. Wiring: - soup_cli/trainer/sft.py: _resolve_mixed_precision helper, batch_size preserved on self, SoupTrainerCallback constructor passes through new spike + grad-accum knobs. - soup_cli/monitoring/callback.py: rich Console import added (was previously module-relative); spike + grad-accum state fields and one-shot helpers. Tests: +15 in tests/test_part_e.py covering the LR-sweep loop with mocked model + optimizer (records, divergence break), mixed-precision resolver across cpu/cuda + auto-flag combinations + qwen2 fp16 quirk on Ampere, spike recovery hint write + attempts increment + disabled no-op, grad-accum advisory one-shot semantics + threshold + cuda-absent + disabled. Known limitations (release notes): - #57 spike recovery is a JSON hint, not in-process optimizer rewind - #59 Phase 2 (live DataLoader rebuild on advisory) deferred Co-Authored-By: Claude Opus 4.7 (1M context) --- soup_cli/commands/train.py | 124 +++++++++--- soup_cli/monitoring/callback.py | 133 +++++++++++++ soup_cli/trainer/sft.py | 70 ++++++- soup_cli/utils/lr_finder.py | 70 +++++++ tests/test_part_e.py | 327 ++++++++++++++++++++++++++++++++ 5 files changed, 697 insertions(+), 27 deletions(-) create mode 100644 tests/test_part_e.py diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index 5a5ab2d..6305ec9 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -153,31 +153,14 @@ def train( except ValueError as exc: console.print(f"[red]Invalid --find-lr range:[/] {exc}") raise typer.Exit(1) from exc - # v0.32.0 ships the LR-sweep schedule + analysis API. The live - # in-process training loop wiring (HF Trainer with custom LR - # callback) is deferred to v0.32.1 — same advisory pattern as - # v0.30.0 --auto-quant. For now we render a stub report so users - # can validate the path containment + plot infrastructure. - console.print( - "[yellow]--find-lr v0.32.0:[/] schedule + analysis API ready; " - "live LR-sweep training loop deferred to v0.32.1. " - "Writing stub report so you can verify the output path." + # v0.33.0 #56: live LR-sweep training loop. Falls back to a + # synthetic curve only when the real loop cannot run (no torch / + # config load failure) so users still get a parseable report. + losses_for_report = _run_live_lr_sweep_or_synth( + config_path, schedule, ) - # Synthetic loss curve: descend through the first 60% of the sweep, - # bottom out, then explode in the tail — mimics a real LR-finder - # output so divergence detection + steepest-gradient logic both - # produce non-trivial values in the stub report. - n = len(schedule) - descend_until = max(1, int(n * 0.6)) - synth_losses = [] - for i in range(n): - if i < descend_until: - synth_losses.append(3.0 - 2.0 * (i / descend_until)) - else: - tail = (i - descend_until) / max(1, n - descend_until) - synth_losses.append(1.0 + 8.0 * tail * tail) try: - save_lr_finder_report(schedule, synth_losses, find_lr_output) + save_lr_finder_report(schedule, losses_for_report, find_lr_output) except ValueError as exc: console.print(f"[red]Invalid --find-lr-output:[/] {exc}") raise typer.Exit(1) from exc @@ -750,3 +733,98 @@ def _resolve_checkpoint(resume: str, output_dir: str, experiment_name: str = Non if checkpoint_path.exists() and checkpoint_path.is_dir(): return str(checkpoint_path) return None + + +def _run_live_lr_sweep_or_synth( + config_path: str, schedule: list[float], +) -> list[float]: + """v0.33.0 #56 — try to run an in-process LR sweep; fall back to a + synthetic curve when prerequisites are missing. + + Falls back when: + - torch / transformers / datasets are not importable + - config load fails + - dataset cannot be tokenized into a small in-memory loader + The fallback curve descends 60% then diverges so the recommended-LR + extraction in :func:`find_optimal_lr` still produces sensible output. + """ + try: + cfg = load_config(config_path) + except Exception as exc: # noqa: BLE001 — fall back rather than abort + console.print( + f"[yellow]--find-lr: config load failed ({exc}); " + f"writing synthetic curve.[/]" + ) + return _synth_lr_curve(len(schedule)) + + try: + return _live_lr_sweep_from_config(cfg, schedule) + except Exception as exc: # noqa: BLE001 — informative fallback + console.print( + f"[yellow]--find-lr: live sweep unavailable ({exc}); " + f"writing synthetic curve.[/]" + ) + return _synth_lr_curve(len(schedule)) + + +def _synth_lr_curve(n: int) -> list[float]: + descend_until = max(1, int(n * 0.6)) + out: list[float] = [] + for i in range(n): + if i < descend_until: + out.append(3.0 - 2.0 * (i / descend_until)) + else: + tail = (i - descend_until) / max(1, n - descend_until) + out.append(1.0 + 8.0 * tail * tail) + return out + + +def _live_lr_sweep_from_config(cfg, schedule: list[float]) -> list[float]: + """Build a tiny in-process loop: load model + tokenizer + a slice of + the train dataset, then call :func:`run_lr_sweep`.""" + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + from soup_cli.data.loader import load_local + from soup_cli.utils.lr_finder import run_lr_sweep + + device = "cuda" if torch.cuda.is_available() else "cpu" + tokenizer = AutoTokenizer.from_pretrained( + cfg.base, trust_remote_code=False, + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + model = AutoModelForCausalLM.from_pretrained( + cfg.base, trust_remote_code=False, + ).to(device) + model.train() + + dataset = load_local(cfg.data.train, cfg.data.format) + rows = list(dataset)[: max(2, len(schedule))] + if not rows: + raise RuntimeError("training dataset is empty") + + def _tokenize(row): + text = row.get("text") or row.get("prompt") or "" + if not text and "messages" in row: + text = " ".join(m.get("content", "") for m in row["messages"]) + enc = tokenizer( + text or " ", return_tensors="pt", truncation=True, + max_length=min(cfg.data.max_length or 256, 256), + padding="max_length", + ) + enc["labels"] = enc["input_ids"].clone() + return {k: v.squeeze(0) for k, v in enc.items()} + + def _batched_loader(): + for row in rows: + tok = _tokenize(row) + yield {k: v.unsqueeze(0) for k, v in tok.items()} + + return run_lr_sweep( + model=model, + dataloader=_batched_loader(), + schedule=schedule, + optimizer_factory=lambda params: torch.optim.AdamW(params, lr=schedule[0]), + device=device, + ) diff --git a/soup_cli/monitoring/callback.py b/soup_cli/monitoring/callback.py index 339e463..b1fc2bb 100644 --- a/soup_cli/monitoring/callback.py +++ b/soup_cli/monitoring/callback.py @@ -5,6 +5,7 @@ from __future__ import annotations import logging from typing import Optional +from rich.console import Console from transformers import ( TrainerCallback, TrainerControl, @@ -15,6 +16,7 @@ from transformers import ( from soup_cli.monitoring.display import TrainingDisplay logger = logging.getLogger(__name__) +console = Console() class SoupTrainerCallback(TrainerCallback): @@ -31,6 +33,14 @@ class SoupTrainerCallback(TrainerCallback): loss_watchdog_threshold: float = 3.0, loss_watchdog_patience: int = 5, eval_gate_config: Optional[object] = None, + spike_recovery: bool = False, + spike_recovery_max_attempts: int = 3, + spike_recovery_lr_decay: float = 0.5, + grad_accum_auto_tune: bool = False, + grad_accum_pressure_threshold: float = 0.9, + grad_accum_total_vram_gb: float = 24.0, + grad_accum_current_steps: int = 1, + grad_accum_current_batch: int = 1, ): self.display = display self.tracker = tracker @@ -43,6 +53,30 @@ class SoupTrainerCallback(TrainerCallback): self._watchdog_patience = loss_watchdog_patience self._watchdog_counter = 0 self._watchdog_fired = False + # v0.33.0 #57 — spike-recovery hint state + self._spike_recovery_enabled = spike_recovery + self._spike_recovery_attempts = 0 + from soup_cli.utils.spike_recovery import SpikeRecoveryStrategy + if spike_recovery: + self._spike_strategy = SpikeRecoveryStrategy( + max_attempts=spike_recovery_max_attempts, + lr_decay=spike_recovery_lr_decay, + ) + else: + self._spike_strategy = None + # v0.33.0 #59 — grad-accum advisory monitor + self._grad_accum_enabled = grad_accum_auto_tune + self._grad_accum_current = max(1, int(grad_accum_current_steps)) + self._grad_accum_batch = max(1, int(grad_accum_current_batch)) + self._grad_accum_advised = False + if grad_accum_auto_tune: + from soup_cli.utils.grad_accum import GradAccumMonitor + self._grad_accum_monitor = GradAccumMonitor( + total_vram_gb=grad_accum_total_vram_gb, + threshold=grad_accum_pressure_threshold, + ) + else: + self._grad_accum_monitor = None # Eval gate state (Part B of v0.26.0) self.eval_gate_config = eval_gate_config # Tests inject these; prod wiring sets them at on_train_begin time. @@ -107,6 +141,16 @@ class SoupTrainerCallback(TrainerCallback): from rich.panel import Panel wc = WatchdogConsole() + + # v0.33.0 #57 — spike recovery hint: write a recovery + # state file the user can resume from. We do NOT mutate + # optimizer state in-place (HF Trainer does not expose a + # safe public API for that mid-loop) but we leave a + # machine-readable hint so a wrapper / re-launch can + # resume with a decayed LR. + if self._spike_strategy is not None: + self._write_spike_recovery_hint(args, loss) + wc.print(Panel( f"[bold red]Loss watchdog triggered![/]\n\n" f"Loss {loss:.4f} exceeded threshold " @@ -121,6 +165,14 @@ class SoupTrainerCallback(TrainerCallback): else: self._watchdog_counter = 0 + # v0.33.0 #59 — grad-accum advisory (one-shot per run) + if ( + self._grad_accum_enabled + and not self._grad_accum_advised + and self._grad_accum_monitor is not None + ): + self._maybe_advise_grad_accum() + # Log to experiment tracker if self.tracker and self.run_id: self.tracker.log_metrics( @@ -246,3 +298,84 @@ class SoupTrainerCallback(TrainerCallback): except Exception as exc: logger.exception("Auto-eval custom failed") console.print(f"[yellow]Auto-eval custom failed: {exc}[/]") + + # ------------------------------------------------------------------ + # v0.33.0 #57 — spike recovery hint + # ------------------------------------------------------------------ + + def _write_spike_recovery_hint(self, args, loss: float) -> None: + """Write a JSON recovery hint next to the run output so a wrapper + script (or `soup train --resume`) can pick up the new LR. + + Best-effort: errors are logged but never crash training. + """ + import json + from pathlib import Path + + if self._spike_strategy is None: + return + attempts = self._spike_recovery_attempts + try: + new_lr = self._spike_strategy.compute_new_lr(args.learning_rate) + except ValueError: + return + recover = self._spike_strategy.should_recover(attempts) + out_dir = Path(self.output_dir or args.output_dir or ".") + try: + out_dir.mkdir(parents=True, exist_ok=True) + hint_path = out_dir / "spike_recovery.json" + hint_path.write_text(json.dumps({ + "attempts": attempts + 1, + "max_attempts": self._spike_strategy.max_attempts, + "loss_at_spike": float(loss), + "previous_lr": float(args.learning_rate), + "recommended_lr": float(new_lr), + "should_recover": recover, + }, indent=2), encoding="utf-8") + self._spike_recovery_attempts = attempts + 1 + console.print( + f"[yellow]Spike recovery hint written:[/] {hint_path} " + f"(recommended_lr={new_lr:.2e}, should_recover={recover})" + ) + except OSError as exc: + logger.warning("Failed to write spike recovery hint: %s", exc) + + # ------------------------------------------------------------------ + # v0.33.0 #59 — grad-accum advisory (Phase 1) + # ------------------------------------------------------------------ + + def _maybe_advise_grad_accum(self) -> None: + """Sample VRAM use; if pressure crosses threshold once, print the + recommended (batch_size, grad_accum_steps) pair. + + Phase 1 is advisory-only. Phase 2 (live DataLoader rebuild) requires + a small upstream TRL change tracked as a known limitation. + """ + if self._grad_accum_advised: + return + try: + import torch + if not torch.cuda.is_available(): + return + used_gb = torch.cuda.max_memory_allocated() / (1024**3) + except Exception: # noqa: BLE001 — VRAM probe is best-effort + return + + if self._grad_accum_monitor is None: + return + self._grad_accum_monitor.observe(used_gb) + if not self._grad_accum_monitor.should_adjust(used_gb): + return + new_batch, new_accum = self._grad_accum_monitor.recommend( + self._grad_accum_batch, self._grad_accum_current, + ) + if new_accum == self._grad_accum_current: + return + self._grad_accum_advised = True + console.print( + f"[yellow]Grad-accum advisory:[/] VRAM pressure crossed " + f"threshold; recommend (batch_size, grad_accum_steps) " + f"({self._grad_accum_batch}, {self._grad_accum_current}) -> " + f"({new_batch}, {new_accum}). " + f"Restart training with the new pair to take effect." + ) diff --git a/soup_cli/trainer/sft.py b/soup_cli/trainer/sft.py index 87a6d01..5afba3c 100644 --- a/soup_cli/trainer/sft.py +++ b/soup_cli/trainer/sft.py @@ -143,6 +143,11 @@ class SFTTrainerWrapper: warmup_steps = int(total_steps * tcfg.warmup_ratio) # --- Training args --- + # v0.33.0 #58: auto_mixed_precision wires pick_mixed_precision() + # into bf16/fp16 kwargs. Default behaviour (bf16 on CUDA) preserved + # when the auto flag is False. + bf16_flag, fp16_flag = self._resolve_mixed_precision(tcfg, cfg.base) + training_kwargs = { "output_dir": str(output_dir), "num_train_epochs": tcfg.epochs, @@ -157,7 +162,8 @@ class SFTTrainerWrapper: "logging_steps": tcfg.logging_steps, "save_steps": tcfg.save_steps, "save_total_limit": 3, - "bf16": self.device == "cuda", + "bf16": bf16_flag, + "fp16": fp16_flag, "report_to": self.report_to, "remove_unused_columns": False, "deepspeed": self.deepspeed_config, @@ -263,6 +269,41 @@ class SFTTrainerWrapper: self.trainer = SFTTrainer(**trainer_kwargs) self._output_dir = str(output_dir) + self._batch_size = batch_size + + def _resolve_mixed_precision(self, tcfg, base_model: str) -> tuple[bool, bool]: + """Return ``(bf16, fp16)`` flags for TrainingArguments. + + - When ``tcfg.auto_mixed_precision`` is True: query GPU compute + capability and call :func:`pick_mixed_precision` to decide. + - Otherwise: preserve legacy default (bf16 on CUDA, no fp16). + """ + if not getattr(tcfg, "auto_mixed_precision", False): + return (self.device == "cuda", False) + + if self.device != "cuda": + return (False, False) + + try: + import torch + + major, minor = torch.cuda.get_device_capability() + cc = float(f"{major}.{minor}") + except (ImportError, RuntimeError, AssertionError, OSError): + return (self.device == "cuda", False) + + from soup_cli.utils.mixed_precision import pick_mixed_precision + + try: + mode = pick_mixed_precision(base_model, cc) + except ValueError: + return (self.device == "cuda", False) + + console.print( + f"[green]Auto mixed-precision picked:[/] {mode} " + f"(model={base_model}, cc={cc})" + ) + return (mode == "bf16", mode == "fp16") def _setup_transformers(self, cfg, tcfg): """Load model via standard transformers + peft pipeline.""" @@ -679,12 +720,33 @@ class SFTTrainerWrapper: if display: from soup_cli.monitoring.callback import SoupTrainerCallback + tcfg_local = self.config.training 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, + output_dir=self._output_dir, + loss_watchdog=tcfg_local.loss_watchdog, + loss_watchdog_threshold=tcfg_local.loss_watchdog_threshold, + loss_watchdog_patience=tcfg_local.loss_watchdog_patience, + spike_recovery=getattr( + tcfg_local, "loss_spike_recovery", False, + ), + spike_recovery_max_attempts=getattr( + tcfg_local, "loss_spike_recovery_max_attempts", 3, + ), + spike_recovery_lr_decay=getattr( + tcfg_local, "loss_spike_recovery_lr_decay", 0.5, + ), + grad_accum_auto_tune=getattr( + tcfg_local, "grad_accum_auto_tune", False, + ), + grad_accum_pressure_threshold=getattr( + tcfg_local, "grad_accum_pressure_threshold", 0.9, + ), + grad_accum_current_steps=getattr( + tcfg_local, "gradient_accumulation_steps", 1, + ), + grad_accum_current_batch=self._batch_size, ) ) diff --git a/soup_cli/utils/lr_finder.py b/soup_cli/utils/lr_finder.py index ee5d555..6ee710d 100644 --- a/soup_cli/utils/lr_finder.py +++ b/soup_cli/utils/lr_finder.py @@ -130,6 +130,76 @@ def _finite_or_reject(values: Sequence[float], label: str) -> list[float]: return cleaned +def run_lr_sweep( + *, model, dataloader, schedule, optimizer_factory, device: str = "cpu", +) -> list[float]: + """Run an in-process LR-sweep training loop (#56, v0.33.0). + + For each LR in ``schedule``, pulls the next batch from ``dataloader``, + runs a forward + backward + optimizer step with that LR, records the + loss. Diverged batches (NaN/Inf loss) terminate the sweep early so the + report's ``diverged_at`` is honest. + + Args: + model: a torch ``nn.Module`` returning a dict with ``loss`` field + (HF causal-LM contract). + dataloader: any iterable producing kwargs dicts for ``model(**batch)``. + schedule: LR sweep from :func:`compute_lr_schedule`. + optimizer_factory: callable ``(params) -> Optimizer`` so we can + instantiate without depending on a specific optimizer here. + device: ``"cpu"`` / ``"cuda"`` / ``"mps"``. + + Returns: + list of per-step losses, length <= ``len(schedule)``. + + Raises: + ValueError: if the schedule is empty. + + Notes: + - We mutate ``param_group["lr"]`` per step (standard LR-finder + pattern, no scheduler interference). + - Loss is captured as a Python float to break the autograd graph. + - The loop is bounded by the schedule length and the dataloader + length — whichever is shorter. + """ + if not schedule: + raise ValueError("schedule must be non-empty") + + import math as _math + + optimizer = optimizer_factory(model.parameters()) + losses: list[float] = [] + + iterator = iter(dataloader) + for lr in schedule: + try: + batch = next(iterator) + except StopIteration: + break + for group in optimizer.param_groups: + group["lr"] = lr + + # Move tensor batch values onto the right device when possible. + # Stays import-free here; ``v.to(device)`` is duck-typed against any + # tensor-like object so we don't need a hard torch dependency. + if isinstance(batch, dict): + batch = { + k: (v.to(device) if hasattr(v, "to") else v) + for k, v in batch.items() + } + + optimizer.zero_grad(set_to_none=True) + out = model(**batch) if isinstance(batch, dict) else model(batch) + loss = out["loss"] if isinstance(out, dict) else out.loss + loss_value = float(loss.detach().item()) if hasattr(loss, "detach") else float(loss) + if not _math.isfinite(loss_value): + break + losses.append(loss_value) + loss.backward() + optimizer.step() + return losses + + def save_lr_finder_report( lrs: Sequence[float], losses: Sequence[float], output_path: Path | str, ) -> None: diff --git a/tests/test_part_e.py b/tests/test_part_e.py new file mode 100644 index 0000000..8c2e1ea --- /dev/null +++ b/tests/test_part_e.py @@ -0,0 +1,327 @@ +"""Part E — v0.32.1 stability live (#56, #57, #58, #59) for v0.33.0. + +Covers: + - #56 run_lr_sweep — in-process LR-sweep loop with mocked model + DataLoader. + - #57 SoupTrainerCallback._write_spike_recovery_hint — writes JSON hint + when watchdog fires and loss_spike_recovery is enabled. + - #58 SFTTrainerWrapper._resolve_mixed_precision — wires + pick_mixed_precision into bf16/fp16 flags; preserves legacy default + when auto flag is False. + - #59 SoupTrainerCallback grad-accum advisory — fires once on threshold + crossing. +""" + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +# --------------------------------------------------------------------------- +# #56 — run_lr_sweep +# --------------------------------------------------------------------------- + + +class TestRunLRSweep: + def test_empty_schedule_rejected(self): + from soup_cli.utils.lr_finder import run_lr_sweep + + with pytest.raises(ValueError, match="schedule must be non-empty"): + run_lr_sweep( + model=MagicMock(), dataloader=iter([]), + schedule=[], optimizer_factory=lambda p: MagicMock(), + ) + + def test_loop_records_loss_per_step(self): + from soup_cli.utils.lr_finder import run_lr_sweep + + # Fake model returning a tensor-like loss + def _fake_loss_value(value): + obj = MagicMock() + obj.detach = MagicMock(return_value=obj) + obj.item = MagicMock(return_value=value) + obj.backward = MagicMock(return_value=None) + return obj + + loss_values = [3.0, 2.0, 1.5, 1.0] + + class FakeModel: + def __init__(self): + self._idx = 0 + + def parameters(self): + return [] + + def __call__(self, **batch): + value = loss_values[self._idx] + self._idx += 1 + return {"loss": _fake_loss_value(value)} + + model = FakeModel() + + # Fake optimizer with mutable param_groups + class FakeOptim: + def __init__(self, _params): + self.param_groups = [{"lr": 0.0}] + + def zero_grad(self, set_to_none: bool = False): # noqa: ARG002 + pass + + def step(self): + pass + + dl = iter([{"input_ids": MagicMock()}] * 4) + schedule = [1e-6, 1e-5, 1e-4, 1e-3] + + losses = run_lr_sweep( + model=model, dataloader=dl, schedule=schedule, + optimizer_factory=FakeOptim, + ) + assert losses == loss_values + + def test_diverged_loss_breaks_loop(self): + from soup_cli.utils.lr_finder import run_lr_sweep + + loss_values = [3.0, float("inf"), 1.0] + + def _wrap(value): + obj = MagicMock() + obj.detach = MagicMock(return_value=obj) + obj.item = MagicMock(return_value=value) + obj.backward = MagicMock(return_value=None) + return obj + + class FakeModel: + def __init__(self): + self._idx = 0 + + def parameters(self): + return [] + + def __call__(self, **batch): + value = loss_values[self._idx] + self._idx += 1 + return {"loss": _wrap(value)} + + class FakeOptim: + def __init__(self, _params): + self.param_groups = [{"lr": 0.0}] + + def zero_grad(self, set_to_none: bool = False): # noqa: ARG002 + pass + + def step(self): + pass + + dl = iter([{"x": MagicMock()}] * 3) + losses = run_lr_sweep( + model=FakeModel(), dataloader=dl, + schedule=[1e-6, 1e-5, 1e-4], + optimizer_factory=FakeOptim, + ) + # Loop terminates after the inf — only the first finite loss kept. + assert losses == [3.0] + + +# --------------------------------------------------------------------------- +# #58 — auto mixed-precision push +# --------------------------------------------------------------------------- + + +class TestResolveMixedPrecision: + def test_auto_flag_off_preserves_legacy(self): + from soup_cli.trainer.sft import SFTTrainerWrapper + + wrapper = SFTTrainerWrapper.__new__(SFTTrainerWrapper) + wrapper.device = "cuda" + tcfg = SimpleNamespace(auto_mixed_precision=False) + bf16, fp16 = wrapper._resolve_mixed_precision(tcfg, "any") + assert bf16 is True + assert fp16 is False + + def test_auto_flag_off_cpu(self): + from soup_cli.trainer.sft import SFTTrainerWrapper + + wrapper = SFTTrainerWrapper.__new__(SFTTrainerWrapper) + wrapper.device = "cpu" + tcfg = SimpleNamespace(auto_mixed_precision=False) + bf16, fp16 = wrapper._resolve_mixed_precision(tcfg, "any") + assert (bf16, fp16) == (False, False) + + def test_auto_flag_cpu_returns_no(self): + from soup_cli.trainer.sft import SFTTrainerWrapper + + wrapper = SFTTrainerWrapper.__new__(SFTTrainerWrapper) + wrapper.device = "cpu" + tcfg = SimpleNamespace(auto_mixed_precision=True) + assert wrapper._resolve_mixed_precision(tcfg, "any") == (False, False) + + def test_auto_flag_picks_bf16_on_ampere(self, monkeypatch): + """Ampere (cc 8.6) + non-quirk model → bf16.""" + import torch + + from soup_cli.trainer.sft import SFTTrainerWrapper + + wrapper = SFTTrainerWrapper.__new__(SFTTrainerWrapper) + wrapper.device = "cuda" + monkeypatch.setattr( + torch.cuda, "get_device_capability", + lambda *_a, **_k: (8, 6), + raising=False, + ) + tcfg = SimpleNamespace(auto_mixed_precision=True) + bf16, fp16 = wrapper._resolve_mixed_precision(tcfg, "neutral-model") + assert (bf16, fp16) == (True, False) + + def test_auto_flag_picks_fp16_for_qwen2_on_ampere(self, monkeypatch): + import torch + + from soup_cli.trainer.sft import SFTTrainerWrapper + + wrapper = SFTTrainerWrapper.__new__(SFTTrainerWrapper) + wrapper.device = "cuda" + monkeypatch.setattr( + torch.cuda, "get_device_capability", + lambda *_a, **_k: (8, 6), + raising=False, + ) + tcfg = SimpleNamespace(auto_mixed_precision=True) + bf16, fp16 = wrapper._resolve_mixed_precision( + tcfg, "Qwen/Qwen2-7B-Instruct", + ) + assert (bf16, fp16) == (False, True) + + +# --------------------------------------------------------------------------- +# #57 — spike recovery hint +# --------------------------------------------------------------------------- + + +def _make_callback(tmp_path, **kwargs): + from soup_cli.monitoring.callback import SoupTrainerCallback + + display = MagicMock() + return SoupTrainerCallback( + display=display, + tracker=None, + run_id="t", + output_dir=str(tmp_path), + **kwargs, + ) + + +class TestSpikeRecoveryHint: + def test_writes_hint_file(self, tmp_path): + cb = _make_callback( + tmp_path, + spike_recovery=True, + spike_recovery_max_attempts=2, + spike_recovery_lr_decay=0.5, + ) + args = SimpleNamespace( + learning_rate=1e-3, output_dir=str(tmp_path), + ) + cb._write_spike_recovery_hint(args, loss=10.0) + hint = tmp_path / "spike_recovery.json" + assert hint.exists() + data = json.loads(hint.read_text(encoding="utf-8")) + assert data["previous_lr"] == pytest.approx(1e-3) + assert data["recommended_lr"] == pytest.approx(5e-4) + assert data["should_recover"] is True + assert data["attempts"] == 1 + + def test_attempts_counter_increments(self, tmp_path): + cb = _make_callback( + tmp_path, + spike_recovery=True, + spike_recovery_max_attempts=3, + spike_recovery_lr_decay=0.5, + ) + args = SimpleNamespace( + learning_rate=1e-3, output_dir=str(tmp_path), + ) + cb._write_spike_recovery_hint(args, loss=10.0) + cb._write_spike_recovery_hint(args, loss=10.0) + data = json.loads((tmp_path / "spike_recovery.json").read_text()) + assert data["attempts"] == 2 + + def test_disabled_when_strategy_not_set(self, tmp_path): + cb = _make_callback(tmp_path, spike_recovery=False) + args = SimpleNamespace( + learning_rate=1e-3, output_dir=str(tmp_path), + ) + cb._write_spike_recovery_hint(args, loss=10.0) + # No hint file written. + assert not (tmp_path / "spike_recovery.json").exists() + + +# --------------------------------------------------------------------------- +# #59 — grad-accum advisory +# --------------------------------------------------------------------------- + + +class TestGradAccumAdvisory: + def test_advise_fires_once_under_pressure(self, tmp_path, monkeypatch, capsys): + cb = _make_callback( + tmp_path, + grad_accum_auto_tune=True, + grad_accum_pressure_threshold=0.5, + grad_accum_total_vram_gb=10.0, + grad_accum_current_steps=1, + grad_accum_current_batch=4, + ) + + # Mock torch presence + memory probe — high pressure (8 GB / 10 GB = 80%). + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = True + fake_torch.cuda.max_memory_allocated.return_value = 8 * (1024**3) + monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) + + cb._maybe_advise_grad_accum() + assert cb._grad_accum_advised is True + + # Second call is a no-op (one-shot). + cb._grad_accum_monitor.observe = MagicMock() + cb._maybe_advise_grad_accum() + # Already advised, so monitor.observe shouldn't be called. + cb._grad_accum_monitor.observe.assert_not_called() + + def test_no_advice_when_under_threshold(self, tmp_path, monkeypatch): + cb = _make_callback( + tmp_path, + grad_accum_auto_tune=True, + grad_accum_pressure_threshold=0.9, + grad_accum_total_vram_gb=10.0, + grad_accum_current_steps=1, + grad_accum_current_batch=4, + ) + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = True + fake_torch.cuda.max_memory_allocated.return_value = 5 * (1024**3) + monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) + + cb._maybe_advise_grad_accum() + assert cb._grad_accum_advised is False + + def test_no_advice_when_disabled(self, tmp_path): + cb = _make_callback(tmp_path, grad_accum_auto_tune=False) + cb._maybe_advise_grad_accum() + assert cb._grad_accum_advised is False + + def test_no_advice_when_cuda_unavailable(self, tmp_path, monkeypatch): + cb = _make_callback( + tmp_path, + grad_accum_auto_tune=True, + grad_accum_pressure_threshold=0.5, + grad_accum_total_vram_gb=10.0, + grad_accum_current_steps=1, + grad_accum_current_batch=4, + ) + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = False + monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch) + + cb._maybe_advise_grad_accum() + assert cb._grad_accum_advised is False