Add LFM2 and LFM2.5 support and testing

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 `&#123;% generation %&#125;` / `&#123;% endgeneration %&#125;` 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`:

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")

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")

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"),
@@ -597,6 +617,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(
@@ -978,6 +999,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(
@@ -1083,10 +1112,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?"},

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",

tests/test_dpo_trainer.py

@@ -194,6 +194,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):

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):
@@ -2457,6 +2465,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",

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):