unsloth/unsloth
Imgyu Kim 14ab6fbfae
BUG: fix _fix_chat_template for ChatML templates missing add_generation_prompt (#4426)
Fixes #4150.

Pre-PR, `_fix_chat_template` only patched templates where a trailing `{{ ... }}` expression followed the last `{% endfor %}`. ChatML templates (Hermes, Magnum, Phi-4, etc.) that end cleanly at `{% endfor %}` with no generation-prompt block were left unchanged, so the outer `fix_chat_template` raised:

```
RuntimeError: Unsloth: The tokenizer `...` does not have a
{% if add_generation_prompt %} for generation purposes.
```

This commonly shows up when a downstream tool (LlamaFactory, Axolotl) re-serializes the tokenizer during LoRA save and strips the generation-prompt block.

This PR adds a second branch to `_fix_chat_template` that fires when:

- the content after the last `{% endfor %}` is empty modulo Jinja `{# ... #}` comments,
- the scrubbed template contains `<|im_start|>` and `<|im_end|>`,
- and the scrubbed template does not already mention `add_generation_prompt`.

The assistant-turn separator is inferred from the template itself (preferring an explicit `'<|im_start|>assistant<sep>'` literal, then the unique `message['role'] + '<sep>'` from role concatenations, then `<|im_sep|>` for Phi-4-mini mixed-separator templates, then `\n`), so Phi-4-style templates are not silently corrupted with the wrong separator.

Verified against the existing chat-template corpus:

- Hermes-3, Magnum-v2, Phi-4-mini, Phi-4 multi-sep, ChatML with trailing whitespace, ChatML with trailing Jinja comment, dot-access `message.role`, split-literal `'<|im_start|>assistant'`: all repaired with the correct assistant prefix.
- Already-fixed ChatML templates: idempotent NOP.
- Trap templates with `<|im_start|>` only inside a Jinja comment: correctly not rewritten.
- Llama-3, Gemma-3, Qwen2.5 (non-ChatML): byte-identical.
- Mistral family (5 models including Mistral-Nemo, Mistral-Small-24B, Mixtral): byte-identical, protected both by the structural guard (no ChatML tokens) and the existing name-based exemption in `load_correct_tokenizer`.
- Qwen family (14 models including Qwen2.5, Qwen3, Qwen3-Coder, QwQ, VL, Math, Qwen3-Guard): byte-identical.

End-to-end reproduction: Hermes-3 LoRA SFT, save with stripped chat_template, reload. Pre-PR code path raises the RuntimeError above. Post-PR reload loads cleanly, patches the template at load time, and `apply_chat_template(add_generation_prompt=True)` produces the correct `<|im_start|>assistant\n` prefix.
2026-04-16 00:21:29 -07:00
..
dataprep Fix raw text paragraph break normalization (#4884) 2026-04-09 04:45:43 -07:00
kernels Add AMD ROCm/HIP support across installer and hardware detection (#4720) 2026-04-10 01:56:12 -07:00
models fix Gemma4 flash attn disable (#5045) 2026-04-15 17:50:48 -05:00
optimizers feat: Implement Q-GaLore optimizer and custom embedding learning rate… (#4511) 2026-03-25 01:03:10 -07:00
registry Revert "[pre-commit.ci] auto fixes from pre-commit.com hooks" 2025-12-01 07:24:58 -08:00
utils Add check to disable xformers on newer GPUs (#4342) 2026-03-16 22:42:38 -07:00
__init__.py Only run ldconfig CUDA-linking recovery when we have permission (#4930) 2026-04-09 00:07:25 -07:00
_auto_install.py Add PyTorch 2.10 and xformers 0.0.34 support (#3985) 2026-02-05 05:56:26 -08:00
chat_templates.py Re-apply #4939: updated models template mappers (#4950) 2026-04-15 07:52:12 -07:00
device_type.py Conditionally enable 4bit on CDNA for bitsandbytes>=v0.49.2 (#4161) 2026-03-07 01:33:40 -08:00
import_fixes.py Also patch accelerate's is_wandb_available for trl callbacks path (#4148) 2026-03-03 08:28:55 -08:00
ollama_template_mappers.py Re-apply #4939: updated models template mappers (#4950) 2026-04-15 07:52:12 -07:00
save.py fix(security): shell injection in GGML export conversion (#4768) 2026-04-02 00:10:43 -07:00
tokenizer_utils.py BUG: fix _fix_chat_template for ChatML templates missing add_generation_prompt (#4426) 2026-04-16 00:21:29 -07:00
trainer.py feat: Implement Q-GaLore optimizer and custom embedding learning rate… (#4511) 2026-03-25 01:03:10 -07:00