NVIDIA-NeMo / NVIDIA-NeMo/RL

get_formatted_message_log fails on system-led conversations for templates that require a user turn

Open
#2,863 1 comment 0 reactions 1 assignee Claimed by @terrykong View on GitHub
bug community-request waiting-on-maintainers
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**, rendering the prefix `message_log[:i+1]` for each turn `i`. Chat templates that **require a user message** (e.g. Qwen3, kanana) raise `TemplateError: No user query found in messages.` when the prefix is a **leading system message alone** (`i == 0`). As a result, **any conversation that starts with a system message fails to tokenize** with such templates.

- **Affected:** `nemo_rl/data/llm_message_utils.py :: get_formatted_message_log` (the per-turn loop calling `apply_chat_template(message_log_strs[:i+1])`).
- **Pre-existing:** reproduces on plain `main`.
- **Impact:** SFT (and any path that builds message logs this way) on data whose conversations begin with a `system` turn crashes for templates that require a user turn. Conversations starting with `user` are unaffected.

## Root cause

For turn `i`, the function renders the prefix:

```python
formatted_message = tokenizer.apply_chat_template(message_log_strs[: i + 1], ...)
```

When `message_log_strs[0]` is a `system` message, the first iteration (`i == 0`) renders `[system]` alone. Templates that assert a user turn exists raise inside the Jinja template.

## Reproduction (Qwen3.6)

```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/Qwen3.6-35B-A3B")

msgs = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "2+2?"},
{"role": "assistant", "content": "4"},
]
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,
)
```

### Actual

```
TemplateError: No user query found in messages.
```

(`tok.apply_chat_template([{"role": "system", ...}])` raises the same error directly.)

### Expected

The conversation tokenizes without error, and the per-turn chunks concatenate to the template's full render. The leading system turn is rendered as the prefix of the first user turn (`<|im_start|>system\n...<|im_end|>\n<|im_start|>user\n...`), as part of the masked, non-assistant context.

## Proposed fix

Defer leading pre-user turns (`i < first_user_msg_id`): when their prefix render fails, emit an empty chunk and let their tokens fold into the first user turn's chunk. The concatenation still equals the full single render, and templates that *can* render a lone system prefix (e.g. Qwen2.5) are unaffected. Template errors at any later turn still propagate.

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.