"""PPO (Proximal Policy Optimization) trainer — wraps trl.PPOTrainer. PPO is the classic RLHF alignment method: generate completions, score them with a reward model (or reward function), then optimize the policy using clipped surrogate objectives with a KL penalty against a frozen reference model. Full RLHF pipeline: SFT → Reward Model → PPO """ 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 PPOTrainerWrapper: """High-level wrapper for PPO training from SoupConfig. PPO generates completions for each prompt, scores them with a reward model or reward function, and optimizes using proximal policy optimization with a KL penalty against a frozen reference model. Supports two reward sources: - reward_model: path/HF ID of a trained reward model (AutoModelForSequenceClassification) - reward_fn: callable reward function (same as GRPO — 'accuracy', 'format', or custom .py) At least one of reward_model or reward_fn must be specified. """ 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.reward_model_instance = None self.reward_fn = None def setup(self, dataset: dict): """Load model, tokenizer, reward model/fn, apply LoRA, create PPO trainer.""" from datasets import Dataset # Import PPOTrainer/PPOConfig — trl >=0.28 moved to trl.experimental ppo_trainer_cls, ppo_config_cls, is_experimental = _import_ppo_classes() # 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 use_unsloth = cfg.backend == "unsloth" # --- Load reward source --- self._setup_reward(cfg, tcfg) 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, ) # PPO needs memory for policy + ref model + reward model → conservative batch_size = max(1, batch_size // 4) console.print(f"[green]Auto batch size (PPO):[/] {batch_size}") # --- Dataset --- train_data = _prepare_ppo_dataset(dataset["train"]) train_ds = Dataset.from_list(train_data) # Tokenize dataset: trl experimental PPOTrainer expects input_ids, # not raw text. Add input_ids/attention_mask from prompt_text. max_len = cfg.data.max_length def _tokenize_ppo(examples): return self.tokenizer( examples["prompt_text"], truncation=True, max_length=max_len, padding=False, ) train_ds = train_ds.map(_tokenize_ppo, batched=True) # --- 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) # --- PPO config --- # Build kwargs, handling trl version differences import inspect ppo_kwargs = { "output_dir": str(output_dir), "per_device_train_batch_size": batch_size, "gradient_accumulation_steps": tcfg.gradient_accumulation_steps, "learning_rate": tcfg.lr, } # FSDP2 — alternative to DeepSpeed if self.fsdp_config: ppo_kwargs.update(self.fsdp_config) ppo_params = inspect.signature(ppo_config_cls).parameters # trl renamed ppo_epochs -> num_ppo_epochs in newer versions if "num_ppo_epochs" in ppo_params: ppo_kwargs["num_ppo_epochs"] = tcfg.ppo_epochs elif "ppo_epochs" in ppo_params: ppo_kwargs["ppo_epochs"] = tcfg.ppo_epochs if "cliprange" in ppo_params: ppo_kwargs["cliprange"] = tcfg.ppo_clip_ratio if "init_kl_coef" in ppo_params: ppo_kwargs["init_kl_coef"] = tcfg.ppo_kl_penalty # Optional params that may not exist in all trl versions if "log_with" in ppo_params: ppo_kwargs["log_with"] = ( self.report_to if self.report_to != "none" else None ) elif "report_to" in ppo_params: ppo_kwargs["report_to"] = self.report_to if "optimize_cuda_cache" in ppo_params: ppo_kwargs["optimize_cuda_cache"] = self.device == "cuda" # CPU support: trl PPOConfig requires use_cpu=True when no CUDA if self.device == "cpu" and "use_cpu" in ppo_params: ppo_kwargs["use_cpu"] = True ppo_config = ppo_config_cls(**ppo_kwargs) # --- Build reward functions list for PPOTrainer --- reward_funcs = [] if self.reward_model_instance is not None: reward_funcs.append(self.reward_model_instance) if self.reward_fn is not None: reward_funcs.append(self.reward_fn) # --- Trainer --- ppo_trainer_params = inspect.signature(ppo_trainer_cls.__init__).parameters # Store for use in train() — experimental API has different .train() signature self._is_experimental = is_experimental if is_experimental: # trl >=0.28 experimental API: PPOTrainer(args, processing_class, # model, ref_model, reward_model, train_dataset, value_model, ...) # ref_model=None is fine (auto-creates from policy model). # reward_model and value_model are required nn.Modules. reward_model_obj = self._get_or_create_reward_model(cfg, tcfg) value_model_obj = self._create_value_model(cfg, tcfg) trainer_kwargs = { "args": ppo_config, "processing_class": self.tokenizer, "model": self.model, "ref_model": None, "reward_model": reward_model_obj, "train_dataset": train_ds, "value_model": value_model_obj, } self._dataset_in_constructor = True self.trainer = ppo_trainer_cls(**trainer_kwargs) elif "args" in ppo_trainer_params: # trl >=0.28 non-experimental (transitional API) trainer_kwargs = { "model": self.model, "args": ppo_config, "processing_class": self.tokenizer, } if "train_dataset" in ppo_trainer_params: trainer_kwargs["train_dataset"] = train_ds elif "dataset" in ppo_trainer_params: trainer_kwargs["dataset"] = train_ds if reward_funcs and "reward_funcs" in ppo_trainer_params: trainer_kwargs["reward_funcs"] = reward_funcs # Pass ref/reward/value models if required positionally if "ref_model" in ppo_trainer_params: trainer_kwargs["ref_model"] = None if "reward_model" in ppo_trainer_params: trainer_kwargs["reward_model"] = self._get_or_create_reward_model( cfg, tcfg ) if "value_model" in ppo_trainer_params: trainer_kwargs["value_model"] = self._create_value_model(cfg, tcfg) self._dataset_in_constructor = ( "train_dataset" in trainer_kwargs or "dataset" in trainer_kwargs ) self.trainer = ppo_trainer_cls(**trainer_kwargs) else: # trl <0.28: PPOTrainer(config=, model=, tokenizer=, dataset=) trainer_kwargs = { "model": self.model, "config": ppo_config, "tokenizer": self.tokenizer, "dataset": train_ds, } self._dataset_in_constructor = True self.trainer = ppo_trainer_cls(**trainer_kwargs) # v0.40.6 #67 — ReLoRA callback (magnitude-prune LoRA every N steps). from soup_cli.utils.peft_wiring import attach_relora_callback attach_relora_callback(self.trainer, tcfg) self._output_dir = str(output_dir) self._train_ds = train_ds self._batch_size = batch_size self._num_epochs = tcfg.epochs self._max_length = cfg.data.max_length def _get_or_create_reward_model(self, cfg, tcfg): """Get existing reward model or create one for trl experimental PPO API. The experimental PPOTrainer requires an nn.Module reward model (not a callable). If we have a loaded reward model instance, use it. Otherwise, load one from the configured reward_model path, or create a fresh AutoModelForSequenceClassification from the base model. """ if self.reward_model_instance is not None: return self.reward_model_instance # If a reward_model path is configured, load it if tcfg.reward_model: return _load_reward_model( tcfg.reward_model, self.device, self.trust_remote_code, tcfg=tcfg, ) # Fallback: create a sequence classification model from the base model from transformers import AutoModelForSequenceClassification console.print( "[yellow]No reward_model path specified. Creating reward model " f"from base model: {cfg.base}[/]" ) reward_model = AutoModelForSequenceClassification.from_pretrained( cfg.base, trust_remote_code=self._trust_remote_code, num_labels=1, device_map="auto" if self.device != "cpu" else None, ) reward_model.eval() return reward_model def _create_value_model(self, cfg, tcfg): """Create a value model for trl experimental PPO API. The value model estimates state values for GAE advantage estimation. We create an AutoModelForSequenceClassification from the base model. """ from transformers import AutoModelForSequenceClassification console.print(f"[dim]Creating value model from: {cfg.base}[/]") value_model = AutoModelForSequenceClassification.from_pretrained( cfg.base, trust_remote_code=self._trust_remote_code, num_labels=1, device_map="auto" if self.device != "cpu" else None, ) return value_model def _setup_reward(self, cfg, tcfg): """Load reward model and/or reward function.""" # Reward model (pre-trained classifier) if tcfg.reward_model: self.reward_model_instance = _load_reward_model( tcfg.reward_model, self.device, self.trust_remote_code, tcfg=tcfg, ) console.print(f"[green]Reward model loaded:[/] {tcfg.reward_model}") # Reward function (callable — reuse GRPO reward functions) if tcfg.reward_fn: from soup_cli.trainer.rewards import load_reward_fn self.reward_fn = load_reward_fn(tcfg.reward_fn) if self.reward_model_instance is None and self.reward_fn is None: console.print( "[yellow]Warning: No reward_model or reward_fn specified. " "Using default reward_fn='accuracy'.[/]" ) from soup_cli.trainer.rewards import load_reward_fn self.reward_fn = load_reward_fn("accuracy") def _setup_transformers(self, cfg: SoupConfig, tcfg) -> 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 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 # Quantization (v0.38.0 Quant Menu — see soup_cli.utils.quant_menu) from soup_cli.utils.quant_menu import build_quantization_config_for_loader quant_config_obj = build_quantization_config_for_loader( tcfg=tcfg, base=cfg.base, console=console, ) console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" model_kwargs = { "trust_remote_code": self._trust_remote_code, "device_map": dev_map, } if quant_config_obj is not None: model_kwargs["quantization_config"] = quant_config_obj self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs) if tcfg.quantization in ("4bit", "8bit", "mxfp4"): 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, ) # v0.40.6 #67 — surgical PEFT patches. from soup_cli.utils.peft_wiring import ( apply_post_lora_patches, apply_pre_lora_patches, ) apply_pre_lora_patches(self.model, cfg.base) self.model = get_peft_model(self.model, lora_config) apply_post_lora_patches(self.model) # QAT — insert fake quantization ops after LoRA. The "fp8" variant # is FP8 training (handled by apply_v028_speed_memory), not int8 QAT. # PPO has its own forward loop, so use_cut_ce will degrade gracefully. 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, tcfg): """Load model via unsloth FastLanguageModel (2-5x faster).""" 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 PPO training loop and return results summary. Supports two trl APIs: - trl >=0.28: uses built-in trainer.train() (OnlineDPOTrainer style) - trl <0.28: manual loop with generate() + step() """ # Detect API: trl >=0.28 PPOTrainer has a .train() from OnlineDPOTrainer has_builtin_train = hasattr(self.trainer, "train") and not hasattr( self.trainer, "step" ) if has_builtin_train: return self._train_builtin(display, tracker, run_id, resume_from_checkpoint) return self._train_manual(display, tracker, run_id) def _train_builtin(self, display, tracker, run_id, resume_from_checkpoint): """Train using trl >=0.28 built-in trainer.train() method.""" start = time.time() # If dataset wasn't accepted by constructor, set it on the trainer if not self._dataset_in_constructor: if hasattr(self.trainer, "train_dataset"): self.trainer.train_dataset = self._train_ds elif hasattr(self.trainer, "dataset"): self.trainer.dataset = self._train_ds 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, ) ) # trl experimental PPOTrainer.train() does not accept resume_from_checkpoint import inspect from soup_cli.utils.v028_features import activation_offloading_context train_params = inspect.signature(self.trainer.train).parameters with activation_offloading_context( self.config.training, self._output_dir, ): if resume_from_checkpoint and "resume_from_checkpoint" in train_params: self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) else: if resume_from_checkpoint: console.print( "[yellow]Warning: This PPOTrainer does not support " "resume_from_checkpoint -- starting from scratch.[/]" ) self.trainer.train() 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 _train_manual(self, display, tracker, run_id): """Train using trl <0.28 manual loop: generate() + step().""" import torch start = time.time() step = 0 all_rewards = [] log_history = [] for epoch in range(self._num_epochs): for batch_idx in range(0, len(self._train_ds), self._batch_size): batch_end = min(batch_idx + self._batch_size, len(self._train_ds)) batch = self._train_ds[batch_idx:batch_end] # Tokenize prompts prompt_texts = batch["prompt_text"] query_tensors = [ self.tokenizer.encode(p, return_tensors="pt").squeeze() for p in prompt_texts ] # Generate completions response_tensors = [] for query in query_tensors: gen_kwargs = { "max_new_tokens": min(256, self._max_length // 2), "do_sample": True, "top_p": 0.9, "temperature": 0.7, } response = self.trainer.generate(query.unsqueeze(0), **gen_kwargs) response_tensors.append(response.squeeze()[len(query):]) # Compute rewards rewards = self._compute_rewards( query_tensors, response_tensors, batch, ) # PPO step stats = self.trainer.step(query_tensors, response_tensors, rewards) step += 1 mean_reward = torch.stack(rewards).mean().item() all_rewards.append(mean_reward) log_entry = { "step": step, "epoch": epoch + 1, "loss": stats.get("ppo/loss/total", 0), "reward": mean_reward, "kl": stats.get("ppo/mean_kl", 0), "lr": stats.get("ppo/learning_rate", self.config.training.lr), } log_history.append(log_entry) # Update display if display and hasattr(display, "update"): display.update( step=step, epoch=epoch + 1, loss=log_entry["loss"], lr=log_entry["lr"], ) # Update tracker if tracker and run_id: tracker.log_metrics( run_id=run_id, step=step, epoch=epoch + 1, loss=log_entry["loss"], lr=log_entry["lr"], ) duration = time.time() - start # Save final model (LoRA adapter) self.model.save_pretrained(self._output_dir) self.tokenizer.save_pretrained(self._output_dir) # Extract metrics losses = [entry["loss"] for entry in log_history 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": losses[0] if losses else 0, "final_loss": losses[-1] if losses else 0, "duration": duration_str, "duration_secs": duration, "output_dir": self._output_dir, "total_steps": step, } def _compute_rewards(self, query_tensors, response_tensors, batch): """Compute rewards using reward model and/or reward function.""" import torch rewards = [] num_samples = len(query_tensors) if self.reward_model_instance is not None: # Score with reward model for query, response in zip(query_tensors, response_tensors): full_ids = torch.cat([query, response]).unsqueeze(0) full_ids = full_ids.to(self.reward_model_instance.device) with torch.no_grad(): output = self.reward_model_instance(full_ids) score = output.logits[:, -1].squeeze().float() rewards.append(score.cpu()) elif self.reward_fn is not None: # Score with reward function completions = [] for response in response_tensors: text = self.tokenizer.decode(response, skip_special_tokens=True) completions.append([{"role": "assistant", "content": text}]) kwargs = {} if "answer" in batch: kwargs["answer"] = batch["answer"] scores = self.reward_fn(completions, **kwargs) rewards = [torch.tensor(score, dtype=torch.float32) for score in scores] else: # Fallback: zero reward rewards = [torch.tensor(0.0) for _ in range(num_samples)] return rewards def _import_ppo_classes(): """Import PPOTrainer and PPOConfig, handling trl version differences. trl >=0.28 moved PPOTrainer to trl.experimental with a new API that requires ref_model, reward_model, train_dataset, and value_model as positional args. The old trl import still works in 0.28 (deprecated) but is removed in 0.29. Returns: (PPOTrainer, PPOConfig, is_experimental): tuple with the classes and a flag indicating whether the experimental API (with required positional args) is used. """ try: from trl.experimental.ppo import PPOConfig, PPOTrainer return PPOTrainer, PPOConfig, True except (ImportError, ModuleNotFoundError): pass from trl import PPOConfig, PPOTrainer return PPOTrainer, PPOConfig, False def _load_reward_model( model_path: str, device: str = "cuda", trust_remote_code: bool = False, tcfg=None, ): """Load a pre-trained reward model for PPO scoring. Reward models are typically AutoModelForSequenceClassification that output a scalar reward score for a given input sequence. v0.40.5 #66: when ``tcfg`` is provided, the reward model is loaded with the same quantization config as the policy model (Quant Menu). This keeps the reward model from silently consuming full-precision VRAM and OOM-ing when the policy is GPTQ/AWQ/HQQ/etc. Pass ``tcfg=None`` for unquantized. """ from transformers import AutoModelForSequenceClassification from soup_cli.utils.trust_remote import ( model_requires_trust_remote_code, resolve_trust_remote_code, ) requires = model_requires_trust_remote_code(model_path) or False resolved = resolve_trust_remote_code( model_path, requested=trust_remote_code, console=console, requires_remote_code=requires, ) console.print(f"[dim]Loading reward model: {model_path}[/]") dev_map = "cpu" if device == "cpu" else "auto" model_kwargs: dict = { "trust_remote_code": resolved, "device_map": dev_map, "num_labels": 1, } if tcfg is not None: from soup_cli.utils.quant_menu import build_quantization_config_for_loader quant_config_obj = build_quantization_config_for_loader( tcfg=tcfg, base=model_path, console=console, ) if quant_config_obj is not None: model_kwargs["quantization_config"] = quant_config_obj reward_model = AutoModelForSequenceClassification.from_pretrained( model_path, **model_kwargs, ) reward_model.eval() return reward_model def _prepare_ppo_dataset(data: list[dict]) -> list[dict]: """Convert dataset rows to PPO format. PPO expects each row to have a 'prompt_text' field (string for tokenization) and optionally an 'answer' field (for reward function scoring). Input can be: - messages format: [{"role": "user", "content": "..."}, ...] - DPO format: {"prompt": "...", "chosen": "...", "rejected": "..."} - prompt field: {"prompt": "..."} - alpaca format: {"instruction": "...", "output": "..."} Returns list of dicts with 'prompt_text' (string) and optional 'answer'. """ prepared = [] for row in data: if "prompt" in row and isinstance(row["prompt"], str): entry = {"prompt_text": row["prompt"]} if "answer" in row: entry["answer"] = row["answer"] prepared.append(entry) elif "messages" in row: messages = row["messages"] # Use user messages as prompt text prompt_parts = [ msg["content"] for msg in messages if msg["role"] in ("system", "user") ] entry = {"prompt_text": " ".join(prompt_parts)} prepared.append(entry) elif "prompt" in row and isinstance(row["prompt"], list): # Message list → join content prompt_parts = [msg.get("content", "") for msg in row["prompt"]] entry = {"prompt_text": " ".join(prompt_parts)} if "answer" in row: entry["answer"] = row["answer"] prepared.append(entry) else: # Alpaca fallback instruction = row.get("instruction", row.get("input", "")) entry = {"prompt_text": str(instruction)} if "output" in row: entry["answer"] = row["output"] prepared.append(entry) return prepared