Add audio-to-audio models: LFM2, LFM2-Audio, Moshi/PersonaPlex - #77
Closed
justinchuby wants to merge 28 commits into
Closed
Add audio-to-audio models: LFM2, LFM2-Audio, Moshi/PersonaPlex#77justinchuby wants to merge 28 commits into
justinchuby wants to merge 28 commits into
Conversation
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 support for LiquidAI’s LFM2 (hybrid ShortConv + attention text LM) and LFM2-Audio (audio-to-audio composite) including a new ShortConv component, a new multi-model audio-to-audio task, hybrid-cache "conv" state handling, and graph-build tests/configs.
Changes:
- Introduces
ShortConvcomponent and adds"conv"as a hybrid-cache layer type (conv_state-only). - Adds LFM2 text model + LFM2-Audio composite model and registers both in the model registry.
- Adds
AudioToAudioTaskto export a 4-modelModelPackage, plus dedicated build-graph tests/configs.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
tests/build_graph_test.py |
Extends hybrid-cache output assertions for "conv" and adds dedicated LFM2-Audio graph build tests. |
tests/_test_configs.py |
Adds tiny LFM2 configs and introduces an (empty) audio-to-audio config bucket into aggregates. |
src/mobius/tasks/_base.py |
Adds "conv" handling to hybrid cache inputs/outputs (conv_state-only). |
src/mobius/tasks/_audio_to_audio.py |
New task that builds 4 sub-model graphs (audio_encoder/embedding/decoder/audio_decoder). |
src/mobius/tasks/__init__.py |
Exports/registers the new "audio-to-audio" task. |
src/mobius/models/lfm2.py |
New LFM2 hybrid conv+attention CausalLM model and HF weight renaming logic. |
src/mobius/models/lfm2_audio.py |
New LFM2-Audio composite model (encoder/embedding/decoder/depthformer) + weight routing/renames. |
src/mobius/models/__init__.py |
Exposes new model classes. |
src/mobius/components/_short_conv.py |
Implements ShortConv gated causal depthwise Conv1d with conv_state caching. |
src/mobius/components/__init__.py |
Exposes ShortConv from the public components API. |
src/mobius/_registry.py |
Registers lfm2 and lfm2_audio and adds sample model IDs / category mapping. |
src/mobius/_configs.py |
Adds Lfm2Config and Lfm2AudioConfig config classes and transformer config handling. |
justinchuby
added a commit
that referenced
this pull request
Apr 1, 2026
- tasks/_base.py: widen StatePair type to include tuple[ir.Value]
for single-state layers (conv/lightning), matching actual 1-tuple
returns from model forward methods
- tasks/_audio_to_audio.py: add explicit parentheses around
'config.audio.num_mel_bins or 128' for clarity
- models/lfm2_audio.py:
* AudioEncoder.forward: add Transpose(perm=[0,2,1]) before
ConformerEncoder (mel is B,n_mels,T; encoder expects B,T,n_mels)
* DepthformerDecoder.forward: build idx_expanded with runtime batch
dim via op.Shape(projected_3d, start=0, end=1) instead of constant 1
* DepthformerDecoder.forward: build position_ids with runtime batch
dim via op.Shape(hidden_states, start=0, end=1) instead of [1,1]
* preprocess_weights: fix tie_word_embeddings head_key from
'lm_head.weight' to 'lfm.lm_head.weight' so the rename function
maps it to 'decoder.lm_head.weight' correctly
- models/moshi.py: remove dead hidden_size/depformer_dim locals
in preprocess_weights (assigned but never used)
- components/_short_conv_test.py: remove unused 'import onnx_ir as ir'
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Apr 2, 2026
- tasks/_base.py: widen StatePair type to include tuple[ir.Value]
for single-state layers (conv/lightning), matching actual 1-tuple
returns from model forward methods
- tasks/_audio_to_audio.py: add explicit parentheses around
'config.audio.num_mel_bins or 128' for clarity
- models/lfm2_audio.py:
* AudioEncoder.forward: add Transpose(perm=[0,2,1]) before
ConformerEncoder (mel is B,n_mels,T; encoder expects B,T,n_mels)
* DepthformerDecoder.forward: build idx_expanded with runtime batch
dim via op.Shape(projected_3d, start=0, end=1) instead of constant 1
* DepthformerDecoder.forward: build position_ids with runtime batch
dim via op.Shape(hidden_states, start=0, end=1) instead of [1,1]
* preprocess_weights: fix tie_word_embeddings head_key from
'lm_head.weight' to 'lfm.lm_head.weight' so the rename function
maps it to 'decoder.lm_head.weight' correctly
- models/moshi.py: remove dead hidden_size/depformer_dim locals
in preprocess_weights (assigned but never used)
- components/_short_conv_test.py: remove unused 'import onnx_ir as ir'
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Comment on lines
+347
to
+352
| # 2. Embedding: text + audio → inputs_embeds (1, S, hidden) | ||
| emb_feeds = { | ||
| "input_ids": np.array([[text_token]], dtype=np.int64), | ||
| "audio_features": audio_features, # (1, T', hidden) | ||
| } | ||
| inputs_embeds = self._embedding.run(None, emb_feeds)[0] # (1, S, hidden) |
There was a problem hiding this comment.
In Lfm2AudioOnnxPipeline.step(), the embedding session is fed an audio_features input, but the exported embedding sub-model (via AudioToAudioTask._build_embedding and _Lfm2AudioEmbedding.forward) only declares input_ids. This will fail in onnxruntime with an invalid input name. Update the example to call the embedding model with input_ids only and then concatenate audio_features (from audio_encoder) with the text embeddings to form inputs_embeds for the decoder (or alternatively change the exported embedding graph to explicitly accept and fuse audio_features).
Suggested change
| # 2. Embedding: text + audio → inputs_embeds (1, S, hidden) | |
| emb_feeds = { | |
| "input_ids": np.array([[text_token]], dtype=np.int64), | |
| "audio_features": audio_features, # (1, T', hidden) | |
| } | |
| inputs_embeds = self._embedding.run(None, emb_feeds)[0] # (1, S, hidden) | |
| # 2. Embedding: text → text_embeds, then fuse with audio_features | |
| emb_input_ids = np.array([[text_token]], dtype=np.int64) | |
| text_embeds = self._embedding.run(None, {"input_ids": emb_input_ids})[0] | |
| # Concatenate along sequence dimension: [audio_features, text_embeds] | |
| inputs_embeds = np.concatenate([audio_features, text_embeds], axis=1) |
Comment on lines
+107
to
+114
| ## Audio-to-Audio | ||
|
|
||
| End-to-end audio language models: audio and text in, audio and text out. | ||
|
|
||
| | Model Type | Module Class | Task | Example HuggingFace Model | | ||
| |---|---|---|---| | ||
| | `lfm2_audio` | `Lfm2AudioModel` | `audio-to-audio` | `LiquidAI/LFM2-Audio-1.5B` | | ||
|
|
There was a problem hiding this comment.
The new Audio-to-Audio section lists only lfm2_audio, but this PR also registers moshi and personaplex as supported audio-to-audio model types. Add rows for moshi / personaplex (and adjust the summary totals if this file is manually maintained), or regenerate the catalog from the registry so the docs reflect the full set of newly supported architectures.
Comment on lines
+10
to
+14
| Typical model split: | ||
| 1. **audio_encoder**: mel/waveform -> audio features (Conformer/encoder) | ||
| 2. **embedding**: text + audio token fusion -> inputs_embeds | ||
| 3. **decoder**: inputs_embeds -> logits + KV cache (hybrid conv+attention LM) | ||
| 4. **audio_decoder**: backbone hidden -> per-codebook logits (depthformer) |
There was a problem hiding this comment.
The module/class docstrings describe the embedding sub-model as fusing text + audio, but AudioToAudioTask._build_embedding() only wires input_ids (text-only). This is correct for the current LFM2-Audio split only if the runtime concatenates audio_encoder outputs with text embeddings before calling decoder. Please clarify the docstring(s) to document that contract explicitly (and note that some subclasses like MoshiTask override embedding inputs).
justinchuby
marked this pull request as draft
April 2, 2026 15:04
New components: - ShortConv: gated causal depthwise Conv1d (LFM2 conv layers) Components: in_proj -> B*x gating -> causal conv -> C gating -> out_proj New model: - Lfm2CausalLMModel: hybrid ShortConv + Attention with SwiGLU MLP Registered as "lfm2" using HybridCausalLMTask Infrastructure: - Lfm2Config with short_conv_kernel/short_conv_bias fields - "conv" layer type support in hybrid cache system (_base.py) - "lfm2" added to attn_qk_norm model list (QK norm via q/k_layernorm) - AudioToAudioTask scaffold for future audio-to-audio models Tests: - 2 test configs: hybrid (conv+attn) and all-attention variants - "conv" layer type assertion in build_graph_test.py - All 2227 tests pass Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Implement LFM2-Audio (LiquidAI/LFM2-Audio-1.5B) as a 4-model ONNX package: - audio_encoder: ConformerEncoder + adapter projection - embedding: text + audio codebook embeddings - decoder: LFM2 hybrid backbone with conv+attention and hybrid cache - audio_decoder: depthformer transformer for codebook-by-codebook generation Key changes: - Lfm2AudioModel with 4 sub-model classes and weight routing - Lfm2AudioConfig extending Lfm2Config with depthformer/codec fields - AudioToAudioTask with hybrid cache support and depthformer audio_decoder - 6 dedicated tests in TestBuildLfm2AudioGraph - Registry entry for lfm2_audio model type Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
- Dynamic codebook head selection via stacked_head_weights parameter indexed by codebook_idx using Gather+MatMul (replaces hardcoded [0]) - preprocess_weights stacks per-codebook weights into stacked tensor - Remove unused audio_features input from embedding model (cleaner I/O) - Fix docstring formatting for D205 lint rule Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
- src/mobius/components/_short_conv_test.py: 18 tests covering ShortConv parameters, prefill (Pad-based causal conv), incremental decode (Concat-based rolling cache), B/C/x gating, kernel sizes - docs/_generate_models.py: add 'Hybrid Conv+Attention' and 'Audio-to-Audio' to _CATEGORY_DESCRIPTIONS and _CATEGORY_ORDER - docs/model-catalog.md: add Hybrid Conv+Attention section (lfm2), Audio-to-Audio section (lfm2_audio), update summary table Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
- Replace duplicated _Lfm2AudioDecoderLayer and _Lfm2AudioConvLayer with aliases to Lfm2AttentionDecoderLayer and Lfm2ConvDecoderLayer from lfm2.py (Lfm2AudioConfig inherits from Lfm2Config, so the constructors accept it directly) - Remove unused ShortConv import from lfm2_audio.py - Add L4+L5 YAML test case for lfm2 (causal-lm/lfm2-1.2b.yaml) - Add L4 YAML test case for lfm2_audio (audio/lfm2-audio-1.5b.yaml) - Add lfm2_audio to model_coverage_test.py _COVERAGE_SKIP (multi-model architecture tested via TestBuildLfm2AudioGraph) Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Adds support for Moshi (kyutai/moshiko-pytorch-bf16) and PersonaPlex (nvidia/personaplex-7b-v1) audio-to-audio models. Architecture (3-model ONNX split via MoshiTask): - embedding: text_emb + 16 per-codebook audio_emb tables - decoder: 32-layer causal transformer (RoPE + SwiGLU MLP + RMSNorm) - audio_decoder: 6-layer depformer depth transformer with per-codebook stacked gating MLPs and head_dim=depformer_dim Key design choices: - No audio_encoder: Moshi/PersonaPlex consume audio as codec token IDs directly (unlike LFM2-Audio which uses a mel-spectrogram conformer). - MoshiTask extends AudioToAudioTask, overriding embedding (adds audio_codes input) and _build_audio_decoder (head_dim = depformer_dim). - preprocess_weights handles: alpha [1,1,H] -> weight [H] reshape, packed in_proj_weight -> split q/k/v, linear_in -> gate/up split, and stacking of per-codebook depformer weights. PersonaPlex-7b-v1 config.json is minimal (only model_type + version); all hyperparameters are inferred from weight shapes with hardcoded v1 defaults in MoshiConfig.from_transformers. Files: - src/mobius/models/moshi.py: new MoshiModel with full architecture - src/mobius/_configs.py: add MoshiConfig (depformer_dim, depformer_layers, depformer_num_heads, num_codebooks, audio_vocab_size) - src/mobius/tasks/_audio_to_audio.py: add MoshiTask - src/mobius/tasks/__init__.py: export MoshiTask, register 'moshi' task - src/mobius/_registry.py: register personaplex + moshi, add test_model_ids - src/mobius/models/__init__.py: export MoshiModel - tests/build_graph_test.py: TestBuildMoshiGraph (6 tests) - tests/synthetic_parity_test.py: skip personaplex/moshi (moshi library) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
examples/moshi_realtime.py demonstrates full-duplex real-time speech with the Moshi ONNX pipeline: - MoshiOnnxPipeline: loads all three ONNX models (embedding, decoder, audio_decoder) and manages KV cache state for streaming inference - MoshiStreamer: full-duplex I/O using sounddevice for simultaneous mic recording and speaker playback in separate threads - EnCodec-based audio tokenization (via moshi library) and detokenization - 12.5 Hz step rate matching Moshi's 80ms frame interval - Graceful ctrl+C handling, resource cleanup, and verbose logging - --export-only flag to export ONNX models without running inference - Dry-run fallback when sounddevice/moshi library is unavailable Requirements: onnxruntime, sounddevice, numpy, moshi (kyutai/moshi) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
- tasks/_base.py: widen StatePair type to include tuple[ir.Value]
for single-state layers (conv/lightning), matching actual 1-tuple
returns from model forward methods
- tasks/_audio_to_audio.py: add explicit parentheses around
'config.audio.num_mel_bins or 128' for clarity
- models/lfm2_audio.py:
* AudioEncoder.forward: add Transpose(perm=[0,2,1]) before
ConformerEncoder (mel is B,n_mels,T; encoder expects B,T,n_mels)
* DepthformerDecoder.forward: build idx_expanded with runtime batch
dim via op.Shape(projected_3d, start=0, end=1) instead of constant 1
* DepthformerDecoder.forward: build position_ids with runtime batch
dim via op.Shape(hidden_states, start=0, end=1) instead of [1,1]
* preprocess_weights: fix tie_word_embeddings head_key from
'lm_head.weight' to 'lfm.lm_head.weight' so the rename function
maps it to 'decoder.lm_head.weight' correctly
- models/moshi.py: remove dead hidden_size/depformer_dim locals
in preprocess_weights (assigned but never used)
- components/_short_conv_test.py: remove unused 'import onnx_ir as ir'
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
examples/lfm2_audio_realtime.py demonstrates real-time speech-to-speech
conversation with the LFM2-Audio ONNX pipeline.
Key differences from moshi_realtime.py (4-model split vs 3):
- audio_encoder: Conformer processes mel spectrogram (B, n_mels, T)
-> audio_features (B, T', hidden) before embedding
- embedding: takes text_ids + audio_features (not raw audio codes)
- decoder: uses hybrid cache -- ShortConv layers carry conv_state,
attention layers carry standard KV pairs; managed via session
input/output name inspection
- audio_decoder: same depthformer structure as Moshi (8 codebooks)
Lfm2AudioOnnxPipeline:
- Inspects ONNX decoder session inputs to discover and initialize
hybrid cache tensors (conv_state vs key/value) automatically
- Updates cache by mapping present.{i}.X outputs back to
past_key_values.{i}.X inputs each step
- Computes log-mel spectrogram from PCM using numpy FFT + triangular
mel filterbank (no extra dependencies)
Lfm2Streamer: real-time duplex I/O via sounddevice (same pattern as
MoshiStreamer); codec via encodec package; --dry-run flag for testing
without audio hardware; --export-only for ONNX export only.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Both models have dedicated build tests in TestBuildMoshiGraph but are not
in _test_configs.py (standard parametrized configs) or testdata/cases/
(YAML test cases). This caused 6 model_coverage_test failures:
- test_all_models_have_test_config_or_skip
- test_model_has_l1_l3_config[moshi/personaplex]
- test_all_models_have_yaml_or_skip
- test_model_has_golden_data[moshi/personaplex]
Fix: add them to _COVERAGE_SKIP under the audio section, matching the
lfm2_audio pattern ('Multi-model audio-to-audio — tested via TestBuildMoshiGraph').
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
The five constants _DEFAULT_NUM_LAYERS, _DEFAULT_NUM_KV_HEADS, _DEFAULT_HEAD_DIM, _DEFAULT_SHORT_CONV_KERNEL, and _DEFAULT_DEPTHFORMER_HEAD_DIM were defined but never referenced. The hybrid decoder cache is initialized dynamically by inspecting ONNX session input shapes (_init_hybrid_cache), so hardcoded layer/head counts are not needed. The depthformer head dim is computed inline from the passed depthformer_dim/heads parameters. Also clarify the remaining _DEFAULT_HIDDEN_SIZE comment to explain it is only used as a fallback when conv_state dims are symbolic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
- testdata/cases/audio/moshiko.yaml: Moshi (kyutai/moshiko-pytorch-bf16) - testdata/cases/audio/personaplex-7b.yaml: PersonaPlex 7B (nvidia/personaplex-7b-v1) Both marked L4 with skip_reason (require real-time audio streaming pipeline, not suitable for CI golden generation). Both models are already exercised via TestBuildMoshiGraph in build_graph_test.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
After rebasing onto main, apply conventions added in main since the PR was originally branched: * Use public onnxscript API (OpBuilder) instead of onnxscript._internal.builder * License headers updated to Microsoft / MIT (was ONNX Project / Apache 2.0) * AudioToAudioTask + MoshiTask refactored to new GraphBuilder/builder.input API and use _make_hybrid_cache_inputs / _make_kv_cache_inputs from _cache_utils (the helpers in _base.py were moved out) * Add model_roles ClassVar on AudioToAudioTask and MoshiTask * Export Lfm2Config, Lfm2AudioConfig, MoshiConfig from mobius._configs * Explicit rope_type="default" for MoshiConfig and the inline depformer/ depthformer ArchitectureConfigs (rope_type now defaults to None so initialize_rope returns None without it) * Add required model_type field to LFM2 / LFM2-Audio / Moshi / PersonaPlex L4 YAML cases (schema now requires it) * Sort __all__ entries (lint fix) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
* lfm2: default hidden_act to 'silu' (LFM2 architecture doesn't set hidden_act in HuggingFace config but always uses SiLU/SwiGLU) * moshi / personaplex: register with task='moshi' (MoshiTask) instead of 'audio-to-audio' — Moshi has no audio_encoder (it consumes codec token IDs, not mel spectrograms) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/audio-to-audio-models
branch
from
June 2, 2026 23:23
c5784c9 to
20de927
Compare
Member
Author
|
Rebased onto main (commit What I didRebase resolution of 14 commits onto current main (158 commits ahead):
Modernization fixes (folded into 2 follow-up commits)
Test results (in conda env
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Two fixes uncovered while smoke-testing the realtime examples against
the actual HuggingFace checkpoints:
1. moshi_realtime.py / lfm2_audio_realtime.py: pass external_data=
to onnx_ir.save so multi-GB decoder protos serialize past the 2 GB
protobuf limit. Without it personaplex-7b's decoder.onnx hits
'Failed to serialize proto' on the very first save.
2. Lfm2AudioConfig.from_transformers: LiquidAI/LFM2-Audio-1.5B ships
a custom nested config (top-level keys: lfm/encoder/depthformer/
preprocessor). The previous flat-config path left hidden_size,
num_hidden_layers, num_attention_heads etc. as their defaults
(None / 0) which then failed downstream. Now we:
* pull backbone fields from config.lfm
* pull depthformer dims from config.depthformer (incl. heads)
* construct AudioConfig from config.encoder + config.preprocessor
so the audio_encoder branch actually receives an AudioConfig
instead of asserting None.
Added 'lfm', 'encoder', 'depthformer', 'preprocessor' to
_dict_to_pretrained_config's nested_config_keys so attribute
access works after the JSON-fallback path.
3. _config_resolver: when config.json lacks model_type, fall back to
the architectures list (Lfm2AudioForConditionalGeneration ->
lfm2_audio) — same idea as the wav2vec2-ForCTC -> mms shim.
After these fixes:
- moshi_realtime.py --export-only works end-to-end on
nvidia/personaplex-7b-v1 (decoder.data 26 GB, audio_decoder.data
5.3 GB, embedding.data 1.06 GB).
- lfm2_audio_realtime.py gets past config+dispatch and starts
applying weights; it still trips on a remaining LFM2-Audio
adapter shape mismatch (HF layout is LayerNorm + Linear512->2048
+ GELU + Linear2048->2048; mobius models it as FCMLP512->2048->512
which is wrong) — left as a follow-up since it's a model-class
rework, not an example bug.
All 20 unit tests for lfm2/moshi/persona/short_conv/audio_to_audio
still pass.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
|
Smoke-tested both realtime examples against the live HF checkpoints. Pushed bdb0864 with two infrastructure fixes that unblock the export path:
Status after fixes:
Remaining LFM2-Audio adapter blocker (not addressed here — needs a model-class rework, not an example fix): the HF All 20 unit tests for the new audio models still pass. |
HF audio_adapter is:
model.0 = LayerNorm(encoder_dim)
model.1 = Linear(encoder_dim, hidden_size, bias=False)
model.2 = GELU
model.3 = Linear(hidden_size, hidden_size, bias=True)
The previous _Lfm2AudioEncoder.adapter used FCMLP(encoder_dim,
hidden_size) which builds Linear(encoder_dim->hidden_size) +
Linear(hidden_size->encoder_dim) — wrong output dim and wrong
parameter count. Replace with a dedicated _Lfm2AudioAdapter module
matching the HF Sequential and update _rename_lfm2_audio_weight
to map model.{0,1,3} to pre_norm/up_proj/out_proj.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The previous _DepthformerLayer / _Lfm2AudioDecoderModule predated
inspection of LiquidAI/LFM2-Audio-1.5B. It assumed
num_heads = depthformer_dim / depthformer_heads, a single shared
embedding_norm, and one Linear-per-codebook head — none of which match
the released checkpoint.
Real geometry per layer (dim=1024):
operator.qkv_proj.weight [1536, 1024] # 32q + 8kv + 8kv heads × hd=32
operator.out_proj.weight [1024, 1024]
operator.bounded_attention.q_layernorm.weight [32] # per-head norm
operator.bounded_attention.k_layernorm.weight [32]
feed_forward.w{1,3}.weight [2816, 1024] # SwiGLU 2*4*d/3
feed_forward.w2.weight [1024, 2816]
And per codebook K:
depth_embeddings.K.embedding.weight [2049, 1024]
depth_embeddings.K.embedding_norm.weight [1024]
depth_embeddings.K.to_logits.weight [2049, 1024]
Changes:
* Lfm2AudioConfig gains depthformer_head_dim (default 32),
depthformer_kv_heads (8) and depthformer_intermediate_size
(auto-derived via block_auto_adjust_ff_dim if absent). The legacy
depthformer_heads field is retained but is no longer the geometry
source of truth.
* _DepthformerLayer now uses GQA with q_heads = dim // head_dim and
kv_heads from the config; activates attn_qk_norm with per-head
head_dim (not full hidden) to match bounded_attention.q/k_layernorm.
* _Lfm2AudioDecoderModule drops the single embedding_norm. Each
codebook now carries its own (embedding, embedding_norm, to_logits)
triple under depth_embeddings.K so HF weights land directly. The
forward gathers per-codebook norm + head from stacked tensors
assembled in preprocess_weights, mirroring the Moshi pattern.
* preprocess_weights splits the fused qkv_proj into q/k/v_proj using
q_rows = num_q*head_dim and kv_rows = num_kv*head_dim chunks, and
stacks per-codebook norm/head tensors.
* AudioToAudioTask._build_audio_decoder now sizes the depthformer KV
cache with depthformer_kv_heads × depthformer_head_dim instead of
num_heads × dim/num_heads, matching the new GQA attention layout.
* tiny test config in tests/build_graph_test.py wires the three new
fields explicitly so all 6 TestBuildLfm2AudioGraph cases still pass.
End-to-end: mobius.build('LiquidAI/LFM2-Audio-1.5B') now produces a
4-model package with the audio_decoder sub-model fully populated
(76/76 initializers). Export via examples/lfm2_audio_realtime.py
--export-only writes 4 .onnx + 4 .data files (audio_decoder.data
≈ 201 MB).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
…point
After the depthformer rewrite, two backbone-side mappings were still
wrong against the LiquidAI/LFM2-Audio-1.5B checkpoint:
1. `lfm.embedding_norm.weight` is HF's post-layers RMSNorm. The
_Lfm2AudioDecoder builds it as `self.norm` so the previous
`decoder.{rest}` generic rule produced `decoder.embedding_norm.weight`
which didn't match the model's `decoder.norm.weight`. Add an
explicit rename.
2. The checkpoint omits `lfm.lm_head.weight` (always tied to
embed_tokens) but ships no top-level `tie_word_embeddings` flag, so
config.tie_word_embeddings was False and the lm_head was never
filled. Force the tie whenever the checkpoint pattern matches
(no lm_head, has embed_tokens) and clone the tensor so the two
sub-models (`embedding.text_embed` and `decoder.lm_head`) each
get their own initializer instead of being deduplicated by
apply_weights' shared-storage detection.
Result for nvidia/personaplex audio_decoder + LFM2-Audio:
- embedding: 1/1 weights, 0 missing
- decoder: 157/157, 0 missing
- audio_decoder: 76/76, 0 missing
- audio_encoder: 9/636 (still broken; needs a NeMo-conformer rewrite
in a follow-up — the mobius ConformerEncoder component uses an
incompatible naming/structure).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The LFM2-Audio checkpoint ships a NeMo-style Conformer encoder
(`conformer.layers.K.{norm_*, feed_forward1, feed_forward2, self_attn,
conv}`), but the previous implementation reused the Sherpa/k2-style
`mobius.components.ConformerEncoder` whose parameter names
(`encoders.K.feed_forward_in.net.0.linear.weight`, ...) do not align.
Only 9 of 687 audio_encoder weights were being loaded, so the encoder
produced nonsense output.
This change adds a private `_NeMoConformerEncoder` plus helpers
(`_NeMoSubsampling`, `_NeMoConformerLayer`, `_NeMoConvBlock`,
`_NeMoRelPosAttention`, `_NeMoFeedForward`, `_NeMoBatchNorm1d`,
`_NeMoConv1d`) **inside lfm2_audio.py** so other models that depend on
the shared component remain untouched. Module/attribute names mirror
the HF checkpoint layout exactly, so weight loading is now a plain
`conformer.* -> audio_encoder.encoder.*` rename.
Highlights of the new encoder:
- `pre_encode` matches the depthwise-separable Conv2d subsampling
stack (Conv -> ReLU -> depthwise+pointwise x2 -> ReLU -> Linear),
reducing time by 8x and projecting freq to d_model.
- Each layer applies pre-LayerNorm + macaron-FF1 (0.5 residual) ->
pre-LayerNorm + relative-position MHA -> pre-LayerNorm + conv module
-> pre-LayerNorm + macaron-FF2 (0.5 residual) -> final LayerNorm.
- The conv module follows NeMo: pointwise(2C) -> GLU(sigmoid gate) ->
depthwise (k=9, groups=C) -> BatchNorm -> SiLU -> pointwise(C). The
BatchNorm runs in inference mode using stored
weight/bias/running_mean/running_var (num_batches_tracked is a
scalar buffer and is skipped during weight loading).
- Relative-position attention is the Transformer-XL formulation used
by NeMo: scores = (Q + u)*K + (Q + v)*R where u/v are learned
per-head biases (`pos_bias_u`, `pos_bias_v`) and R is a learned
projection (`linear_pos`, no bias) of a sinusoidal pos embedding
of length 2T-1. The standard rel-shift trick folds the
[B,H,T,2T-1] BD matrix down to [B,H,T,T].
`_rename_lfm2_audio_weight` also now skips `conformer.preprocessor.*`
(host-side mel-extractor buffers) and `*.num_batches_tracked` (no
ONNX home).
Results on LiquidAI/LFM2-Audio-1.5B:
audio_encoder: 687/687 initializers filled (was 9/687)
embedding: 1/1
decoder: 157/157
audio_decoder: 76/76
All 58 lfm2/moshi/persona/short_conv/audio_to_audio tests pass; the
full non-integration suite (2811 tests) is green; the lfm2-audio
real-time export emits all 4 ONNX + 4 .data files end-to-end.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Five issues that surfaced when actually running the example end-to-end on the LiquidAI/LFM2-Audio-1.5B checkpoint: 1. mobius.build() defaulted to BF16 (matching the HF checkpoint), but ORT's CPU EP doesn't have a BF16 Conv kernel for the Conformer subsampling stage. Force dtype="float32" in the example. 2. The temp-dir save in --dry-run mode called onnx_ir.save() without external_data and tripped the 2 GB protobuf limit on the decoder. Mirror the --save-to path and pass external_data=. 3. The embedding sub-model only takes input_ids (text); the example was passing audio_features too, which ORT rejected as an unknown input. Splice audio_features + text_embeds in Python instead (the audio side is already projected to hidden_size by the audio_encoder's adapter). 4. The depthformer KV-cache shapes were initialized from the *legacy* (and wrong) Lfm2AudioConfig.depthformer_heads=16 assumption. The real depthformer geometry is 32 query heads / 8 KV heads / head_dim=32, matching the post-rewrite config fields depthformer_head_dim=32 and depthformer_kv_heads=8. 5. Removed the unused _DEFAULT_DEPTHFORMER_HEADS constant. Verified end-to-end on testdata/652-129742-0006.flac (a ~9 s LibriSpeech English clip): the pipeline runs without error, emits text tokens within vocab range and 8 codebook tokens per 80 ms frame that vary across frames. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
|
LFM2-Audio end-to-end works on real audio now. Pushed 1ad7b78 + 7893103. Tested with Text tokens land in the 65 536 vocab range; each frame emits 8 codebook tokens in [0, 2049); outputs vary across frames (not constant), so the model is doing real work end-to-end. What was wrong with the original PRWhen I started smoke-testing the example, the LFM2-Audio model class had never been validated against the actual HF checkpoint. Every major sub-module had a structural mismatch:
Per-sub-model weight load status (final)Export sizes (float32)Tests still green
Caveats / follow-ups
|
Member
Author
audio_decoder (depthformer) numerical parity ✅Validated Setup: synthetic single-step inputs, fp32 CPU, empty past KV cache.
Per-codebook logits diff (shape (2049,) each):
Per-layer present KV cache diff (cb=0, shape (1, 8, 1, 32) each): | layer | max |K| diff | max |V| diff | KV cache layout Verdict: Confirms the depthformer rewrite is numerically correct:
HF state_dict loaded into |
The LFM2 HuggingFace checkpoint stores the post-layers RMSNorm as `model.embedding_norm.weight`, but the text backbone exposed it as `self.norm`, leaving `model.norm.weight` unbound at build time (MISSING initializer) and the actual HF tensor unmapped. Rename the attribute to `embedding_norm` so weight names line up with HF directly, no preprocess rename needed. This matches what the LFM2-Audio decoder already does for the same tensor (see `_Lfm2AudioDecoder.embedding_norm`). Verified by greedy-decoding LFM2-1.2B end-to-end via ORT: 'Hello, my name is' -> 'Alex, and I'm here to help you with any questions or tasks you have in mind. How can I assist you today?' Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
LFM2-1.2B text-only smoke test — PASS ✅Build: Success after fix (commit Bug found & fixed: The HF checkpoint stores the post-layers RMSNorm as Per-sub-model weight load: Greedy decode (CPU EP, fp32, 30 tokens): Prompt:
Coherent English: Yes — natural assistant-style continuation, terminates cleanly at Result: Hybrid conv+attention backbone ( |
Member
Author
audio_encoder parity validation — FAIL (bug found, stopped per instructions)Setup
Element-wise diff vs HF (
|
| metric | value |
|---|---|
max_abs_diff |
4.945e-01 |
mean_abs_diff |
1.825e-02 |
median_abs_diff |
1.451e-02 |
np.testing.assert_allclose(..., atol=1e-3, rtol=1e-3) → FAIL.
First-5 elements at [0,0,:5]:
- HF:
[-0.00649, -0.07804, 0.05648, -0.09988, 0.26178] - ONNX:
[ 0.01316, -0.10752, 0.04475, -0.12681, 0.26010]
Root-cause: missing audio_adapter.model.1.bias
HF audio_adapter is:
Sequential(
(0): LayerNorm(512, eps=1e-5, bias=True)
(1): Linear(512 -> 2048, bias=True) # bias mean|.|=0.0295, max|.|=0.254
(2): GELU(approximate='none')
(3): Linear(2048 -> 2048, bias=True)
)
But src/mobius/models/lfm2_audio.py (_Lfm2AudioAdapter, ~line 631) declares:
self.up_proj = Linear(encoder_dim, hidden_size, bias=False) # ← should be bias=Trueand the docstring claim "audio_adapter.model.1.weight (no bias in HF)" is incorrect — the HF checkpoint does carry a non-trivial bias on this layer. preprocess_weights (line ~1272) doesn't drop it explicitly, so the bias tensor is silently unmapped during weight loading.
Isolation
| comparison | max_abs | mean_abs |
|---|---|---|
| HF (with bias) vs ONNX | 4.945e-01 | 1.825e-02 |
| HF-simulated without bias on layer 1 vs ONNX | 1.174e-01 | 8.896e-03 |
| HF (with bias) vs HF-simulated (no bias) | 4.954e-01 | 1.500e-02 |
So fixing the missing bias collapses most — but not all — of the divergence (max 0.49 → 0.117, mean 1.8e-2 → 8.9e-3). A second, smaller divergence remains, almost certainly inside the conformer encoder (since the simulated adapter exactly mirrors mobius). That residual is not reproducibility noise — running ONNX twice is bit-identical.
Recommended fix
_Lfm2AudioAdapter.up_proj:bias=False→bias=True.- Drop the incorrect docstring/comments in
_Lfm2AudioEncoderandpreprocess_weightsthat say HF layer 1 has no bias. - After applying (1), re-run this parity script and drill into the conformer (e.g. compare encoder pre-adapter output by re-exporting with adapter disabled) to identify the residual ~0.1 max-abs source.
Reproduction script: scripts/compare_audio_encoder.py (added in workdir, not committed).
Status: blocked — surfacing the bug rather than modifying model code, per the task instructions.
The HF `audio_adapter.model.1` is `Linear(512, 2048, bias=True)` with a real non-zero bias (mean|b|≈0.030, max|b|≈0.254). The earlier rewrite declared it `bias=False" based on an incomplete weight dump and silently dropped the bias during weight load. End-to-end HF↔ONNX parity max abs diff drops from 0.49 → 0.12 with the bias restored. The remaining ~0.12 residual lives inside the conformer encoder itself (per encoder-only isolation attempt) and likely comes from one of: - relative-position attention math (BD shift-trick alignment) - BatchNorm eps / running-stats convention - LayerNorm eps mismatch (HF nn.LayerNorm uses 1e-5 by default) - subsampling Conv2d output indexing Filing the residual as a follow-up; this commit closes the larger correctness gap that broke generation quality at the adapter boundary. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
Moshi / PersonaPlex end-to-end validation — PASSReplicated the LFM2-Audio-depth validation for Fixes (already landed in 6e0876e on
|
….2e-6)
The mobius _NeMoSubsampling ran a vanilla Conv2d stack over the input
mel, while HF's ConvSubsampling wraps the same convs in a
MaskedConvSequential that re-zeros padded time frames before every layer.
Conv2d has a bias, so any all-zero padded frame becomes a non-zero
(bias-only) row after the first conv, and that contamination is mixed
into neighbouring valid frames by subsequent stride-2 convolutions.
Attention then smears the corruption across every output frame, giving
~0.12 max-abs-diff vs the HF reference on a 1-second test clip.
Fix: derive the validity mask from the input itself — the mel
preprocessor writes padded frames as exact zeros after normalization, so
ReduceMax(|mel|, dim=freq) > 0 recovers it — and apply it before every
Conv2d layer. After each stride-2/k=3/pad=1 step, recompute the per-batch
valid length with L_next = (L - 1) // 2 + 1 (same formula HF uses) and
re-broadcast a fresh arange(T_new) < L_next mask. This avoids changing
the audio_encoder ONNX I/O contract (no new feat_lengths input).
Bisection notes (see scripts/bisect_conformer.py): a PyTorch re-impl of
the mobius math, loaded with HF weights, matches HF *block-for-block at
0.0 abs diff* once fed the same layer input. The entire residual lived
in pre_encode (max=1.6e+2 at the masked frame, then propagated via
attention). All other Conformer sub-blocks (rel-pos attention BD shift,
GLU, BatchNorm, macaron FF residuals, norm_out) were already correct.
Verification:
scripts/compare_audio_encoder.py:
before: max=1.15e-1 mean=8.5e-3
after: max=1.22e-6 mean=1.6e-7 PASS @ atol=1e-3
tests/build_graph_test.py -k 'lfm2 or moshi or persona or short_conv or audio_to_audio':
20 passed
Weight remap: pre_encode.conv.{0,2,3,5,6} (HF Sequential indices with
ReLU at 1/4/7) -> pre_encode.conv_{0,2,3,5,6} (named attributes); the
ReLU-only indices are now dropped in _rename_lfm2_audio_weight since the
ReLUs are inlined into the masking flow.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
Conformer parity fix: 1.15e-1 → 1.22e-6 ✅Commit: Root cause
BisectionA PyTorch re-impl of the mobius math (
So the rel-pos attention BD-shift, GLU direction, BatchNorm running stats, macaron FF 0.5 scaling, per-layer FixDerive the validity mask from the input itself — the mel preprocessor writes padded frames as exact zeros after normalization, so VerificationWeight-name remap: Instrumentation kept in |
Member
Author
Fleet-mode validation completeDispatched 5 parallel validation agents. All passed. Two real bugs caught + fixed.
Per-component final parityBoth well inside fp32 tolerance. All test suites still green
Validation scripts shipped
These can be reused for regression testing. |
These were used to chase down the LFM2-Audio numerical-parity bugs (adapter up_proj bias, subsampling padded-frame masking) — fixes are captured in commits 6e0876e and 49c0a49. The scripts hardcoded checkpoint paths and weren't wired into pytest, so they have no ongoing regression value. The production L1 unit tests in tests/build_graph_test.py cover the same graph-construction surface. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
The previous 'too large for CI golden data generation' skip_reason predated the model actually working — now that the LFM2 backbone is fully wired up (per commits 387ea45 and the audio-to-audio rebase), generate_golden.py produces clean L4 top-K logits + L5 generated text in ~15 s on this machine. Both: TestL4CheckpointVerified::test_prefill_argmax_matches_golden[text-generation/lfm2-1.2b] TestL5GenerationE2E::test_generation_matches_golden[text-generation/lfm2-1.2b] pass against the ONNX export. L5 generated text from prompt 'Here is my poem:' is coherent English ('In the quiet of the night, under a blanket of stars, A lone owl hoot'). LFM2-Audio / Moshi / PersonaPlex still need a custom audio-to-audio test scaffold (their pipelines are multi-model + custom audio processing); leaving those skipped for a follow-up. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
A focused, microphone-free example specifically for testing
nvidia/personaplex-7b-v1. Reuses MoshiOnnxPipeline from
moshi_realtime.py for the actual ORT plumbing but drops the streaming
audio I/O complexity:
- Loads an audio file (any sample rate; resampled to 24 kHz with
librosa) and Mimi-encodes per 80 ms frame.
- Falls back to all-zero codes when the Mimi codec isn't installed
so the script still smoke-tests every ONNX sub-model.
- --synthetic flag uses deterministic non-zero random codes,
confirming the pipeline produces frame-varied output (the
all-zero-input case correctly produces constant output, which is
not as informative).
- --save-to builds + exports from HuggingFace; --onnx-dir reloads.
- Best-effort text decode via the HF tokenizer at the end.
Verified end-to-end on testdata/652-129742-0006.flac (without Mimi:
constant outputs as expected for silence-input) and with --synthetic
(per-step varied outputs). KV caches (32 decoder + 6 depformer layer
pairs) update correctly across steps.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.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
Add three new audio/speech model architectures and the
AudioToAudioTaskscaffold.New Models
LFM2 (
lfm2)Hybrid conv+attention language model with ShortConv, sliding window attention, and linear attention.
Lfm2Configwithshort_conv_kernel,short_conv_bias,layer_typesLFM2-Audio (
lfm2_audio)Audio-to-audio composite model producing 4 ONNX sub-models:
audio_encoderembeddingdecoderaudio_decoderDRY implementation: decoder/conv layers reuse classes from
lfm2.pyvia aliases.Moshi/PersonaPlex (
moshi,personaplex)Full-duplex speech model with 3-model split:
embeddingdecoderaudio_decoderNew Task
AudioToAudioTask: Multi-model task for audio-to-audio pipelines with auto-detection of hybrid vs standard cache fromconfig.layer_typesNew Components
ShortConv: Gated causal depthwise Conv1d with streaming state managementExamples
examples/lfm2_audio_realtime.py: Real-time LFM2-Audio streaming demoexamples/moshi_realtime.py: Real-time Moshi/PersonaPlex streaming demo (deep-reviewed with 13 bug fixes: Mimi codec API, thread safety, onnx_ir.save, etc.)Testing
build_graph_test.pylfm2-1.2b.yaml,lfm2-audio-1.5b.yamlmoshi,personaplex,lfm2_audioin_COVERAGE_SKIP