Mark pre-templated samples with chat_template_preapplied, and tokenize them with the loop NeMo-RL already has

NVIDIA-NeMo/RL #4295 (split out of #4108). Permalinks pin the PR head 4fcc987. Every value in the figure comes from running the proposed code with the Qwen2.5-0.5B-Instruct tokenizer.

Before
get_formatted_message_log always applies the tokenizer's chat template. Energon sources whose text is already templated would be templated twice.
This PR
Adds skip_chat_template=True, backed by a separate 100-line tokenizer. Nothing sets the flag yet.
Status
Proposed — not implemented. Keep the feature, name it for what it means, and run it through the existing loop. The code below was run against the real module, but it is not in the PR yet.

Two pages on this PR. Page 2 shows the bug in the PR's own tokenizer, which this design deletes.

Some Energon sources in an SFT blend hand over messages whose text already has the chat template in it. NeMo-RL must not apply the template again, but it still needs it for every other source in the blend. The blend marks such a source with the subflavor skip_chat_template: true, the same key Megatron-LM reads. Subflavors are set per dataset entry in the blend, so one run can mix pre-templated and normal sources. The cooker turns the subflavor into a typed field on the sample, and the tokenizing loop skips just the one step that applies the template. Why not policy.tokenizer.chat_template: null: that turns the template off for the whole run, not for one source.

One blend, two sources, one loop

Normal source Pre-templated source raw Energon sample (blend YAML, Megatron-LM's key) __subflavors__ = {} raw Energon sample (blend YAML, Megatron-LM's key) __subflavors__ = {skip_chat_template: true} cook_conversation (+1 line) cooker returns a CanonicalSFTSample (NeMo-RL) messages=[{user: "A: "}, {assistant: "Paris"}] chat_template_preapplied=False (default) cooker returns a CanonicalSFTSample (NeMo-RL) messages=[{user: "<|im_start|>user\nA: <|im_end|>\n <|im_start|>assistant\n"}, {assistant: "Paris<|im_end|>\n"}] chat_template_preapplied=True HFMultimodalSFTProcessorAdapter.encode, generic_sft.py L211 (+1 line), same for both get_formatted_message_log(messages, ..., chat_template_preapplied=sample.chat_template_preapplied) build turns 0..i: apply the chat template apply_chat_template, L677 i=0 → "…system…<|im_start|>user\nA: <|im_end|>\n" i=1 → (same) + "…assistant\nParis<|im_end|>\n" build turns 0..i: join the text as it is the one new branch in the loop i=0 → "<|im_start|>user\nA: …<|im_start|>assistant\n" i=1 → (same) + "Paris<|im_end|>\n" keep only the text turn i added, L682, same for both assistant chunk → "<|im_start|>assistant\nParis<|im_end|>\n" assistant chunk → "Paris<|im_end|>\n" tokenize each turn on its own, L750, same for both user 29 ids […, 25, 220, 151645, 198] assistant 6 ids […, 198, 59604, 151645, 198] user 11 ids [151644, 872, 198, 32, 25, 220, …] assistant 3 ids [59604, 151645, 198] add_loss_mask_to_message_log, L165: loss on the assistant turn trains on "<|im_start|>assistant\nParis<|im_end|>\n" trains on "Paris<|im_end|>\n" the only step that differs set by the cooker from the dataset's skip_chat_template subflavor Which turn holds "<|im_start|>assistant\n" follows the cooker's messages: here the cooker put it in the user turn, so it is not trained on. The PR's helper, given these same messages, splits them the same way. The tokenizer's own chat template is still used for the normal source, and is not changed.

The change — proposed, not in the PR

Where it lands: RL only, four files. llm_message_utils.py (+39 / −113, the helper deleted), plus one field in types.py one line in generic_sft.py, and one line in the generic cooker, generic.py. Ruff format and ruff check pass.

