diff --git a/docs/source/chat_templates.md b/docs/source/chat_templates.md index 499c0d8a814..d3aba5050b7 100644 --- a/docs/source/chat_templates.md +++ b/docs/source/chat_templates.md @@ -20,7 +20,7 @@ TRL ships patched templates under [`trl/chat_templates/`](https://github.com/hug ## Supported model families -TRL stores reference copies of the original templates so it can identify supported models at init and swap in a training template when needed. The following families are recognized: Cohere, Cohere2, DeepSeek-V3, Gemma, Gemma3, GLM-4-MoE, GPT-OSS, Idefics3, Llama 3 / 3.1 / 3.2, Llava-Next, Nemotron 3 (Nano, Super, Ultra), Phi-3, Phi-3.5, Qwen2-VL, Qwen2.5, Qwen2.5-VL, Qwen3 (including the Instruct-2507 variant), Qwen3-VL, Qwen3.5, Qwen3.6. +TRL stores reference copies of the original templates so it can identify supported models at init and swap in a training template when needed. The following families are recognized: Cohere, Cohere2, DeepSeek-V3, Gemma, Gemma3, GLM-4-MoE, GPT-OSS, Idefics3, LFM2, LFM2.5, Llama 3 / 3.1 / 3.2, Llava-Next, Nemotron 3 (Nano, Super, Ultra), Phi-3, Phi-3.5, Qwen2-VL, Qwen2.5, Qwen2.5-VL, Qwen3 (including the Instruct-2507 variant), Qwen3-VL, Qwen3.5, Qwen3.6. ## Training templates @@ -114,6 +114,12 @@ Patched Idefics3 template. Diff vs `idefics3.jinja`: Split the assistant message into its own branch so the `{% generation %}` / `{% endgeneration %}` markers wrap the assistant content. This enables `return_assistant_tokens_mask=True` to produce correct masks for SFT assistant-only loss. +### `lfm2_training.jinja` + +Patched LFM2 template. Diff vs `lfm2.jinja`: + +Split the unified message output line into role-specific branches, so the `<|im_start|>assistant\n` prompt cue sits outside the generation block (it is not generated by the model), while the assistant's content and `<|im_end|>\n` (which the model must learn to produce and to stop on) sit inside. Wrap the assistant content with `{% generation %}` / `{% endgeneration %}` so that `return_assistant_tokens_mask=True` produces correct masks for SFT assistant-only loss. + ### `llama3_training.jinja` Patched Llama 3 template. Diff vs `llama3.jinja`: diff --git a/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm.py b/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm.py new file mode 100644 index 00000000000..70b1ce92ca8 --- /dev/null +++ b/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm.py @@ -0,0 +1,68 @@ +# Copyright 2020-2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from transformers import AutoTokenizer, GenerationConfig, Lfm2Config, Lfm2ForCausalLM + +from .._common import ( + check_dtype_pattern, + check_transformers_version, + init_weights_tiny_model, + print_config_diff, + push_to_hub, + smoke_test, +) + + +check_transformers_version() + +MODEL_ID = "LiquidAI/LFM2-1.2B" + +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) +generation_config = GenerationConfig.from_pretrained(MODEL_ID) +config = Lfm2Config( + vocab_size=65536, + hidden_size=8, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + intermediate_size=32, + # LFM2 interleaves short convolution layers with full attention layers; `full_attn_idxs` selects which layers get + # attention, and the rest get a convolution. Keep one of each so both layer types are exercised. + full_attn_idxs=[1], + # The MLP derives its actual dim from `intermediate_size`, rounded *up* to a multiple of `block_multiple_of`. At + # the reference value (256), any tiny `intermediate_size` would still yield a 256-wide MLP, dwarfing the rest of + # the model. Scale it down instead of disabling `block_auto_adjust_ff_dim`, so the MLP keeps the same code path as + # the reference. + block_multiple_of=8, + # Non-size fields kept aligned with the reference so the tiny config only differs in what we scale down. + max_position_embeddings=128000, + norm_eps=1e-05, + rope_theta=1000000.0, + conv_bias=False, + conv_L_cache=3, + block_auto_adjust_ff_dim=True, + block_ffn_dim_multiplier=1.0, + bos_token_id=1, + # The reference tokenizer's EOS is <|im_end|> (7), not <|endoftext|> (2) which `Lfm2Config` defaults to. + eos_token_id=7, + pad_token_id=0, + use_cache=False, +) +model = Lfm2ForCausalLM(config).to(dtype=torch.bfloat16) +init_weights_tiny_model(model) +smoke_test(model, tokenizer) +check_dtype_pattern(MODEL_ID, model) +print_config_diff(MODEL_ID, model) +push_to_hub(model, tokenizer, generation_config, "tiny") diff --git a/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm_2_5.py b/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm_2_5.py new file mode 100644 index 00000000000..5dfb1bc1daf --- /dev/null +++ b/scripts/generate_tiny_models/for_causal_lm/lfm2_for_causal_lm_2_5.py @@ -0,0 +1,67 @@ +# Copyright 2020-2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from transformers import AutoTokenizer, GenerationConfig, Lfm2Config, Lfm2ForCausalLM + +from .._common import ( + check_dtype_pattern, + check_transformers_version, + init_weights_tiny_model, + print_config_diff, + push_to_hub, + smoke_test, +) + + +# LFM2.5 ships a `TokenizersBackend` tokenizer, which was introduced in transformers 5.0.0. This is above TRL's +# transformers floor, so unlike the other tiny models this one can't be loaded by the floor CI job, and the tests +# using it are skipped there. +check_transformers_version("5.0.0") + +MODEL_ID = "LiquidAI/LFM2.5-230M" + +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) +generation_config = GenerationConfig.from_pretrained(MODEL_ID) +config = Lfm2Config( + vocab_size=65536, + hidden_size=8, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + intermediate_size=32, + # LFM2 interleaves short convolution layers with full attention layers. The reference spells the pattern out with + # `layer_types` rather than `full_attn_idxs`; keep one of each so both layer types are exercised. + layer_types=["conv", "full_attention"], + # Non-size fields kept aligned with the reference so the tiny config only differs in what we scale down. + max_position_embeddings=128000, + norm_eps=1e-05, + rope_parameters={"rope_theta": 1000000.0, "rope_type": "default"}, + conv_bias=False, + conv_L_cache=3, + block_auto_adjust_ff_dim=False, + block_multiple_of=256, + block_ffn_dim_multiplier=1.0, + bos_token_id=1, + # The reference tokenizer's EOS is <|im_end|> (7), not <|endoftext|> (2) which `Lfm2Config` defaults to. + eos_token_id=7, + pad_token_id=0, + use_cache=False, +) +model = Lfm2ForCausalLM(config).to(dtype=torch.bfloat16) +init_weights_tiny_model(model) +smoke_test(model, tokenizer) +check_dtype_pattern(MODEL_ID, model) +print_config_diff(MODEL_ID, model) +push_to_hub(model, tokenizer, generation_config, "tiny", "2.5") diff --git a/tests/test_chat_template_utils.py b/tests/test_chat_template_utils.py index 8b17d3cd20d..5557e7110a3 100644 --- a/tests/test_chat_template_utils.py +++ b/tests/test_chat_template_utils.py @@ -135,6 +135,14 @@ class TestAddResponseSchema: strict=True, ), ), + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + id="lfm2-2.5", + marks=pytest.mark.skipif( + not _SUPPORTS_RESPONSE_TEMPLATE, + reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13", + ), + ), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"), pytest.param( @@ -253,6 +261,14 @@ class TestSupportsToolCalling: ), ), pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"), + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + id="lfm2-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"), pytest.param( @@ -361,6 +377,10 @@ def test_supports_tool_calling(self, model_id): pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3", id="phi3"), pytest.param("trl-internal-testing/tiny-Phi3ForCausalLM-3.5", id="phi3.5"), # Renders tool message content as plain text but drops assistant tool_calls + # LFM2 renders `tools` into the system prompt and wraps tool message content in + # <|tool_response_start|> / <|tool_response_end|>, but never reads `tool_calls`: the model is trained to + # emit <|tool_call_start|> / <|tool_call_end|> as plain text inside `content`. + pytest.param("trl-internal-testing/tiny-Lfm2ForCausalLM", id="lfm2"), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3", id="llama3"), pytest.param("trl-internal-testing/tiny-Qwen2VLForConditionalGeneration", id="qwen2_vl"), pytest.param("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration", id="qwen2.5_vl"), @@ -608,6 +628,7 @@ def test_template_error_returns_false(self): pytest.param( "trl-internal-testing/tiny-Idefics3ForConditionalGeneration", id="idefics3", marks=require_vision ), + pytest.param("trl-internal-testing/tiny-Lfm2ForCausalLM", id="lfm2"), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3", id="llama3"), pytest.param("trl-internal-testing/tiny-LlavaForConditionalGeneration", id="llava", marks=require_vision), pytest.param( @@ -989,6 +1010,14 @@ def test_assistant_masks_multi_turn(self, tokenizer_name, request): [ pytest.param("trl-internal-testing/tiny-Glm4MoeForCausalLM", id="glm4moe"), pytest.param("trl-internal-testing/tiny-GptOssForCausalLM", id="gptoss"), + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + id="lfm2-2.5", + marks=pytest.mark.skipif( + not _SUPPORTS_RESPONSE_TEMPLATE, + reason="LFM2.5 only ships a new-style response template, which requires transformers>=5.13", + ), + ), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.1", id="llama3.1"), pytest.param("trl-internal-testing/tiny-LlamaForCausalLM-3.2", id="llama3.2"), pytest.param( @@ -1094,10 +1123,13 @@ def test_parse_response_with_reasoning_content(self, model_name): pytest.skip("gpt-oss thinking/content separation requires the response_template parser (>= 5.13).") processing_class = self._load(model_name) - # gpt-oss uses the `thinking` field name (matching its harmony chat template) rather - # than the `reasoning_content` convention used by other models. + # gpt-oss and LFM2.5 use the `thinking` field name (matching their own chat templates, which read it back off + # the message) rather than the `reasoning_content` convention used by other models. reasoning_field = ( - "thinking" if model_name == "trl-internal-testing/tiny-GptOssForCausalLM" else "reasoning_content" + "thinking" + if model_name + in ("trl-internal-testing/tiny-GptOssForCausalLM", "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5") + else "reasoning_content" ) messages = [ {"role": "user", "content": "What is 3*4?"}, diff --git a/tests/test_data_utils.py b/tests/test_data_utils.py index 6ebddf8034c..5f510fd738f 100644 --- a/tests/test_data_utils.py +++ b/tests/test_data_utils.py @@ -537,6 +537,14 @@ class TestApplyChatTemplate(TrlTestCase): reason="GLM4 tokenizer requires transformers>=5.0.0", ), ), + "trl-internal-testing/tiny-Lfm2ForCausalLM", + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), "trl-internal-testing/tiny-LlamaForCausalLM-3.1", "trl-internal-testing/tiny-LlamaForCausalLM-3.2", "trl-internal-testing/tiny-LlamaForCausalLM-3", diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py index 2ce14e8db84..28a031803c0 100644 --- a/tests/test_dpo_trainer.py +++ b/tests/test_dpo_trainer.py @@ -195,6 +195,14 @@ class TestDPOTrainer(TrlTestCase): reason="Olmo 3 requires transformers>=4.57.0", ), ), + "trl-internal-testing/tiny-Lfm2ForCausalLM", + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), ], ) def test_train(self, model_id): diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py index 62a4ac854ae..54c3f3a3b83 100644 --- a/tests/test_sft_trainer.py +++ b/tests/test_sft_trainer.py @@ -294,6 +294,14 @@ def test_init_with_training_arguments(self): reason="Olmo 3 requires transformers>=4.57.0", ), ), + "trl-internal-testing/tiny-Lfm2ForCausalLM", + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), ], ) def test_train(self, model_id): @@ -2462,6 +2470,14 @@ def test_train_offloading(self, model_name, packing): "trl-internal-testing/tiny-GemmaForCausalLM", "trl-internal-testing/tiny-Glm4MoeForCausalLM", "trl-internal-testing/tiny-GptOssForCausalLM", + "trl-internal-testing/tiny-Lfm2ForCausalLM", + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), "trl-internal-testing/tiny-LlamaForCausalLM-3.1", "trl-internal-testing/tiny-LlamaForCausalLM-3.2", "trl-internal-testing/tiny-LlamaForCausalLM-3", diff --git a/tests/test_utils.py b/tests/test_utils.py index a39dba2052a..9f9b1289f68 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1282,6 +1282,14 @@ def forward(self, input_ids, attention_mask=None, labels=None, **kwargs): "trl-internal-testing/tiny-GemmaForCausalLM", "trl-internal-testing/tiny-Glm4MoeForCausalLM", "trl-internal-testing/tiny-GptOssForCausalLM", + "trl-internal-testing/tiny-Lfm2ForCausalLM", + pytest.param( + "trl-internal-testing/tiny-Lfm2ForCausalLM-2.5", + marks=pytest.mark.skipif( + Version(transformers.__version__) < Version("5.0.0"), + reason="LFM2.5 tokenizer requires transformers>=5.0.0", + ), + ), "trl-internal-testing/tiny-LlamaForCausalLM-3.1", "trl-internal-testing/tiny-LlamaForCausalLM-3.2", "trl-internal-testing/tiny-LlamaForCausalLM-3", @@ -1420,7 +1428,13 @@ def test_forward(self, model_id, temperature): @pytest.mark.parametrize("model_id", _CHUNKED_LM_HEAD_MODEL_IDS) @pytest.mark.parametrize("temperature", [1.0, 0.7]) def test_backward(self, model_id, temperature): - model_ref = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16).to(torch_device) + # Run in float32, not bfloat16. For models that tie their embeddings, `lm_head.weight.grad` is the sum of the + # LM-head projection and the input-embedding lookup, so its absolute error is set by the larger of the two + # rather than by the element's own value. In bfloat16 that error (~1 ULP at the gradient's scale) swamps the + # assertion: every untied model lands on the same ~1.6e-2 bf16 quantum, and the tied ones range up to ~4e-1, + # which no useful tolerance can separate from a real bug. float32 drops the worst case to ~3e-5 and makes the + # test sensitive to the thing it is meant to check: that the chunked gradient matches the reference. + model_ref = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.float32).to(torch_device) model_chunked = copy.deepcopy(model_ref) B, S, chunk_size = 2, 8, 32 @@ -1442,7 +1456,7 @@ def test_backward(self, model_id, temperature): out["log_probs"].sum().backward() chunked_grad = model_chunked.lm_head.weight.grad.clone() - torch.testing.assert_close(chunked_grad, ref_grad, atol=5e-2, rtol=5e-2) + torch.testing.assert_close(chunked_grad, ref_grad, atol=1e-3, rtol=1e-3) class TestComputeFlopsPerToken(TrlTestCase): diff --git a/trl/chat_template_utils.py b/trl/chat_template_utils.py index 180c3c762aa..5f3a1c08c49 100644 --- a/trl/chat_template_utils.py +++ b/trl/chat_template_utils.py @@ -446,6 +446,56 @@ def clone_chat_template( }, } +# An LFM2.5 argument value is either a single-quoted string, a JSON list/mapping, or a bare literal running up to the +# next `,` or `)`. Lists and mappings contain those very separators, so the pattern has to consume balanced brackets, +# which a regex can only do up to a fixed nesting depth. Three levels covers realistic tool arguments; deeper values +# truncate. +_LFM2_2_5_ATOM = r"[^\[\]{}'\"]|'[^']*'|\"[^\"]*\"" +_LFM2_2_5_NESTED = _LFM2_2_5_ATOM +for _ in range(3): + _LFM2_2_5_NESTED = rf"{_LFM2_2_5_ATOM}|[\[{{](?:{_LFM2_2_5_NESTED})*[\]}}]" +_LFM2_2_5_ARG_VALUE = rf"'[^']*'|[\[{{](?:{_LFM2_2_5_NESTED})*[\]}}]|[^,)]*" + +lfm2_2_5_template = { + # LFM2.5 renders every tool call of a turn into a single `<|tool_call_start|>[...]<|tool_call_end|>` block, as a + # comma-separated list of Python-style calls: `[get_weather(city='Paris', days=3), ping()]`. The `open_pattern` + # therefore matches either the opening bracket or the comma separating two calls, and `close_pattern` matches the + # closing paren plus, for the last call, the trailing `]<|tool_call_end|>`. + # + # Argument values are rendered by the template's `format_arg_value` macro: strings single-quoted, mappings as JSON, + # everything else via Jinja's `string` filter. `string_delims` turns the single-quoted strings back into JSON + # strings so they parse. Note that this round-trip is lossy for Python literals the `string` filter emits but JSON + # cannot express: `flag=False` parses back as the string `"False"`, not `False`. That's inherent to the upstream + # template, not something we can recover here. + # + # The thinking field is named `thinking` (not `reasoning_content` as in other families) to match what the LFM2.5 + # template itself reads back off the message. + "defaults": {"role": "assistant"}, + "start_anchor": "<|im_start|>assistant\n", + "fields": { + "thinking": { + "open": "", + "close_pattern": r"\s*", + "content": "text", + }, + "tool_calls": { + "open_pattern": r"(?:<\|tool_call_start\|>\[|,\s*)(?P\w+)\(", + "close_pattern": r"\)(?:\]<\|tool_call_end\|>)?\s*", + "repeats": True, + "content": "xml-inline", + "content_args": { + "tag_pattern": rf"(?P\w+)=(?P{_LFM2_2_5_ARG_VALUE})", + "value_parser": {"name": "json", "args": {"string_delims": [["'", "'"]], "allow_non_json": True}}, + }, + "transform": {"type": "function", "function": {"name": "{name}", "arguments": "{content}"}}, + }, + "content": { + "close_pattern": r"<\|im_end\|>\s*", + "content": "text", + }, + }, +} + nemotron_3_template = { # Nemotron 3 always prefills the assistant turn with the `` opener (and a trailing `\n` when thinking is # enabled, or `` when it is not), so the start anchor consumes it. Reasoning content is therefore not framed @@ -497,6 +547,10 @@ def clone_chat_template( idefics3_chat_template = (_CHAT_TEMPLATES_DIR / "idefics3.jinja").read_text(encoding="utf-8") +lfm2_chat_template = (_CHAT_TEMPLATES_DIR / "lfm2.jinja").read_text(encoding="utf-8") + +lfm2_2_5_chat_template = (_CHAT_TEMPLATES_DIR / "lfm2_2_5.jinja").read_text(encoding="utf-8") + llama3_chat_template = (_CHAT_TEMPLATES_DIR / "llama3.jinja").read_text(encoding="utf-8") llama3_1_chat_template = (_CHAT_TEMPLATES_DIR / "llama3_1.jinja").read_text(encoding="utf-8") @@ -602,6 +656,10 @@ def add_response_schema(processing_class: ProcessingClassT) -> ProcessingClassT: nemotron_3_ultra_chat_template, ]: schema, template = qwen3_5_schema, nemotron_3_template + elif chat_template == lfm2_2_5_chat_template: + # Only the new-style template; the legacy schema is on its way out, so it isn't worth adding for a family whose + # tokenizer already requires transformers >= 5.0.0. + schema, template = None, lfm2_2_5_template else: raise ValueError( "Unrecognized chat template, failed to add response schema. Please manually set the response schema on " @@ -614,6 +672,12 @@ def add_response_schema(processing_class: ProcessingClassT) -> ProcessingClassT: # attribute reflects the parser actually in play. if _SUPPORTS_RESPONSE_TEMPLATE: tokenizer.response_template = template + elif schema is None: + raise ValueError( + f"Response parsing for this chat template is only provided as a new-style `response_template`, which " + f"requires transformers >= 5.13 (installed: {transformers.__version__}). Please upgrade transformers, or " + f"manually set a legacy `response_schema` on the tokenizer or processor." + ) else: tokenizer.response_schema = schema return processing_class @@ -835,6 +899,8 @@ def is_chat_template_stop_token_trained( idefics3_training_chat_template = (_CHAT_TEMPLATES_DIR / "idefics3_training.jinja").read_text(encoding="utf-8") +lfm2_training_chat_template = (_CHAT_TEMPLATES_DIR / "lfm2_training.jinja").read_text(encoding="utf-8") + llama3_training_chat_template = (_CHAT_TEMPLATES_DIR / "llama3_training.jinja").read_text(encoding="utf-8") llava_next_training_chat_template = (_CHAT_TEMPLATES_DIR / "llava_next_training.jinja").read_text(encoding="utf-8") @@ -887,9 +953,9 @@ def get_training_chat_template( Returns a patched chat template that is prefix-preserving and includes `{%% generation %%}` / `{%% endgeneration %%}` markers for assistant-only loss masking. Returns `None` if the template already satisfies both requirements. - Currently Cohere, Cohere 2, DeepSeek-V3, Gemma, Gemma 2, Gemma 3, GLM-4-MoE, GPT-OSS, Idefics3, LLaMA 3, Phi-3, - Phi-3.5, Qwen2-VL, Qwen2.5, Qwen2.5-VL, Qwen3 (including the Instruct-2507 variant), Qwen3-VL, Qwen3.5, and Qwen3.6 - are supported. + Currently Cohere, Cohere 2, DeepSeek-V3, Gemma, Gemma 2, Gemma 3, GLM-4-MoE, GPT-OSS, Idefics3, LFM2, LLaMA 3, + Phi-3, Phi-3.5, Qwen2-VL, Qwen2.5, Qwen2.5-VL, Qwen3 (including the Instruct-2507 variant), Qwen3-VL, Qwen3.5, and + Qwen3.6 are supported. Args: processing_class (`PreTrainedTokenizerBase` or `ProcessorMixin`): @@ -982,6 +1048,9 @@ def get_training_chat_template( if processing_class.chat_template == idefics3_chat_template: return idefics3_training_chat_template + if processing_class.chat_template == lfm2_chat_template: + return lfm2_training_chat_template + if processing_class.chat_template == llama3_chat_template: return llama3_training_chat_template diff --git a/trl/chat_templates/README.md b/trl/chat_templates/README.md index 685284a0009..09b94b5a26a 100644 --- a/trl/chat_templates/README.md +++ b/trl/chat_templates/README.md @@ -45,6 +45,14 @@ Original GPT-OSS chat template. Original Idefics3 chat template (as shipped by `HuggingFaceM4/Idefics3-8B-Llama3`). Does not support tool calling. +### `lfm2.jinja` + +Original LFM2 chat template (as shipped by `LiquidAI/LFM2-*` checkpoints). ChatML-style. Renders `tools` into the system prompt wrapped in `<|tool_list_start|>` / `<|tool_list_end|>`, and wraps `tool` message content in `<|tool_response_start|>` / `<|tool_response_end|>`. It never reads `message['tool_calls']` though, so assistant tool calls are silently dropped: the model is trained to emit `<|tool_call_start|>` / `<|tool_call_end|>` as plain text inside `content`. + +### `lfm2_2_5.jinja` + +Original LFM2.5 chat template (as shipped by `LiquidAI/LFM2.5-230M` and the other checkpoints in that generation). Unlike `lfm2.jinja`, it renders assistant `tool_calls` — as a single `<|tool_call_start|>[name(key=value, ...)]<|tool_call_end|>` block holding a comma-separated list of Python-style calls — supports a `` block (read off `message.thinking`), and already carries `{% generation %}` markers, so no training patch is needed. Response parsing uses `lfm2_2_5_template`. + ### `llama3.jinja` Original Llama 3 chat template. @@ -162,6 +170,12 @@ Patched Idefics3 template. Diff vs `idefics3.jinja`: Split the assistant message into its own branch so the `{% generation %}` / `{% endgeneration %}` markers wrap the assistant content. This enables `return_assistant_tokens_mask=True` to produce correct masks for SFT assistant-only loss. +### `lfm2_training.jinja` + +Patched LFM2 template. Diff vs `lfm2.jinja`: + +Split the unified message output line into role-specific branches, so the `<|im_start|>assistant\n` prompt cue sits outside the generation block (it is not generated by the model), while the assistant's content and `<|im_end|>\n` (which the model must learn to produce and to stop on) sit inside. Wrap the assistant content with `{% generation %}` / `{% endgeneration %}` so that `return_assistant_tokens_mask=True` produces correct masks for SFT assistant-only loss. + ### `llama3_training.jinja` Patched Llama 3 template. Diff vs `llama3.jinja`: diff --git a/trl/chat_templates/lfm2.jinja b/trl/chat_templates/lfm2.jinja new file mode 100644 index 00000000000..ed66b6de622 --- /dev/null +++ b/trl/chat_templates/lfm2.jinja @@ -0,0 +1,37 @@ +{{- bos_token -}} +{%- set system_prompt = "" -%} +{%- set ns = namespace(system_prompt="") -%} +{%- if messages[0]["role"] == "system" -%} + {%- set ns.system_prompt = messages[0]["content"] -%} + {%- set messages = messages[1:] -%} +{%- endif -%} +{%- if tools -%} + {%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: <|tool_list_start|>[" -%} + {%- for tool in tools -%} + {%- if tool is not string -%} + {%- set tool = tool | tojson -%} + {%- endif -%} + {%- set ns.system_prompt = ns.system_prompt + tool -%} + {%- if not loop.last -%} + {%- set ns.system_prompt = ns.system_prompt + ", " -%} + {%- endif -%} + {%- endfor -%} + {%- set ns.system_prompt = ns.system_prompt + "]<|tool_list_end|>" -%} +{%- endif -%} +{%- if ns.system_prompt -%} + {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}} +{%- endif -%} +{%- for message in messages -%} + {{- "<|im_start|>" + message["role"] + "\n" -}} + {%- set content = message["content"] -%} + {%- if content is not string -%} + {%- set content = content | tojson -%} + {%- endif -%} + {%- if message["role"] == "tool" -%} + {%- set content = "<|tool_response_start|>" + content + "<|tool_response_end|>" -%} + {%- endif -%} + {{- content + "<|im_end|>\n" -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|im_start|>assistant\n" -}} +{%- endif -%} \ No newline at end of file diff --git a/trl/chat_templates/lfm2_2_5.jinja b/trl/chat_templates/lfm2_2_5.jinja new file mode 100644 index 00000000000..8bca4a545e9 --- /dev/null +++ b/trl/chat_templates/lfm2_2_5.jinja @@ -0,0 +1,115 @@ +{{- bos_token -}} +{%- set preserve_thinking = preserve_thinking | default(false) -%} + +{%- macro format_arg_value(arg_value) -%} + {%- if arg_value is string -%} + {{- "'" + arg_value + "'" -}} + {%- elif arg_value is mapping -%} + {{- arg_value | tojson -}} + {%- else -%} + {{- arg_value | string -}} + {%- endif -%} +{%- endmacro -%} + +{%- macro parse_content(content) -%} + {%- if content is string -%} + {{- content -}} + {%- else -%} + {%- set _ns = namespace(result="") -%} + {%- for item in content -%} + {%- if item["type"] == "image" -%} + {%- set _ns.result = _ns.result + "" -%} + {%- elif item["type"] == "text" -%} + {%- set _ns.result = _ns.result + item["text"] -%} + {%- else -%} + {%- set _ns.result = _ns.result + item | tojson -%} + {%- endif -%} + {%- endfor -%} + {{- _ns.result -}} + {%- endif -%} +{%- endmacro -%} + +{%- macro render_tool_calls(tool_calls) -%} + {%- set tool_calls_ns = namespace(tool_calls=[]) -%} + {%- for tool_call in tool_calls -%} + {%- set func_name = tool_call["function"]["name"] -%} + {%- set func_args = tool_call["function"]["arguments"] -%} + {%- set args_ns = namespace(arg_strings=[]) -%} + {%- for arg_name, arg_value in func_args.items() -%} + {%- set args_ns.arg_strings = args_ns.arg_strings + [arg_name + "=" + format_arg_value(arg_value)] -%} + {%- endfor -%} + {%- set tool_calls_ns.tool_calls = tool_calls_ns.tool_calls + [func_name + "(" + (args_ns.arg_strings | join(", ")) + ")"] -%} + {%- endfor -%} + {{- "<|tool_call_start|>[" + (tool_calls_ns.tool_calls | join(", ")) + "]<|tool_call_end|>" -}} +{%- endmacro -%} + +{%- set ns = namespace(system_prompt="", last_user_index=-1) -%} +{%- if messages[0]["role"] == "system" -%} + {%- if messages[0].get("content") -%} + {%- set ns.system_prompt = parse_content(messages[0]["content"]) -%} + {%- endif -%} + {%- set messages = messages[1:] -%} +{%- endif -%} +{%- if tools -%} + {%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: [" -%} + {%- for tool in tools -%} + {%- if tool is not string -%} + {%- set tool = tool | tojson -%} + {%- endif -%} + {%- set ns.system_prompt = ns.system_prompt + tool -%} + {%- if not loop.last -%} + {%- set ns.system_prompt = ns.system_prompt + ", " -%} + {%- endif -%} + {%- endfor -%} + {%- set ns.system_prompt = ns.system_prompt + "]" -%} +{%- endif -%} +{%- if ns.system_prompt -%} + {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}} +{%- endif -%} +{%- for message in messages -%} + {%- if message["role"] == "user" -%} + {%- set ns.last_user_index = loop.index0 -%} + {%- endif -%} +{%- endfor -%} +{%- for message in messages -%} + {{- "<|im_start|>" + message.role + "\n" -}} + {%- if message.role == "assistant" -%} + {%- generation -%} + {%- if message.thinking is defined and (preserve_thinking or loop.index0 > ns.last_user_index) -%} + {{- "" + message.thinking + "" -}} + {%- endif -%} + {%- set _cfm_tag = "CONTINUE_FINAL_MESSAGE_TAG " -%} + {%- set _has_cfm = false -%} + {%- if message.content is defined -%} + {%- set content = parse_content(message.content) -%} + {%- if not (preserve_thinking or loop.index0 > ns.last_user_index) -%} + {%- if "" in content -%} + {%- set content = content.split("")[-1] | trim -%} + {%- endif -%} + {%- endif -%} + {%- if message.tool_calls is defined and content.endswith(_cfm_tag) -%} + {%- set _has_cfm = true -%} + {%- set _trunc_len = (content | length) - (_cfm_tag | length) -%} + {{- content[:_trunc_len] -}} + {%- else -%} + {{- content -}} + {%- endif -%} + {%- endif -%} + {%- if message.tool_calls is defined -%} + {{- render_tool_calls(message.tool_calls) -}} + {%- endif -%} + {%- if _has_cfm -%} + {{- _cfm_tag -}} + {%- endif -%} + {{- "<|im_end|>\n" -}} + {%- endgeneration -%} + {%- else %} + {%- if message.get("content") -%} + {{- parse_content(message["content"]) -}} + {%- endif -%} + {{- "<|im_end|>\n" -}} + {%- endif %} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|im_start|>assistant\n" -}} +{%- endif -%} \ No newline at end of file diff --git a/trl/chat_templates/lfm2_training.jinja b/trl/chat_templates/lfm2_training.jinja new file mode 100644 index 00000000000..e15ffddc12d --- /dev/null +++ b/trl/chat_templates/lfm2_training.jinja @@ -0,0 +1,41 @@ +{{- bos_token -}} +{%- set system_prompt = "" -%} +{%- set ns = namespace(system_prompt="") -%} +{%- if messages[0]["role"] == "system" -%} + {%- set ns.system_prompt = messages[0]["content"] -%} + {%- set messages = messages[1:] -%} +{%- endif -%} +{%- if tools -%} + {%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: <|tool_list_start|>[" -%} + {%- for tool in tools -%} + {%- if tool is not string -%} + {%- set tool = tool | tojson -%} + {%- endif -%} + {%- set ns.system_prompt = ns.system_prompt + tool -%} + {%- if not loop.last -%} + {%- set ns.system_prompt = ns.system_prompt + ", " -%} + {%- endif -%} + {%- endfor -%} + {%- set ns.system_prompt = ns.system_prompt + "]<|tool_list_end|>" -%} +{%- endif -%} +{%- if ns.system_prompt -%} + {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}} +{%- endif -%} +{%- for message in messages -%} + {{- "<|im_start|>" + message["role"] + "\n" -}} + {%- set content = message["content"] -%} + {%- if content is not string -%} + {%- set content = content | tojson -%} + {%- endif -%} + {%- if message["role"] == "tool" -%} + {%- set content = "<|tool_response_start|>" + content + "<|tool_response_end|>" -%} + {%- endif -%} + {%- if message["role"] == "assistant" -%} + {% generation %}{{- content + "<|im_end|>\n" -}}{% endgeneration %} + {%- else -%} + {{- content + "<|im_end|>\n" -}} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|im_start|>assistant\n" -}} +{%- endif -%} \ No newline at end of file