Skip unnecessary no_grad forward pass for IS correction in on-policy vLLM training
- 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
Assessment
This issue has not been assessed yet.