mirror of https://github.com/razor-ai/soup.git
464 lines
13 KiB
Python
464 lines
13 KiB
Python
"""Pydantic schemas for soup.yaml config — single source of truth."""
|
||
|
||
import re
|
||
from typing import List, Literal, Optional, Union
|
||
|
||
from pydantic import BaseModel, Field, field_validator
|
||
|
||
|
||
class LoraConfig(BaseModel):
|
||
r: int = Field(default=64, description="LoRA rank")
|
||
alpha: int = Field(default=16, description="LoRA alpha")
|
||
dropout: float = Field(default=0.05, description="LoRA dropout")
|
||
target_modules: Union[str, List[str]] = Field(
|
||
default="auto",
|
||
description="Target modules for LoRA. 'auto' = let peft decide.",
|
||
)
|
||
use_dora: bool = Field(
|
||
default=False,
|
||
description="Enable DoRA (Weight-Decomposed Low-Rank Adaptation)",
|
||
)
|
||
|
||
|
||
class DataConfig(BaseModel):
|
||
train: str = Field(..., description="Path to training data or HF dataset name")
|
||
format: Literal[
|
||
"alpaca", "sharegpt", "chatml", "dpo", "kto", "llava", "sharegpt4v", "auto"
|
||
] = Field(
|
||
default="auto",
|
||
description="Data format",
|
||
)
|
||
val_split: float = Field(default=0.1, ge=0.0, le=0.5, description="Validation split ratio")
|
||
max_length: int = Field(default=2048, description="Max sequence length in tokens")
|
||
image_dir: Optional[str] = Field(
|
||
default=None,
|
||
description="Base directory for resolving relative image paths in vision datasets",
|
||
)
|
||
|
||
|
||
class TrainingConfig(BaseModel):
|
||
epochs: int = Field(default=3, ge=1, description="Number of training epochs")
|
||
lr: float = Field(default=2e-5, gt=0, description="Learning rate")
|
||
batch_size: Union[int, Literal["auto"]] = Field(
|
||
default="auto",
|
||
description="Batch size. 'auto' = find max that fits in memory.",
|
||
)
|
||
gradient_accumulation_steps: int = Field(default=4, ge=1)
|
||
warmup_ratio: float = Field(default=0.03, ge=0.0, le=0.5)
|
||
weight_decay: float = Field(default=0.01, ge=0.0)
|
||
max_grad_norm: float = Field(default=1.0, gt=0)
|
||
lora: LoraConfig = Field(default_factory=LoraConfig)
|
||
quantization: Literal["4bit", "8bit", "none"] = Field(
|
||
default="4bit",
|
||
description="Quantization: 4bit (QLoRA), 8bit, or none (full precision)",
|
||
)
|
||
quantization_aware: bool = Field(
|
||
default=False,
|
||
description="Enable Quantization-Aware Training (QAT) for better post-quantization quality",
|
||
)
|
||
optimizer: str = Field(default="adamw_torch", description="Optimizer name")
|
||
scheduler: str = Field(default="cosine", description="LR scheduler type")
|
||
save_steps: int = Field(default=100, description="Save checkpoint every N steps")
|
||
logging_steps: int = Field(default=10, description="Log metrics every N steps")
|
||
# DPO-specific
|
||
dpo_beta: float = Field(
|
||
default=0.1, gt=0, description="DPO beta — KL penalty coefficient"
|
||
)
|
||
# KTO-specific
|
||
kto_beta: float = Field(
|
||
default=0.1, gt=0, description="KTO beta — KL penalty coefficient"
|
||
)
|
||
# ORPO-specific
|
||
orpo_beta: float = Field(
|
||
default=0.1, gt=0, description="ORPO beta — odds ratio weight"
|
||
)
|
||
# SimPO-specific
|
||
simpo_gamma: float = Field(
|
||
default=0.5, ge=0, description="SimPO gamma — reward margin term"
|
||
)
|
||
cpo_alpha: float = Field(
|
||
default=1.0, gt=0, description="CPO/SimPO alpha — NLL loss weight"
|
||
)
|
||
# IPO-specific (uses DPO trainer with loss_type='ipo')
|
||
ipo_tau: float = Field(
|
||
default=0.1, gt=0, description="IPO tau — regularization strength"
|
||
)
|
||
# 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",
|
||
)
|
||
# PPO-specific
|
||
ppo_epochs: int = Field(
|
||
default=4, ge=1, description="Number of PPO optimization epochs per batch"
|
||
)
|
||
ppo_clip_ratio: float = Field(
|
||
default=0.2, gt=0, le=1.0, description="PPO clipping range for policy ratio"
|
||
)
|
||
ppo_kl_penalty: float = Field(
|
||
default=0.05, ge=0, description="KL divergence penalty coefficient for PPO"
|
||
)
|
||
reward_model: Optional[str] = Field(
|
||
default=None,
|
||
description="Path or HF ID of a trained reward model for PPO",
|
||
)
|
||
# LoRA+ — different learning rates for A and B matrices
|
||
loraplus_lr_ratio: Optional[float] = Field(
|
||
default=None,
|
||
gt=0,
|
||
description="LoRA+ lr ratio: lr_B = lr × ratio. None = disabled (standard LoRA).",
|
||
)
|
||
# GaLore — memory-efficient full-parameter training
|
||
use_galore: bool = Field(
|
||
default=False,
|
||
description="Enable GaLore (Gradient Low-Rank Projection) for memory-efficient training",
|
||
)
|
||
galore_rank: int = Field(
|
||
default=128, ge=1, description="GaLore projection rank"
|
||
)
|
||
galore_update_proj_gap: int = Field(
|
||
default=200, ge=1, description="GaLore projection update interval (steps)"
|
||
)
|
||
galore_scale: float = Field(
|
||
default=0.25, gt=0, description="GaLore gradient scaling factor"
|
||
)
|
||
|
||
|
||
class SoupConfig(BaseModel):
|
||
"""Root config for soup.yaml."""
|
||
|
||
base: str = Field(..., description="Base model name or path (HF model ID)")
|
||
task: Literal[
|
||
"sft", "dpo", "grpo", "ppo", "reward_model", "kto", "orpo", "simpo", "ipo"
|
||
] = Field(
|
||
default="sft", description="Training task type"
|
||
)
|
||
modality: Literal["text", "vision"] = Field(
|
||
default="text",
|
||
description="Training modality: text (default) or vision (multimodal)",
|
||
)
|
||
backend: Literal["transformers", "unsloth"] = Field(
|
||
default="transformers",
|
||
description="Training backend: transformers (default) or unsloth (2-5x faster)",
|
||
)
|
||
data: DataConfig
|
||
training: TrainingConfig = Field(default_factory=TrainingConfig)
|
||
output: str = Field(default="./output", description="Output directory for trained model")
|
||
experiment_name: Optional[str] = Field(default=None, description="Experiment name for tracking")
|
||
|
||
@field_validator("experiment_name")
|
||
@classmethod
|
||
def experiment_name_safe(cls, value: Optional[str]) -> Optional[str]:
|
||
"""Disallow path separators and null bytes in experiment_name."""
|
||
if value is None:
|
||
return value
|
||
if re.search(r'[/\\:\x00]', value):
|
||
raise ValueError(
|
||
"experiment_name must not contain path separators (/ \\ :) or null bytes"
|
||
)
|
||
return value
|
||
|
||
|
||
# --- Built-in templates ---
|
||
|
||
TEMPLATES: dict[str, str] = {
|
||
"chat": """# Soup template: Chat Assistant
|
||
# Fine-tune a model for conversational chat
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: sft
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/train.jsonl
|
||
format: alpaca
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 2e-5
|
||
batch_size: auto
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
|
||
output: ./output
|
||
""",
|
||
"code": """# Soup template: Code Model
|
||
# Fine-tune a model for code generation / completion
|
||
|
||
base: codellama/CodeLlama-7b-Instruct-hf
|
||
task: sft
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/code_train.jsonl
|
||
format: alpaca
|
||
val_split: 0.1
|
||
max_length: 4096
|
||
|
||
training:
|
||
epochs: 2
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
lora:
|
||
r: 128
|
||
alpha: 32
|
||
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
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
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
|
||
""",
|
||
"vision": """# Soup template: Vision / Multimodal
|
||
# Fine-tune a vision-language model for image understanding
|
||
|
||
base: meta-llama/Llama-3.2-11B-Vision-Instruct
|
||
task: sft
|
||
modality: vision
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/vision_train.jsonl
|
||
format: llava
|
||
image_dir: ./data/images
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
|
||
output: ./output
|
||
""",
|
||
"medical": """# Soup template: Medical / Domain Expert
|
||
# Fine-tune a model with domain-specific knowledge
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: sft
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/medical_train.jsonl
|
||
format: alpaca
|
||
val_split: 0.15
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 5
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 8
|
||
lora:
|
||
r: 128
|
||
alpha: 32
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
|
||
output: ./output
|
||
""",
|
||
"kto": """# Soup template: KTO (Kahneman-Tversky Optimization)
|
||
# Align a model using unpaired preference data (no need for chosen+rejected pairs)
|
||
#
|
||
# Data format (JSONL):
|
||
# {"prompt": "What is 2+2?", "completion": "4", "label": true}
|
||
# {"prompt": "What is 2+2?", "completion": "Fish", "label": false}
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: kto
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/kto_train.jsonl
|
||
format: kto
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 4
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
kto_beta: 0.1
|
||
|
||
output: ./output
|
||
""",
|
||
"orpo": """# Soup template: ORPO (Odds Ratio Preference Optimization)
|
||
# Align a model without a reference model — simpler than DPO
|
||
#
|
||
# Data format (JSONL):
|
||
# {"prompt": "What is 2+2?", "chosen": "4", "rejected": "Fish"}
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: orpo
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/preference_train.jsonl
|
||
format: dpo
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 4
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
orpo_beta: 0.1
|
||
|
||
output: ./output
|
||
""",
|
||
"simpo": """# Soup template: SimPO (Simple Preference Optimization)
|
||
# Reference-free preference alignment with length-normalized rewards
|
||
#
|
||
# Data format (JSONL):
|
||
# {"prompt": "What is 2+2?", "chosen": "4", "rejected": "Fish"}
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: simpo
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/preference_train.jsonl
|
||
format: dpo
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 4
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
simpo_gamma: 0.5
|
||
cpo_alpha: 1.0
|
||
|
||
output: ./output
|
||
""",
|
||
"ipo": """# Soup template: IPO (Identity Preference Optimization)
|
||
# A theoretically grounded variant of DPO with stronger regularization
|
||
#
|
||
# Data format (JSONL):
|
||
# {"prompt": "What is 2+2?", "chosen": "4", "rejected": "Fish"}
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: ipo
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/preference_train.jsonl
|
||
format: dpo
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 3
|
||
lr: 1e-5
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 4
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
ipo_tau: 0.1
|
||
|
||
output: ./output
|
||
""",
|
||
"rlhf": """# Soup template: Full RLHF Pipeline (SFT + Reward Model + PPO)
|
||
# Three-stage training: 1) SFT warmup, 2) Reward model, 3) PPO alignment
|
||
#
|
||
# Usage:
|
||
# Step 1: soup train --config soup_sft.yaml # SFT warmup
|
||
# Step 2: soup train --config soup_rm.yaml # Train reward model
|
||
# Step 3: soup train --config soup_ppo.yaml # PPO with reward model
|
||
#
|
||
# This template generates the PPO config (step 3).
|
||
# For steps 1-2, use: soup init --template chat (SFT) and edit task to reward_model.
|
||
|
||
base: meta-llama/Llama-3.1-8B-Instruct
|
||
task: ppo
|
||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||
|
||
data:
|
||
train: ./data/prompts.jsonl
|
||
format: chatml
|
||
val_split: 0.1
|
||
max_length: 2048
|
||
|
||
training:
|
||
epochs: 1
|
||
lr: 1e-6
|
||
batch_size: auto
|
||
gradient_accumulation_steps: 4
|
||
lora:
|
||
r: 64
|
||
alpha: 16
|
||
target_modules: auto
|
||
quantization: 4bit
|
||
reward_model: ./output_rm
|
||
ppo_epochs: 4
|
||
ppo_clip_ratio: 0.2
|
||
ppo_kl_penalty: 0.05
|
||
|
||
output: ./output_ppo
|
||
""",
|
||
}
|