Gemma 4 checkpoints are rejected for assistant-only loss: The chat template is not training-compatible

Open Beginner friendly
  • Is something wrong?
#8,095 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
72/100
Issue type
Feature
Clarity
Clearly specified
Activity status
Active
Tech stack
python
Domain
machine-learning

Research direction

Start with get_training_chat_template in the TRL chat template utilities, and look at how the existing training templates such as the one from #7967 are stored and registered. The work is to add training variants of gemma4_v5.jinja and gemma4_v10.jinja that carry {% generation %} markers, since the issue says the rest of each template is already prefix-preserving. Done means get_training_chat_template accepts the Gemma 4 E2B/E4B and 31B/12B checkpoints, and return_assistant_tokens_mask=True returns a non-zero mask.

Written by the indexing model from the issue text.

Description

Reproduction

TRL stores the Gemma 4 chat templates (gemma4.jinja to gemma4_v10.jinja) for response parsing, but has no training version of any of them. Gemma 4 templates have no {% generation %} markers, so get_training_chat_template rejects every Gemma 4 checkpoint and SFTTrainer with assistant_only_loss=True fails at init.

from transformers import AutoTokenizer
from trl.chat_template_utils import get_training_chat_template

tokenizer = AutoTokenizer.from_pretrained("google/gemma-4-E4B-it")
get_training_chat_template(tokenizer)

outputs:

ValueError: The chat template is not training-compatible (missing prefix-preservation or `{% generation %}` markers) and patching is not supported for this template. Please manually modify the chat template for training.

Same for google/gemma-4-31B-it and google/gemma-4-12B-it. return_assistant_tokens_mask=True returns an all-zero mask on all three.

Tested with TRL main (1.18.0.dev0) and transformers 5.22.1.

How did you hit this?

I hit this fine-tuning Gemma 4 E4B on multi-turn agent trajectories (reasoning plus tool calls), and ended up computing the assistant spans myself.

Proposal

Add training templates for the two current revisions, the ones users download today:

  • gemma4_v5.jinja (google/gemma-4-E2B-it, google/gemma-4-E4B-it)
  • gemma4_v10.jinja (google/gemma-4-31B-it, google/gemma-4-12B-it)

Both templates are already prefix-preserving, including with reasoning and tool calls in the history, so only the {% generation %} markers are needed.

I'd be happy to open the PR, following #7967. Could you assign this to me?

System Info
  • Platform: macOS-15.7.7-arm64-arm-64bit
  • Python version: 3.12.13
  • TRL version: 1.18.0.dev0
  • PyTorch version: 2.15.0
  • accelerator(s): MPS
  • Transformers version: 5.22.1
  • Accelerate version: 1.15.0
  • Accelerate config: not found
  • Datasets version: 5.1.0
  • HF Hub version: 2.4.0
  • bitsandbytes version: not installed
  • DeepSpeed version: not installed
  • Liger-Kernel version: not installed
  • PEFT version: not installed
  • vLLM version: not installed
Checklist
  • I have checked that my issue isn't already filed (see open issues)
  • I have included my system information
  • Any code provided is minimal, complete, and reproducible (more on MREs)
  • Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
  • Any traceback provided is complete
Dominant language
Python
Stars
19.7k
Forks
3.1k
Avg merge
2d 13h
Merged PRs (30d)
283

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from huggingface/trl

All issues in huggingface/trl

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.