mirror of https://github.com/razor-ai/soup.git
454 lines
15 KiB
Python
454 lines
15 KiB
Python
"""soup sweep — hyperparameter search over training configs."""
|
||
|
||
import itertools
|
||
import random
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import typer
|
||
from rich.console import Console
|
||
from rich.panel import Panel
|
||
from rich.table import Table
|
||
|
||
from soup_cli.config.loader import load_config
|
||
|
||
console = Console()
|
||
|
||
|
||
def sweep(
|
||
config: str = typer.Option(
|
||
"soup.yaml",
|
||
"--config",
|
||
"-c",
|
||
help="Path to base soup.yaml config file",
|
||
),
|
||
param: list[str] = typer.Option(
|
||
...,
|
||
"--param",
|
||
"-p",
|
||
help="Parameter to sweep: key=val1,val2,val3 (e.g., lr=1e-5,2e-5,5e-5)",
|
||
),
|
||
strategy: str = typer.Option(
|
||
"grid",
|
||
"--strategy",
|
||
"-s",
|
||
help="Search strategy: grid, random",
|
||
),
|
||
max_runs: Optional[int] = typer.Option(
|
||
None,
|
||
"--max-runs",
|
||
help="Max number of runs (useful for random strategy)",
|
||
),
|
||
name: Optional[str] = typer.Option(
|
||
None,
|
||
"--name",
|
||
"-n",
|
||
help="Sweep experiment name prefix",
|
||
),
|
||
dry_run: bool = typer.Option(
|
||
False,
|
||
"--dry-run",
|
||
help="Show planned runs without executing",
|
||
),
|
||
early_stop: float = typer.Option(
|
||
None,
|
||
"--early-stop",
|
||
help="Stop early if run's loss exceeds best loss by this factor (e.g. 1.5 = 50% worse)",
|
||
),
|
||
yes: bool = typer.Option(
|
||
False,
|
||
"--yes",
|
||
"-y",
|
||
help="Skip confirmation prompt",
|
||
),
|
||
):
|
||
"""Run hyperparameter sweep: grid or random search over training parameters."""
|
||
config_path = Path(config)
|
||
if not config_path.exists():
|
||
console.print(f"[red]Config not found: {config_path}[/]")
|
||
raise typer.Exit(1)
|
||
|
||
if strategy not in ("grid", "random"):
|
||
console.print(f"[red]Invalid strategy: {strategy}. Must be grid or random.[/]")
|
||
raise typer.Exit(1)
|
||
|
||
# Parse sweep parameters
|
||
sweep_params = _parse_sweep_params(param)
|
||
if not sweep_params:
|
||
console.print("[red]No valid sweep parameters provided.[/]")
|
||
raise typer.Exit(1)
|
||
|
||
# Generate parameter combinations
|
||
combinations = _generate_combinations(sweep_params, strategy, max_runs)
|
||
|
||
console.print(
|
||
Panel(
|
||
f"Config: [bold]{config_path}[/]\n"
|
||
f"Strategy: [bold]{strategy}[/]\n"
|
||
f"Params: [bold]{', '.join(sweep_params.keys())}[/]\n"
|
||
f"Runs: [bold]{len(combinations)}[/]",
|
||
title="Sweep Plan",
|
||
)
|
||
)
|
||
|
||
# Show parameter table
|
||
param_table = Table(title="Parameter Grid")
|
||
param_table.add_column("Run", style="bold")
|
||
for key in sweep_params:
|
||
param_table.add_column(key)
|
||
|
||
for idx, combo in enumerate(combinations):
|
||
row_values = [str(combo[key]) for key in sweep_params]
|
||
param_table.add_row(f"#{idx + 1}", *row_values)
|
||
|
||
console.print(param_table)
|
||
|
||
if dry_run:
|
||
console.print("[yellow]Dry run - no training will be executed.[/]")
|
||
raise typer.Exit()
|
||
|
||
if not yes:
|
||
if not typer.confirm(f"Start {len(combinations)} training run(s)?", default=True):
|
||
console.print("[yellow]Cancelled.[/]")
|
||
raise typer.Exit()
|
||
|
||
# Execute sweep
|
||
base_cfg = load_config(config_path)
|
||
results = []
|
||
best_loss = float("inf")
|
||
skipped = 0
|
||
|
||
for idx, combo in enumerate(combinations):
|
||
run_name = f"{name or 'sweep'}_{idx + 1}"
|
||
console.print(f"\n[bold]--- Run {idx + 1}/{len(combinations)}: {run_name} ---[/]")
|
||
|
||
for key, val in combo.items():
|
||
console.print(f" {key} = {val}")
|
||
|
||
try:
|
||
result = _run_single(base_cfg, combo, run_name, config_path)
|
||
final_loss = result.get("final_loss", 0)
|
||
results.append({
|
||
"name": run_name,
|
||
"params": combo,
|
||
"run_id": result.get("run_id", ""),
|
||
"final_loss": final_loss,
|
||
"duration": result.get("duration", ""),
|
||
"status": "completed",
|
||
})
|
||
|
||
# Update best loss and check early stopping for remaining runs
|
||
if final_loss and final_loss < best_loss:
|
||
best_loss = final_loss
|
||
|
||
if early_stop and final_loss and best_loss < float("inf"):
|
||
if final_loss > best_loss * early_stop:
|
||
console.print(
|
||
f"[yellow]Loss {final_loss:.4f} exceeds threshold "
|
||
f"({best_loss:.4f} × {early_stop} = {best_loss * early_stop:.4f})[/]"
|
||
)
|
||
except Exception as exc:
|
||
console.print(f"[red]Run {run_name} failed: {exc}[/]")
|
||
results.append({
|
||
"name": run_name,
|
||
"params": combo,
|
||
"run_id": "",
|
||
"final_loss": 0,
|
||
"duration": "",
|
||
"status": "failed",
|
||
})
|
||
|
||
# Early stopping: skip remaining runs if too many are poor
|
||
if early_stop and len(results) >= 2:
|
||
completed = [r for r in results if r["status"] == "completed" and r["final_loss"]]
|
||
if completed:
|
||
recent = completed[-1]
|
||
if recent["final_loss"] > best_loss * early_stop:
|
||
remaining = len(combinations) - idx - 1
|
||
if remaining > 0:
|
||
skipped = remaining
|
||
console.print(
|
||
f"[yellow]Early stopping: skipping {remaining} remaining run(s). "
|
||
f"Last loss {recent['final_loss']:.4f} exceeded threshold.[/]"
|
||
)
|
||
break
|
||
|
||
# Summary table
|
||
_display_summary(results, sweep_params)
|
||
|
||
if skipped:
|
||
console.print(f"\n[yellow]Early stopping: {skipped} run(s) skipped.[/]")
|
||
|
||
|
||
def _parse_sweep_params(params: list[str]) -> dict[str, list]:
|
||
"""Parse sweep parameter strings into a dict of {key: [values]}."""
|
||
result = {}
|
||
for param_str in params:
|
||
if "=" not in param_str:
|
||
console.print(f"[yellow]Skipping invalid param: {param_str} (missing '=')[/]")
|
||
continue
|
||
|
||
key, values_str = param_str.split("=", 1)
|
||
key = key.strip()
|
||
values = []
|
||
|
||
for val in values_str.split(","):
|
||
val = val.strip()
|
||
values.append(_parse_value(val))
|
||
|
||
if values:
|
||
result[key] = values
|
||
|
||
return result
|
||
|
||
|
||
def _parse_value(val: str):
|
||
"""Parse a string value into the appropriate Python type."""
|
||
# Bool
|
||
if val.lower() in ("true", "false"):
|
||
return val.lower() == "true"
|
||
# None
|
||
if val.lower() == "none":
|
||
return None
|
||
# Int
|
||
try:
|
||
return int(val)
|
||
except ValueError:
|
||
pass
|
||
# Float (including scientific notation)
|
||
try:
|
||
return float(val)
|
||
except ValueError:
|
||
pass
|
||
# String
|
||
return val
|
||
|
||
|
||
def _generate_combinations(
|
||
sweep_params: dict[str, list],
|
||
strategy: str,
|
||
max_runs: Optional[int],
|
||
) -> list[dict]:
|
||
"""Generate parameter combinations based on strategy."""
|
||
keys = list(sweep_params.keys())
|
||
value_lists = [sweep_params[k] for k in keys]
|
||
|
||
if strategy == "grid":
|
||
combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
|
||
elif strategy == "random":
|
||
total_possible = 1
|
||
for vals in value_lists:
|
||
total_possible *= len(vals)
|
||
|
||
num_runs = max_runs or min(total_possible, 10)
|
||
num_runs = min(num_runs, total_possible)
|
||
|
||
if num_runs >= total_possible:
|
||
# Just do all of them
|
||
combos = [dict(zip(keys, vals)) for vals in itertools.product(*value_lists)]
|
||
else:
|
||
seen = set()
|
||
combos = []
|
||
while len(combos) < num_runs:
|
||
vals = tuple(random.choice(vals_list) for vals_list in value_lists)
|
||
if vals not in seen:
|
||
seen.add(vals)
|
||
combos.append(dict(zip(keys, vals)))
|
||
else:
|
||
combos = []
|
||
|
||
if max_runs and len(combos) > max_runs:
|
||
combos = combos[:max_runs]
|
||
|
||
return combos
|
||
|
||
|
||
def _set_nested_param(config_dict: dict, key: str, value) -> dict:
|
||
"""Set a nested parameter in a config dict using dot notation.
|
||
|
||
Supports keys like: lr, lora.r, training.epochs, etc.
|
||
Maps common short names to their full paths.
|
||
"""
|
||
# Short name mappings
|
||
shortcuts = {
|
||
"lr": "training.lr",
|
||
"epochs": "training.epochs",
|
||
"batch_size": "training.batch_size",
|
||
"lora_r": "training.lora.r",
|
||
"lora_alpha": "training.lora.alpha",
|
||
"lora_dropout": "training.lora.dropout",
|
||
"quantization": "training.quantization",
|
||
"warmup_ratio": "training.warmup_ratio",
|
||
"weight_decay": "training.weight_decay",
|
||
"gradient_accumulation_steps": "training.gradient_accumulation_steps",
|
||
"max_grad_norm": "training.max_grad_norm",
|
||
"optimizer": "training.optimizer",
|
||
"scheduler": "training.scheduler",
|
||
"val_split": "data.val_split",
|
||
"max_length": "data.max_length",
|
||
"dpo_beta": "training.dpo_beta",
|
||
"kto_beta": "training.kto_beta",
|
||
"grpo_beta": "training.grpo_beta",
|
||
"num_generations": "training.num_generations",
|
||
"reward_fn": "training.reward_fn",
|
||
"ppo_epochs": "training.ppo_epochs",
|
||
"ppo_clip_ratio": "training.ppo_clip_ratio",
|
||
"ppo_kl_penalty": "training.ppo_kl_penalty",
|
||
"reward_model": "training.reward_model",
|
||
"orpo_beta": "training.orpo_beta",
|
||
"simpo_gamma": "training.simpo_gamma",
|
||
"cpo_alpha": "training.cpo_alpha",
|
||
"ipo_tau": "training.ipo_tau",
|
||
"loraplus_lr_ratio": "training.loraplus_lr_ratio",
|
||
"use_dora": "training.lora.use_dora",
|
||
"use_galore": "training.use_galore",
|
||
"galore_rank": "training.galore_rank",
|
||
"backend": "backend",
|
||
}
|
||
|
||
full_key = shortcuts.get(key, key)
|
||
parts = full_key.split(".")
|
||
|
||
obj = config_dict
|
||
for part in parts[:-1]:
|
||
if part not in obj:
|
||
obj[part] = {}
|
||
obj = obj[part]
|
||
obj[parts[-1]] = value
|
||
|
||
return config_dict
|
||
|
||
|
||
def _run_single(base_cfg, params: dict, run_name: str, config_path: Path) -> dict:
|
||
"""Run a single training with modified parameters."""
|
||
from soup_cli.config.schema import SoupConfig
|
||
from soup_cli.data.loader import load_dataset
|
||
from soup_cli.experiment.tracker import ExperimentTracker
|
||
from soup_cli.monitoring.display import TrainingDisplay
|
||
from soup_cli.trainer.sft import SFTTrainerWrapper
|
||
from soup_cli.utils.gpu import detect_device, get_gpu_info
|
||
|
||
# Deep copy and modify config
|
||
config_dict = base_cfg.model_dump()
|
||
for key, val in params.items():
|
||
_set_nested_param(config_dict, key, val)
|
||
|
||
# Override experiment name
|
||
config_dict["experiment_name"] = run_name
|
||
cfg = SoupConfig(**config_dict)
|
||
|
||
# Detect hardware
|
||
device, device_name = detect_device()
|
||
gpu_info = get_gpu_info()
|
||
|
||
# Load data
|
||
dataset = load_dataset(cfg.data)
|
||
console.print(f"[dim]Loaded {len(dataset['train'])} train samples[/]")
|
||
|
||
# Start tracking
|
||
tracker = ExperimentTracker()
|
||
run_id = tracker.start_run(
|
||
config_dict=cfg.model_dump(),
|
||
device=device,
|
||
device_name=device_name,
|
||
gpu_info=gpu_info,
|
||
experiment_name=run_name,
|
||
)
|
||
|
||
# Build trainer
|
||
if cfg.task == "dpo":
|
||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||
|
||
trainer_wrapper = DPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "kto":
|
||
from soup_cli.trainer.kto import KTOTrainerWrapper
|
||
|
||
trainer_wrapper = KTOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "grpo":
|
||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||
|
||
trainer_wrapper = GRPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "ppo":
|
||
from soup_cli.trainer.ppo import PPOTrainerWrapper
|
||
|
||
trainer_wrapper = PPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "orpo":
|
||
from soup_cli.trainer.orpo import ORPOTrainerWrapper
|
||
|
||
trainer_wrapper = ORPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "simpo":
|
||
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
||
|
||
trainer_wrapper = SimPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "ipo":
|
||
from soup_cli.trainer.ipo import IPOTrainerWrapper
|
||
|
||
trainer_wrapper = IPOTrainerWrapper(cfg, device=device)
|
||
elif cfg.task == "reward_model":
|
||
from soup_cli.trainer.reward_model import RewardModelTrainerWrapper
|
||
|
||
trainer_wrapper = RewardModelTrainerWrapper(cfg, device=device)
|
||
else:
|
||
trainer_wrapper = SFTTrainerWrapper(cfg, device=device)
|
||
trainer_wrapper.setup(dataset)
|
||
|
||
# Train
|
||
display = TrainingDisplay(cfg, device_name=device_name)
|
||
try:
|
||
result = trainer_wrapper.train(display=display, tracker=tracker, run_id=run_id)
|
||
tracker.finish_run(
|
||
run_id=run_id,
|
||
initial_loss=result["initial_loss"],
|
||
final_loss=result["final_loss"],
|
||
total_steps=result["total_steps"],
|
||
duration_secs=result["duration_secs"],
|
||
output_dir=result["output_dir"],
|
||
)
|
||
result["run_id"] = run_id
|
||
return result
|
||
except Exception:
|
||
tracker.fail_run(run_id)
|
||
raise
|
||
|
||
|
||
def _display_summary(results: list[dict], sweep_params: dict[str, list]):
|
||
"""Display sweep results summary table."""
|
||
table = Table(title="Sweep Results")
|
||
table.add_column("Run", style="bold")
|
||
for key in sweep_params:
|
||
table.add_column(key)
|
||
table.add_column("Final Loss", justify="right", style="green")
|
||
table.add_column("Duration", justify="right")
|
||
table.add_column("Status")
|
||
|
||
# Sort by final loss (best first)
|
||
sorted_results = sorted(results, key=lambda r: r.get("final_loss", float("inf")))
|
||
|
||
for idx, res in enumerate(sorted_results):
|
||
status_style = "green" if res["status"] == "completed" else "red"
|
||
param_vals = [str(res["params"].get(k, "")) for k in sweep_params]
|
||
loss_str = f"{res['final_loss']:.4f}" if res["final_loss"] else "-"
|
||
best_marker = " [bold yellow]*[/]" if idx == 0 and res["status"] == "completed" else ""
|
||
table.add_row(
|
||
res["name"],
|
||
*param_vals,
|
||
f"{loss_str}{best_marker}",
|
||
res.get("duration", "-"),
|
||
f"[{status_style}]{res['status']}[/]",
|
||
)
|
||
|
||
console.print(table)
|
||
|
||
# Best run
|
||
completed = [r for r in sorted_results if r["status"] == "completed"]
|
||
if completed:
|
||
best = completed[0]
|
||
console.print(
|
||
f"\n[bold green]Best run:[/] {best['name']} "
|
||
f"(loss: {best['final_loss']:.4f})"
|
||
)
|
||
for key, val in best["params"].items():
|
||
console.print(f" {key} = {val}")
|
||
if best.get("run_id"):
|
||
console.print(f"\n[dim]View details: soup runs show {best['run_id']}[/]")
|