soup/soup_cli/config/schema.py

773 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""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)",
)
use_rslora: bool = Field(
default=False,
description="Enable rank-stabilized LoRA scaling (better for high ranks)",
)
class DataConfig(BaseModel):
train: str = Field(..., description="Path to training data or HF dataset name")
format: Literal[
"alpaca", "sharegpt", "chatml", "dpo", "kto", "llava", "sharegpt4v",
"plaintext", "embedding", "audio", "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, ge=64, le=1048576,
description="Max sequence length in tokens",
)
image_dir: Optional[str] = Field(
default=None,
description="Base directory for resolving relative image paths in vision datasets",
)
audio_dir: Optional[str] = Field(
default=None,
description="Base directory for resolving relative audio paths in audio 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"
)
# MoE-specific
moe_lora: bool = Field(
default=False,
description="Enable MoE-aware LoRA (ScatterMoE) — applies LoRA to expert FFN layers",
)
moe_aux_loss_coeff: float = Field(
default=0.01,
ge=0,
description="Auxiliary load-balancing loss coefficient for MoE models",
)
# Performance — Liger Kernel (fused operations)
use_liger: bool = Field(
default=False,
description="Enable Liger Kernel fused operations (20-60% memory savings, 20-40% speedup)",
)
# Performance — FlashAttention
use_flash_attn: bool = Field(
default=False,
description="Enable FlashAttention (auto-detects v2/v3/v4 for faster attention)",
)
# Performance — Ring FlashAttention (sequence parallelism)
use_ring_attention: bool = Field(
default=False,
description="Enable Ring FlashAttention for sequence parallelism across GPUs",
)
# Long-context — RoPE scaling
rope_scaling_type: Optional[Literal["linear", "dynamic", "yarn", "longrope"]] = Field(
default=None,
description="RoPE scaling method for long-context: linear, dynamic, yarn, longrope",
)
gradient_checkpointing: bool = Field(
default=False,
description="Enable gradient checkpointing for memory savings on long sequences",
)
# Embedding-specific
embedding_loss: Literal["contrastive", "triplet", "cosine"] = Field(
default="contrastive",
description="Loss function for embedding training: contrastive, triplet, or cosine",
)
embedding_margin: float = Field(
default=0.5, gt=0,
description="Margin for contrastive/triplet loss (higher = stricter separation)",
)
embedding_pooling: Literal["mean", "cls", "last"] = Field(
default="mean",
description="Pooling strategy for sentence embeddings: mean, cls, or last token",
)
embedding_temperature: float = Field(
default=0.05, gt=0,
description="Temperature for contrastive (InfoNCE) loss — lower = stricter similarity",
)
# Curriculum learning — sort dataset by difficulty
curriculum: bool = Field(
default=False,
description="Enable curriculum learning (sort dataset by difficulty, easy → hard)",
)
curriculum_metric: Literal["length", "perplexity", "loss"] = Field(
default="length",
description="Metric for curriculum difficulty: length, perplexity, or loss",
)
curriculum_buckets: int = Field(
default=4, ge=1, le=20,
description="Number of difficulty stages for curriculum learning",
)
# Loss watchdog — auto-stop on loss spikes
loss_watchdog: bool = Field(
default=False,
description="Enable loss spike detection (auto-stop if loss exceeds threshold)",
)
loss_watchdog_threshold: float = Field(
default=3.0,
gt=0,
le=100.0,
description="Stop training if loss exceeds this threshold",
)
loss_watchdog_patience: int = Field(
default=5,
ge=1,
le=1000,
description="Consecutive high-loss steps before stopping",
)
# Freeze training — freeze bottom layers for parameter-efficient training
freeze_layers: Optional[int] = Field(
default=None,
ge=1,
le=1000,
description="Freeze first N layers (from bottom). Train only remaining layers.",
)
freeze_ratio: Optional[float] = Field(
default=None,
gt=0.0,
lt=1.0,
description="Freeze this fraction of layers (0.75 = freeze 75% from bottom).",
)
# Sample packing — pack multiple short samples into one sequence
packing: bool = Field(
default=False,
description="Pack multiple short samples into one sequence for faster training",
)
# NEFTune — noisy embeddings for better fine-tuning
neftune_alpha: Optional[float] = Field(
default=None,
ge=0.0,
le=50.0,
description="NEFTune noise alpha (0-50). Adds noise to embeddings for better chat quality.",
)
class EvalConfig(BaseModel):
"""Evaluation configuration for auto-eval after training."""
auto_eval: bool = Field(
default=False,
description="Run evaluation automatically after training completes",
)
benchmarks: Optional[List[str]] = Field(
default=None,
description="lm-evaluation-harness benchmark names to run",
)
custom_tasks: Optional[str] = Field(
default=None,
description="Path to custom eval JSONL file",
)
judge: Optional[dict] = Field(
default=None,
description="LLM-as-a-judge config: model, rubric, provider",
)
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",
"pretrain", "embedding",
] = Field(
default="sft", description="Training task type"
)
modality: Literal["text", "vision", "audio"] = Field(
default="text",
description="Training modality: text (default), vision (multimodal), or audio",
)
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")
eval: Optional[EvalConfig] = Field(
default=None,
description="Evaluation configuration for auto-eval after training",
)
@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
""",
"pretrain": """# Soup template: Continued Pre-training
# Continue pre-training a model on raw text data (domain adaptation)
#
# Data format (JSONL):
# {"text": "Your raw text document here..."}
#
# Or plain .txt files (one document per line or entire file as one document).
base: meta-llama/Llama-3.1-8B
task: pretrain
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
data:
train: ./data/corpus.jsonl
format: plaintext
val_split: 0.05
max_length: 4096
training:
epochs: 1
lr: 1e-5
batch_size: auto
gradient_accumulation_steps: 8
lora:
r: 64
alpha: 16
target_modules: auto
quantization: 4bit
output: ./output_pretrain
""",
"moe": """# Soup template: MoE (Mixture of Experts) Fine-tuning
# Fine-tune a Mixture of Experts model with ScatterMoE LoRA
#
# Supported MoE models: Qwen3-30B-A3B, Mixtral-8x7B, DeepSeek-V3, etc.
base: Qwen/Qwen3-30B-A3B
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: 1e-5
batch_size: auto
gradient_accumulation_steps: 8
lora:
r: 64
alpha: 16
target_modules: auto
quantization: 4bit
moe_lora: true
moe_aux_loss_coeff: 0.01
output: ./output
""",
"longcontext": """# Soup template: Long-Context Fine-tuning (128k+)
# Extend model context window for long-document understanding
#
# Uses RoPE scaling + gradient checkpointing + FlashAttention for 128k tokens.
# Optionally enable Liger Kernel for additional memory savings.
base: meta-llama/Llama-3.1-8B-Instruct
task: sft
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
data:
train: ./data/long_context_train.jsonl
format: alpaca
val_split: 0.05
max_length: 131072
training:
epochs: 1
lr: 5e-6
batch_size: 1
gradient_accumulation_steps: 16
lora:
r: 64
alpha: 16
target_modules: auto
quantization: 4bit
gradient_checkpointing: true
rope_scaling_type: dynamic
use_flash_attn: true
# use_liger: true # pip install 'soup-cli[liger]' for fused ops
# use_ring_attention: true # Multi-GPU sequence parallelism
output: ./output_longctx
""",
"embedding": """# Soup template: Embedding Model Fine-tuning
# Fine-tune a sentence embedding model (BGE, E5, GTE, etc.)
#
# Data format (JSONL) — contrastive pairs:
# {"anchor": "What is Python?", "positive": "Python is a programming language."}
#
# Data format (JSONL) — triplets:
# {"anchor": "query", "positive": "relevant doc", "negative": "unrelated doc"}
base: BAAI/bge-base-en-v1.5
task: embedding
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
data:
train: ./data/embedding_train.jsonl
format: embedding
val_split: 0.1
max_length: 512
training:
epochs: 3
lr: 2e-5
batch_size: auto
gradient_accumulation_steps: 4
lora:
r: 64
alpha: 16
target_modules: auto
quantization: none
embedding_loss: contrastive
embedding_margin: 0.5
embedding_pooling: mean
output: ./output_embedding
""",
"audio": """# Soup template: Audio / Speech
# Fine-tune an audio-language model for speech understanding
#
# Supported models: Qwen2-Audio, Whisper (via transformers)
#
# Data format (JSONL):
# {"audio": "path/to/audio.wav", "messages": [
# {"role": "user", "content": "Transcribe."},
# {"role": "assistant", "content": "Hello world."}]}
base: Qwen/Qwen2-Audio-7B-Instruct
task: sft
modality: audio
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
data:
train: ./data/audio_train.jsonl
format: audio
audio_dir: ./data/audio
val_split: 0.1
max_length: 2048
training:
epochs: 3
lr: 1e-5
batch_size: auto
gradient_accumulation_steps: 8
lora:
r: 64
alpha: 16
target_modules: auto
quantization: 4bit
output: ./output_audio
""",
"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
""",
}