feat: add Llama 3 training chat template with generation markers#5493
Conversation
Adds `llama3_training.jinja` — a training-compatible variant of the Llama 3
chat template that wraps assistant output in `{% generation %}` /
`{% endgeneration %}` markers, enabling `return_assistant_tokens_mask=True`
for assistant-only loss masking in SFT training.
- `trl/chat_templates/llama3.jinja`: original Llama 3 chat template (source of truth)
- `trl/chat_templates/llama3_training.jinja`: training variant with generation markers
- `trl/chat_template_utils.py`: load templates and add Llama 3 branch in `get_training_chat_template()`
- `tests/test_chat_template_utils.py`: add `trl-internal-testing/tiny-LlamaForCausalLM-3` to `TestGetTrainingChatTemplate`
Closes part of huggingface#5471.
Built by Rudrendu Paul, developed with Claude Code
|
@codex review |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Codex Review: Didn't find any major issues. 👍 ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 5efca64. Configure here.
|
|
||
| '+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|> | ||
|
|
||
| ' }}{% endif %} No newline at end of file |
There was a problem hiding this comment.
Template newline encoding mismatch breaks tokenizer comparison
High Severity
The llama3.jinja file uses actual newline characters (line breaks) in the Jinja string literals, while the real Llama 3 tokenizer's chat_template stores \n as literal two-character escape sequences (backslash+n), matching how the JSON tokenizer_config.json round-trips them. This mirrors how qwen3.jinja correctly uses literal \n escape characters. Because Path.read_text() reads actual newlines (chr(10)) while the tokenizer attribute has backslash+n (chr(92)+chr(110)), the equality check tokenizer.chat_template == llama3_chat_template will be False, causing get_training_chat_template() to raise a ValueError instead of returning the training template.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 5efca64. Configure here.
What does this PR do?
Adds a training-compatible variant of the Llama 3 chat template with
{% generation %}/{% endgeneration %}markers, enablingreturn_assistant_tokens_mask=Truefor assistant-only loss masking in SFT training.Contributes to #5471 (Tracking: add
{% generation %}chat templates for common model families).Files changed:
trl/chat_templates/llama3.jinja— original Llama 3 chat template (source of truth for the exact-match comparison inget_training_chat_template())trl/chat_templates/llama3_training.jinja— training variant: same text output as original, with{% generation %}/{% endgeneration %}wrapping assistant output (<|eot_id|>included inside the block)trl/chat_template_utils.py— loads both templates; adds Llama 3 branch before Qwen3 inget_training_chat_template(); updates docstringtests/test_chat_template_utils.py— addstrl-internal-testing/tiny-LlamaForCausalLM-3toTestGetTrainingChatTemplateparametrizeTemplate design: The Llama 3 template is already prefix-preserving (no conditional logic that changes earlier messages when new messages are appended), so no structural prefix-preservation fix is needed — only the
{% generation %}markers are added. The training template restructures the original single-pass loop to handle the assistant role separately so markers can wrap the generated tokens only (content +<|eot_id|>), while the role header and BOS token remain outside.Before submitting
{% generation %}chat templates for common model families #5471AI writing disclosure
Who can review?
@qgallouedec (opened #5471)
Note
Low Risk
Low risk: adds an opt-in training chat template for LLaMA 3 and a new selection branch in
get_training_chat_template, with tests asserting rendered output stays identical while enabling assistant-token masking.Overview
Adds LLaMA 3 support to
get_training_chat_template()by introducing a canonicalllama3.jinjatemplate and a newllama3_training.jinjavariant that wraps assistant content in{% generation %}/{% endgeneration %}forreturn_assistant_tokens_mask=True.Updates template loading/dispatch in
trl/chat_template_utils.pyand extendsTestGetTrainingChatTemplateto cover a LLaMA 3 tokenizer, ensuring the patched template remains prefix-preserving and preserves rendered text output.Reviewed by Cursor Bugbot for commit 5efca64. Bugbot is set up for automated code reviews on this repo. Configure here.