get_formatted_message_log duplicates assistant content in multi-turn with reasoning chat templates
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
`nemo_rl.data.llm_message_utils.get_formatted_message_log` tokenizes a conversation **turn by turn** using an incremental **string diff** (`get_first_index_that_differs`) of each turn's render against the previously rendered prefix. This assumes every turn's rendered string is a *monotonic prefix extension* of the previous turn's render.
That assumption is violated by **reasoning chat templates that re-render history differently from the last turn** (e.g. Qwen3, DeepSeek-R1, kanana): they keep the `…` block for the *current* assistant turn but **strip it from past assistant turns**. When this happens, the previous assistant turn's content is **duplicated** into the following turn's token chunk, producing a corrupted training sequence.
- **Affected:** `nemo_rl/data/llm_message_utils.py :: get_formatted_message_log` (the per-turn loop using `get_first_index_that_differs(prev_formatted_message, formatted_message)`).
- **Pre-existing:** reproduces on plain `main`.
- **Impact:** SFT (and any path that builds message logs this way) on **multi-turn** data with a reasoning template produces sequences where each prior assistant answer appears twice. Single-turn is unaffected.
## Root cause
For turn `i`:
```python
formatted_message = tokenizer.apply_chat_template(message_log_strs[: i + 1], ...)
prev_len = get_first_index_that_differs(prev_formatted_message, formatted_message)
message_chunk = formatted_message[prev_len:]
prev_formatted_message = formatted_message
```
This is correct only when `prev_formatted_message` is a prefix of `formatted_message`. With a reasoning template, once an assistant turn becomes history its `` is stripped, so the two renders diverge *before* that turn's content and `get_first_index_that_differs` returns an offset that re-includes the previous answer in this turn's chunk.
## Reproduction (model-agnostic)
```python
from transformers import AutoTokenizer
from nemo_rl.data.interfaces import TaskDataSpec
from nemo_rl.data.llm_message_utils import get_formatted_message_log
tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
# Keep only for the last assistant turn, strip from history.
tok.chat_template = (
"{%- for m in messages -%}"
"{%- if m['role']=='user' -%}<|user|>{{ m['content'] }}"
"{%- elif m['role']=='assistant' -%}"
"{%- if loop.last -%}<|assistant|>{{ m['content'] }}"
"{%- else -%}<|assistant|>{{ m['content'] }}{%- endif -%}"
"{%- endif -%}{%- endfor -%}"
)
msgs = [
{"role": "user", "content": "q1"},
{"role": "assistant", "content": "ANSWER1"},
{"role": "user", "content": "q2"},
{"role": "assistant", "content": "ANSWER2"},
]
mlog = get_formatted_message_log(
[dict(m) for m in msgs], tok, TaskDataSpec(task_name="repro"),
add_bos_token=False, add_eos_token=False, add_generation_prompt=False, tools=None,
)
full = "".join(tok.decode(m["token_ids"].tolist()) for m in mlog)
print("ANSWER1 count:", full.count("ANSWER1")) # -> 2 (should be 1)
```
### Actual
```
'<|user|>q1<|assistant|>ANSWER1ANSWER1<|user|>q2<|assistant|>ANSWER2'
ANSWER1 count: 2
```
### Expected
```
'<|user|>q1<|assistant|>ANSWER1<|user|>q2<|assistant|>ANSWER2'
ANSWER1 count: 1
```
Also confirmed on the real `Qwen/Qwen3` tokenizer (think block in content) and the kanana-2 reasoning tokenizer across all `thinking_mode` values.
## Why keep the per-turn `` (not just render once)
NeMo-RL's own multi-turn rollout builds context by concatenating each turn's `token_ids` verbatim (no chat-template re-application), so past `` blocks are preserved at inference. The default SFT loss also trains on all assistant turns (`only_unmask_final=False`). The fix should therefore keep each assistant turn's `` and only remove the duplication, rather than collapsing history to its inference-stripped form.
## Proposed fix
Detect the non-monotonic case (the previous render is not a prefix of the current one) and re-anchor the turn's chunk at this turn's opening delimiter, derived from the chat template (e.g. `<|im_start|>`, `<|start_header_id|>`), instead of slicing from the string diff. Fall back to stripping the chunk's overlap with the already-emitted text when no delimiter can be derived. This keeps each assistant turn's `` block and leaves the monotonic (normal-template) path unchanged.
Contributor guide
Assessment
This issue has not been assessed yet.