MakazhanAlpamys/Soup v0.53.11

v0.53.11 — GRPO Plus finish + preference live

on GitHub

What's New

Five deep TRL trainer subclassing tasks lifted from the v0.50.0 / v0.49.0 / v0.40.1 deferred stubs.

• Live GRPO objective variants. apply_variant_loss ships real math kernels for gspo (group-stabilised importance ratio), dapo (decoupled asymmetric clip), dr_grpo (no length-norm), bnpo (length-normalised PPO), two_sided (symmetric grpo_delta clip), rft (rejection-sampling fine-tuning). New _GRPOTrainerVariant subclass (built via make_grpo_trainer_variant lru_cache factory with case-insensitive normalisation) overrides compute_loss to route through the kernel; reads inputs FIRST and only calls super() as fallback (no double forward pass).
• task='prm' Process Reward Model trainer. New PRMTrainerWrapper loads AutoModelForCausalLM + nn.Linear(hidden, 1) reward head; make_prm_trainer_class(HF Trainer) factory subclasses any HF Trainer and overrides compute_loss to gather hidden states at step boundaries, project to scalars, and MSE against per-step labels. _prepare_prm_dataset tokenises with reserved-token truncation; dataset wrapped in datasets.Dataset.from_list for full HF Trainer compatibility.
• GRPOStabilityCallback live. Inherits transformers.TrainerCallback (lazy). EMA reference-model update fires post-step with strict=True + fallback-to-strict=False-with-WARNING on key mismatch (silent corruption defence). Replay buffer, TIS alert counter, and ema_alpha surfaced via state.log_history.
• LongLoRA forward override live. New shift_heads_for_s2 math kernel rolls the second half of attention heads by group_size // 2 along the sequence dim. LongLoRAForwardOverride context manager monkey-patches every Llama*Attention / Mistral*Attention / Qwen*Attention / Phi*Attention module's forward with idempotent install (_soup_longlora_patched marker prevents re-entry double-wrap), 256-char class name cap, restore on __exit__ AND on exception.
• True weighted-sum preference combine. attach_weighted_preference_combine reads policy_chosen_logps / policy_rejected_logps / reference_chosen_logps / reference_rejected_logps from TRL inputs and computes each requested loss via compute_dpo_term / compute_ipo_term / compute_simpo_term / compute_orpo_term, then combines via combine_losses(terms, weights). Explicit None checks on trainer attrs (no or on possibly-tensor); DEBUG log on per-term skip.
• +75 net new tests (8330 → 8400). Four review agents ran (python / code / security / tdd); 25 findings fixed: 10 HIGH (double-forward GRPO variant, non-idempotent LongLoRA install, strict=False EMA load_state_dict, silent term-skip, wrong logits_to_keep key, or on tensors, missing TrainerCallback inherit, missing future annotations, dead import, bool bypass on batch_size) + 8 MEDIUM + 7 LOW.

Install / Upgrade

```bash
pip install --upgrade soup-cli
```

Or via Docker:

```bash
docker pull ghcr.io/makazhanalpamys/soup:v0.53.11
```

Security

• + "update_ema" + now calls + "load_state_dict(strict=True)" + and falls back to + "strict=False" + with a WARNING on key mismatch — closes the silent-corruption vector flagged by the security review (adversarial checkpoint could silently degrade EMA without diagnostic).
• + "LongLoRAForwardOverride._install" + caps class-name length at 256 chars before regex match (DoS defence).
• + "attach_weighted_preference_combine" + uses explicit + "is None" + checks instead of + "or" + on TRL trainer attributes (defends against zero-element tensor truth-value RuntimeError).
• + "make_grpo_trainer_variant" + normalises variant name BEFORE the lru_cache boundary so case variations share one class entry (prevents cache eviction attack).

Known Limitations

1. LongLoRA S² shift applied to attention OUTPUT, not Q/K projections — the output-side shift is cheaper and works with FA-v2 but is a documented approximation; the canonical input-side shift requires per-arch attention internals refactoring (deferred).
2. PRM dataset prep does not learn tokenizer-aware step boundaries — concatenates prompt + each completion sequentially and records the last token index as the boundary. Non-standard step delimiters need pre-splitting.
3. GRPO variant compute_loss override has TRL-version coupling — defensively probes + "per_token_logps" + / + "old_per_token_logps" + / + "ref_per_token_logps" + attribute names. A TRL refactor renaming all of these will trigger the fallback-to-original-loss path silently (only DEBUG logged).
4. EMA update reads full state_dict — large policy models will see per-step memory spikes from + "state_dict()" + materialisation. For 70B+ scale, operators should set + "ref_model_ema_alpha=None" + and rely on the legacy fixed-ref path.
5. GRPO + "async_grpo_prefetch" + / + "defer_rerolling" + / + "mask_truncated_completions" + / + "off_policy_mask_threshold" + — schema-validated and stored as callback fields, but the actual rollout-loop wiring lives inside TRL's + "GRPOTrainer.training_step" + (not exposed via TrainerCallback events). Full integration requires a TRL fork or PR upstream.
6. True weighted-sum requires TRL trainer to expose + "policy_chosen_logps" + etc on inputs — TRL ≥0.7 exposes them. Older TRL versions fall back to the v0.40.1 primary-loss scaling silently.

🤖 Generated with Claude Code