How one token that crosses two turns leaves the assistant's answer with zero tokens

NVIDIA-NeMo/RL #4295. All permalinks pin the PR head 4fcc987. Token ids measured with the Qwen2.5-0.5B-Instruct tokenizer; Llama-3.1, Llama-2 and Nemotron-Nano-9B-v2 give the same split.

Before
The chat-template path tokenizes each turn on its own, so this cannot happen there.
This PR
Adds skip_chat_template=True for messages whose text is already templated. Its helper tokenizes the joined text once and cuts it into turns.
Status
Open — not fixed in this PR, and also present in Megatron-LM (same code, see the box below). The recommended design deletes this helper, and the bug with it. The fix at the bottom applies only if the helper stays.

Two pages on this PR. This one: why an answer can end up with no tokens.

A tokenizer cuts text into pieces called tokens, each with an id. SFT needs those ids split back into turns, because only the assistant's tokens are trained on. The PR's helper tokenizes all turns joined together, then has to work out where turn 1 ends in that token list. That breaks when one token covers the end of turn 1 and the start of turn 2. The figure walks one real example through each step.

"A: " then "Paris", step by step

the two turns each box is one character user turn: "A: " assistant turn: "Paris" A : ␣ P a r i s where turn 1 ends ① tokenize "A: Paris" once L519 A32 :25 ␣Paris12095 one token, half in each turn 3 tokens in total ② can't cut inside a token, so count turn 1 on its own fallback, L544-L550 A32 :25 ␣220 tokenize "A: " alone: 3 tokens, so "turn 1 = the first 3 tokens of ①" ③ cut ①'s list after 3 L552-L564 A32 :25 ␣Paris12095 user turn = tokens 1-3: all of them [ ] empty assistant turn = token 4 onward: ① has only 3, so nothing is left ④ what reaches training add_loss_mask_to_message_log, L165 message token_ids loss mask user [32, 25, 12095] [0, 0, 0] assistant [] [] the model reads "A: Paris" but learns from no token of the answer fix: tokenize each turn on its own (as the chat-template path does) L750 A32 :25 ␣220 Paris59604 user [32, 25, 220] [0, 0, 0] assistant [59604] [1] trains on "Paris" token in the wrong turn turn with no tokens each turn keeps its own tokens where turn 1 ends The number under each piece is its token id (Qwen2.5 tokenizer). ␣ = a space. No error or warning is raised at any step.

The same fallback can also repeat tokens. Token counts of growing prefixes do not always grow (Qwen2.5: turns ["\n ", " \n", "xx \n", "a"] give 5 tokens where the full text has 4). It is also the only path for tokenizers that cannot report character ranges, and for the PR's own unit test, whose fake tokenizer cannot be called directly. Turns that end on a chat-template marker like <|im_end|>\n split correctly.

Same code, same bug, in Megatron-LM. The PR's helper follows Megatron-LM's raw-conversation tokenizer step for step. Links pin Megatron-LM 6a36609, the commit NeMo-RL uses (through its Megatron-Bridge submodule 1f8873bb).

stepthis PR — _tokenize_preformatted_message_logMegatron-LM — MultimodalTokenizer
find where each turn ends in the token list; give up if a token crosses 
if token_index < len(offsets) and offsets[token_index][0] < boundary:
    cumulative_lengths = []
    break
llm_message_utils.py L530-L536		 
if start < boundary:
    return None
_tokenize_text_with_offsets L617-L631
fall back: count the tokens of each growing prefix 
for turn in raw_turns:
    cumulative_text += turn
    cumulative_lengths.append(len(_encode_raw(cumulative_text)))
L544-L550		 
for turn, ... in zip(rendered_turns, ...):
    cumulative_text += turn["content"]
    cumulative_lengths.append(
        len(self._encode_with_markers(cumulative_text, ...)))
_tokenize_raw_conversation_slow L644-L652
cut the full token list at those counts 
new_message["token_ids"] = token_ids[start:end]
L552-L564		 
if turn["role"].lower() == "assistant":
    target[start:end] = tokens[start:end]
L661-L666
run on ["A: ", "Paris"] (Qwen2.5) 
user      [32, 25, 12095]
assistant []
 
tokens  [32, 25, 12095]
trained []   (target all IGNORE_INDEX)

Megatron-LM's own docstring calls the fallback deliberate: "In that case callers must use the cumulative-prefix path to preserve legacy behavior." Megatron-LM turns this path on per dataset with the same subflavor, skip_chat_template, so Megatron-LM SFT runs on such data lose the same answers. The Megatron-LM result above comes from running its _tokenize_raw_conversation methods, copied unchanged from that commit, with the Qwen2.5 tokenizer.

Longer conversations: one token lost at every crossing

