v0.53.11 — GRPO variants live, PRM live, LongLoRA live, weighted preference

The capstone of the v0.50 / v0.49 / v0.40 deferred-stub debt. Five deep TRL trainer-subclassing items lifted.

Six GRPO objective variants live

apply_variant_loss ships real math kernels for every non-standard variant:

  • gspo — replaced in v0.75.0 and no longer what this line used to describe. It was a group-stabilised, column-centered importance ratio; it is now the published Group Sequence Policy Optimization objective (Qwen Team, arXiv:2507.18071): a length-normalised sequence importance ratio with sequence-level surrogate clipping, masked tokens held at exactly 0.0 gradient, invariant to padding and to its placement, and invariant to batch permutation. An existing gspo config will not reproduce prior runs. The heuristic it replaced also centred the per-token log-ratio across the batch before applying the completion mask, so one padding token shifted the loss and gradient of every unmasked row sharing its column
  • dapo — decoupled asymmetric clip
  • dr_grpo — no length normalization
  • bnpo — length-normalised PPO
  • two_sided — symmetric grpo_delta clip
  • rft — rejection-sampling fine-tuning

(Changed in v0.76.0, #1232 and #1249: none of the six variants applied grpo_beta through v0.75.2. Each now adds a per-token KL penalty weighted by grpo_beta (default 0.1), so an existing variant config will not reproduce a prior run, and grpo_beta: 0 runs KL-free. The variant kernels also clamp the importance ratio before exponentiation (#1609). See The KL penalty.)

Subclassing is done via make_grpo_trainer_variant (an lru_cache factory over _GRPOTrainerVariant) that overrides compute_loss to route through the kernel, with a defensive fallback to the original loss if TRL renames input attributes.

task: prm Process Reward Model — live

PRMTrainerWrapper loads AutoModelForCausalLM + nn.Linear(hidden, 1) as the reward head. make_prm_trainer_class(HF Trainer) factory subclasses Trainer and overrides compute_loss to:

  1. Gather hidden states at step-boundary tokens
  2. Project to scalars through the linear head
  3. MSE-loss against per-step labels

_prepare_prm_dataset tokenizes {prompt, completions, labels} with reserved-token truncation; _build_collator pads with -1 sentinel for missing step positions.

yaml
task: prm
base: deepseek-ai/deepseek-math-7b-base
data:
  train: prm_steps.jsonl
  format: prm   # {prompt, completions, labels}

(Changed in v0.76.0, #1235, #1564, #1220: when the device was given as cuda this trainer loaded the base and the head in bf16 with no fp32 copy for the optimizer; since v0.74.0 that made the first step fail outside DeepSpeed. It now loads fp32 on every device string and trains on fp32 master weights, which upstream estimates at about 16 bytes per parameter before activations; without DeepSpeed the saved checkpoint is fp32. It also reads optimizer, scheduler, warmup_ratio, weight_decay, max_grad_norm and gradient_accumulation_steps now, which it did not, and --resume continues the run instead of restarting at step 0. These are rows of the v0.76.0 re-run table.)

GRPOStabilityCallback (attached, mostly record-only)

The callback is attached to the GRPO trainer, but what it does at v0.76.0 is narrower than the v0.53.11 notes described (the file differs from v0.75.2 only by the gradient watchdog added in v0.76.0, described below). What the code contains, and what reaches a run:

  • An EMA reference-model update, (1-α) * ref + α * policy per parameter. It only acts when the callback is handed a separate reference model, which the Hugging Face callback handler does not do, and a LoRA run carries none, so ref_model_ema_alpha changes nothing in practice (read from the code, not run on a GPU)
  • A bounded-deque replay buffer, which the trainer never feeds
  • A TIS (truncated importance sampling) alert counter, which nothing calls
  • ema_alpha, replay_size and tis_alerts entries (for whichever of those keys you set) appended to state.log_history

GRPO Plus has the full accounting for the eight keys.

(Updated in v0.76.0: the eight keys are still mostly record-only. What changed is that every task: grpo run now carries a gradient watchdog that zeroes a non-finite gradient before the optimizer step, so that step updates nothing, whether or not you set any of those keys.)

LongLoRA S² forward override live

The v0.49.0 schema is now backed by a runtime. shift_heads_for_s2 rolls the second half of attention heads by group_size // 2 along the sequence axis (LongLoRA paper §3.2). LongLoRAForwardOverride is a context manager that monkey-patches every Llama / Mistral / Qwen / Phi attention module's forward, restoring on exit. Idempotent __del__ cleanup and best-effort exception swallow ensure training never crashes from a shape mismatch.

(Superseded in v0.76.0, #1240: training.use_longlora: true is refused at config load. Upstream's message says the implementation leaks future tokens and is not S² shifted sparse attention, and that it stays refused until real S² attention exists. To extend the context, set training.rope_scaling_type (linear, dynamic, yarn or llama3) and train plain LoRA. The current account is under LongLoRA S² shifted-sparse attention.)

True weighted-sum preference combine

attach_weighted_preference_combine now reads the four logprob tensors from TRL inputs (policy_chosen_logps / policy_rejected_logps / reference_chosen_logps / reference_rejected_logps) and computes each requested loss via the matching kernel:

  • compute_dpo_term
  • compute_ipo_term
  • compute_simpo_term
  • compute_orpo_term

Then combines via combine_losses(terms, weights). BCO mixed with paired losses is still rejected at validation time. Through v0.75.2 there was a defence-in-depth fallback to the v0.40.1 primary-loss scaling when TRL renames the logps attributes; it was removed in v0.76.0 (see the note below the example). At v0.75.2 that fallback was the path to expect on TRL 0.29, whose DPO, ORPO and CPO trainers do not put those tensors in the batch (read from the source, not run on a GPU); see Preference Variety.

yaml
task: preference
training:
  preference_loss_weights:
    dpo: 0.7
    simpo: 0.3

(Changed in v0.76.0, #1425, #1513 and #1556: the blend above no longer trains. Through v0.75.2 it printed Multi-objective preference loss: 0.70·dpo + 0.30·simpo at setup and then trained DPO alone, scaled by 0.7. Now a blend that cannot be computed stops the run at the first step, and only a blend that names simpo and orpo trains, for example preference_loss_weights: {simpo: 0.6, orpo: 0.4}. See Multi-objective preference loss.)

Tests

  • 8,330 → 8,400 (+75) in test_v05311.py (54 initial + 21 review-fix coverage gaps from python / code / security / tdd review agents).
  • Math kernels are real and unit-tested. Full GPU smoke runs documented in the v0.53.11.1 follow-up plan.

See also

Soup is free and Apache-2.0. If it saved you a training run, starring the repo costs nothing and helps most. You can also fund the GPU time behind the work a 4 GB laptop cannot reach.