Support hard and soft reasoning budget controls for RL rollouts
- 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
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