NVIDIA-NeMo / NVIDIA-NeMo/RL

[Proporsal]Stop injecting the chat template into `messages_log_batch[i]["content"]`; store raw text instead

Open
#854 4 comments 0 reactions 1 assignee Claimed by @SahilJain314 View on GitHub
external x-SBIntuitions
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**TL;DR**

* **Problem:** In GRPO, Environment input `messages_log_batch[i]["content"]`contains **chat-template–applied text(ex: ` {"content": "<|user|>Who are you<|assistant|>"}`), which prevents the Environment from accessing the original prompt.

* I want to use the original prompt for reward modeling, as in the LLM-as-a-judge case.
* **Proposal:**

* Change `content` to store **chat-template–free (raw) text**.
* Keep **`token_ids`** as the chat-templated representation for training/inference .

## Background (current behavior)

* In NeMo/RL’s GRPO pipeline, the [Data Processor](https://github.com/NVIDIA-NeMo/RL/blob/b6269f70f40fd33ca5adfd1f9aa0cac4619581d4/nemo_rl/data/processors.py#L68) **applies the chat template to each message**, stores that templated string in `content`, and generates **`token_ids`** from it to add to the `message_log`.
* Consequently, `messages_log_batch` (the conversation history passed to `Environments`) contains **templated strings** like `<|user|>Who are you<|assistant|>` in `content`.

## Problems

* It is difficult to **faithfully reconstruct the original prompt** from chat-templated text.

* Many use cases require the **original prompt (raw text)** for reward and next-step computation:
* LLM-as-a-Judge (reward from prompt + response)
* Multi-turn environments that generate responses via another LLM
* etc.
* Having the chat template embedded in `content` is unintuitive and harder for humans to read.

## Proposal

Change the Data Processor to behave as follows for each `message_log` entry:

* **`content`**: **raw text before chat template** (the source of truth for humans/Environments)
* **`token_ids`**: token IDs produced **after** applying the chat template (used for training/inference)
* **`content_templated` (optional)**: the **chat-templated text** corresponding to `token_ids` (helpful for debugging/inspection)

### Sample (Before / After)

**Before**

```json
{
"role": "user",
"content": "<|user|>Who are you<|assistant|>",
"token_ids": [ ... ]
}
```

**After (proposed)**

```json
{
"role": "user",
"content": "Who are you?",
"content_templated": "<|user|>Who are you<|assistant|>", // optional
"token_ids": [ ... ] // training/inference unchanged
}
```

## Impact

My understanding is that GRPO (generation/training) **primarily reads `token_ids`** and does not reference `content`, so switching `content` to raw text should **not change training behavior**.

However, **DPO/SFT** paths (e.g., `get_formatted_message_log()`) currently **expect templated text in `content`**, which we should keep in mind.

## Alternative

We could introduce a separate `raw_content` field that holds the chat-template–free string. This minimizes disruption to existing code, but continuing to place templated text in `content` is not common practice. For that reason, I prefer the original proposal (raw in `content`).

## Related issue

* **#682:** I think adopting this proposal would make that issue easier to resolve.

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.