mirror of https://github.com/razor-ai/soup.git
346 lines
13 KiB
Python
346 lines
13 KiB
Python
"""BCO (Binary Classifier Optimization) trainer — wraps trl.BCOTrainer.
|
|
|
|
v0.40.0 Part A. Mirrors the ORPO / SimPO / IPO wrapper pattern.
|
|
|
|
Data format (same as DPO): {"prompt": ..., "chosen": ..., "rejected": ...}.
|
|
Each row is internally split into two rows for TRL's BCOTrainer
|
|
(``{"prompt", "completion", "label"}``): one chosen-as-completion with
|
|
``label=True`` and one rejected-as-completion with ``label=False``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, 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
|
|
|
|
if TYPE_CHECKING:
|
|
from soup_cli.config.schema import TrainingConfig # noqa: F401
|
|
|
|
console = Console()
|
|
|
|
|
|
def _split_dpo_rows_to_bco(rows: list[dict]) -> list[dict]:
|
|
"""Convert DPO-shaped rows to TRL BCO unpaired format.
|
|
|
|
Each input row produces two output rows: one with ``label=True``
|
|
(chosen) and one with ``label=False`` (rejected). Rows missing any
|
|
required field are skipped; the count is emitted at DEBUG so production
|
|
silent-degradation is inspectable (mirrors v0.33.0 #47 CrossDocCollator
|
|
DEBUG-on-fallback policy).
|
|
"""
|
|
out: list[dict] = []
|
|
skipped = 0
|
|
for row in rows:
|
|
prompt = row.get("prompt")
|
|
chosen = row.get("chosen")
|
|
rejected = row.get("rejected")
|
|
if prompt is None or chosen is None or rejected is None:
|
|
skipped += 1
|
|
continue
|
|
out.append({"prompt": prompt, "completion": chosen, "label": True})
|
|
out.append({"prompt": prompt, "completion": rejected, "label": False})
|
|
if skipped:
|
|
import logging
|
|
|
|
logging.getLogger(__name__).debug(
|
|
"BCO: skipped %d row(s) missing prompt/chosen/rejected.", skipped,
|
|
)
|
|
return out
|
|
|
|
|
|
class BCOTrainerWrapper:
|
|
"""High-level wrapper for BCO training from SoupConfig.
|
|
|
|
BCO uses a binary classifier objective (chosen vs rejected). Soup
|
|
accepts the DPO data format and adapts it to TRL's BCOTrainer
|
|
unpaired ``{prompt, completion, label}`` schema.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: SoupConfig,
|
|
device: str = "cuda",
|
|
report_to: str = "none",
|
|
deepspeed_config: Optional[str] = None,
|
|
fsdp_config: Optional[dict] = None,
|
|
trust_remote_code: bool = False,
|
|
):
|
|
self.config = config
|
|
self.device = device
|
|
self.report_to = report_to
|
|
self.deepspeed_config = deepspeed_config
|
|
self.fsdp_config = fsdp_config
|
|
self.trust_remote_code = trust_remote_code
|
|
from soup_cli.utils.trust_remote import (
|
|
model_requires_trust_remote_code,
|
|
resolve_trust_remote_code,
|
|
)
|
|
|
|
requires = model_requires_trust_remote_code(config.base) or False
|
|
self._trust_remote_code = resolve_trust_remote_code(
|
|
config.base,
|
|
requested=trust_remote_code,
|
|
console=console,
|
|
requires_remote_code=requires,
|
|
)
|
|
self.model = None
|
|
self.tokenizer = None
|
|
self.trainer = None
|
|
self._output_dir: Optional[str] = None
|
|
|
|
def setup(self, dataset: dict) -> None:
|
|
"""Load model, tokenizer, apply LoRA, create BCO trainer."""
|
|
from datasets import Dataset
|
|
from trl import BCOConfig, BCOTrainer
|
|
|
|
from soup_cli.trainer.sft import _enable_hf_transfer_progress
|
|
|
|
_enable_hf_transfer_progress()
|
|
|
|
cfg = self.config
|
|
tcfg = cfg.training
|
|
use_unsloth = cfg.backend == "unsloth"
|
|
|
|
if use_unsloth:
|
|
self._setup_unsloth(cfg, tcfg)
|
|
else:
|
|
self._setup_transformers(cfg, tcfg)
|
|
|
|
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,
|
|
)
|
|
# BCO sees ~2x rows per sample (chosen + rejected) → halve.
|
|
batch_size = max(1, batch_size // 2)
|
|
console.print(f"[green]Auto batch size (BCO):[/] {batch_size}")
|
|
|
|
# --- Dataset (split DPO rows into BCO unpaired) ---
|
|
train_rows = _split_dpo_rows_to_bco(dataset["train"])
|
|
if not train_rows:
|
|
raise ValueError(
|
|
"BCO training: dataset['train'] produced no usable rows. "
|
|
"Each row must contain 'prompt', 'chosen', and 'rejected'."
|
|
)
|
|
train_ds = Dataset.from_list(train_rows)
|
|
eval_ds = None
|
|
if "val" in dataset and dataset["val"]:
|
|
val_rows = _split_dpo_rows_to_bco(dataset["val"])
|
|
if val_rows:
|
|
eval_ds = Dataset.from_list(val_rows)
|
|
|
|
# --- 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 ---
|
|
total_steps = (
|
|
math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps)
|
|
* tcfg.epochs
|
|
)
|
|
warmup_steps = int(total_steps * tcfg.warmup_ratio)
|
|
|
|
# --- BCO config ---
|
|
bco_config = BCOConfig(
|
|
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,
|
|
**(self.fsdp_config or {}),
|
|
beta=tcfg.bco_beta,
|
|
max_length=cfg.data.max_length,
|
|
max_prompt_length=cfg.data.max_length // 2,
|
|
**(
|
|
{"neftune_noise_alpha": tcfg.neftune_alpha}
|
|
if tcfg.neftune_alpha is not None
|
|
else {}
|
|
),
|
|
)
|
|
|
|
# --- Trainer ---
|
|
self.trainer = BCOTrainer(
|
|
model=self.model,
|
|
args=bco_config,
|
|
train_dataset=train_ds,
|
|
eval_dataset=eval_ds,
|
|
processing_class=self.tokenizer,
|
|
)
|
|
|
|
self._output_dir = str(output_dir)
|
|
|
|
def _setup_transformers(self, cfg: SoupConfig, tcfg: "TrainingConfig") -> None:
|
|
"""Load model via standard transformers + peft pipeline."""
|
|
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
|
|
|
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
cfg.base, trust_remote_code=self._trust_remote_code
|
|
)
|
|
if self.tokenizer.pad_token is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
bnb_config = None
|
|
if tcfg.quantization == "4bit":
|
|
from soup_cli.utils.gpu import get_compute_dtype
|
|
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=get_compute_dtype(),
|
|
bnb_4bit_use_double_quant=True,
|
|
)
|
|
elif tcfg.quantization == "8bit":
|
|
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
|
|
|
|
console.print(f"[dim]Loading model: {cfg.base}[/]")
|
|
dev_map = "cpu" if self.device == "cpu" else "auto"
|
|
model_kwargs = {
|
|
"trust_remote_code": self._trust_remote_code, "device_map": dev_map,
|
|
}
|
|
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)
|
|
|
|
target_modules = tcfg.lora.target_modules
|
|
if target_modules == "auto":
|
|
target_modules = None
|
|
|
|
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",
|
|
use_dora=tcfg.lora.use_dora,
|
|
use_rslora=tcfg.lora.use_rslora,
|
|
)
|
|
self.model = get_peft_model(self.model, lora_config)
|
|
|
|
# QAT — int8 only; "fp8" handled by apply_v028_speed_memory below.
|
|
if tcfg.quantization_aware and tcfg.quantization_aware != "fp8":
|
|
from soup_cli.utils.qat import prepare_model_for_qat
|
|
|
|
self.model = prepare_model_for_qat(self.model)
|
|
|
|
# v0.35.0 #60 — multi-trainer wiring of v0.28.0 speed/memory features.
|
|
from soup_cli.utils.v028_features import apply_v028_speed_memory
|
|
apply_v028_speed_memory(
|
|
model=self.model, tcfg=tcfg, base_model=cfg.base,
|
|
console=console, device=self.device, backend=cfg.backend,
|
|
)
|
|
|
|
def _setup_unsloth(self, cfg: SoupConfig, tcfg: "TrainingConfig") -> None:
|
|
"""Load model via unsloth FastLanguageModel."""
|
|
from soup_cli.utils.unsloth import load_model_and_tokenizer
|
|
|
|
console.print(f"[dim]Loading model via [bold]unsloth[/]: {cfg.base}[/]")
|
|
self.model, self.tokenizer = load_model_and_tokenizer(
|
|
model_name=cfg.base,
|
|
max_seq_length=cfg.data.max_length,
|
|
quantization=tcfg.quantization,
|
|
lora_r=tcfg.lora.r,
|
|
lora_alpha=tcfg.lora.alpha,
|
|
lora_dropout=tcfg.lora.dropout,
|
|
target_modules=tcfg.lora.target_modules,
|
|
)
|
|
if self.tokenizer.pad_token is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
def train(
|
|
self,
|
|
display: Optional[object] = None,
|
|
tracker: Optional[object] = None,
|
|
run_id: str = "",
|
|
resume_from_checkpoint: Optional[str] = None,
|
|
) -> dict:
|
|
"""Run BCO training and return results summary."""
|
|
if self.trainer is None:
|
|
raise RuntimeError(
|
|
"BCOTrainerWrapper.train() called before setup(). "
|
|
"Call setup(dataset) first."
|
|
)
|
|
start = time.time()
|
|
|
|
if display:
|
|
from soup_cli.monitoring.callback import SoupTrainerCallback
|
|
|
|
self.trainer.add_callback(
|
|
SoupTrainerCallback(
|
|
display, tracker=tracker, run_id=run_id,
|
|
loss_watchdog=self.config.training.loss_watchdog,
|
|
loss_watchdog_threshold=self.config.training.loss_watchdog_threshold,
|
|
loss_watchdog_patience=self.config.training.loss_watchdog_patience,
|
|
)
|
|
)
|
|
|
|
from soup_cli.utils.v028_features import activation_offloading_context
|
|
|
|
with activation_offloading_context(
|
|
self.config.training, self._output_dir,
|
|
):
|
|
self.trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
|
duration = time.time() - start
|
|
|
|
self.trainer.save_model(self._output_dir)
|
|
self.tokenizer.save_pretrained(self._output_dir)
|
|
|
|
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,
|
|
}
|