Add BitNet b1.58 ternary-weight model support - #83
Draft
justinchuby wants to merge 13 commits into
Draft
Conversation
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>
Add BitNet b1.58 causal language model with ternary {-1, 0, +1} weights.
The architecture is Llama-like with sub-layer RMSNorms (attn_sub_norm
before o_proj, ffn_sub_norm before down_proj) and squared ReLU activation.
Implementation:
- BitNetAttention: Attention subclass with attn_sub_norm
- BitNetMLP: MLP subclass with ffn_sub_norm
- BitNetCausalLMModel: CausalLMModel subclass with preprocess_weights
that unpacks uint8-packed 2-bit ternary values and applies weight_scale
- Uses standard Linear ops (not MatMulNBits) for initial correctness;
2-bit quantized kernels can be added as a future optimization
Also adds:
- YAML golden test case (with skip_reason pending golden JSON)
- Skill file documenting the quantized/low-bit model pattern
- Test config, registry entry, and test_model_id
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Documents patterns for adding quantized/low-bit weight models to mobius: - linear_class injection pattern and when to use it vs standard Linear - Weight packing/unpacking conventions (ternary uint8 → float) - Sub-layer norm pattern (RMSNorm before projections) - Testing quantized models with tiny configs - BitNet implementation as reference example Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
This reverts commit f3da0ce. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
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 new architecture support to mobius for (1) BitNet b1.58 ternary-weight CausalLMs and (2) LiquidAI/Moshi-style hybrid conv+attention and audio-to-audio model pipelines, including new task wiring, components, registry/config integration, tests, and docs updates.
Changes:
- Add
BitNetCausalLMModelwith ternary weight unpacking + BitNet-specific sub-layer RMSNorms and ReLU² activation. - Introduce LFM2 (ShortConv + attention) and audio-to-audio task support (LFM2-Audio + Moshi/PersonaPlex), including
ShortConvcomponent + unit tests. - Update registry/config parsing, test coverage harness, and documentation scaffolding for new categories/model types.
Reviewed changes
Copilot reviewed 23 out of 23 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
tests/synthetic_parity_test.py |
Adds skip reasons for moshi/personaplex synthetic parity. |
tests/model_coverage_test.py |
Marks lfm2_audio as covered by dedicated build-graph tests. |
tests/build_graph_test.py |
Extends hybrid-cache assertions for conv layers; adds dedicated build tests for LFM2-Audio and Moshi splits. |
tests/_test_configs.py |
Adds tiny BitNet config and LFM2 configs to test matrix; introduces placeholder audio-to-audio config list. |
testdata/cases/causal-lm/lfm2-1.2b.yaml |
Adds skipped golden case for LFM2. |
testdata/cases/causal-lm/bitnet-b1.58-2b.yaml |
Adds skipped golden case for BitNet. |
testdata/cases/audio/lfm2-audio-1.5b.yaml |
Adds skipped golden case for LFM2-Audio. |
src/mobius/tasks/_base.py |
Extends hybrid cache input/output support to single-state conv layers. |
src/mobius/tasks/_audio_to_audio.py |
New multi-model task wiring for audio-to-audio pipelines + Moshi variant. |
src/mobius/tasks/__init__.py |
Exports/registers AudioToAudioTask and MoshiTask. |
src/mobius/models/moshi.py |
Adds Moshi/PersonaPlex 3-model split modules + preprocessing for HF weight layouts. |
src/mobius/models/lfm2.py |
Adds LFM2 hybrid ShortConv+attention CausalLM + HF key renames. |
src/mobius/models/lfm2_audio.py |
Adds LFM2-Audio 4-model split (audio encoder/embedding/decoder/audio decoder) + weight routing. |
src/mobius/models/bitnet.py |
Adds BitNet model with ternary unpacking and sub-layer norms. |
src/mobius/models/__init__.py |
Exposes new model classes. |
src/mobius/components/_short_conv.py |
Implements ShortConv gated causal depthwise Conv1d block with cached state. |
src/mobius/components/_short_conv_test.py |
Adds unit tests for ShortConv shapes/ops/state handling. |
src/mobius/components/__init__.py |
Exports ShortConv from public components API. |
src/mobius/_registry.py |
Registers bitnet, lfm2, lfm2_audio, moshi, personaplex and test model IDs/variants. |
src/mobius/_configs.py |
Adds Lfm2Config, Lfm2AudioConfig, MoshiConfig config parsing. |
examples/moshi_realtime.py |
Adds an example streaming pipeline for Moshi/PersonaPlex ONNX sub-models. |
docs/model-catalog.md |
Updates catalog counts/categories and adds sections for LFM2 + audio-to-audio. |
docs/_generate_models.py |
Adds category descriptions/order for new categories. |
Comment on lines
+59
to
+63
| # Q/K/V projections — (batch, seq_len, num_heads * head_dim) | ||
| query_states = self.q_proj(op, hidden_states) | ||
| key_states = self.k_proj(op, hidden_states) | ||
| value_states = self.v_proj(op, hidden_states) | ||
|
|
There was a problem hiding this comment.
BitNetAttention.forward reimplements Attention.forward but omits the optional Q/K normalization path (q_norm/k_norm) when config.attn_qk_norm is enabled. This silently changes behavior vs the base Attention class for configs that use QK-norm; consider either reusing the base forward logic and inserting attn_sub_norm before o_proj, or duplicating the q_norm/k_norm handling from Attention.forward.
Comment on lines
+149
to
+177
| def preprocess_weights( | ||
| self, state_dict: dict[str, torch.Tensor] | ||
| ) -> dict[str, torch.Tensor]: | ||
| """Unpack ternary weights and apply per-tensor weight_scale. | ||
|
|
||
| HuggingFace packs ternary {-1, 0, +1} values as 2-bit unsigned | ||
| integers, 4 values per uint8 byte. This method: | ||
| 1. Collects ``*.weight_scale`` tensors (per-tensor scaling factors). | ||
| 2. For each packed ``*.weight`` with a matching scale, unpacks the | ||
| uint8 → int8 ternary values and multiplies by the scale. | ||
| 3. Drops the consumed ``weight_scale`` keys from the state dict. | ||
| """ | ||
| # Collect per-tensor weight scales | ||
| weight_scales: dict[str, torch.Tensor] = {} | ||
| for key, value in state_dict.items(): | ||
| if key.endswith(".weight_scale"): | ||
| prefix = key[: -len(".weight_scale")] | ||
| weight_scales[prefix] = value | ||
|
|
||
| new_state_dict: dict[str, torch.Tensor] = {} | ||
| for key, value in state_dict.items(): | ||
| if key.endswith(".weight_scale"): | ||
| continue # consumed by the matching .weight key | ||
| if key.endswith(".weight") and key[: -len(".weight")] in weight_scales: | ||
| scale = weight_scales[key[: -len(".weight")]] | ||
| if value.dtype == torch.uint8: | ||
| value = _unpack_ternary_weights(value) | ||
| value = value.float() * scale.float() | ||
| new_state_dict[key] = value |
There was a problem hiding this comment.
The ternary weight unpacking in preprocess_weights is untested. Since this logic is model-specific and easy to regress (bit packing order, shape expansion, scale application), add a focused unit test that validates _unpack_ternary_weights and the weight_scale multiplication against a small hand-constructed packed tensor.
Comment on lines
+133
to
+152
| def __init__(self, config: Lfm2AudioConfig): | ||
| super().__init__() | ||
| self.text_embed = Embedding(config.vocab_size, config.hidden_size, config.pad_token_id) | ||
| # Audio codebook embedding: codebooks * audio_vocab_size entries | ||
| audio_vocab = config.audio_vocab_size * config.num_codebooks | ||
| self.audio_embed = Embedding(audio_vocab, config.hidden_size) | ||
|
|
||
| def forward( | ||
| self, | ||
| op: builder.OpBuilder, | ||
| input_ids: ir.Value, | ||
| ): | ||
| """Forward: text_ids -> inputs_embeds. | ||
|
|
||
| Returns text embeddings only. Audio codebook embeddings are | ||
| handled separately by the audio_decoder's depth_embeddings. | ||
| Runtime assembles text + audio at the sequence level. | ||
| """ | ||
| return self.text_embed(op, input_ids) | ||
|
|
There was a problem hiding this comment.
_Lfm2AudioEmbedding defines self.audio_embed but forward() never uses it, so its weights will never be realized in the exported ONNX graph and will show up as "not applied" during weight application. Either wire audio_embed into the embedding computation (if the model needs discrete audio token embeddings) or remove it (and the corresponding weight rename) to avoid dead parameters and confusing logs.
Comment on lines
+331
to
+345
| # Per-codebook logit projection | ||
| # Each codebook has its own embedding + tied head | ||
| self.depth_embeddings = nn.ModuleList([]) | ||
| for _ in range(config.num_codebooks): | ||
| self.depth_embeddings.append( | ||
| Linear(depthformer_dim, config.audio_vocab_size, bias=False) | ||
| ) | ||
|
|
||
| # Stacked head weights for dynamic codebook selection via Gather. | ||
| # Shape: (num_codebooks, audio_vocab_size, depthformer_dim) | ||
| # In preprocess_weights, this is assembled from per-codebook weights. | ||
| self.stacked_head_weights = nn.Parameter( | ||
| [config.num_codebooks, config.audio_vocab_size, depthformer_dim] | ||
| ) | ||
|
|
There was a problem hiding this comment.
_Lfm2AudioDecoderModule keeps a per-codebook depth_embeddings ModuleList, but forward() selects weights exclusively from stacked_head_weights. This leaves depth_embeddings.* weights unused in the ONNX graph (and likely unmapped at apply time). Consider deleting the per-codebook weights from the preprocessed state_dict after stacking, or removing depth_embeddings entirely if it’s only present for HF name compatibility, to reduce log noise and memory footprint.
Comment on lines
+3
to
5
| **mobius** supports 275 registered model types across 12 categories. | ||
| This catalog lists every supported architecture with its module class, task type, | ||
| and example HuggingFace model IDs. |
There was a problem hiding this comment.
docs/model-catalog.md is updated to reflect new totals/categories, but it still doesn’t mention newly added model types from this PR (e.g., bitnet, moshi, personaplex) anywhere in the catalog tables. If this page is meant to list every supported model_type, add rows for these entries (or regenerate the page from the registry to keep it consistent).
- Fix __all__ sort order in models/__init__.py (BitNetCausalLMModel now alphabetical) - Add input validation in _unpack_ternary_weights (ndim check) - Add bitnet_test.py with 9 unit tests for _unpack_ternary_weights: shape, ternary value set, known byte round-trip, edge cases (all-zero, all-one, all-two bytes), dtype, invalid input, multi-column Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Address Copilot reviewer comment: BitNetAttention.forward was missing the optional Q/K normalization path (q_norm/k_norm) that the base Attention class supports. Added the full q_norm/k_norm handling (both full and per-head modes) between Q/K/V projections and RoPE, matching the base Attention.forward logic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <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
Adds support for BitNet b1.58 models with ternary {-1,0,+1} weights.
Changes
src/mobius/models/bitnet.py: BitNetCausalLMModel with BitNetAttention (attn_sub_norm) and BitNetMLP (ffn_sub_norm)bitnetArchitecture
BitNet b1.58 is Llama-like with these key differences:
Testing
All existing tests pass + 4 new BitNet build graph tests.
Note
Initial implementation uses float Linear (weights unpacked to float during preprocessing). MatMulNBits bits=2 optimization is future work pending ORT kernel support.