ppo.py

"""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,
    resolve_device_map,
    resolve_frozen_base_load_dtype,
)
from soup_cli.utils.mixed_precision import align_trainable_dtype_for_fp16
from soup_cli.utils.seeding import apply_training_seed, training_seed_kwargs

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

        # #353: seed before the model and any adapter are built.
        apply_training_seed(tcfg)

        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,
            **training_seed_kwargs(tcfg),
        }

        # 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)

        # v0.71.26 — reward-hack mitigation buffer parity with GRPO. When a
        # detector / mitigation mode is set, capture the callable reward fns'
        # rewards + completions into a shared RLSignalBuffer so the mitigation
        # callback can observe the step (BETA — the on-GPU proof is GRPO-only).
        # nn.Module reward models are skipped (their call shape differs).
        self._rl_buffer = None
        from soup_cli.utils.peft_wiring import rl_callbacks_need_buffer

        if rl_callbacks_need_buffer(tcfg) and reward_funcs:
            # v0.71.26 Stage 3 — reward-shaping shim over callable reward fns,
            # BEFORE the buffer capture (no-op when shaping is off).
            from soup_cli.utils.reward_hack_control import apply_reward_shaping
            from soup_cli.utils.rl_signal_buffer import (
                RLSignalBuffer,
                wrap_reward_funcs,
            )

            self._rl_buffer = RLSignalBuffer()
            reward_funcs = [
                wrap_reward_funcs(
                    apply_reward_shaping(fn, tcfg), self._rl_buffer
                )
                if callable(fn) and not hasattr(fn, "forward")
                else fn
                for fn in reward_funcs
            ]

        # --- 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,
                    reward_funcs_supplied=bool(reward_funcs)
                    and "reward_funcs" in ppo_trainer_params,
                )
            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)

        # #359 - the same exposure #336 fixed in sft.py: with LoRA the
        # no-decay optimizer group is empty, DeepSpeed drops it, and the LR
        # scheduler keeps two base_lrs until torch's strict zip raises at the
        # first step. The guard prunes inside create_optimizer, i.e. before
        # the scheduler is built. No-op for full fine-tuning, and only under
        # DeepSpeed so the ordinary path keeps its own optimizer.
        if self.deepspeed_config:
            from soup_cli.utils.deepspeed import attach_empty_param_group_guard

            attach_empty_param_group_guard(self.trainer)

        # v0.71.26 — reward-hack mitigation / echo-trap / RL-checkpoint callbacks
        # (PPO parity with GRPO; kl_coef mutation for the controller).
        from soup_cli.utils.peft_wiring import attach_rl_callbacks

        attach_rl_callbacks(
            self.trainer,
            tcfg,
            buffer=self._rl_buffer,
            tokenizer=self.tokenizer,
            output_dir=str(output_dir),
            task="ppo",
        )

        # v0.40.6 #67 — ReLoRA callback (magnitude-prune LoRA every N steps).
        from soup_cli.utils.peft_wiring import (
            attach_curriculum_callback,
            attach_plugin_callback,
            attach_relora_callback,
        )
        attach_relora_callback(self.trainer, tcfg)
        # v0.53.5 #114/#115 — dynamic curriculum live callback.
        attach_curriculum_callback(self.trainer, tcfg, str(output_dir), console)
        # v0.53.6 #101 — Soup plugin TrainerCallback.
        attach_plugin_callback(self.trainer, console)

        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, *, reward_funcs_supplied=False):
        """Get the reward model for trl's PPO API (an nn.Module, not a callable).

        Uses a loaded instance, else loads the configured ``reward_model`` path.
        If neither exists, PPO has NO real reward signal: a fresh
        ``AutoModelForSequenceClassification`` has a randomly-initialised
        regression head, so optimising the policy against it trains toward
        noise (silent, GPU-hours wasted). A ``reward_fn`` (callable) cannot fill
        this nn.Module slot on trl's PPO API. Fail loudly instead — unless the
        trainer is separately consuming ``reward_funcs`` (a newer trl API), in
        which case the reward model is unused and a neutral head is fine.
        """
        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,
            )

        if not reward_funcs_supplied:
            raise RuntimeError(
                "PPO requires a trained reward_model (set training.reward_model "
                "to a path or HF id). trl's PPO reward-model slot needs an "
                "nn.Module; a reward_fn cannot drive it, and creating one from "
                "the base model would train the policy against a randomly-"
                "initialised reward head. Use `task: grpo` for reward-function-"
                "based RL instead."
            )

        # reward_funcs handles the reward signal on this trl API; the reward
        # model is a required-but-unused formality (neutral head is fine).
        from transformers import AutoModelForSequenceClassification

        console.print(
            "[dim]reward_funcs active; creating a neutral reward-model "
            f"placeholder from base model: {cfg.base}[/]"
        )
        reward_model = AutoModelForSequenceClassification.from_pretrained(
            cfg.base,
            trust_remote_code=self._trust_remote_code,
            num_labels=1,
            device_map=resolve_device_map(self.device) 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=resolve_device_map(self.device) 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, verifiable_domain=tcfg.verifiable_domain
            )

        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 = resolve_device_map(self.device)
        model_kwargs = {
            "trust_remote_code": self._trust_remote_code, "device_map": dev_map,
            "torch_dtype": resolve_frozen_base_load_dtype(self.device),
        }
        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)

        from soup_cli.utils.peft_wiring import resolve_lora_target_modules

        target_modules = resolve_lora_target_modules(self.model, tcfg.lora.target_modules)

        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,
                    eval_gate_config=self.config.training.eval_gate,
                )
            )

        # 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:
                align_trainable_dtype_for_fp16(
                    self.trainer.model,
                    fp16=getattr(self.trainer.args, "fp16", False),
                    bf16=getattr(self.trainer.args, "bf16", False),
                )
                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.[/]"
                    )
                align_trainable_dtype_for_fp16(
                    self.trainer.model,
                    fp16=getattr(self.trainer.args, "fp16", False),
                    bf16=getattr(self.trainer.args, "bf16", False),
                )
                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 and ``training.quantize_reward_model``
    is set, the reward model is loaded with the same quantization config as
    the policy model (Quant Menu), so it doesn't silently consume
    full-precision VRAM and OOM when the policy is GPTQ/AWQ/HQQ/etc. v0.53.0
    made this opt-in via that flag (default False) rather than unconditional,
    matching ``double_quant_on``'s #321 precedent of honouring a configured
    flag over a hardcoded default. Pass ``tcfg=None`` for unquantized
    regardless of the flag.
    """
    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 = resolve_device_map(device)
    model_kwargs: dict = {
        "trust_remote_code": resolved,
        "device_map": dev_map,
        "num_labels": 1,
    }
    if tcfg is not None and tcfg.quantize_reward_model:
        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