Skip to content

feat: add Qwen3-8B template - #47

Open
Neonkraft wants to merge 10 commits into
mainfrom
feat/qwen3-sft-template
Open

feat: add Qwen3-8B template#47
Neonkraft wants to merge 10 commits into
mainfrom
feat/qwen3-sft-template

Conversation

@Neonkraft

Copy link
Copy Markdown
Collaborator

Summary

Adds qwen3-8b.jinja to the chat-template registry directory: Qwen3-8B's upstream template with {% generation %} markers spliced in so it can drive assistant_only_loss=True.

The upstream template concatenates the turn header and the message body into a single emission ('<|im_start|>' + message.role + '\n' + content), so the markers can't simply wrap the assistant branch. The header emission is split out and left outside the markers; {% generation %} opens after it and closes around <|im_end|>, with the trailing \n outside. Result: assistant content, reasoning trace and tool calls are in the loss, prompt tokens and the turn separator are not.

Type of change

  • Bug fix
  • New feature
  • Refactor
  • Performance
  • Documentation
  • Maintenance

Validation

Rendered output compared byte-for-byte against the pristine upstream template across 20 combinations — 5 conversation shapes (single-turn reasoning, multi-turn, no-think, tool call, system+tools) × add_generation_prompt on/off × with/without tools — all identical, confirming the restructuring is loss-mask-only and changes no inference behaviour.

@Neonkraft
Neonkraft requested a review from KonstiNik August 6, 2026 15:32

@KonstiNik KonstiNik left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, thanks for taking care of it!

  • Found one aspect worth noting.
  • Adding it to the tests would also be great I think.
  • Just curious about the naming – are there differences across qwen model sizes?

Comment on lines +51 to +52
{%- else %}
{{- content }}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the history content? If so, it is part of the training here, but I think we only want to train on the last generation and have the rest as context. So I think we want to exclude this part, maybe by doing sth like this:

{%- set assistant_body %}
    …existing think/content logic, tool_calls loop, '<|im_end|>'…
{%- endset %}
{%- if loop.index0 > ns.last_query_index %}
    {%- generation %}{{- assistant_body -}}{%- endgeneration %}
{%- else %}
    {{- assistant_body }}
{%- endif %}

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. Yes, the else block is the history content. I've fixed it using your using suggestion.

Qwen3 drops <think> from assistant turns at or before the last user query, so history rendered without a reasoning trace was being trained on.
Qwen3 drops <think> from assistant turns at or before the
last user query, so history rendered without a reasoning
trace was being trained on.
Vendors the pristine Qwen/Qwen3-8B chat_template as a fixture and diffs rendered output across the 20-combination matrix.
@Neonkraft

Copy link
Copy Markdown
Collaborator Author

Adding it to the tests would also be great I think.

Done. Also added a test to ensure that messages rendered with the new template are byte-equivalent to those rendered with the upstream template.

Just curious about the naming – are there differences across qwen model sizes?

Fair point. All Qwen3 models except Qwen3-{size}-{Instruct/Thinking}-* seem to have the same template, so I've renamed the template to qwen3.jinja.

@Neonkraft Neonkraft added the enhancement New feature or request label Aug 7, 2026
@Neonkraft

Copy link
Copy Markdown
Collaborator Author

I think we only want to train on the last generation and have the rest as context.

Do we have a strong reason for this? The only explanation I can think of is that the non-final assistant turns wouldn't have the reasoning traces, so the model would be learning to produce a mix of reasoning as well as non-reasoning outputs. On the other hand, wouldn't we be losing a lot of training signal in case of multi-turn conversations?

@KonstiNik

Copy link
Copy Markdown
Collaborator

AFAIK, this is how the Qwen template handles it.
The fix you're referring to would then be done in the data: a sample with multi-turn reasoning would be unrolled into several data samples, stopping at the different turns. So it would create more samples. What are your thoughts on that?

@KonstiNik

Copy link
Copy Markdown
Collaborator

Nice, this looks right to me now. I checked the rendering on the normal shapes and it does what it should.

One thing worth handling before this goes in, and I think it should go into this PR.

Gating the markers makes zero-span rows possible. If a conversation has no assistant turn after the last real user message, there is no generation region and the mask comes out empty. TRL checks for exactly that and raises inside dataset.map:

You're using assistant_only_loss=True, but at least one example has no assistant tokens. This usually means the tokenizer's chat template doesn't generate assistant masks — it may be missing the {% generation %} keyword. […]

So one bad row kills tokenization for the whole mixture, before a single training step, with an error that points at the template instead of the data.

It splits in two, and only half of it is about Qwen:

  • no assistant turn at all ([user], [system, user]) – empty mask under the olmo3-*-sft templates too, I checked, since there is simply nothing to wrap. Pre-existing hole: _sft_row_filter only checks len(messages) > 0, and open-perfectblend has rows like this.
  • an assistant turn, but none after the last real user query – typically ending on a user turn, and also anything with no user turn at all, since ns.last_query_index then falls back to the final index. Qwen-only: trains fine under OLMo, where the dangling user turn is just unsupervised context.

So the solution is part global, part template-aware:

  • require ≥1 assistant turn in _sft_row_filter unconditionally (drop the rows that don't fulfill). That half is a bug fix for OLMo too
  • for the Qwen-only: a capability flag on the template in the registry set only for qwen3, with _sft_row_filter conditioned on it. build_sft_trainer already reads config.data.chat_template just above where it passes the filter, so it's a closure over one boolean
  • log how many rows the filter dropped per dataset. loader.py filters silently today, so a heavily filtered dataset looks the same as a small one – and the weight is applied to the surviving rows, so a silent drop quietly shrinks that dataset's share of the mix. Pre-existing, so happy for it to be a separate PR.

Also worth a zero-span test – the current masking tests cover marker presence and history exclusion, which is why this didn't show up.

One edge case for completeness, though I doubt it shows up in practice: a conversation ending on a tool turn whose assistant caller sits before the last user query ([user, assistant, user, tool]) is also empty. Ordinary tool trajectories are fine. Only worth mentioning because it means the check can't just be "does it end on an assistant turn" – [user, assistant, tool] doesn't either, and that one is fine.

What do you think?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants