google

/

diffusiongemma-26B-A4B-it

like 1.18k

Follow

Google 64.2k

Problem

chat_template.jinja has no {% generation %}...{% endgeneration %} markers, so
apply_chat_template(..., return_assistant_tokens_mask=True) cannot return an assistant mask.
SFT frameworks then fall back to prefix-length heuristics to locate the supervised region, and
those heuristics place the boundary before the <|turn>model\n header rather than after it.

The model is therefore trained to re-emit the turn header as the first thing in its response.
Because model is an ordinary (non-special) token, it survives detokenization and shows up as a
stray model\n at the start of generations. (When the base channel prior wins instead, it
surfaces as thought\n.)

Concretely, with NVIDIA NeMo AutoModel's ChatDataset — the framework used by the official
DiffusionGemma SFT guide —
on one openai/gsm8k row (transformers 4.57.6):

The 3 wrongly-supervised tokens are <|turn> (105), model (4368), \n (107).

Filed upstream as NVIDIA-NeMo/Automodel#3352; the framework side should also assert this
boundary, but shipping generation markers here fixes it for every framework at once.

Fix

One semantic hunk. The assistant content and its turn close are captured into variables and
emitted together inside a single generation block:

-            {{- captured_content -}}
             {%- set has_content = captured_content | trim | length > 0 -%}

         {#- Forward-scan ... -#}   (unchanged)

+        {%- set captured_close -%}
         {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
             {{- '<|tool_response>' -}}
         {%- elif continues_into_next -%}
         {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%}
             {{- '<turn|>\n' -}}
         {%- endif -%}
+        {%- endset -%}
+
+        {%- if role == 'model' -%}
+            {% generation %}{{- captured_content -}}{{- captured_close -}}{% endgeneration %}
+        {%- else -%}
+            {{- captured_content -}}{{- captured_close -}}
+        {%- endif -%}

Two deliberate choices:

• The turn close stays inside the generation block, so termination remains supervised.
  DeepMind's own SFT pipeline does the same — in
  hackable_diffusion_adapter
  the prompt template ends with <|turn>model\n and the response template is "{text}<turn|>",
  i.e. the header belongs to the prompt and the closer belongs to the supervised response. This
  PR makes the chat template agree with that boundary.
• The forward-scan block is hoisted above the content emission so both can sit in one
  generation block. It only assigns variables (next_nt, ns_tr_out) and depends solely on
  loop_messages, so the rendered text is unchanged.

Verification

• rendered text is byte-identical to the current template (all render modes)
• supervised-region start now equals the inference prompt length exactly
• inference prompt length unchanged
• {% generation %} detected by the standard regex, so return_assistant_tokens_mask=True
  returns assistant_masks instead of falling back

Notes

• Orthogonal to #25. That PR moves the {%- endif -%} around ns.prev_non_tool_role; this
  one inserts a generation block just above it. The two are textually adjacent but semantically
  independent and can be merged in either order.
• The same gap exists in google/gemma-4-31B-it and google/gemma-4-E4B-it (also 0 generation
  markers). I've limited this PR to DiffusionGemma since it's the variant with a documented SFT
  recipe, but the same patch shape applies to the whole Gemma 4 template family if you'd prefer
  to fix it consistently.