Upload training/train_grpo.py with huggingface_hub

training/train_grpo.py

@@ -782,19 +782,23 @@ def train(cfg: TrainConfig) -> None:
                     # so on the first opt step the ratio == 1; the clip becomes
                     # active only across multiple opt steps per batch. We still
                     # apply it for stability when group_size is large.
-                    log_ratio = cur_lp - ref_lp.detach()  # tiny KL surrogate
-                    ratio = torch.exp(cur_lp.detach() - ref_lp.detach())
-                    unclipped = ratio * adv
-                    clipped = torch.clamp(
+                    # Stable PPO surrogate: clamp log-ratio before exp to avoid
+                    # overflow/underflow from very large policy deltas.
+                    log_ratio = (cur_lp - ref_lp.detach()).clamp(-20.0, 20.0)
+                    ratio = torch.exp(log_ratio)
+                    clipped_ratio = torch.clamp(
                         ratio, 1.0 - cfg.ppo_clip, 1.0 + cfg.ppo_clip
-                    ) * adv
-                    pg_term = -torch.min(unclipped, clipped) * cur_lp / (cur_lp.detach().abs() + 1e-6)
-                    # Equivalent to -adv * log_pi (REINFORCE-style) when ratio~1,
-                    # but bounded for stability.
+                    )
+                    adv_t = torch.tensor(float(adv), device=device, dtype=cur_lp.dtype)
+                    pg_term = -torch.min(ratio * adv_t, clipped_ratio * adv_t)
 
-                    kl_term = cfg.kl_coef * (cur_lp - ref_lp.detach()).pow(2).mean()
+                    kl_term = cfg.kl_coef * (cur_lp - ref_lp.detach()).pow(2)
 
                     loss = pg_term + kl_term
+                    if not torch.isfinite(loss):
+                        # Skip pathological pairs instead of poisoning optimizer
+                        # state with inf/nan gradients.
+                        continue
                     loss.backward()
 
                     total_loss_val += float(loss.detach().item())