mirror of https://github.com/razor-ai/soup.git
Add GRPO reasoning training (Phase 4) — v0.4.2
- New `task: grpo` for Group Relative Policy Optimization (DeepSeek-R1 style) - Reward function API: built-in (accuracy, format) + custom Python files - Config fields: grpo_beta, num_generations, reward_fn - `soup init --template reasoning` template - GRPOTrainerWrapper with auto batch size scaling for multi-generation - Sweep shortcuts for grpo_beta, num_generations, reward_fn - 42 new tests (371 total), lint clean Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
0ad2b182b9
commit
5296b0019b
|
|
@ -43,7 +43,9 @@ soup train --config soup.yaml
|
|||
|
||||
**Data pipeline:** `data/loader.py` handles local files (JSONL/JSON/CSV/Parquet) and HuggingFace datasets. `data/formats.py` auto-detects and normalizes alpaca/sharegpt/chatml formats into a unified `{"messages": [...]}` structure. Also supports reverse conversion via `messages_to_format()`.
|
||||
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`) and `trainer/dpo.py` (`DPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. Both trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`.
|
||||
**Trainer:** `trainer/sft.py` (`SFTTrainerWrapper`), `trainer/dpo.py` (`DPOTrainerWrapper`), and `trainer/grpo.py` (`GRPOTrainerWrapper`) wrap HuggingFace's SFTTrainer/DPOTrainer/GRPOTrainer with auto quantization (BitsAndBytes), LoRA (PEFT), and batch size estimation. Heavy ML imports are lazy (inside methods) so CLI stays fast for non-training commands. All trainers enable Rich progress bars for HuggingFace Hub model downloads via `_enable_hf_transfer_progress()`.
|
||||
|
||||
**GRPO (Group Relative Policy Optimization):** `trainer/grpo.py` implements reasoning model training (DeepSeek-R1 style). Generates multiple completions per prompt, scores them with reward functions, and optimizes using group-relative advantages. `trainer/rewards.py` provides built-in reward functions (`accuracy` — checks final answer, `format` — checks `<think>` blocks) and supports custom rewards via Python files. Config: `task: grpo`, `grpo_beta`, `num_generations`, `reward_fn`.
|
||||
|
||||
**Monitoring:** `monitoring/callback.py` is a HuggingFace `TrainerCallback` that streams metrics to `monitoring/display.py` (Rich Live panel at 2Hz) and optionally to the experiment tracker.
|
||||
|
||||
|
|
@ -141,4 +143,5 @@ Test suite lives in `tests/`:
|
|||
| `test_errors.py` | Friendly error messages, --verbose flag, error mapping |
|
||||
| `test_doctor.py` | `soup doctor` command, version checking, dependency table |
|
||||
| `test_quickstart.py` | `soup quickstart` demo, data/config creation, --dry-run |
|
||||
| `test_grpo.py` | GRPO config, rewards, data prep, template, sweep shortcuts |
|
||||
| `test_progress.py` | Rich download progress bar, `_enable_hf_transfer_progress` |
|
||||
|
|
|
|||
59
README.md
59
README.md
|
|
@ -63,9 +63,10 @@ pip install git+https://github.com/MakazhanAlpamys/Soup.git
|
|||
soup init
|
||||
|
||||
# Or use a template
|
||||
soup init --template chat # conversational fine-tune
|
||||
soup init --template code # code generation
|
||||
soup init --template medical # domain expert
|
||||
soup init --template chat # conversational fine-tune
|
||||
soup init --template code # code generation
|
||||
soup init --template medical # domain expert
|
||||
soup init --template reasoning # GRPO reasoning training
|
||||
```
|
||||
|
||||
### 3. Train
|
||||
|
|
@ -142,6 +143,55 @@ training:
|
|||
quantization: 4bit
|
||||
```
|
||||
|
||||
## GRPO Training (Reasoning)
|
||||
|
||||
Train reasoning models with Group Relative Policy Optimization (DeepSeek-R1 style):
|
||||
|
||||
```yaml
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: grpo
|
||||
|
||||
data:
|
||||
train: ./data/reasoning_train.jsonl
|
||||
format: sharegpt
|
||||
max_length: 4096
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 1e-5
|
||||
grpo_beta: 0.1
|
||||
num_generations: 4
|
||||
reward_fn: accuracy # or 'format', or path to custom .py
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
quantization: 4bit
|
||||
```
|
||||
|
||||
```bash
|
||||
# Create a reasoning config
|
||||
soup init --template reasoning
|
||||
|
||||
# Train
|
||||
soup train --config soup.yaml
|
||||
```
|
||||
|
||||
**Built-in reward functions:**
|
||||
- `accuracy` — checks if the final answer matches expected (supports `####` and `\boxed{}` formats)
|
||||
- `format` — checks for structured `<think>...</think>` reasoning blocks
|
||||
|
||||
**Custom reward functions** — point to a Python file:
|
||||
```python
|
||||
# my_reward.py
|
||||
def reward_fn(completions, **kwargs):
|
||||
"""Score each completion. Return list of floats."""
|
||||
return [1.0 if "correct" in c[-1]["content"] else 0.0 for c in completions]
|
||||
```
|
||||
```yaml
|
||||
training:
|
||||
reward_fn: ./my_reward.py
|
||||
```
|
||||
|
||||
## Chat with your model
|
||||
|
||||
```bash
|
||||
|
|
@ -463,6 +513,7 @@ soup eval --model ./output --benchmarks mmlu --run-id run_20260223_143052_a1b2
|
|||
| LoRA / QLoRA fine-tuning | ✅ |
|
||||
| SFT (Supervised Fine-Tune) | ✅ |
|
||||
| DPO (Direct Preference Optimization) | ✅ |
|
||||
| GRPO (Reasoning / DeepSeek-R1 style) | ✅ |
|
||||
| Auto batch size | ✅ |
|
||||
| Auto GPU detection (CUDA/MPS/CPU) | ✅ |
|
||||
| Live terminal dashboard | ✅ |
|
||||
|
|
@ -493,7 +544,7 @@ soup eval --model ./output --benchmarks mmlu --run-id run_20260223_143052_a1b2
|
|||
## All Commands
|
||||
|
||||
```
|
||||
soup init [--template chat|code|medical] Create config
|
||||
soup init [--template chat|code|medical|reasoning] Create config
|
||||
soup train --config soup.yaml Start training
|
||||
soup chat --model ./output Interactive chat
|
||||
soup push --model ./output --repo user/name Upload to HuggingFace
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.4.1"
|
||||
version = "0.4.2"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.4.1"
|
||||
__version__ = "0.4.2"
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ def init(
|
|||
None,
|
||||
"--template",
|
||||
"-t",
|
||||
help="Use a template: chat, code, medical",
|
||||
help="Use a template: chat, code, medical, reasoning",
|
||||
),
|
||||
output: str = typer.Option(
|
||||
"soup.yaml",
|
||||
|
|
@ -62,7 +62,7 @@ def _interactive_wizard() -> str:
|
|||
"Base model",
|
||||
default="meta-llama/Llama-3.1-8B-Instruct",
|
||||
)
|
||||
task = Prompt.ask("Task", choices=["sft", "dpo"], default="sft")
|
||||
task = Prompt.ask("Task", choices=["sft", "dpo", "grpo"], default="sft")
|
||||
data_path = Prompt.ask("Training data path", default="./data/train.jsonl")
|
||||
data_format = Prompt.ask(
|
||||
"Data format", choices=["alpaca", "sharegpt", "chatml"], default="alpaca",
|
||||
|
|
@ -72,6 +72,18 @@ def _interactive_wizard() -> str:
|
|||
|
||||
quantization = "4bit" if use_qlora == "yes" else "none"
|
||||
|
||||
grpo_block = ""
|
||||
if task == "grpo":
|
||||
reward_fn = Prompt.ask(
|
||||
"Reward function", choices=["accuracy", "format", "custom"], default="accuracy",
|
||||
)
|
||||
if reward_fn == "custom":
|
||||
reward_fn = Prompt.ask("Path to reward .py file", default="./reward.py")
|
||||
grpo_block = f""" grpo_beta: 0.1
|
||||
num_generations: 4
|
||||
reward_fn: {reward_fn}
|
||||
"""
|
||||
|
||||
return f"""# Soup training config
|
||||
# Docs: https://github.com/MakazhanAlpamys/Soup
|
||||
|
||||
|
|
@ -92,6 +104,6 @@ training:
|
|||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: {quantization}
|
||||
|
||||
{grpo_block}
|
||||
output: ./output
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -287,6 +287,9 @@ def _set_nested_param(config_dict: dict, key: str, value) -> dict:
|
|||
"val_split": "data.val_split",
|
||||
"max_length": "data.max_length",
|
||||
"dpo_beta": "training.dpo_beta",
|
||||
"grpo_beta": "training.grpo_beta",
|
||||
"num_generations": "training.num_generations",
|
||||
"reward_fn": "training.reward_fn",
|
||||
}
|
||||
|
||||
full_key = shortcuts.get(key, key)
|
||||
|
|
|
|||
|
|
@ -154,6 +154,12 @@ def train(
|
|||
trainer_wrapper = DPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
elif cfg.task == "grpo":
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = GRPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
else:
|
||||
trainer_wrapper = SFTTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
|
|
|
|||
|
|
@ -49,13 +49,24 @@ class TrainingConfig(BaseModel):
|
|||
dpo_beta: float = Field(
|
||||
default=0.1, gt=0, description="DPO beta — KL penalty coefficient"
|
||||
)
|
||||
# GRPO-specific
|
||||
grpo_beta: float = Field(
|
||||
default=0.1, gt=0, description="GRPO beta — KL penalty coefficient"
|
||||
)
|
||||
num_generations: int = Field(
|
||||
default=4, ge=2, description="Number of generations per prompt for GRPO"
|
||||
)
|
||||
reward_fn: Optional[str] = Field(
|
||||
default="accuracy",
|
||||
description="Reward function: 'accuracy', 'format', or path to custom .py file",
|
||||
)
|
||||
|
||||
|
||||
class SoupConfig(BaseModel):
|
||||
"""Root config for soup.yaml."""
|
||||
|
||||
base: str = Field(..., description="Base model name or path (HF model ID)")
|
||||
task: Literal["sft", "dpo"] = Field(default="sft", description="Training task type")
|
||||
task: Literal["sft", "dpo", "grpo"] = Field(default="sft", description="Training task type")
|
||||
data: DataConfig
|
||||
training: TrainingConfig = Field(default_factory=TrainingConfig)
|
||||
output: str = Field(default="./output", description="Output directory for trained model")
|
||||
|
|
@ -111,6 +122,34 @@ training:
|
|||
target_modules: auto
|
||||
quantization: 4bit
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"reasoning": """# Soup template: Reasoning / GRPO
|
||||
# Fine-tune a model for chain-of-thought reasoning with GRPO
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: grpo
|
||||
|
||||
data:
|
||||
train: ./data/reasoning_train.jsonl
|
||||
format: sharegpt
|
||||
val_split: 0.1
|
||||
max_length: 4096
|
||||
|
||||
training:
|
||||
epochs: 3
|
||||
lr: 1e-5
|
||||
batch_size: auto
|
||||
gradient_accumulation_steps: 8
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
grpo_beta: 0.1
|
||||
num_generations: 4
|
||||
reward_fn: accuracy
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"medical": """# Soup template: Medical / Domain Expert
|
||||
|
|
|
|||
|
|
@ -0,0 +1,277 @@
|
|||
"""GRPO (Group Relative Policy Optimization) trainer — wraps trl.GRPOTrainer."""
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import 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
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
class GRPOTrainerWrapper:
|
||||
"""High-level wrapper for GRPO training from SoupConfig.
|
||||
|
||||
GRPO generates multiple completions per prompt, scores them with a reward
|
||||
function, and optimizes using group-relative advantages. This is the approach
|
||||
used by DeepSeek-R1 for reasoning model training.
|
||||
|
||||
Data format: same as SFT (messages with prompt/response) or DPO-style prompts.
|
||||
The reward_fn in config determines how completions are scored.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: SoupConfig,
|
||||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
||||
def setup(self, dataset: dict):
|
||||
"""Load model, tokenizer, apply LoRA, create GRPO trainer."""
|
||||
from datasets import Dataset
|
||||
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
BitsAndBytesConfig,
|
||||
)
|
||||
from trl import GRPOConfig, GRPOTrainer
|
||||
|
||||
# Enable Rich progress bar for HuggingFace downloads
|
||||
from soup_cli.trainer.sft import _enable_hf_transfer_progress
|
||||
|
||||
_enable_hf_transfer_progress()
|
||||
|
||||
cfg = self.config
|
||||
tcfg = cfg.training
|
||||
|
||||
# --- Load reward function ---
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
reward_fn = load_reward_fn(tcfg.reward_fn)
|
||||
|
||||
# --- Tokenizer ---
|
||||
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
|
||||
|
||||
# --- Quantization ---
|
||||
bnb_config = None
|
||||
if tcfg.quantization == "4bit":
|
||||
import torch
|
||||
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
elif tcfg.quantization == "8bit":
|
||||
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
# --- Model ---
|
||||
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
||||
model_kwargs = {"trust_remote_code": True, "device_map": "auto"}
|
||||
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)
|
||||
|
||||
# --- LoRA ---
|
||||
target_modules = tcfg.lora.target_modules
|
||||
if target_modules == "auto":
|
||||
target_modules = None # peft will auto-detect
|
||||
|
||||
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",
|
||||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
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,
|
||||
)
|
||||
# GRPO generates N completions per prompt → more memory
|
||||
batch_size = max(1, batch_size // tcfg.num_generations)
|
||||
console.print(f"[green]Auto batch size (GRPO):[/] {batch_size}")
|
||||
|
||||
# --- Dataset ---
|
||||
# GRPO expects prompts — extract from messages or use prompt field
|
||||
train_data = _prepare_grpo_dataset(dataset["train"])
|
||||
train_ds = Dataset.from_list(train_data)
|
||||
eval_ds = None
|
||||
if "val" in dataset and dataset["val"]:
|
||||
eval_data = _prepare_grpo_dataset(dataset["val"])
|
||||
eval_ds = Dataset.from_list(eval_data)
|
||||
|
||||
# --- 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 ---
|
||||
import math
|
||||
|
||||
total_steps = (
|
||||
math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps)
|
||||
* tcfg.epochs
|
||||
)
|
||||
warmup_steps = int(total_steps * tcfg.warmup_ratio)
|
||||
|
||||
# --- GRPO config ---
|
||||
grpo_config = GRPOConfig(
|
||||
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,
|
||||
beta=tcfg.grpo_beta,
|
||||
num_generations=tcfg.num_generations,
|
||||
max_completion_length=cfg.data.max_length,
|
||||
)
|
||||
|
||||
# --- Trainer ---
|
||||
self.trainer = GRPOTrainer(
|
||||
model=self.model,
|
||||
args=grpo_config,
|
||||
train_dataset=train_ds,
|
||||
eval_dataset=eval_ds,
|
||||
reward_funcs=reward_fn,
|
||||
processing_class=self.tokenizer,
|
||||
)
|
||||
|
||||
self._output_dir = str(output_dir)
|
||||
|
||||
def train(
|
||||
self,
|
||||
display: Optional[object] = None,
|
||||
tracker: Optional[object] = None,
|
||||
run_id: str = "",
|
||||
resume_from_checkpoint: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Run GRPO training and return results summary."""
|
||||
start = time.time()
|
||||
|
||||
# Add callback for live display and experiment tracking
|
||||
if display:
|
||||
from soup_cli.monitoring.callback import SoupTrainerCallback
|
||||
|
||||
self.trainer.add_callback(
|
||||
SoupTrainerCallback(display, tracker=tracker, run_id=run_id)
|
||||
)
|
||||
|
||||
self.trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
||||
duration = time.time() - start
|
||||
|
||||
# Save final model (LoRA adapter)
|
||||
self.trainer.save_model(self._output_dir)
|
||||
self.tokenizer.save_pretrained(self._output_dir)
|
||||
|
||||
# Extract metrics
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
def _prepare_grpo_dataset(data: list[dict]) -> list[dict]:
|
||||
"""Convert dataset rows to GRPO format.
|
||||
|
||||
GRPO expects each row to have a 'prompt' field (list of messages or string).
|
||||
Input can be:
|
||||
- messages format: [{"role": "user", "content": "..."}, ...]
|
||||
- DPO format: {"prompt": "...", "chosen": "...", "rejected": "..."}
|
||||
- prompt field: {"prompt": "..."}
|
||||
|
||||
Returns list of dicts with 'prompt' as a message list for chat models.
|
||||
"""
|
||||
prepared = []
|
||||
for row in data:
|
||||
if "prompt" in row and isinstance(row["prompt"], str):
|
||||
# DPO or plain prompt format — convert to message list
|
||||
entry = {"prompt": [{"role": "user", "content": row["prompt"]}]}
|
||||
# Preserve 'answer' field if present (for accuracy reward)
|
||||
if "answer" in row:
|
||||
entry["answer"] = row["answer"]
|
||||
prepared.append(entry)
|
||||
elif "messages" in row:
|
||||
# Messages format — use the user message(s) as prompt
|
||||
messages = row["messages"]
|
||||
prompt_msgs = [msg for msg in messages if msg["role"] != "assistant"]
|
||||
entry = {"prompt": prompt_msgs}
|
||||
prepared.append(entry)
|
||||
elif "prompt" in row and isinstance(row["prompt"], list):
|
||||
# Already in message list format
|
||||
entry = {"prompt": row["prompt"]}
|
||||
if "answer" in row:
|
||||
entry["answer"] = row["answer"]
|
||||
prepared.append(entry)
|
||||
else:
|
||||
# Fallback: treat any 'instruction' field as prompt
|
||||
instruction = row.get("instruction", row.get("input", ""))
|
||||
entry = {"prompt": [{"role": "user", "content": str(instruction)}]}
|
||||
if "output" in row:
|
||||
entry["answer"] = row["output"]
|
||||
prepared.append(entry)
|
||||
return prepared
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
"""Reward functions for GRPO training.
|
||||
|
||||
Built-in reward functions:
|
||||
- accuracy: checks if the model answer matches the expected answer
|
||||
- format: checks if the response follows a structured format (e.g., <think>...</think>)
|
||||
|
||||
Custom reward functions can be loaded from a Python file with a
|
||||
`reward_fn(completions, **kwargs)` callable.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
def accuracy_reward(completions: list[list[dict]], **kwargs) -> list[float]:
|
||||
"""Reward based on whether the final answer matches the expected answer.
|
||||
|
||||
Looks for the answer after the last '####' or in a \\boxed{} block.
|
||||
Falls back to checking if the expected answer appears anywhere in the response.
|
||||
|
||||
Args:
|
||||
completions: list of message lists, each containing a completion with 'content'.
|
||||
**kwargs: must contain 'answer' — the expected answer for each prompt.
|
||||
|
||||
Returns:
|
||||
List of float rewards (1.0 for correct, 0.0 for incorrect).
|
||||
"""
|
||||
answers = kwargs.get("answer", [])
|
||||
rewards = []
|
||||
for completion, expected in zip(completions, answers):
|
||||
content = completion[-1]["content"] if completion else ""
|
||||
predicted = _extract_answer(content)
|
||||
if predicted is not None and predicted.strip() == str(expected).strip():
|
||||
rewards.append(1.0)
|
||||
elif str(expected).strip().lower() in content.lower():
|
||||
rewards.append(0.5)
|
||||
else:
|
||||
rewards.append(0.0)
|
||||
return rewards
|
||||
|
||||
|
||||
def format_reward(completions: list[list[dict]], **kwargs) -> list[float]:
|
||||
"""Reward based on whether the response follows a structured reasoning format.
|
||||
|
||||
Checks for:
|
||||
- <think>...</think> block (chain-of-thought)
|
||||
- A final answer section after the thinking block
|
||||
|
||||
Args:
|
||||
completions: list of message lists.
|
||||
**kwargs: unused.
|
||||
|
||||
Returns:
|
||||
List of float rewards (0.0 to 1.0).
|
||||
"""
|
||||
rewards = []
|
||||
for completion in completions:
|
||||
content = completion[-1]["content"] if completion else ""
|
||||
score = 0.0
|
||||
# Check for <think> block
|
||||
if re.search(r"<think>.*?</think>", content, re.DOTALL):
|
||||
score += 0.5
|
||||
# Check for content after </think>
|
||||
after_think = re.split(r"</think>", content)
|
||||
if len(after_think) > 1 and after_think[-1].strip():
|
||||
score += 0.5
|
||||
rewards.append(score)
|
||||
return rewards
|
||||
|
||||
|
||||
def _extract_answer(text: str) -> str | None:
|
||||
"""Extract the final answer from model output.
|
||||
|
||||
Supports:
|
||||
- #### <answer> format (GSM8K style)
|
||||
- \\boxed{<answer>} format (math style)
|
||||
"""
|
||||
# Try #### format
|
||||
parts = text.split("####")
|
||||
if len(parts) > 1:
|
||||
return parts[-1].strip()
|
||||
# Try \\boxed{} format
|
||||
match = re.search(r"\\boxed\{([^}]+)\}", text)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
return None
|
||||
|
||||
|
||||
# Registry of built-in reward functions
|
||||
BUILTIN_REWARDS: dict[str, callable] = {
|
||||
"accuracy": accuracy_reward,
|
||||
"format": format_reward,
|
||||
}
|
||||
|
||||
|
||||
def load_reward_fn(reward_fn_spec: str) -> callable:
|
||||
"""Load a reward function by name or from a custom Python file.
|
||||
|
||||
Args:
|
||||
reward_fn_spec: Either a built-in name ('accuracy', 'format') or
|
||||
a path to a .py file containing a `reward_fn` callable.
|
||||
|
||||
Returns:
|
||||
A callable reward function with signature:
|
||||
(completions: list[list[dict]], **kwargs) -> list[float]
|
||||
"""
|
||||
# Built-in reward function
|
||||
if reward_fn_spec in BUILTIN_REWARDS:
|
||||
console.print(f"[dim]Using built-in reward function: {reward_fn_spec}[/]")
|
||||
return BUILTIN_REWARDS[reward_fn_spec]
|
||||
|
||||
# Custom Python file
|
||||
reward_path = Path(reward_fn_spec)
|
||||
if reward_path.exists() and reward_path.suffix == ".py":
|
||||
console.print(f"[dim]Loading custom reward function from: {reward_path}[/]")
|
||||
spec = importlib.util.spec_from_file_location("custom_reward", reward_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if not hasattr(module, "reward_fn"):
|
||||
raise ValueError(
|
||||
f"Custom reward file {reward_path} must define a 'reward_fn' callable.\n"
|
||||
f"Example:\n"
|
||||
f" def reward_fn(completions, **kwargs):\n"
|
||||
f" return [1.0] * len(completions)"
|
||||
)
|
||||
return module.reward_fn
|
||||
|
||||
raise ValueError(
|
||||
f"Unknown reward function: '{reward_fn_spec}'\n"
|
||||
f"Options: {', '.join(BUILTIN_REWARDS.keys())} or path to a .py file"
|
||||
)
|
||||
|
|
@ -0,0 +1,426 @@
|
|||
"""Tests for GRPO training — config, rewards, data preparation, template."""
|
||||
|
||||
import textwrap
|
||||
|
||||
import pytest
|
||||
|
||||
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
||||
|
||||
# ─── Config Tests ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGRPOConfig:
|
||||
"""Test GRPO task config validation."""
|
||||
|
||||
def test_grpo_task_accepted(self):
|
||||
"""GRPO task should be a valid task type."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.task == "grpo"
|
||||
|
||||
def test_grpo_beta_default(self):
|
||||
"""grpo_beta should default to 0.1."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.training.grpo_beta == 0.1
|
||||
|
||||
def test_grpo_beta_custom(self):
|
||||
"""Custom grpo_beta should be accepted."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"grpo_beta": 0.04},
|
||||
)
|
||||
assert cfg.training.grpo_beta == pytest.approx(0.04)
|
||||
|
||||
def test_grpo_beta_must_be_positive(self):
|
||||
"""grpo_beta must be > 0."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"grpo_beta": 0},
|
||||
)
|
||||
|
||||
def test_num_generations_default(self):
|
||||
"""num_generations should default to 4."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.training.num_generations == 4
|
||||
|
||||
def test_num_generations_custom(self):
|
||||
"""Custom num_generations should be accepted."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"num_generations": 8},
|
||||
)
|
||||
assert cfg.training.num_generations == 8
|
||||
|
||||
def test_num_generations_minimum(self):
|
||||
"""num_generations must be >= 2."""
|
||||
with pytest.raises(Exception):
|
||||
SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"num_generations": 1},
|
||||
)
|
||||
|
||||
def test_reward_fn_default(self):
|
||||
"""reward_fn should default to 'accuracy'."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
assert cfg.training.reward_fn == "accuracy"
|
||||
|
||||
def test_reward_fn_custom_path(self):
|
||||
"""reward_fn should accept a custom file path."""
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"reward_fn": "./my_reward.py"},
|
||||
)
|
||||
assert cfg.training.reward_fn == "./my_reward.py"
|
||||
|
||||
def test_grpo_full_config(self):
|
||||
"""Full GRPO config should validate correctly."""
|
||||
cfg = SoupConfig(
|
||||
base="meta-llama/Llama-3.1-8B-Instruct",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl", "format": "sharegpt", "max_length": 4096},
|
||||
training={
|
||||
"epochs": 3,
|
||||
"lr": 1e-5,
|
||||
"grpo_beta": 0.1,
|
||||
"num_generations": 4,
|
||||
"reward_fn": "format",
|
||||
"lora": {"r": 64, "alpha": 16},
|
||||
"quantization": "4bit",
|
||||
},
|
||||
)
|
||||
assert cfg.task == "grpo"
|
||||
assert cfg.training.reward_fn == "format"
|
||||
assert cfg.training.num_generations == 4
|
||||
assert cfg.data.max_length == 4096
|
||||
|
||||
|
||||
# ─── Reward Function Tests ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestAccuracyReward:
|
||||
"""Test the accuracy reward function."""
|
||||
|
||||
def test_exact_match(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [[{"role": "assistant", "content": "The answer is #### 42"}]]
|
||||
rewards = accuracy_reward(completions, answer=["42"])
|
||||
assert rewards == [1.0]
|
||||
|
||||
def test_boxed_match(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [[{"role": "assistant", "content": "So \\boxed{42} is the result"}]]
|
||||
rewards = accuracy_reward(completions, answer=["42"])
|
||||
assert rewards == [1.0]
|
||||
|
||||
def test_partial_match(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [[{"role": "assistant", "content": "The answer is 42 degrees"}]]
|
||||
rewards = accuracy_reward(completions, answer=["42"])
|
||||
assert rewards == [0.5]
|
||||
|
||||
def test_no_match(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [[{"role": "assistant", "content": "I don't know"}]]
|
||||
rewards = accuracy_reward(completions, answer=["42"])
|
||||
assert rewards == [0.0]
|
||||
|
||||
def test_multiple_completions(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [
|
||||
[{"role": "assistant", "content": "#### 42"}],
|
||||
[{"role": "assistant", "content": "Wrong answer"}],
|
||||
[{"role": "assistant", "content": "The answer is 42"}],
|
||||
]
|
||||
rewards = accuracy_reward(completions, answer=["42", "42", "42"])
|
||||
assert rewards == [1.0, 0.0, 0.5]
|
||||
|
||||
def test_empty_completion(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward
|
||||
|
||||
completions = [[]]
|
||||
rewards = accuracy_reward(completions, answer=["42"])
|
||||
assert rewards == [0.0]
|
||||
|
||||
|
||||
class TestFormatReward:
|
||||
"""Test the format reward function."""
|
||||
|
||||
def test_perfect_format(self):
|
||||
from soup_cli.trainer.rewards import format_reward
|
||||
|
||||
content = "<think>Let me think step by step...</think>\nThe answer is 42."
|
||||
completions = [[{"role": "assistant", "content": content}]]
|
||||
rewards = format_reward(completions)
|
||||
assert rewards == [1.0]
|
||||
|
||||
def test_think_only(self):
|
||||
from soup_cli.trainer.rewards import format_reward
|
||||
|
||||
content = "<think>Thinking...</think>"
|
||||
completions = [[{"role": "assistant", "content": content}]]
|
||||
rewards = format_reward(completions)
|
||||
assert rewards == [0.5]
|
||||
|
||||
def test_no_format(self):
|
||||
from soup_cli.trainer.rewards import format_reward
|
||||
|
||||
completions = [[{"role": "assistant", "content": "Just a plain answer"}]]
|
||||
rewards = format_reward(completions)
|
||||
assert rewards == [0.0]
|
||||
|
||||
def test_multiple_completions(self):
|
||||
from soup_cli.trainer.rewards import format_reward
|
||||
|
||||
completions = [
|
||||
[{"role": "assistant", "content": "<think>A</think>\nB"}],
|
||||
[{"role": "assistant", "content": "No format"}],
|
||||
]
|
||||
rewards = format_reward(completions)
|
||||
assert rewards == [1.0, 0.0]
|
||||
|
||||
|
||||
class TestExtractAnswer:
|
||||
"""Test answer extraction from model output."""
|
||||
|
||||
def test_hash_format(self):
|
||||
from soup_cli.trainer.rewards import _extract_answer
|
||||
|
||||
assert _extract_answer("Some work\n#### 42") == "42"
|
||||
|
||||
def test_boxed_format(self):
|
||||
from soup_cli.trainer.rewards import _extract_answer
|
||||
|
||||
assert _extract_answer("So \\boxed{42} is the answer") == "42"
|
||||
|
||||
def test_no_answer(self):
|
||||
from soup_cli.trainer.rewards import _extract_answer
|
||||
|
||||
assert _extract_answer("Just plain text") is None
|
||||
|
||||
def test_multiple_hashes(self):
|
||||
from soup_cli.trainer.rewards import _extract_answer
|
||||
|
||||
assert _extract_answer("#### step\n#### 42") == "42"
|
||||
|
||||
|
||||
class TestLoadRewardFn:
|
||||
"""Test reward function loading."""
|
||||
|
||||
def test_load_builtin_accuracy(self):
|
||||
from soup_cli.trainer.rewards import accuracy_reward, load_reward_fn
|
||||
|
||||
fn = load_reward_fn("accuracy")
|
||||
assert fn is accuracy_reward
|
||||
|
||||
def test_load_builtin_format(self):
|
||||
from soup_cli.trainer.rewards import format_reward, load_reward_fn
|
||||
|
||||
fn = load_reward_fn("format")
|
||||
assert fn is format_reward
|
||||
|
||||
def test_load_custom_file(self, tmp_path):
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
custom_file = tmp_path / "my_reward.py"
|
||||
custom_file.write_text(textwrap.dedent("""\
|
||||
def reward_fn(completions, **kwargs):
|
||||
return [1.0] * len(completions)
|
||||
"""))
|
||||
fn = load_reward_fn(str(custom_file))
|
||||
result = fn([[{"content": "test"}]])
|
||||
assert result == [1.0]
|
||||
|
||||
def test_load_custom_file_missing_fn(self, tmp_path):
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
custom_file = tmp_path / "bad_reward.py"
|
||||
custom_file.write_text("x = 1\n")
|
||||
with pytest.raises(ValueError, match="must define a 'reward_fn'"):
|
||||
load_reward_fn(str(custom_file))
|
||||
|
||||
def test_load_unknown_name(self):
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown reward function"):
|
||||
load_reward_fn("nonexistent")
|
||||
|
||||
|
||||
# ─── Data Preparation Tests ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestPrepareGRPODataset:
|
||||
"""Test GRPO dataset preparation."""
|
||||
|
||||
def test_from_prompt_string(self):
|
||||
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
||||
|
||||
data = [{"prompt": "What is 2+2?", "answer": "4"}]
|
||||
result = _prepare_grpo_dataset(data)
|
||||
assert len(result) == 1
|
||||
assert result[0]["prompt"] == [{"role": "user", "content": "What is 2+2?"}]
|
||||
assert result[0]["answer"] == "4"
|
||||
|
||||
def test_from_messages(self):
|
||||
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
||||
|
||||
data = [
|
||||
{
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi!"},
|
||||
]
|
||||
}
|
||||
]
|
||||
result = _prepare_grpo_dataset(data)
|
||||
assert len(result) == 1
|
||||
# Should only include non-assistant messages as prompt
|
||||
assert len(result[0]["prompt"]) == 2
|
||||
assert result[0]["prompt"][0]["role"] == "system"
|
||||
assert result[0]["prompt"][1]["role"] == "user"
|
||||
|
||||
def test_from_prompt_message_list(self):
|
||||
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
||||
|
||||
data = [
|
||||
{
|
||||
"prompt": [{"role": "user", "content": "What is 2+2?"}],
|
||||
"answer": "4",
|
||||
}
|
||||
]
|
||||
result = _prepare_grpo_dataset(data)
|
||||
assert result[0]["prompt"] == [{"role": "user", "content": "What is 2+2?"}]
|
||||
assert result[0]["answer"] == "4"
|
||||
|
||||
def test_from_alpaca_format(self):
|
||||
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
||||
|
||||
data = [{"instruction": "Translate hello", "input": "", "output": "hola"}]
|
||||
result = _prepare_grpo_dataset(data)
|
||||
assert result[0]["prompt"] == [{"role": "user", "content": "Translate hello"}]
|
||||
assert result[0]["answer"] == "hola"
|
||||
|
||||
def test_multiple_rows(self):
|
||||
from soup_cli.trainer.grpo import _prepare_grpo_dataset
|
||||
|
||||
data = [
|
||||
{"prompt": "Q1", "answer": "A1"},
|
||||
{"prompt": "Q2", "answer": "A2"},
|
||||
{"prompt": "Q3", "answer": "A3"},
|
||||
]
|
||||
result = _prepare_grpo_dataset(data)
|
||||
assert len(result) == 3
|
||||
|
||||
|
||||
# ─── Template Tests ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestReasoningTemplate:
|
||||
"""Test the reasoning/GRPO template."""
|
||||
|
||||
def test_reasoning_template_exists(self):
|
||||
assert "reasoning" in TEMPLATES
|
||||
|
||||
def test_reasoning_template_valid_yaml(self):
|
||||
import yaml
|
||||
|
||||
config = yaml.safe_load(TEMPLATES["reasoning"])
|
||||
assert config["task"] == "grpo"
|
||||
assert config["training"]["grpo_beta"] == 0.1
|
||||
assert config["training"]["num_generations"] == 4
|
||||
assert config["training"]["reward_fn"] == "accuracy"
|
||||
|
||||
def test_reasoning_template_valid_config(self):
|
||||
import yaml
|
||||
|
||||
raw = yaml.safe_load(TEMPLATES["reasoning"])
|
||||
cfg = SoupConfig(**raw)
|
||||
assert cfg.task == "grpo"
|
||||
assert cfg.training.grpo_beta == 0.1
|
||||
|
||||
|
||||
# ─── Train Command Routing Tests ─────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGRPOTrainRouting:
|
||||
"""Test that train command routes to GRPO trainer."""
|
||||
|
||||
def test_grpo_import_exists(self):
|
||||
"""GRPOTrainerWrapper should be importable."""
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
assert GRPOTrainerWrapper is not None
|
||||
|
||||
def test_grpo_wrapper_init(self):
|
||||
"""GRPOTrainerWrapper should initialize without error."""
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
cfg = SoupConfig(
|
||||
base="some-model",
|
||||
task="grpo",
|
||||
data={"train": "./data.jsonl"},
|
||||
)
|
||||
wrapper = GRPOTrainerWrapper(cfg, device="cpu")
|
||||
assert wrapper.config.task == "grpo"
|
||||
assert wrapper.device == "cpu"
|
||||
assert wrapper.model is None
|
||||
assert wrapper.trainer is None
|
||||
|
||||
|
||||
# ─── Sweep Shortcut Tests ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestGRPOSweepParams:
|
||||
"""Test GRPO parameter shortcuts in sweep."""
|
||||
|
||||
def test_grpo_beta_shortcut(self):
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"grpo_beta": 0.1}}
|
||||
_set_nested_param(config, "grpo_beta", 0.04)
|
||||
assert config["training"]["grpo_beta"] == 0.04
|
||||
|
||||
def test_num_generations_shortcut(self):
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"num_generations": 4}}
|
||||
_set_nested_param(config, "num_generations", 8)
|
||||
assert config["training"]["num_generations"] == 8
|
||||
|
||||
def test_reward_fn_shortcut(self):
|
||||
from soup_cli.commands.sweep import _set_nested_param
|
||||
|
||||
config = {"training": {"reward_fn": "accuracy"}}
|
||||
_set_nested_param(config, "reward_fn", "format")
|
||||
assert config["training"]["reward_fn"] == "format"
|
||||
Loading…
Reference in New Issue