feat(grpo_trainer.py): Variational Sequence-Level Soft Policy Optimization (VESPO)

[casinca]

What does this PR do?

This PR implements the VESPO loss, resolve #5196

Official implementation: https://github.com/FloyedShen/VESPO/blob/main/recipe/vespo/code/core_algos.py
Paper: https://huggingface.co/papers/2602.10693

Note:

• The paper and the official implementation can have different variable names, to make things clearer:

  • c1 = k = α
  • c2 = lambda

• Docstrings/comments are a mix of official impl and my writing.

 

Alternative options:

• Currently VESPO has 4 hparams k_pos, lambda_pos, k_neg, lambda_neg but I could reduce with 2 tuples of 2 floats eg: lambdas (pos, neg) if it's better.
• Original impl also returns for metrics w_seq. I can include it in metrics, but this would force me to return a tuple in get_gamma_weights or remove @staticmethod. Not sure here what's the preference.

 

For efficiency, the TRL VESPO implementation is slightly different than the official one. It's ~25% faster per call on gpu, and tested for equivalence.

With importance_sampling_ratio:
-----------------------------------------------------------------
B x T         TRL_VESPO (ms)    OG_VESPO (ms)     Faster
-----------------------------------------------------------------
8 x 128         0.4290          0.5301          TRL_VESPO (1.24x)
16 x 256        0.4281          0.5302          TRL_VESPO (1.24x)
32 x 512        0.4283          0.5299          TRL_VESPO (1.24x)
64 x 512        0.4284          0.5294          TRL_VESPO (1.24x)
128 x 512       0.4286          0.5322          TRL_VESPO (1.24x)
32 x 1024       0.4473          0.5313          TRL_VESPO (1.19x)
64 x 1024       0.4285          0.5360          TRL_VESPO (1.25x)
128 x 1024      0.4240          0.5203          TRL_VESPO (1.23x)
-----------------------------------------------------------------

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

---

Note

Medium Risk
Adds a new loss formulation to GRPOTrainer that changes how gradients are scaled and how vLLM importance-sampling correction is applied, so training behavior can shift for users selecting loss_type="vespo". Scope is contained to a new loss branch plus config/docs/tests, but it touches core loss computation.

Overview
Adds VESPO (loss_type="vespo") to GRPOTrainer, including a new get_gamma_weights helper that computes detached, advantage-sign-dependent Gamma weights from sequence-level importance ratios (optionally incorporating vLLM TIS/MIS correction in log space) and uses them in the loss.

Extends GRPOConfig with four new hyperparameters (vespo_k_pos, vespo_lambda_pos, vespo_k_neg, vespo_lambda_neg) and updates validation/behavior so VESPO warns about importance_sampling_level usage, restricts vLLM correction modes to token_truncate/token_mask, and skips the generic per-token vLLM correction multiplier for VESPO.

Updates docs (paper_index.md) to add a VESPO paper entry/config snippet and expands the GRPO loss-type test matrix to include vespo.

^{Written by Cursor Bugbot for commit c5d0d50. This will update automatically on new commits. Configure here.}

[casinca]

I owe some better explanations to facilitate the review concerning importing math for lower_clamp = math.log(1e-8) in get_gamma_weights

From the original implementation below, the author is recomputing the log_w_seq from w_seq but we already have the log from seq_log_ratio_clampled. The only difference for recomputing, is the range of the min clamp being reduced to min=1e-8.

 

In order to avoid a 2nd log op in TRL, I'm directly clamping in logspace log_w_seq = torch.clamp(seq_log_ratio, lower_clamp, 20.0) once. Which ends up being the same.

This is solely to follow the original implementation, otherwise I'm not really sure if reducing from $e^{-20}$ to $e^{-18.42}$ (ie $e^{log(1e-8)}$) is important. I had opened an issue in OP for this: FloyedShen/VESPO#6

If keeping the original logic and importing math is problematic, alternative could be to hardcode log(1e-8) or retrieved from a tensor torch.log(torch.tensor(1e-8))

[casinca]

hey @FloyedShen , feel free to share any thoughts on this TRL implementation.

[FloyedShen]

Hi @casinca, sorry for the late reply — been swamped with work lately 😅

Thanks for this great contribution! I'm one of the authors of VESPO. I took some time to go through your implementation and ran a quick validation experiment comparing VESPO against the GRPO baseline.

Setup: Qwen3-4B-Base, DAPO-Math-17k, 8×H20 GPUs, 8 generations, 8 steps per generation, lr=1e-6, max_completion_length=8192.

The results look solid and are consistent with what we see in our verl-based experiments — VESPO shows improved reward over GRPO as training progresses, with notably more stable gradient norms and better entropy retention throughout training. Everything checks out on my end!

Full training logs: https://wandb.ai/brain-cog/vespo_trl

Thanks again for the clean implementation and the thorough benchmarking 👍

[casinca]

Np, thanks for the feedback and for taking the time to test, appreciate it.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Thanks, clean implementation!