NVIDIA-NeMo / NVIDIA-NeMo/RL

Support hard and soft reasoning budget controls for RL rollouts

Open
#2,946 1 comment 0 reactions 0 assignees View on GitHub
community-request waiting-on-customer
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Context

Nemotron-3-Nano-Omni style reasoning models expose two related budget controls that matter for RL rollouts:

- a hard vLLM limit, `thinking_token_budget`, which can force reasoning to close by injecting the configured reasoning end token; and
- soft prompt-level controls such as HF chat template `reasoning_budget` / `reasoning_grace`, which guide the model's own reasoning length rather than imposing a hard stop.

For GRPO/GSPO vision rollouts, missing or mismatched reasoning-budget handling can make many samples never reach the answer segment after ``. Reward functions that evaluate only the post-thinking answer can then see many zero-reward samples, which can make training appear stalled even when the model and reward are otherwise reasonable.

## Hard Budget

The hard-limit path is vLLM's `thinking_token_budget`. NeMo-RL can expose this as a generation config option and forward it to vLLM `SamplingParams`.
https://docs.vllm.ai/en/latest/features/reasoning_outputs/#thinking-budget-control

Related PR: https://github.com/NVIDIA-NeMo/RL/pull/2947

When using `thinking_token_budget`, vLLM also requires `reasoning_config`, for example:

```yaml
policy:
generation:
thinking_token_budget: 1024
vllm_kwargs:
reasoning_config:
reasoning_parser: nemotron_v3
reasoning_start_str:
reasoning_end_str:
```

## Soft Budget

Some model chat templates also reference `reasoning_budget` / `reasoning_grace` as soft controls. This is different from vLLM's hard token budget: the template text conditions the model to choose a target reasoning length, but it does not itself guarantee that `` is emitted at a fixed token boundary.

One possible NeMo-RL interface would be to support these values through `policy.tokenizer.chat_template_kwargs`, but that can be awkward when the desired budget varies per training example.

### Current Workaround

For per-sample soft budgets during training with Nano 3 Omni, a practical workaround is to add the budget instruction directly to the dataset prompt, for example:

```text
\n\n{thinking token budget: }
```
https://huggingface.co/nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16/blob/main/chat_template.jinja#L245

This keeps the budget attached to the data example and avoids needing a global `chat_template_kwargs` value for all samples in a run.

### Open Question

Should NeMo-RL add first-class config support for model-specific soft reasoning budget controls in chat template kwargs, or should that remain dataset/template-specific guidance while NeMo-RL only exposes vLLM's hard `thinking_token_budget`?

Contributor guide

Open the contributing guide

Research direction

Start by tracing policy.generation into vLLM SamplingParams, then inspect policy.tokenizer.chat_template_kwargs and the referenced chat_template.jinja budget handling. Compare the hard thinking_token_budget path with related PR 2947 and decide whether soft per-sample controls belong in configuration or dataset prompts. Done means the supported interface and behavior for hard and soft budgets are clearly established.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.