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