Add LFM2 and LFM2.5 support and testing#6428
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| 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": "<think>", | ||
| "close_pattern": r"</think>\s*", | ||
| "content": "text", | ||
| }, | ||
| "tool_calls": { | ||
| "open_pattern": r"(?:<\|tool_call_start\|>\[|,\s*)(?P<name>\w+)\(", | ||
| "close_pattern": r"\)(?:\]<\|tool_call_end\|>)?\s*", | ||
| "repeats": True, | ||
| "content": "xml-inline", | ||
| "content_args": { | ||
| "tag_pattern": r"(?P<key>\w+)=(?P<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", | ||
| }, | ||
| }, | ||
| } |
kashif
left a comment
kashif
left a comment
There was a problem hiding this comment.
curious if the CI passes
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e29c92649
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "repeats": True, | ||
| "content": "xml-inline", | ||
| "content_args": { | ||
| "tag_pattern": r"(?P<key>\w+)=(?P<value>'[^']*'|\{.*?\}|[^,)]*)", |
There was a problem hiding this comment.
When an LFM2.5 tool call has an array argument or a nested mapping, this regex stops the value at the first comma or first }. The companion template can render valid calls such as foo(items=[1, 2]) or foo(cfg={"a":{"b":1},"c":2}), but response parsing turns those into truncated strings like "[1" or partial JSON before the tool is executed. This silently corrupts tool arguments for GRPO/response parsing; the value parser needs to consume balanced list/object values rather than splitting on commas/braces.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch, reproduced: foo(items=[1, 2]) parsed back as {"items": "[1"}, and foo(cfg={"a": {"b": 1}, "c": 2}) as truncated JSON. The value pattern had no branch for [, and its \{.*?\} branch was non-greedy. Fixed by matching balanced brackets (quote-aware, up to three nesting levels, the regex-fixed-depth cap, noted in the comment). Lists, nested mappings, and lists-of-objects now round-trip.
Not sure if it's the best solution, but it seems to work.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c4b6d94. Configure here.
| 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) |
There was a problem hiding this comment.
Missing tied-embedding backward xfail
Medium Severity
This PR adds LFM2 models to _CHUNKED_LM_HEAD_MODEL_IDS and tightens test_backward, but the PR notes tied-embedding models should xfail at temperature != 1.0. No xfail or skip on tie_word_embeddings appears, so CI can fail for LFM2 and other tied tiny models at temperature=0.7.
Reviewed by Cursor Bugbot for commit c4b6d94. Configure here.
2 participants
LFM2 sees roughly 1000 runs per week as per our telemetry, but the architecture/template isn't covered by our test suite at all. This adds both, for the two generations that behave differently:
trl-internal-testing/tiny-Lfm2ForCausalLM(fromLiquidAI/LFM2-1.2B), landed in transformers 4.54.0, below our floor.trl-internal-testing/tiny-Lfm2ForCausalLM-2.5(fromLiquidAI/LFM2.5-230M), gated ontransformers>=5.0.0, which itsTokenizersBackendtokenizer requires.Both are wired into the usual parametrized lists (
test_data_utils,test_utils,test_sft_trainer,test_dpo_trainer,test_chat_template_utils).Chat templates
The two generations differ in a way worth stating plainly:
toolsand thetoolrole, but never readsmessage["tool_calls"]: an assistant tool call renders as an empty turn. So it goes in the does-not-support tool-calling list, and gets alfm2_training.jinjapatch since its template has no{% generation %}markers.tool_calls, as<|tool_call_start|>[name(key=value, ...)]<|tool_call_end|>. It needs no training patch (its template already ships generation markers and is prefix-preserving), but it does get anlfm2_2_5_templateresponse template so tool calls can be parsed back. Only the new-style template is provided; the legacyresponse_schemais on its way out and isn't worth adding for a new family.One known limitation, inherent to the upstream template: argument values are rendered with Jinja's
stringfilter, soflag=Falseparses back as the string"False". Strings, numbers, and nested JSON round-trip fine.Adding LFM2 to
_CHUNKED_LM_HEAD_MODEL_IDSsurfaced a pre-existing issue. For models that tie their embeddings, the chunked LM head'sgrad_hiddenreacheslm_head.weightthrough the embedding, where the chunked path (fp32 accumulation) and the reference (whose.float()backward casts to bf16) round differently. Both end up several percent from the true fp32 gradient.It is not LFM2-specific, the behaviour splits cleanly on
tie_word_embeddingsacross five architectures (~1e-6 fp32 relative error tied vs ~1e-11 untied). Gemma and Cohere are tied too, and stay inside tolerance only because their gradients are ~40x smaller.Consequently,
test_backwardis markedxfail(strict=True)for LFM2 attemperature != 1.0.Follow-ups?
get_training_chat_template. Highest-value next step.Note
Low Risk
Changes are additive (templates, parsers, tests, docs) with no changes to core training loops; the chunked-LM-head test tweak only affects test numerics.
Overview
Adds first-class support for Liquid LFM2 and LFM2.5 in TRL’s chat-template stack and CI, including tiny hub fixtures and broad test parametrization (SFT, DPO, data utils, template utilities).
LFM2 gets bundled reference templates (
lfm2.jinja,lfm2_training.jinja) and automatic swapping viaget_training_chat_templatefor assistant-only loss. Its template renders tools/tool responses but ignores structuredtool_calls, so it is classified as not supporting standard tool calling. LFM2.5 shipslfm2_2_5.jinja(already has{% generation %}markers) plus a newlfm2_2_5_templateresponse parser for Python-style<|tool_call_start|>[fn(a=1)]<|tool_call_end|>blocks,thinkingfields, and multi-call turns;add_response_schemawires this only as new-styleresponse_template(transformers ≥ 5.13).Docs and
trl/chat_templates/README.mddocument both families.test_utils.test_backwardfor chunked LM head now runs in float32 with tighter tolerance so tied-embedding bf16 noise does not mask real gradient bugs when LFM2 joins the model list.Reviewed by Cursor Bugbot for commit e6cefec. Bugbot is set up for automated code reviews on this repo. Configure here.