Spaces:
Sleeping
Sleeping
| """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 | |
| ) | |