NVIDIA-NeMo / NVIDIA-NeMo/RL

`get_formatted_message_log` tokenizes each turn in isolation, adding a spurious leading space to assistant tokens with SentencePiece tokenizers

Open
#2,844 0 comments 0 reactions 1 assignee Claimed by @ashors1 View on GitHub
accuracy bug community-request waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

`nemo_rl.data.llm_message_utils.get_formatted_message_log` builds each sample's `token_ids` turn by turn:

1. render the chat template on the growing prefix `messages[:i+1]` (`tokenize=False`),
2. extract the current turn's substring via a string diff (`get_first_index_that_differs`),
3. tokenize that substring on its own: `tokenizer(text=message_chunk, add_special_tokens=False)`.

(main: `nemo_rl/data/llm_message_utils.py` lines \~530, \~535, \~541, \~603.)

Step 3 assumes that tokenizing a chunk on its own gives the same tokens as tokenizing it in context. That holds for byte-level BPE tokenizers (GPT-2, Qwen2, Llama-3) but not for SentencePiece tokenizers, which prepend a word-boundary marker `▁` (a space) to the first token of every string they encode (via a `Metaspace` pre-tokenizer with `prepend_scheme: "first"`, or a `Prepend`/`add_dummy_prefix` normalizer; this covers Llama, Mistral, TinyLlama, Phi-3, Zephyr, etc.).

As a result the assistant response's first token is tokenized as `▁The` instead of `The`. Assistant tokens are the supervised targets, so the model learns to emit a leading space right after the `<|assistant|>\n` generation prompt. At inference the prompt is tokenized as a whole, where the token after `\n` has no leading space, so this is a train/inference mismatch: after SFT the model generates `" The ..."` (token `▁The`) at the start of every answer.

**Steps/Code to reproduce bug**

Self-contained on `TinyLlama/TinyLlama-1.1B-Chat-v1.0` (public, ungated, no model or GPU needed):

```python
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

msgs = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 6 times 7?"},
{"role": "assistant", "content": "The answer is 42."},
]

# (a) What get_formatted_message_log does: render prefix[:i+1], diff, tokenize the chunk on its own.
def first_diff(a, b):
i = 0
while i < len(a) and i < len(b) and a[i] == b[i]:
i += 1
return i

prev = ""
for i, m in enumerate(msgs):
add_gen = m["role"] in ("user", "tool") # data.add_generation_prompt=true
cur = tok.apply_chat_template(msgs[: i + 1], add_generation_prompt=add_gen, tokenize=False)
chunk = cur[first_diff(prev, cur):]
if m["role"] == "assistant":
print("assistant chunk:", repr(chunk))
print("first tokens :", tok.convert_ids_to_tokens(tok(chunk, add_special_tokens=False)["input_ids"])[:3])
# -> assistant chunk: 'The answer is 42.\n'
# -> first tokens : ['▁The', '▁answer', '▁is'] # spurious leading space
prev = cur

# (b) Joint tokenization, what inference actually sees:
full = tok.apply_chat_template(msgs, tokenize=False)
gen_prompt = tok.apply_chat_template(msgs[:2], add_generation_prompt=True, tokenize=False)
joint = tok(full, add_special_tokens=False)["input_ids"]
plen = len(tok(gen_prompt, add_special_tokens=False)["input_ids"])
print("token after '<|assistant|>\\n' in joint encoding:", tok.convert_ids_to_tokens([joint[plen]]))
# -> ['The'] # 'The', not '▁The'
```

Through the actual code path (needs the repo's Python 3.12 env, e.g. `uv run python`):

```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("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
spec = TaskDataSpec(task_name="demo")
msgs = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 6 times 7?"},
{"role": "assistant", "content": "The answer is 42."},
]
out = get_formatted_message_log(
msgs, tok, spec,
add_bos_token=True, add_eos_token=True, add_generation_prompt=True,
)
print(tok.convert_ids_to_tokens(out[-1]["token_ids"])[:3])
# -> ['▁The', ...] # the supervised first token of the answer carries a leading space
```

Originally observed on a 1B model trained with a custom 128k SentencePiece tokenizer (Metaspace `prepend_scheme="first"`, Qwen-style `<|im_start|>` template): greedy generation from a prompt ending in `<|im_start|>assistant\n` yields first token `▁The` on every answer.

**Expected behavior**

Training tokens should match the tokens produced by tokenizing the whole conversation (what inference sees). The first token of the assistant answer after `<|assistant|>\n` should be `The`, not `▁The`. In short: `concat(per_turn_token_ids) == tokenize(whole_rendered_conversation)`, with the loss mask covering the assistant spans.

**Additional context**

Suggested fix: tokenize the fully rendered conversation once and derive per-turn token spans by character offsets:

* render the whole conversation, tokenize it once with `return_offsets_mapping=True`,
* render each assistant span (the `{% generation %}` region) to get its character range,
* map character ranges to token ranges to build the loss mask.

This gives tokens identical to inference for any tokenizer family. HF's `apply_chat_template(..., return_assistant_tokens_mask=True)`, which uses the `{% generation %}` markers already present in these templates, is a related building block.

Related, a second symptom of the same per-turn strategy: reasoning templates that re-render history differently from the current turn (e.g. Qwen3, which keeps `...` only for the last assistant turn) break the string diff and duplicate prior assistant content into the next chunk. That facet is tracked in NVIDIA-NeMo/RL#2821 / PR [#2822]() (fixed there via string overlap-stripping, which does not address the leading-space artifact in this issue). See also NVIDIA-NeMo/RL#38 ("apply chat template per message?", closed) for the original concern. A single joint-tokenization-with-offsets fix would cover both.

Affected paths: `get_formatted_message_log` is used by SFT (`examples/run_sft.py`), DPO/RM (`examples/run_dpo.py`, `examples/run_rm.py`), and the reward-model environment (`nemo_rl/environments/reward_model_environment.py`).

Environment: reproduces on current `main`, transformers 4.57.x. Reproducible on the public `TinyLlama/TinyLlama-1.1B-Chat-v1.0` and on any SentencePiece-family tokenizer.

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.