Add LiquidAI LFM2.5 hybrid model support - #462
Open
justinchuby wants to merge 5 commits into
Open
Conversation
Implement the hybrid double-gated short-convolution and QK-normalized GQA architecture, including recurrent cache and ORT GenAI wiring. Add synthetic and real-weight parity, pinned L4/L5 goldens, CLI/runtime coverage, and dtype validation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds first-class support for LiquidAI’s LFM2.5 hybrid Conv+GQA architecture to mobius, including its hybrid cache contract and ORT GenAI packaging metadata, plus multi-tier validation coverage and goldens.
Changes:
- Introduces an
lfm2model implementation with mixed short-convolution and full-attention layers, plus weight key remapping for upstream naming differences. - Adds a reusable double-gated depthwise short convolution component with a full-kernel recurrent cache (aligned to ORT GenAI’s LFM2 cache contract).
- Extends test/config/registry + ORT GenAI config generation to recognize LFM2 hybrid cache I/O and validate deterministic generation.
Reviewed changes
Copilot reviewed 20 out of 20 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Adds HF config translation for lfm2 and filters conv-specific config keys for synthetic parity. |
| tests/ort_genai_test.py | Adds an integration test that exports LFM2.5 and validates deterministic ORT GenAI generation. |
| tests/integration_test.py | Extends feed construction to handle per-layer conv_state inputs/outputs for conv layers. |
| tests/build_graph_test.py | Verifies hybrid-cache output naming for conv layers (present.{i}.conv_state). |
| tests/_test_configs.py | Adds a representative tiny lfm2 config for graph-build/unit coverage. |
| testdata/golden/causal-lm/lfm2_5-230m.json | Adds pinned golden logits/top-k snapshot for LFM2.5-230M. |
| testdata/golden/causal-lm/lfm2_5-230m_generation.json | Adds pinned deterministic 20-token generation golden for LFM2.5-230M. |
| testdata/cases/causal-lm/lfm2_5-230m.yaml | Adds an L4/L5 real-checkpoint case pinned to a specific HF revision. |
| src/mobius/tasks/_cache_utils.py | Updates hybrid cache input shapes for conv layers to store a full kernel-wide window. |
| src/mobius/models/lfm2.py | Implements Lfm2CausalLMModel + mixed conv/attention decoder layers and projection key remapping. |
| src/mobius/models/init.py | Exports Lfm2CausalLMModel. |
| src/mobius/integrations/ort_genai/genai_config.py | Adds LFM2-specific decoder metadata (layer_types/conv cache) and disables share-buffer mode. |
| src/mobius/integrations/ort_genai/genai_config_test.py | Tests that LFM2 genai_config declares hybrid conv cache fields correctly. |
| src/mobius/integrations/ort_genai/auto_export.py | Maps mobius lfm2 to ORT GenAI model type lfm2. |
| src/mobius/components/_short_conv.py | Adds GatedShortConv implementing the double-gated depthwise causal short-conv with recurrent state. |
| src/mobius/components/_short_conv_test.py | Adds a unit test to ensure the short-conv component builds a stateful graph and exposes expected params. |
| src/mobius/components/init.py | Exports GatedShortConv via the public components API. |
| src/mobius/_testing/torch_reference.py | Extends HF reference execution to support opaque hybrid Cache objects for conv/hybrid models. |
| src/mobius/_registry.py | Registers lfm2 model type, test model id, and dashboard family/category mappings. |
| src/mobius/_configs/_base.py | Adds short-conv config fields and extracts them from HF configs (conv_L_cache/conv_bias), plus lfm2 defaults. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 20 out of 20 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
testdata/cases/causal-lm/lfm2_5-230m.yaml:3
- This SHA does not actually pin the L4/L5 run:
GoldenTestCase.revisionis parsed, buttests/e2e_golden_test.py::_build_model_packagecallsbuild(case.model_id, ...)without using it, andbuild()has no revision parameter. The golden tests therefore load the repository's current default revision and can drift from these committed logits/tokens. Resolve this revision to a local snapshot before building (or thread a revision argument throughbuild) so the advertised pinned coverage is real.
revision: "13a53837c4906b4f7405932532ba85d182bb013b"
tests/ort_genai_test.py:233
- This deterministic exact-token test resolves both weights and tokenizer artifacts from the Hub's current default revision, so an upstream update can invalidate the hard-coded sequence. Use the same pinned checkpoint revision as the golden case and pass its local snapshot to both
buildandwrite_ort_genai_config.
model_id = "LiquidAI/LFM2.5-230M"
pkg = build(model_id, dtype="f32", load_weights=True)
output_dir = str(tmp_path / "lfm2")
pkg.save(output_dir)
write_ort_genai_config(pkg, output_dir, hf_model_id=model_id)
Comment on lines
+713
to
+714
| short_conv_kernel=getattr(config, "conv_L_cache", 3), | ||
| short_conv_bias=getattr(config, "conv_bias", False), |
Compute LFM2 RMS variance in float32 before casting back for gamma, matching Transformers for decoder, final, and QK norms. This restores deterministic fp16 CUDA generation parity. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Use only the K-1 cached values required by causal convolution so Conv emits exactly the current token span. Replace dynamic length arithmetic with equivalent negative slices and equal splitting, removing redundant shape, subtraction, slicing, and constant nodes while preserving the K-wide ORT GenAI cache contract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Use the public INT64_MAX component constant for all LFM2 short-convolution slices, removing duplicated magic bounds identified during the optimized-graph audit. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Member
Author
|
Can you set stash type for the RMS norm op? We would prefer fused ops. |
Keep config-only changes conservative, but let new model implementations use import-graph affected-model selection. This prevents LFM2 PRs from running the entire GPU golden catalog while preserving run-all behavior for standalone test config and unmapped task changes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
lfm2support forLiquidAI/LFM2.5-230M, following upstream TransformersUpstream references:
Final validation
Hardware: NVIDIA RTX A1000 (8 GB), driver 573.44, CUDA 12.8. Restored global ONNX Runtime GPU 1.28.0 reports TensorRT/CUDA/CPU; actual LFM2 sessions report
['CUDAExecutionProvider', 'CPUExecutionProvider'].Restored-ORT-1.28 parity harness (
Once upon a time):All CUDA outputs were finite. CUDA-resident Hugging Face reference validation after the fp32 RMSNorm fix also produced exact 20/20 generation for fp32, fp16, and bf16.
Optimized graph audit
LinearandINT64_MAX, uses equal Split, retains only required current-mask/K-1-history/K-wide-state slices, and has no duplicated dynamic length arithmeticCI scoping fix
The first two L4/L5 attempts were externally terminated mid-pytest because
tests/_test_configs.pyforcedrun_all=true, causing each GPU job to run the entire golden catalog.detect_affected_modelsnow keeps standalone test-config and unmapped-task changes conservative (run_all=true) while allowing a real model implementation plus its test config to use import-graph scoping. The full PR diff now resolves to:{"affected": ["lfm2"], "run_all": false}This is covered by 61 detector tests, including config-only, model+config, and unmapped-task+config cases.
Quality checklist
lintrunner -aand targeted testsRemaining infrastructure limitations