3a00ff4d33
- Complete Fastify gateway with 8-stage pipeline - Circuit breaker (opossum) per model tier - Rate limiting per caller - Ban list validation (EN/DE/auto-detected) - TIP validator (SFF-8024, part numbers, wavelengths) - Prometheus metrics - pg-boss async queue - PostgreSQL audit log + review queue - 9 prompt templates (TIP, LinkedIn, ShieldX) - Learning engine scaffolding - Auto-learning: ban-list, few-shot, routing, prompt optimizer
287 lines
8.2 KiB
Python
287 lines
8.2 KiB
Python
"""
|
|
dpo_trainer.py - DPO (Direct Preference Optimization) fine-tuning.
|
|
|
|
Turns human-edited outputs into (chosen, rejected) preference pairs
|
|
and trains a policy model to prefer the human-edited versions.
|
|
|
|
MPS limitations apply identically to trainer.py:
|
|
- float32, no fp16/bf16, gradient_checkpointing=False,
|
|
dataloader_num_workers=0, device_map not used.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from datasets import Dataset
|
|
from peft import LoraConfig, TaskType, get_peft_model
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
|
|
from trl import DPOTrainer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Preference pair construction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_preference_pairs(corpus_rows: list[dict]) -> list[dict]:
|
|
"""
|
|
Build (prompt, chosen, rejected) triples from human-edited corpus rows.
|
|
|
|
Filters:
|
|
- human_edited must be True
|
|
- edited_output must be non-empty and differ from output_text
|
|
|
|
Returns a list of dicts with keys: prompt, chosen, rejected, task_type.
|
|
Immutable input — corpus_rows is never mutated.
|
|
"""
|
|
pairs: list[dict] = []
|
|
skipped = 0
|
|
|
|
for row in corpus_rows:
|
|
if not row.get("human_edited"):
|
|
skipped += 1
|
|
continue
|
|
|
|
prompt = (row.get("input_text") or "").strip()
|
|
chosen = (row.get("edited_output") or "").strip()
|
|
rejected = (row.get("output_text") or "").strip()
|
|
|
|
if not prompt or not chosen or not rejected:
|
|
skipped += 1
|
|
continue
|
|
|
|
if chosen == rejected:
|
|
skipped += 1
|
|
continue
|
|
|
|
pairs.append(
|
|
{
|
|
"prompt": prompt,
|
|
"chosen": chosen,
|
|
"rejected": rejected,
|
|
"task_type": row.get("task_type", "general"),
|
|
}
|
|
)
|
|
|
|
logger.info(
|
|
"build_preference_pairs: %d valid pairs, %d skipped", len(pairs), skipped
|
|
)
|
|
return pairs
|
|
|
|
|
|
def _pairs_to_dataset(pairs: list[dict]) -> Dataset:
|
|
"""
|
|
Convert preference pair dicts to a HuggingFace Dataset.
|
|
DPOTrainer expects columns: prompt, chosen, rejected.
|
|
"""
|
|
records = [
|
|
{
|
|
"prompt": p["prompt"],
|
|
"chosen": p["chosen"],
|
|
"rejected": p["rejected"],
|
|
}
|
|
for p in pairs
|
|
]
|
|
return Dataset.from_list(records)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Device selection (mirrors trainer.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _select_device() -> str:
|
|
if torch.backends.mps.is_available() and torch.backends.mps.is_built():
|
|
return "mps"
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
return "cpu"
|
|
|
|
|
|
def _load_model_and_tokenizer(
|
|
base_model_path: str, device: str
|
|
) -> tuple:
|
|
logger.info("DPO: loading tokenizer from %s", base_model_path)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
base_model_path,
|
|
trust_remote_code=True,
|
|
padding_side="right",
|
|
)
|
|
if tokenizer.pad_token is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
logger.info("DPO: loading model from %s on device=%s", base_model_path, device)
|
|
|
|
if device == "cuda":
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
base_model_path,
|
|
torch_dtype=torch.bfloat16,
|
|
device_map="auto",
|
|
trust_remote_code=True,
|
|
)
|
|
else:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
base_model_path,
|
|
torch_dtype=torch.float32,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
model.config.use_cache = False
|
|
return model, tokenizer
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main DPO training entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def run_dpo_training(
|
|
base_model_path: str,
|
|
pairs: list[dict],
|
|
output_dir: str,
|
|
beta: float = 0.1,
|
|
num_epochs: int = 1,
|
|
batch_size: int = 1,
|
|
gradient_accumulation_steps: int = 4,
|
|
learning_rate: float = 5e-5,
|
|
max_length: int = 2048,
|
|
max_prompt_length: int = 512,
|
|
lora_r: int = 16,
|
|
lora_alpha: int = 32,
|
|
lora_dropout: float = 0.05,
|
|
) -> dict:
|
|
"""
|
|
Run DPO preference learning.
|
|
|
|
Uses LoRA adapters on the base model to keep memory footprint small.
|
|
The reference model is the frozen base model; the policy model trains
|
|
on top of the LoRA adapter.
|
|
|
|
Returns:
|
|
{
|
|
"train_loss": float,
|
|
"adapter_path": str,
|
|
"n_pairs": int,
|
|
"device": str,
|
|
"beta": float,
|
|
}
|
|
|
|
Raises on fatal errors.
|
|
"""
|
|
if len(pairs) < 10:
|
|
raise ValueError(
|
|
f"Insufficient preference pairs: need >= 10, got {len(pairs)}"
|
|
)
|
|
|
|
device = _select_device()
|
|
logger.info(
|
|
"run_dpo_training: device=%s beta=%.2f pairs=%d output=%s",
|
|
device,
|
|
beta,
|
|
len(pairs),
|
|
output_dir,
|
|
)
|
|
|
|
dataset = _pairs_to_dataset(pairs)
|
|
|
|
# Split: 90% train, 10% eval
|
|
split = dataset.train_test_split(test_size=0.1, seed=42)
|
|
train_dataset = split["train"]
|
|
eval_dataset = split["test"]
|
|
|
|
model, tokenizer = _load_model_and_tokenizer(base_model_path, device)
|
|
|
|
# Separate reference model (frozen copy of base)
|
|
ref_model, _ = _load_model_and_tokenizer(base_model_path, device)
|
|
for param in ref_model.parameters():
|
|
param.requires_grad = False
|
|
|
|
# LoRA adapter on the policy model
|
|
lora_config = LoraConfig(
|
|
r=lora_r,
|
|
lora_alpha=lora_alpha,
|
|
lora_dropout=lora_dropout,
|
|
bias="none",
|
|
task_type=TaskType.CAUSAL_LM,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
inference_mode=False,
|
|
)
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
|
|
if device in ("mps", "cpu"):
|
|
model = model.to(device)
|
|
ref_model = ref_model.to(device)
|
|
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
training_args = TrainingArguments(
|
|
output_dir=str(output_path),
|
|
num_train_epochs=num_epochs,
|
|
per_device_train_batch_size=batch_size,
|
|
per_device_eval_batch_size=batch_size,
|
|
gradient_accumulation_steps=gradient_accumulation_steps,
|
|
learning_rate=learning_rate,
|
|
warmup_ratio=0.1,
|
|
eval_strategy="steps",
|
|
eval_steps=25,
|
|
save_strategy="steps",
|
|
save_steps=50,
|
|
load_best_model_at_end=True,
|
|
metric_for_best_model="eval_loss",
|
|
greater_is_better=False,
|
|
logging_steps=5,
|
|
report_to="none",
|
|
dataloader_num_workers=0,
|
|
fp16=False,
|
|
bf16=False,
|
|
optim="adamw_torch",
|
|
gradient_checkpointing=False,
|
|
remove_unused_columns=False,
|
|
label_names=["labels"],
|
|
)
|
|
|
|
trainer = DPOTrainer(
|
|
model=model,
|
|
ref_model=ref_model,
|
|
args=training_args,
|
|
beta=beta,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
tokenizer=tokenizer,
|
|
max_length=max_length,
|
|
max_prompt_length=max_prompt_length,
|
|
)
|
|
|
|
logger.info(
|
|
"Starting DPO training: %d train pairs, %d eval pairs",
|
|
len(train_dataset),
|
|
len(eval_dataset),
|
|
)
|
|
|
|
train_result = trainer.train()
|
|
|
|
eval_metrics = trainer.evaluate()
|
|
logger.info("DPO eval metrics: %s", eval_metrics)
|
|
|
|
adapter_path = str(output_path / "dpo_adapter")
|
|
model.save_pretrained(adapter_path)
|
|
tokenizer.save_pretrained(adapter_path)
|
|
logger.info("Saved DPO adapter to %s", adapter_path)
|
|
|
|
return {
|
|
"train_loss": round(train_result.training_loss, 4),
|
|
"eval_loss": round(eval_metrics.get("eval_loss", -1.0), 4),
|
|
"train_runtime": round(train_result.metrics.get("train_runtime", 0.0), 1),
|
|
"n_pairs": len(pairs),
|
|
"train_pairs": len(train_dataset),
|
|
"eval_pairs": len(eval_dataset),
|
|
"adapter_path": adapter_path,
|
|
"device": device,
|
|
"beta": beta,
|
|
}
|