The damage does not spread to everything after it. At each turn boundary that falls inside a token, that one token moves into the turn before it; the cut lines up again at the next clean boundary (a newline, an <|im_end|>). But it repeats at every boundary that crosses. In question-and-answer text where every answer starts right after "A: ", every answer loses its first word:

tokens each turn receives from the PR's helper (Qwen2.5); number = token id turn 1 · user'Q: France?\nA: 'Q48:25␣France9625?\n5267A32:25␣Paris12095turn 2 · assistant'Paris\n'\n198should be Paris \nturn 3 · user'Q: Spain?\nA: 'Q48:25␣Spain17689?\n5267A32:25␣Madrid24081turn 4 · assistant'Madrid\n'\n198should be Mad rid \nturn 5 · user'Q: Italy?\nA: 'Q48:25␣Italy15344?\n5267A32:25␣Rome21718turn 6 · assistant'Rome\n'\n198should be R ome \n assistant tokens trained on: 3 of 8. None of them is an answer word. belongs to the next turn, moved here assistant token that is trained on user token (not trained on)

Longer answers lose only that first word: "Paris is the capital.\n" trains on " is the capital.\n". Turns that end on a clean boundary, such as answers on their own line or chat-template turns ending in <|im_end|>\n, come out exactly right. All values above are from running the PR's helper with the Qwen2.5 tokenizer.

Below is the same conversation as the trainer receives it: all turns joined into one token sequence, with the loss mask under each token (1 = trained on, 0 = only read). The mask comes from NeMo-RL's own add_loss_mask_to_message_log, which sets 1 on assistant tokens. The fix row is the proposed chat_template_preapplied path, run on the same turns. The two sequences have different lengths (24 vs 29) because one encode of the joined text merges each space into the next word.

PR as written: what reaches trainingQ0:0␣France0?\n0A0:0␣Paris0\n1Q0:0␣Spain0?\n0A0:0␣Madrid0\n1Q0:0␣Italy0?\n0A0:0␣Rome0\n124 tokens; mask 1 on 3: \n \n \nafter the fix (proposed): each turn tokenized on its ownQ0:0␣France0?\n0A0:0␣0Paris1\n1Q0:0␣Spain0?\n0A0:0␣0Mad1rid1\n1Q0:0␣Italy0?\n0A0:0␣0R1ome1\n129 tokens; mask 1 on 8: Paris \n Mad rid \n R ome \nmask 1: trained onmask 0: read, not trained onanswer word stuck in the user turn

The fix if the helper stays — proposed, not in the PR

Where the fix lands: RL only, one function. _tokenize_preformatted_message_log, L494-L565. Nothing below is implemented yet. It was run locally on the cases above, and the PR's existing test still passes.

@@ llm_message_utils.py L494-L565 · _tokenize_preformatted_message_log raw_tokenizer = getattr(tokenizer, "tokenizer", tokenizer) - def _encode_raw(text: str) -> torch.Tensor: ... # L496-L507 - cumulative_text = "".join(raw_turns) # L509 - try: # L518-L542 - encoded = raw_tokenizer(cumulative_text, ..., return_offsets_mapping=True) - ... walk the character ranges, find each turn's end ... - except (KeyError, NotImplementedError, TypeError, ValueError): - pass - if token_ids is None: # L544-L550 - ... len(_encode_raw(prefix)) for each prefix ... - for message, turn, end in zip(message_log, raw_turns, cumulative_lengths, strict=True): - new_message["token_ids"] = token_ids[start:end] + # Tokenize each turn on its own, like the chat-template path does, so a + # token can never span two turns and be given to the wrong one. + raw_message_log: LLMMessageLogType = [] + for message, turn in zip(message_log, raw_turns, strict=True): + new_message = message.copy() + new_message["content"] = turn + new_message["token_ids"] = torch.as_tensor( + raw_tokenizer.encode(turn, add_special_tokens=False), dtype=torch.long + ) + raw_message_log.append(new_message) return raw_message_log

The - lines are shortened with ...; the full block is in the review's suggestion block. Each turn is tokenized on its own, so a token can never belong to two turns, and every turn decodes back to its own text. This is also what generation sees: the prompt is tokenized alone, then the model writes the answer. It deletes the broad try/except and the fallback together. For chat-template turns the ids come out the same as before.

🔴 review comment — llm_message_utils.py:494-565

So what. Pre-templated data in a plain-text format, where a turn ends in a space or half a word, hits this. Every answer that starts in the middle of a token loses that token to the turn before it, and a one-word answer loses all of it. Nothing raises, and training looks normal while it learns less than it should. The recommended design removes the helper, so the bug goes with it; if the helper stays, tokenizing each turn on its own fixes it.