mirror of https://github.com/razor-ai/soup.git
484 lines
19 KiB
Python
484 lines
19 KiB
Python
"""HuggingFace Trainer callback that feeds metrics to our display and tracker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from rich.console import Console
|
|
from transformers import (
|
|
TrainerCallback,
|
|
TrainerControl,
|
|
TrainerState,
|
|
TrainingArguments,
|
|
)
|
|
|
|
from soup_cli.monitoring.display import TrainingDisplay
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
|
|
class SoupTrainerCallback(TrainerCallback):
|
|
"""Bridges HF Trainer events to Soup's Rich live display and experiment tracker."""
|
|
|
|
def __init__(
|
|
self,
|
|
display: TrainingDisplay,
|
|
tracker: Optional[object] = None,
|
|
run_id: str = "",
|
|
eval_config: Optional[object] = None,
|
|
output_dir: str = "",
|
|
loss_watchdog: bool = False,
|
|
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
|
|
self.run_id = run_id
|
|
self.eval_config = eval_config
|
|
self.output_dir = output_dir
|
|
# Loss watchdog state
|
|
self._watchdog_enabled = loss_watchdog
|
|
self._watchdog_threshold = loss_watchdog_threshold
|
|
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.
|
|
self._gate_suite: Optional[object] = None
|
|
self._gate_generate_fn: Optional[object] = None
|
|
self._gate_baseline: Optional[dict] = None
|
|
# Injectable for tests; default is the real implementation.
|
|
self._gate_run_fn = None
|
|
|
|
def on_train_begin(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
self.display.start(total_steps=state.max_steps)
|
|
|
|
def on_step_end(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
"""v0.53.10 #156 — peek at the active batch for tool-calling rows.
|
|
|
|
HF Trainer does not pass ``inputs`` to ``on_step_end`` by default, so
|
|
this hook can only see a batch when the trainer subclass explicitly
|
|
threads ``inputs=`` via ``kwargs``. When present, route any
|
|
``tool_calls`` entries through the global tool-output buffer so the
|
|
Web UI's Tool Outputs panel populates in real time.
|
|
|
|
Best-effort: any exception inside the probe MUST NEVER take down
|
|
training. Lazy buffer access — zero overhead when ``tool_calls``
|
|
is absent from the batch.
|
|
"""
|
|
inputs = kwargs.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
return
|
|
tool_calls = inputs.get("tool_calls")
|
|
if not tool_calls:
|
|
return
|
|
# stdlib import outside the try-block (python-review MEDIUM fix:
|
|
# ``time`` is impossible to ImportError, so don't pretend it can).
|
|
import time
|
|
|
|
try:
|
|
from soup_cli.utils.tool_outputs import get_global_tool_buffer
|
|
|
|
count = 0
|
|
if isinstance(tool_calls, (list, tuple)):
|
|
count = len(tool_calls)
|
|
elif isinstance(tool_calls, (int, float)) and not isinstance(
|
|
tool_calls, bool
|
|
):
|
|
count = int(tool_calls)
|
|
if count <= 0:
|
|
return
|
|
get_global_tool_buffer().record_call(
|
|
name="sft_batch",
|
|
started_ts=time.time(),
|
|
duration_ms=0.0,
|
|
success=True,
|
|
output_preview=f"step {state.global_step}: {count} tool_calls",
|
|
)
|
|
except Exception: # noqa: BLE001 — best-effort, must never crash training
|
|
return
|
|
|
|
def on_log(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, logs=None, **kwargs,
|
|
):
|
|
if logs is None:
|
|
return
|
|
|
|
# Try to get GPU memory
|
|
gpu_mem = ""
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
used = torch.cuda.memory_allocated() / (1024**3)
|
|
total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
|
|
gpu_mem = f"{used:.1f}/{total:.1f} GB"
|
|
except Exception:
|
|
pass
|
|
|
|
step = state.global_step
|
|
epoch = state.epoch or 0
|
|
loss = logs.get("loss", 0.0)
|
|
lr = logs.get("learning_rate", 0.0)
|
|
grad_norm = logs.get("grad_norm", 0.0)
|
|
speed = logs.get("train_steps_per_second", 0.0)
|
|
|
|
self.display.update(
|
|
step=step,
|
|
epoch=epoch,
|
|
loss=loss,
|
|
lr=lr,
|
|
grad_norm=grad_norm,
|
|
speed=speed,
|
|
gpu_mem=gpu_mem,
|
|
)
|
|
|
|
# v0.53.9 #94 — push to the global SSE buffer so the live Web UI
|
|
# dashboard sees per-step metrics in real time. Best-effort: any
|
|
# exception inside the push must NEVER take down training.
|
|
try:
|
|
from soup_cli.utils.sse_train_stream import TrainEvent
|
|
from soup_cli.utils.train_event_buffer import push_train_event
|
|
|
|
push_train_event(
|
|
TrainEvent(
|
|
type="metric",
|
|
step=int(step) if step is not None else None,
|
|
epoch=float(epoch) if epoch is not None else None,
|
|
loss=float(loss) if loss else None,
|
|
lr=float(lr) if lr else None,
|
|
grad_norm=float(grad_norm) if grad_norm else None,
|
|
)
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# v0.53.9 #100 — tool-call observation for SFT runs whose batch logs
|
|
# surface a `tool_calls` count. Best-effort; the SFT trainer emits
|
|
# the field only when the data format is `tool-calling`.
|
|
if "tool_calls" in (logs or {}):
|
|
try:
|
|
import time
|
|
|
|
from soup_cli.utils.tool_outputs import get_global_tool_buffer
|
|
|
|
tool_count = logs.get("tool_calls")
|
|
if isinstance(tool_count, (int, float)) and not isinstance(
|
|
tool_count, bool
|
|
) and tool_count > 0:
|
|
get_global_tool_buffer().record_call(
|
|
name="batch",
|
|
started_ts=time.time(),
|
|
duration_ms=0.0,
|
|
success=True,
|
|
output_preview=f"observed {int(tool_count)} tool_calls at step {step}",
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# Loss watchdog — detect loss spikes and auto-stop
|
|
if self._watchdog_enabled and not self._watchdog_fired and "loss" in logs:
|
|
if loss > self._watchdog_threshold:
|
|
self._watchdog_counter += 1
|
|
if self._watchdog_counter >= self._watchdog_patience:
|
|
self._watchdog_fired = True
|
|
# Stop Live display before printing panel
|
|
self.display.stop()
|
|
|
|
from rich.console import Console as WatchdogConsole
|
|
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 "
|
|
f"{self._watchdog_threshold} for "
|
|
f"{self._watchdog_counter} consecutive steps "
|
|
f"(patience={self._watchdog_patience}).\n\n"
|
|
f"Training will stop.",
|
|
title="Loss Watchdog",
|
|
border_style="red",
|
|
))
|
|
control.should_training_stop = True
|
|
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(
|
|
run_id=self.run_id,
|
|
step=step,
|
|
epoch=epoch,
|
|
loss=loss,
|
|
lr=lr,
|
|
grad_norm=grad_norm,
|
|
speed=speed,
|
|
gpu_mem=gpu_mem,
|
|
)
|
|
|
|
def on_epoch_end(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
"""Run the eval gate, if configured."""
|
|
self._run_eval_gate(state, control)
|
|
|
|
def _run_eval_gate(self, state: TrainerState, control: TrainerControl) -> None:
|
|
"""Execute the eval gate and react per ``on_regression`` policy."""
|
|
cfg = self.eval_gate_config
|
|
if cfg is None or not getattr(cfg, "enabled", False):
|
|
return
|
|
suite = self._gate_suite
|
|
if suite is None:
|
|
return
|
|
every_n = int(getattr(cfg, "every_n_epochs", 1) or 1)
|
|
current_epoch = int(state.epoch or 0)
|
|
if current_epoch <= 0 or current_epoch % every_n != 0:
|
|
return
|
|
|
|
from soup_cli.eval.gate import run_gate as _default_run_gate
|
|
|
|
run_fn = self._gate_run_fn if self._gate_run_fn is not None else _default_run_gate
|
|
|
|
baseline = self._gate_baseline or {}
|
|
threshold = float(getattr(cfg, "regression_threshold", 0.05))
|
|
generate_fn = self._gate_generate_fn or (lambda prompt: "")
|
|
on_reg = getattr(cfg, "on_regression", "stop")
|
|
try:
|
|
result = run_fn(
|
|
suite, generate_fn=generate_fn, baseline=baseline,
|
|
regression_threshold=threshold,
|
|
)
|
|
except (ValueError, FileNotFoundError, OSError) as exc:
|
|
# Structured errors — log and react per policy. 'stop' means
|
|
# the user wants safety; treat errors as regressions for stop.
|
|
logger.warning("eval gate failed to execute: %s", exc)
|
|
if on_reg == "stop":
|
|
control.should_training_stop = True
|
|
return
|
|
except Exception as exc: # unexpected — fail safe under 'stop'
|
|
logger.exception("eval gate raised unexpected error: %s", exc)
|
|
if on_reg == "stop":
|
|
control.should_training_stop = True
|
|
return
|
|
|
|
if not result.passed:
|
|
if on_reg == "stop":
|
|
control.should_training_stop = True
|
|
logger.warning(
|
|
"eval gate FAILED (%d task(s) regressed); stopping training",
|
|
sum(1 for r in result.task_results if not r.passed),
|
|
)
|
|
elif on_reg == "warn":
|
|
logger.warning(
|
|
"eval gate FAILED (%d task(s) regressed); continuing per policy",
|
|
sum(1 for r in result.task_results if not r.passed),
|
|
)
|
|
# on_reg == "continue": silent per user policy
|
|
|
|
def on_train_end(
|
|
self, args: TrainingArguments, state: TrainerState,
|
|
control: TrainerControl, **kwargs,
|
|
):
|
|
self.display.stop()
|
|
self._run_auto_eval()
|
|
|
|
def _run_auto_eval(self) -> None:
|
|
"""Run automatic evaluation if configured."""
|
|
if self.eval_config is None:
|
|
return
|
|
if not getattr(self.eval_config, "auto_eval", False):
|
|
return
|
|
if not self.output_dir:
|
|
return
|
|
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
console.print("\n[bold blue]Running auto-eval...[/]")
|
|
|
|
benchmarks = getattr(self.eval_config, "benchmarks", None) or []
|
|
custom_tasks = getattr(self.eval_config, "custom_tasks", None)
|
|
|
|
# Run standard benchmarks
|
|
if benchmarks:
|
|
try:
|
|
from soup_cli.commands.eval import benchmark
|
|
benchmark(
|
|
model=self.output_dir,
|
|
benchmarks=",".join(benchmarks),
|
|
num_fewshot=None,
|
|
batch_size=8,
|
|
run_id=self.run_id,
|
|
device=None,
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("Auto-eval benchmark failed")
|
|
console.print(f"[yellow]Auto-eval benchmark failed: {exc}[/]")
|
|
|
|
# Run custom eval
|
|
if custom_tasks:
|
|
try:
|
|
from soup_cli.commands.eval import custom
|
|
custom(
|
|
tasks=custom_tasks,
|
|
model=self.output_dir,
|
|
run_id=self.run_id,
|
|
)
|
|
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 ".")
|
|
# Containment: never write outside cwd (matches v0.32.0 lr-finder /
|
|
# autopilot / cans policy). Caller's output_dir is already validated
|
|
# by SoupConfig but args.output_dir from raw HF TrainingArguments is
|
|
# not — defence-in-depth.
|
|
from soup_cli.utils.paths import is_under_cwd
|
|
if not is_under_cwd(out_dir):
|
|
logger.warning(
|
|
"spike_recovery hint skipped: output_dir %s is outside cwd",
|
|
out_dir,
|
|
)
|
|
return
|
|
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."
|
|
)
|