@@ llm_message_utils.py L465-L565 · the PR's helper, replaced by a small one -def _tokenize_preformatted_message_log(message_log, tokenizer, *, add_generation_prompt): - ... join turns, tokenize once with character ranges, cut, fall back ... # 101 lines +def _join_text_parts(message: dict[str, Any]) -> dict[str, Any]: + """Return the message with its text-only content joined into one string.""" + ... # str content passes through; a list of {"type": "text"} parts is joined; + ... # anything else raises "chat_template_preapplied=True requires text-only turns." @@ llm_message_utils.py L577 · the argument - skip_chat_template: bool = False, + chat_template_preapplied: bool = False, @@ llm_message_utils.py L588-L589 · docstring - skip_chat_template: Tokenize already-formatted turns directly. BOS and EOS - insertion are also skipped because those tokens must already be present. + chat_template_preapplied: Set when each message's content already has the chat + template applied upstream (for example by an Energon cooker), so it must + not be applied again. The text must use the same chat template the + tokenizer applies at generation time, and already contain any BOS/EOS: + add_bos_token and add_eos_token are ignored. Turns are tokenized as-is, + one at a time. Default: False @@ llm_message_utils.py L595-L600 · before the loop - if skip_chat_template: - return _tokenize_preformatted_message_log( - message_log, tokenizer, add_generation_prompt=add_generation_prompt, - ) + if chat_template_preapplied: + # The template is already in the text, so anything that would change the + # text before templating (a task prompt, tools, a generation prompt) can't + # be applied any more. + if add_generation_prompt or tools is not None or task_data_spec.prompt: + raise ValueError( + "chat_template_preapplied=True does not support add_generation_prompt, " + "tools, or a task_data_spec prompt: the chat template was already applied" + ) + # The text is final, so BOS/EOS must already be in it too. + add_bos_token = add_eos_token = False + message_log = [_join_text_parts(message) for message in message_log] @@ llm_message_utils.py L677-L679 · inside the loop - formatted_message: str = tokenizer.apply_chat_template( # type: ignore - message_log_strs[: i + 1], **template_kwargs - ) + if chat_template_preapplied: + # The turns already hold templated text, so join them as they are + # instead of applying the chat template again. + formatted_message = "".join(m["content"] for m in message_log_strs[: i + 1]) + else: + formatted_message = tokenizer.apply_chat_template( # type: ignore + message_log_strs[: i + 1], **template_kwargs + ) @@ energon/multimodal/types.py L59 · CanonicalSFTSample tools: list[dict[str, Any]] | None + # True when the source's text already has the chat template applied + # (the blend's skip_chat_template subflavor). + chat_template_preapplied: bool = False @@ energon/multimodal/cookers/generic.py L135 · cook_conversation tools=deepcopy(tools), + # "skip_chat_template" is Megatron-LM's subflavor key for pre-templated + # sources; keep it so one blend YAML means the same in both frameworks. + chat_template_preapplied=bool( + (sample.get("__subflavors__") or {}).get("skip_chat_template", False) + ), @@ task_encoders/generic_sft.py L220 · HFMultimodalSFTProcessorAdapter.encode tools=sample.tools, + chat_template_preapplied=sample.chat_template_preapplied, )

A blend marks a pre-templated dataset with subflavors: {skip_chat_template: true}, exactly as it would for Megatron-LM; nothing else in the blend changes. Run through cook_conversation and the adapter, such a sample gives the right-hand column above, and a sample without the subflavor keeps the chat template. _join_text_parts keeps the PR's support for content given as a list of text parts, which the adapter's _normalize_messages can produce.

🟢 review comment — :577 (rename)   review comment — :595 (use the loop)   review comment — :465 (wiring)   review comment — :588 (BOS in mixed blends)   review comment — :589 (data question)

Things that behave differently from the PR's helper

So what. The feature stays and gets a name that says what it means. The PR's separate tokenizer, and the bug in it (page 2), go away: pre-templated text goes through the same per-turn tokenizing as every other source, and a cooker opts in with one field.