training/pm_ops_trainer.py

"""PMOpsGRPOTrainer — GRPOTrainer subclass with direct reward injection.

Problem being solved
--------------------
TRL 1.2.0 + Unsloth PatchFastRL strips custom rollout_func output keys
before calling reward_funcs. The 'reward' key never arrives in kwargs, so
reward_func always returns 0.0 and GRPO sees constant reward → zero gradient.

Root cause: TRL builds reward_kwargs from the *dataset* batch columns (only
'prompt'). rollout_func extra keys are not stored in the dataset, so they
never reach _calculate_rewards via inputs.

Fix
---
Wrap rollout_func inside __init__ using a closure that captures self directly
(standard Python closure semantics — no forward reference tricks needed).
Each time rollout_func is called during training, the wrapper stores the
'reward' list into self._rollout_reward_cache. _calculate_rewards then injects
those cached values as a tensor, bypassing the broken kwargs path entirely.

Usage
-----
    from training.pm_ops_trainer import PMOpsGRPOTrainer

    trainer = PMOpsGRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=reward_func,   # kept as fallback; not called on inject path
        train_dataset=dataset,
        args=grpo_config,
        rollout_func=rollout_func,  # must return 'reward': list[float] in output
    )

rollout_func contract
---------------------
rollout_func(prompts, trainer=None) must return a dict containing at minimum:
    {
        "prompt_ids":    list[list[int]],
        "completion_ids": list[list[int]],
        "logprobs":      list[list[float]],
        "reward":        list[float],   # one combined float per episode
    }
"""

import torch
from trl.trainer.grpo_trainer import GRPOTrainer


class PMOpsGRPOTrainer(GRPOTrainer):
    """Drop-in GRPOTrainer with closure-based reward caching.

    Wraps rollout_func at construction time to intercept the 'reward' list
    on every call. _calculate_rewards injects these rewards directly as a
    tensor — no kwargs, no inputs lookup.
    Falls back to standard GRPOTrainer behaviour when cache is empty.
    """

    def __init__(self, *args, rollout_func=None, **kwargs):
        # Reward cache populated by the rollout wrapper, consumed by _calculate_rewards
        self._rollout_reward_cache: list[float] = []
        self._rollout_capture_calls = 0
        self._warned_rollout_bypass = False

        wrapped_rollout_func = None
        if rollout_func is not None:
            original_rollout = rollout_func

            # self is captured directly by closure; accept flexible call signatures
            # because cloud runtimes may invoke rollout_func with positional/keyword variations.
            def _capturing_rollout(prompts, trainer=None, *rollout_args, **rollout_kwargs):
                rollout_kwargs.setdefault("trainer", trainer)
                result = original_rollout(prompts, *rollout_args, **rollout_kwargs)

                raw_rewards = result.get("reward", result.get("rewards", []))
                if raw_rewards is None:
                    rewards = []
                elif isinstance(raw_rewards, torch.Tensor):
                    rewards = [float(r) for r in raw_rewards.detach().cpu().flatten().tolist()]
                elif isinstance(raw_rewards, (int, float)):
                    rewards = [float(raw_rewards)]
                else:
                    rewards = [float(r) for r in list(raw_rewards)]

                self._rollout_reward_cache = rewards
                self._rollout_capture_calls += 1
                print(f"[PMOpsGRPOTrainer] captured {len(rewards)} rewards "
                      f"(mean={sum(rewards)/len(rewards):.3f})" if rewards else
                      "[PMOpsGRPOTrainer] WARNING: rollout returned no rewards")
                return result

            wrapped_rollout_func = _capturing_rollout
            rollout_func = wrapped_rollout_func

        super().__init__(*args, rollout_func=rollout_func, **kwargs)

        # Keep wrapper bound explicitly in case an upstream patch reassigns rollout_func.
        if wrapped_rollout_func is not None:
            self.rollout_func = wrapped_rollout_func

    def _calculate_rewards(
        self,
        inputs,
        prompts,
        completions,
        completion_ids_list,
    ):
        """Inject cached rewards when available; fall back to reward_funcs otherwise."""
        def _coerce_rewards(raw):
            if raw is None:
                return []
            if isinstance(raw, torch.Tensor):
                return [float(r) for r in raw.detach().cpu().flatten().tolist()]
            if isinstance(raw, (int, float)):
                return [float(raw)]
            return [float(r) for r in list(raw)]

        cache = self._rollout_reward_cache
        n = len(completions)

        # Some TRL variants forward rollout extra_fields into `inputs` directly.
        if not cache:
            input_rewards: list[float] = []
            for row in inputs:
                if isinstance(row, dict) and "reward" in row:
                    input_rewards.append(float(row["reward"]))
                else:
                    input_rewards = []
                    break
            if input_rewards:
                cache = input_rewards

        # Cloud fallback: some patched runtimes skip the normal rollout capture path
        # before calling _calculate_rewards. Actively invoke rollout_func here.
        # Pass ALL n prompts (including repeated ones for num_generations > 1) so
        # the rollout_func can generate distinct rewards per generation — NOT tile-mod.
        if not cache and self.rollout_func is not None and prompts:
            try:
                print(f"[PMOpsGRPOTrainer] cache empty — probing rollout_func for {n} rewards")
                out = self.rollout_func(list(prompts), trainer=self)
                cache = self._rollout_reward_cache
                if not cache and isinstance(out, dict):
                    cache = _coerce_rewards(out.get("reward", out.get("rewards", [])))
                    self._rollout_reward_cache = cache
            except Exception as exc:
                print(f"[PMOpsGRPOTrainer] rollout probe failed: {exc!r}")

        if cache:
            if len(cache) == n:
                rewards_list = cache
            elif len(cache) > n:
                rewards_list = cache[:n]
            else:
                # Still short — extend with mean rather than tile-mod so we don't
                # duplicate rewards for the same prompt (tile-mod → zero advantage).
                mean_r = sum(cache) / len(cache)
                rewards_list = list(cache) + [mean_r] * (n - len(cache))
                print(f"[PMOpsGRPOTrainer] WARNING: cache has {len(cache)} rewards for n={n}; "
                      f"padding with mean={mean_r:.3f}. Consider matching num_generations.")

            device = self.accelerator.device
            rewards = torch.tensor(
                [float(r) for r in rewards_list],
                dtype=torch.float32,
                device=device,
            ).unsqueeze(1)  # [batch_size, 1]

            std = rewards.std().item() if n > 1 else 0.0
            self.log({"reward/injected_mean": rewards.mean().item(),
                      "reward/injected_std": std})
            self._rollout_reward_cache = []  # consume cache
            print(f"[PMOpsGRPOTrainer] injected {n} rewards "
                  f"mean={rewards.mean().item():.3f} std={std:.3f}")
            return rewards

        if self.rollout_func is not None and self._rollout_capture_calls == 0 and not self._warned_rollout_bypass:
            print(
                "[PMOpsGRPOTrainer] WARNING: rollout_func was never called before reward calculation. "
                "This cloud runtime is likely bypassing rollout_func (TRL/Unsloth mismatch), so "
                "reward_funcs only receive prompts/completion_ids/trainer_state and no 'reward' key."
            )
            self._warned_rollout_bypass = True

        # Fallback: standard TRL reward_funcs path
        print("[PMOpsGRPOTrainer] cache empty — falling back to reward_funcs")
        return super()._calculate_rewards(
            inputs, prompts, completions, completion_ids_list
        )