huggingface / huggingface/trl

Skip unnecessary no_grad forward pass for IS correction in on-policy vLLM training

Open
#5,412 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

# Feature request

## Problem

When using vLLM with `vllm_importance_sampling_correction=True`, GRPOTrainer always computes `old_per_token_logps` via an extra `no_grad` forward pass, even in on-policy settings (`num_iterations=1`, aligned steps). This adds an unnecessary full model forward pass to each training step.

On-policy, `old_per_token_logps == per_token_logps.detach()`, so this extra forward pass is redundant. The existing fallback in `_compute_loss` already handles this:

```python
# _compute_loss already falls back to per_token_logps.detach() when old_per_token_logps is None
old_per_token_logps = inputs.get("old_per_token_logps")
old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps
```

## Current behavior

```python
# _generate_and_score_completions
if self.args.gradient_accumulation_steps % generate_every != 0 or (
self.use_vllm and self.vllm_importance_sampling_correction # always True with vLLM + IS
):
old_per_token_logps, _ = self._get_per_token_logps_and_entropies(...) # extra forward pass
```

The condition forces a no_grad forward pass whenever vLLM IS correction is enabled, regardless of on/off-policy.

## Proposed behavior

- Only force the extra forward pass when using liger kernel (which needs pre-computed IS ratio since it computes logprobs internally) or when off-policy.
- For the non-liger path, compute IS correction inline in `_compute_loss` using `old_per_token_logps` (which falls back to `per_token_logps.detach()` on-policy).
- This eliminates an unnecessary full model forward pass per training step in on-policy + non-liger settings with no change in training behavior.
- Changes are in `trl/trainer/grpo_trainer.py` only (RLOO trainer does not have IS correction logic).

## Example

```python
# On-policy settings where this optimization applies:
args = GRPOConfig(
use_vllm=True,
vllm_mode="colocate",
vllm_importance_sampling_correction=True,
num_iterations=1, # on-policy
# gradient_accumulation_steps is a multiple of steps_per_generation (aligned)
)
# Before: extra no_grad forward pass every step
# After: no_grad forward pass skipped, IS correction computed inline in _compute_loss
```

---

# Motivation

When using vLLM with `vllm_importance_sampling_correction=True` in on-policy settings, GRPOTrainer performs an unnecessary `no_grad` forward pass to compute `old_per_token_logps` every training step. On-policy, this value equals `per_token_logps.detach()`, which is already available in `_compute_loss`. This wastes a full model forward pass per step.

---

# Your contribution

I'm going to submit a PR that implements the proposed change. It modifies only `trl/trainer/grpo_trainer.py`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.