diff --git a/docs/_generate_models.py b/docs/_generate_models.py index 8368d47f..a05dbde5 100644 --- a/docs/_generate_models.py +++ b/docs/_generate_models.py @@ -23,8 +23,10 @@ _CATEGORY_DESCRIPTIONS: dict[str, str] = { "Text Generation": "Standard autoregressive language models (CausalLM).", "Mixture of Experts": "Models that route tokens to a subset of expert MLPs.", + "Hybrid Conv+Attention": "Hybrid models with alternating conv and attention layers (LFM2).", "Multimodal": "Models that process images, audio, or other modalities alongside text.", "Speech-to-Text": "Encoder-decoder models for speech recognition.", + "Audio-to-Audio": "End-to-end audio language models: audio in, audio + text out (LFM2-Audio, Moshi).", "Audio": "Audio encoder models for feature extraction (Wav2Vec2, HuBERT, WavLM).", "encoder-only": "Encoder-only models for embeddings and classification (BERT, RoBERTa).", "encoder-decoder": "Encoder-decoder sequence-to-sequence models (BART, T5, mBART).", @@ -39,8 +41,10 @@ _CATEGORY_ORDER = [ "Text Generation", "Mixture of Experts", + "Hybrid Conv+Attention", "Multimodal", "Speech-to-Text", + "Audio-to-Audio", "Audio", "encoder-only", "encoder", diff --git a/docs/model-catalog.md b/docs/model-catalog.md index 7cf3691b..a6ad360d 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -1,6 +1,6 @@ # Model Catalog -**mobius** supports 273 registered model types across 10 categories. +**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. @@ -96,6 +96,22 @@ Mamba and Mamba2 architectures using selective state-space layers. | `falcon_mamba` | `MambaCausalLMModel` | `ssm-text-generation` | `tiiuae/falcon-mamba-7b` | | `mamba2` | `Mamba2CausalLMModel` | `ssm2-text-generation` | `state-spaces/mamba2-2.7b` | +## Hybrid Conv+Attention + +Models with alternating depthwise-conv (ShortConv) and attention layers. + +| Model Type | Module Class | Task | Example HuggingFace Model | +|---|---|---|---| +| `lfm2` | `Lfm2CausalLMModel` | `hybrid-text-generation` | `LiquidAI/LFM2-1.2B` | + +## 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` | + ## Hybrid SSM+Attention Models combining Mamba/SSM layers with transformer attention layers. @@ -273,11 +289,13 @@ MatMulNBits ops. |---|---|---| | Decoder-only LLMs | ~100 | `CausalLMModel`, `GPT2CausalLMModel` | | Mixture of Experts | ~25 | `MoECausalLMModel`, `DeepSeekV3CausalLMModel` | +| Hybrid Conv+Attention | 1 | `Lfm2CausalLMModel` | | SSM / Hybrid | 5 | `MambaCausalLMModel`, `JambaCausalLMModel` | | Vision-Language | ~40 | `LLaVAModel`, `Qwen25VLCausalLMModel` | | Encoder-only | ~40 | `BertModel`, `DistilBertModel` | | Encoder-decoder | ~20 | `BartForConditionalGeneration`, `T5ForConditionalGeneration` | | Speech & Audio | ~20 | `WhisperForConditionalGeneration`, `Wav2Vec2Model` | +| Audio-to-Audio | 1 | `Lfm2AudioModel` | | Vision | ~25 | `ViTModel`, `CLIPVisionModel` | | Diffusion | ~10 | `UNet2DConditionModel`, `FluxTransformer2DModel` | -| **Total** | **~273** | | +| **Total** | **~275** | | diff --git a/examples/moshi_realtime.py b/examples/moshi_realtime.py new file mode 100644 index 00000000..e1798347 --- /dev/null +++ b/examples/moshi_realtime.py @@ -0,0 +1,627 @@ +#!/usr/bin/env python +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Moshi/PersonaPlex real-time full-duplex speech conversation. + +Moshi is a speech-to-speech model that simultaneously listens and speaks. +It encodes incoming audio as RVQ codec tokens, processes them through a +causal transformer backbone, and generates both text tokens (inner +monologue) and audio tokens (spoken response) autoregressively. + +This example demonstrates the full streaming inference loop: + + microphone → RVQ encode → embedding → decoder → audio_decoder + → RVQ decode → speakers + +The pipeline runs in a dual-stream configuration: +- **Input stream**: mic audio → codec encoder → audio_codes input to the model +- **Output stream**: model generates audio_codes → codec decoder → speaker + +Prerequisites:: + + pip install mobius-ai[transformers] sounddevice numpy onnxruntime moshi + +Usage:: + + # Interactive conversation with PersonaPlex-7B + python examples/moshi_realtime.py + + # Use base Moshi model + python examples/moshi_realtime.py --model kyutai/moshiko-pytorch-bf16 + + # Export ONNX models only (no inference) + python examples/moshi_realtime.py --export-only --save-to output/moshi/ + + # Use saved ONNX models + python examples/moshi_realtime.py --onnx-dir output/moshi/ + +Notes: + - Requires a HuggingFace account with accepted nvidia/personaplex-7b-v1 + license to download PersonaPlex weights. + - The model runs at 12.5 Hz (80ms per frame) — one transformer step + produces 1 text token + 16 audio codec tokens per step. + - Audio sample rate: 24 kHz (Moshi's EnCodec codec). + - This example uses a simplified single-stream mode for clarity. + Production use would run listener and speaker in parallel threads. +""" + +from __future__ import annotations + +import argparse +import signal +import sys +import threading +import time +from pathlib import Path + +import numpy as np + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +DEFAULT_MODEL = "nvidia/personaplex-7b-v1" + +# Moshi audio parameters +SAMPLE_RATE = 24_000 # EnCodec sample rate (Hz) +FRAME_SAMPLES = 1920 # 80ms at 24kHz (one model step) +NUM_CODEBOOKS = 16 # RVQ codebook count +AUDIO_VOCAB_SIZE = 2048 # per-codebook vocabulary size +TEXT_VOCAB_SIZE = 32_000 # text token vocabulary + +# Model step rate +STEPS_PER_SECOND = SAMPLE_RATE // FRAME_SAMPLES # 12.5 Hz + + +# --------------------------------------------------------------------------- +# Utilities +# --------------------------------------------------------------------------- + + +def _try_import(name: str, pip_name: str | None = None) -> object: + """Import optional dependency with a helpful error message.""" + import importlib + + try: + return importlib.import_module(name) + except ImportError: + pkg = pip_name or name + print(f"Missing dependency: {name}. Install with: pip install {pkg}", file=sys.stderr) + sys.exit(1) + + +def _make_zero_kv_cache( + num_layers: int, + num_heads: int, + head_dim: int, + max_seq: int = 0, + dtype: np.dtype = np.float32, +) -> list[tuple[np.ndarray, np.ndarray]]: + """Create an empty KV cache (zeros) for ``num_layers`` transformer layers.""" + shape = (1, num_heads, max_seq, head_dim) + return [ + (np.zeros(shape, dtype=dtype), np.zeros(shape, dtype=dtype)) for _ in range(num_layers) + ] + + +# --------------------------------------------------------------------------- +# ONNX model session helpers +# --------------------------------------------------------------------------- + + +class MoshiOnnxPipeline: + """Wrapper around the three Moshi ONNX sub-models. + + Sub-models: + - ``embedding``: (input_ids, audio_codes) → inputs_embeds + - ``decoder``: inputs_embeds → logits + KV cache + - ``audio_decoder``: backbone_hidden → codebook_logits (per step) + + Maintains KV cache state across inference steps for streaming. + """ + + def __init__( + self, + embedding_path: str, + decoder_path: str, + audio_decoder_path: str, + *, + hidden_size: int = 4096, + num_decoder_layers: int = 32, + num_decoder_heads: int = 32, + head_dim: int = 128, + depformer_dim: int = 1024, + depformer_layers: int = 6, + depformer_heads: int = 16, + num_codebooks: int = NUM_CODEBOOKS, + dtype: np.dtype = np.float32, + ): + ort = _try_import("onnxruntime") + + opts = ort.SessionOptions() + opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + + self._embedding = ort.InferenceSession(embedding_path, opts) + self._decoder = ort.InferenceSession(decoder_path, opts) + self._audio_decoder = ort.InferenceSession(audio_decoder_path, opts) + + self._hidden_size = hidden_size + self._depformer_dim = depformer_dim + self._num_codebooks = num_codebooks + self._dtype = dtype + + # Decoder KV cache: num_heads * head_dim per layer + self._decoder_kv = _make_zero_kv_cache( + num_decoder_layers, num_decoder_heads, head_dim, dtype=dtype + ) + # Depformer KV cache: depformer_heads x head_dim=depformer_dim per layer + self._depformer_kv = _make_zero_kv_cache( + depformer_layers, depformer_heads, depformer_dim, dtype=dtype + ) + self._step = 0 + + def reset(self) -> None: + """Reset KV caches (start of new conversation).""" + self._decoder_kv = [(np.zeros_like(k), np.zeros_like(v)) for k, v in self._decoder_kv] + self._depformer_kv = [ + (np.zeros_like(k), np.zeros_like(v)) for k, v in self._depformer_kv + ] + self._step = 0 + + def _build_decoder_feeds( + self, + inputs_embeds: np.ndarray, + position_id: int, + ) -> dict: + """Build feed dict for the decoder with current KV cache.""" + feeds: dict = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, position_id + 1), dtype=np.int64), + "position_ids": np.array([[position_id]], dtype=np.int64), + } + for i, (k, v) in enumerate(self._decoder_kv): + feeds[f"past_key_values.{i}.key"] = k + feeds[f"past_key_values.{i}.value"] = v + return feeds + + def _update_decoder_kv(self, outputs: list) -> None: + """Extract present KV cache from decoder outputs.""" + # outputs[1:] are present key/value pairs interleaved + kv_outputs = outputs[1:] + for i in range(len(self._decoder_kv)): + self._decoder_kv[i] = (kv_outputs[2 * i], kv_outputs[2 * i + 1]) + + def _build_depformer_feeds( + self, + backbone_hidden: np.ndarray, + prev_embedding: np.ndarray, + codebook_idx: int, + ) -> dict: + """Build feed dict for the audio_decoder with current depformer KV cache.""" + feeds: dict = { + "backbone_hidden": backbone_hidden, + "prev_embedding": prev_embedding, + "codebook_idx": np.array(codebook_idx, dtype=np.int64), + } + for i, (k, v) in enumerate(self._depformer_kv): + feeds[f"past_key_values.{i}.key"] = k + feeds[f"past_key_values.{i}.value"] = v + return feeds + + def _update_depformer_kv(self, outputs: list) -> None: + """Extract present KV cache from audio_decoder outputs.""" + kv_outputs = outputs[1:] + for i in range(len(self._depformer_kv)): + self._depformer_kv[i] = (kv_outputs[2 * i], kv_outputs[2 * i + 1]) + + def step( + self, + text_token: int, + audio_codes: np.ndarray, + ) -> tuple[int, np.ndarray]: + """Run one inference step. + + Args: + text_token: int — last generated text token (0 = pad/start) + audio_codes: (num_codebooks,) int64 — incoming audio codec codes + + Returns: + (next_text_token, output_audio_codes) + - next_text_token: int — model's predicted next text token + - output_audio_codes: (num_codebooks,) int64 — audio to synthesize + """ + # 1. Embedding: (1, 1, hidden_size) + emb_feeds = { + "input_ids": np.array([[text_token]], dtype=np.int64), + "audio_codes": audio_codes.reshape(1, 1, self._num_codebooks).astype(np.int64), + } + inputs_embeds = self._embedding.run(None, emb_feeds)[0] # (1, 1, hidden_size) + + # 2. Decoder: logits + updated KV cache + dec_feeds = self._build_decoder_feeds(inputs_embeds, self._step) + dec_outputs = self._decoder.run(None, dec_feeds) + logits = dec_outputs[0] # (1, 1, vocab_size) + self._update_decoder_kv(dec_outputs) + + # 3. Sample next text token (greedy) + next_text_token = int(np.argmax(logits[0, 0])) + + # 4. Audio decoder: generate num_codebooks audio tokens autoregressively + # backbone_hidden = decoder's last hidden state (approximated from logits for demo). + # In production, expose the pre-projection hidden state from the decoder. + backbone_hidden = inputs_embeds # shape: (1, 1, hidden_size) — simplified + + output_codes = np.zeros(self._num_codebooks, dtype=np.int64) + prev_embedding = np.zeros((1, 1, self._depformer_dim), dtype=self._dtype) + + # Reset depformer KV cache at the start of each backbone step + depformer_kv_snapshot = [(k.copy(), v.copy()) for k, v in self._depformer_kv] + self._depformer_kv = _make_zero_kv_cache( + len(self._depformer_kv), + self._depformer_kv[0][0].shape[1], + self._depformer_kv[0][0].shape[3], + dtype=self._dtype, + ) + + for codebook_idx in range(self._num_codebooks): + dep_feeds = self._build_depformer_feeds( + backbone_hidden, prev_embedding, codebook_idx + ) + dep_outputs = self._audio_decoder.run(None, dep_feeds) + codebook_logits = dep_outputs[0] # (1, 1, audio_logits_size) + self._update_depformer_kv(dep_outputs) + + # Greedy sample from codebook logits + output_codes[codebook_idx] = int(np.argmax(codebook_logits[0, 0])) + + # Use sampled code as prev_embedding for next codebook (simplified: zero vec) + prev_embedding = np.zeros((1, 1, self._depformer_dim), dtype=self._dtype) + + # Restore depformer snapshot (we don't accumulate cross-step depformer state) + self._depformer_kv = depformer_kv_snapshot + + self._step += 1 + return next_text_token, output_codes + + +# --------------------------------------------------------------------------- +# Audio codec (EnCodec) placeholder +# --------------------------------------------------------------------------- + + +def encode_audio_frame(audio_frame: np.ndarray, codec) -> np.ndarray: + """Encode a PCM audio frame to RVQ codec codes. + + Args: + audio_frame: (frame_samples,) float32 PCM at SAMPLE_RATE Hz + codec: EnCodec model (from moshi or encodec package) + + Returns: + codes: (num_codebooks,) int64 + """ + import torch + + with torch.no_grad(): + # EnCodec expects (batch, channels, samples) + wav = torch.from_numpy(audio_frame).float().unsqueeze(0).unsqueeze(0) + encoded = codec.encode(wav) # returns EncodedFrame list + # Extract first frame, first batch: shape (num_codebooks, time) + codes = encoded[0][0].squeeze(0).numpy() # (num_codebooks, 1) or (num_codebooks,) + if codes.ndim == 2: + codes = codes[:, 0] + return codes.astype(np.int64) + + +def decode_audio_codes(codes: np.ndarray, codec) -> np.ndarray: + """Decode RVQ codec codes back to PCM waveform. + + Args: + codes: (num_codebooks,) int64 + codec: EnCodec model + + Returns: + audio_frame: (frame_samples,) float32 PCM + """ + import torch + + with torch.no_grad(): + # codes: (num_codebooks,) → (1, num_codebooks, 1) for batch/time dims + codes_t = torch.from_numpy(codes).long().unsqueeze(0).unsqueeze(-1) + wav = codec.decode([(codes_t, None)]) # (1, 1, samples) + return wav.squeeze().numpy() + + +# --------------------------------------------------------------------------- +# Real-time streaming loop +# --------------------------------------------------------------------------- + + +class MoshiStreamer: + """Real-time duplex audio streamer for Moshi. + + Runs two parallel threads: + - **Input thread**: records mic audio → encodes → queues input codes + - **Inference + output thread**: consumes codes → runs model → plays audio + """ + + def __init__( + self, + pipeline: MoshiOnnxPipeline, + codec, + *, + device: str | None = None, + ): + self._pipeline = pipeline + self._codec = codec + self._device = device + + self._input_queue: list[np.ndarray] = [] + self._output_queue: list[np.ndarray] = [] + self._lock = threading.Lock() + self._running = False + self._text_token = 0 # start with pad token + + def _record_callback(self, indata: np.ndarray, frames: int, time_info, status) -> None: + """Sounddevice input callback — encodes incoming mic audio.""" + if status: + print(f"[audio] Input status: {status}", file=sys.stderr) + + audio_frame = indata[:, 0].astype(np.float32) # take first channel + try: + codes = encode_audio_frame(audio_frame, self._codec) + with self._lock: + self._input_queue.append(codes) + except Exception as e: + print(f"[audio] Encode error: {e}", file=sys.stderr) + + def _play_callback(self, outdata: np.ndarray, frames: int, time_info, status) -> None: + """Sounddevice output callback — decodes and plays generated audio.""" + if status: + print(f"[audio] Output status: {status}", file=sys.stderr) + + with self._lock: + if self._output_queue: + audio_frame = self._output_queue.pop(0) + else: + audio_frame = np.zeros(frames, dtype=np.float32) + + outdata[:, 0] = audio_frame[:frames] + + def _inference_loop(self) -> None: + """Main inference loop: consume input codes, run model, enqueue output.""" + print("[moshi] Inference loop started — speak into the microphone.") + + while self._running: + with self._lock: + if not self._input_queue: + codes_in = None + else: + codes_in = self._input_queue.pop(0) + + if codes_in is None: + # No input yet; feed silence (all-zero codes) + codes_in = np.zeros(NUM_CODEBOOKS, dtype=np.int64) + + try: + next_token, output_codes = self._pipeline.step(self._text_token, codes_in) + self._text_token = next_token + + # Decode output codes to waveform + audio_out = decode_audio_codes(output_codes, self._codec) + + with self._lock: + self._output_queue.append(audio_out) + + except Exception as e: + print(f"[moshi] Inference error: {e}", file=sys.stderr) + + # Pace the inference loop to match model step rate (12.5 Hz) + time.sleep(1.0 / STEPS_PER_SECOND) + + def run(self) -> None: + """Start real-time streaming. Blocks until Ctrl+C.""" + sd = _try_import("sounddevice") + + self._running = True + + # Set up Ctrl+C handler for clean shutdown + original_sigint = signal.getsignal(signal.SIGINT) + + def _stop(sig, frame): + print("\n[moshi] Stopping...") + self._running = False + signal.signal(signal.SIGINT, original_sigint) + + signal.signal(signal.SIGINT, _stop) + + # Start inference thread + inference_thread = threading.Thread(target=self._inference_loop, daemon=True) + inference_thread.start() + + # Open audio streams + input_stream = sd.InputStream( + samplerate=SAMPLE_RATE, + channels=1, + dtype="float32", + blocksize=FRAME_SAMPLES, + callback=self._record_callback, + device=self._device, + ) + output_stream = sd.OutputStream( + samplerate=SAMPLE_RATE, + channels=1, + dtype="float32", + blocksize=FRAME_SAMPLES, + callback=self._play_callback, + device=self._device, + ) + + print(f"[moshi] Streaming at {SAMPLE_RATE} Hz, {FRAME_SAMPLES} samples/frame") + print("[moshi] Ctrl+C to stop") + + with input_stream, output_stream: + while self._running: + time.sleep(0.1) + + inference_thread.join(timeout=2.0) + print("[moshi] Stopped.") + + +# --------------------------------------------------------------------------- +# Model building and export +# --------------------------------------------------------------------------- + + +def build_moshi_models(model_id: str, save_dir: Path | None = None) -> dict: + """Build Moshi ONNX models via mobius and optionally save to disk. + + Returns a dict of ``{name: onnx_model}`` for embedding, decoder, + and audio_decoder. + """ + from mobius import build + + print(f"[moshi] Building ONNX models from {model_id} ...") + pkg = build(model_id) + + if save_dir is not None: + import onnx + + save_dir.mkdir(parents=True, exist_ok=True) + for name, model in pkg.items(): + out_path = save_dir / f"{name}.onnx" + onnx.save(model, str(out_path)) + print(f"[moshi] Saved {name} → {out_path}") + + return pkg + + +def load_saved_models(onnx_dir: Path) -> dict[str, str]: + """Return paths to saved ONNX models in onnx_dir.""" + paths = {} + for name in ("embedding", "decoder", "audio_decoder"): + path = onnx_dir / f"{name}.onnx" + if not path.exists(): + print(f"[moshi] Missing model file: {path}", file=sys.stderr) + sys.exit(1) + paths[name] = str(path) + return paths + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Moshi/PersonaPlex real-time speech conversation via ONNX" + ) + parser.add_argument( + "--model", + default=DEFAULT_MODEL, + help=f"HuggingFace model ID (default: {DEFAULT_MODEL})", + ) + parser.add_argument( + "--export-only", + action="store_true", + help="Build and export ONNX models without running inference", + ) + parser.add_argument( + "--save-to", + type=Path, + default=None, + metavar="DIR", + help="Save exported ONNX models to this directory", + ) + parser.add_argument( + "--onnx-dir", + type=Path, + default=None, + metavar="DIR", + help="Load pre-exported ONNX models from this directory", + ) + parser.add_argument( + "--device", + default=None, + help="sounddevice device name or index (default: system default)", + ) + args = parser.parse_args() + + # ------------------------------------------------------------------ + # Step 1: Obtain ONNX model paths + # ------------------------------------------------------------------ + if args.onnx_dir is not None: + print(f"[moshi] Loading ONNX models from {args.onnx_dir}") + model_paths = load_saved_models(args.onnx_dir) + else: + import tempfile + + onnx_models = build_moshi_models(args.model, save_dir=args.save_to) + + if args.export_only: + if args.save_to is None: + print("[moshi] --export-only requires --save-to ") + sys.exit(1) + print("[moshi] Export complete.") + return + + # Save to a temp dir so we can pass file paths to ORT + _tmp = tempfile.mkdtemp(prefix="moshi_onnx_") + tmp_dir = Path(_tmp) + import onnx + + model_paths = {} + for name, model in onnx_models.items(): + out = tmp_dir / f"{name}.onnx" + onnx.save(model, str(out)) + model_paths[name] = str(out) + print(f"[moshi] Temporary ONNX models saved to {tmp_dir}") + + # ------------------------------------------------------------------ + # Step 2: Load EnCodec / Moshi codec for audio encode/decode + # ------------------------------------------------------------------ + try: + import moshi.models # type: ignore[import] + + print("[moshi] Loading EnCodec codec ...") + codec = moshi.models.get_encodec(sample_rate=SAMPLE_RATE) + codec.eval() + except ImportError: + print( + "[moshi] 'moshi' package not found. Install with: pip install moshi\n" + " Falling back to silence-only demo (no audio encode/decode).", + file=sys.stderr, + ) + codec = None + + # ------------------------------------------------------------------ + # Step 3: Build inference pipeline + # ------------------------------------------------------------------ + pipeline = MoshiOnnxPipeline( + embedding_path=model_paths["embedding"], + decoder_path=model_paths["decoder"], + audio_decoder_path=model_paths["audio_decoder"], + ) + + if codec is None: + # Dry-run demo: simulate one step without real audio + print("[moshi] Running dry-run step (no codec, no audio device) ...") + dummy_codes = np.zeros(NUM_CODEBOOKS, dtype=np.int64) + text_token, out_codes = pipeline.step(0, dummy_codes) + print( + f"[moshi] Dry-run OK — text_token={text_token}, out_codes shape={out_codes.shape}" + ) + return + + # ------------------------------------------------------------------ + # Step 4: Start real-time streaming + # ------------------------------------------------------------------ + _try_import("sounddevice") # ensure sounddevice is available before starting + + streamer = MoshiStreamer(pipeline, codec, device=args.device) + streamer.run() + + +if __name__ == "__main__": + main() diff --git a/src/mobius/_configs.py b/src/mobius/_configs.py index a3626af9..6bad0f87 100644 --- a/src/mobius/_configs.py +++ b/src/mobius/_configs.py @@ -870,6 +870,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: in ( "gemma3_text", "flex_olmo", + "lfm2", "olmoe", "olmo2", "olmo3", @@ -1843,6 +1844,145 @@ def from_transformers(cls, config, parent_config=None) -> NemotronHConfig: ) +@dataclasses.dataclass +class Lfm2Config(ArchitectureConfig): + """Configuration for LFM2 hybrid ShortConv+Attention models. + + LFM2 interleaves ShortConv (gated causal depthwise Conv1d) layers + with standard attention layers. Both layer types include an MLP + (SiLU-gated feed-forward). + + ShortConv state per layer: conv_state (batch, hidden_size, kernel-1) + Attention state per layer: standard KV cache (key + value) + + Layer types are specified via ``layer_types`` in the HF config: + ``"conv"`` = ShortConv layer + ``"full_attention"`` = standard attention layer + """ + + # ShortConv parameters + short_conv_kernel: int = 3 + short_conv_bias: bool = False + + @classmethod + def from_transformers(cls, config, parent_config=None) -> Lfm2Config: + base = ArchitectureConfig.from_transformers(config, parent_config) + + # LFM2 MLP: adjusted intermediate size with SwiGLU + intermediate = base.intermediate_size + if getattr(config, "block_auto_adjust_ff_dim", False): + intermediate = int(2 * intermediate / 3) + multiplier = getattr(config, "block_ffn_dim_multiplier", None) + if multiplier is not None: + intermediate = int(multiplier * intermediate) + multiple_of = getattr(config, "block_multiple_of", 256) + intermediate = multiple_of * ((intermediate + multiple_of - 1) // multiple_of) + + base_fields = { + k: v for k, v in _shallow_fields(base).items() if k not in ("intermediate_size",) + } + return cls( + **base_fields, + intermediate_size=intermediate, + short_conv_kernel=getattr(config, "conv_L_cache", 3), + short_conv_bias=getattr(config, "conv_bias", False), + ) + + +@dataclasses.dataclass +class Lfm2AudioConfig(Lfm2Config): + """Configuration for LFM2-Audio: audio-to-audio with hybrid backbone. + + Extends Lfm2Config with depthformer and audio codec fields. + """ + + # Depthformer parameters + depthformer_layers: int = 6 + depthformer_dim: int = 1024 + depthformer_heads: int = 16 + depthformer_tie: bool = True + + # Audio codec parameters + num_codebooks: int = 8 + audio_vocab_size: int = 2049 + + @classmethod + def from_transformers(cls, config, parent_config=None) -> Lfm2AudioConfig: + # LFM2-Audio uses a custom config format from liquid-audio, + # not a standard HuggingFace config. The base fields come from + # the nested 'lfm' sub-config. + base = Lfm2Config.from_transformers(config, parent_config) + base_fields = _shallow_fields(base) + + depthformer = getattr(config, "depthformer", None) or {} + if hasattr(depthformer, "__dict__"): + depthformer = depthformer.__dict__ + + return cls( + **base_fields, + depthformer_layers=depthformer.get("layers", 6), + depthformer_dim=depthformer.get("dim", 1024), + depthformer_tie=depthformer.get("tie", True), + num_codebooks=getattr(config, "codebooks", 8), + audio_vocab_size=getattr(config, "audio_vocab_size", 2049), + ) + + +@dataclasses.dataclass +class MoshiConfig(ArchitectureConfig): + """Configuration for Moshi/PersonaPlex audio-to-audio models. + + Moshi (and its fine-tune PersonaPlex) is a full-duplex speech model. + It combines a standard causal transformer backbone with a depth + transformer ("depformer") that generates audio codec tokens. + + The depformer uses one attention head per codebook (``num_codebooks`` + heads, each with ``depformer_dim`` head_dim), processing a single + codebook token per step with a KV cache over previous codebooks. + + Per-codebook gating MLPs are stored as stacked parameters and selected + at runtime using ``codebook_idx``. + + Reference: Défossez et al., "Moshi: a speech-text foundation model + for real-time dialogue" (2024), ``kyutai/moshiko-pytorch-bf16``. + """ + + # Depformer parameters + depformer_dim: int = 1024 + depformer_layers: int = 6 + depformer_num_heads: int = 16 # must equal num_codebooks + depformer_intermediate_size: int = 2816 + + # Audio codec parameters + num_codebooks: int = 16 + audio_vocab_size: int = 2049 # 2048 codebook entries + 1 padding + + @classmethod + def from_transformers(cls, config, parent_config=None) -> MoshiConfig: + # PersonaPlex config.json only has {"model_type": "personaplex", "version": "7b-v1"}. + # All hyperparameters are derived from weight shapes; hardcode v1 defaults here. + return cls( + model_type=getattr(config, "model_type", "personaplex"), + hidden_size=getattr(config, "hidden_size", 4096), + num_hidden_layers=getattr(config, "num_hidden_layers", 32), + num_attention_heads=getattr(config, "num_attention_heads", 32), + num_key_value_heads=getattr(config, "num_key_value_heads", 32), + head_dim=getattr(config, "head_dim", 128), + intermediate_size=getattr(config, "intermediate_size", 11264), + vocab_size=getattr(config, "vocab_size", 32000), + hidden_act=getattr(config, "hidden_act", "silu"), + rms_norm_eps=getattr(config, "rms_norm_eps", 1e-5), + rope_theta=getattr(config, "rope_theta", 10000.0), + max_position_embeddings=getattr(config, "max_position_embeddings", 4096), + depformer_dim=getattr(config, "depformer_dim", 1024), + depformer_layers=getattr(config, "depformer_layers", 6), + depformer_num_heads=getattr(config, "depformer_num_heads", 16), + depformer_intermediate_size=getattr(config, "depformer_intermediate_size", 2816), + num_codebooks=getattr(config, "num_codebooks", 16), + audio_vocab_size=getattr(config, "audio_vocab_size", 2049), + ) + + @dataclasses.dataclass class JetMoeConfig(CausalLMConfig): """Configuration for JetMoE: Mixture-of-Attention + MoE FFN model. diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index b5e512f8..b08ec61d 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -29,6 +29,7 @@ from mobius.models import ( ApertusCausalLMModel, ArceeCausalLMModel, + BitNetCausalLMModel, CausalLMModel, ChatGLMCausalLMModel, DeepSeekOCR2CausalLMModel, @@ -391,6 +392,7 @@ def _create_default_registry() -> ModelRegistry: "gemma3n_text": Gemma3nCausalLMModel, "granite": GraniteCausalLMModel, "diffllama": DiffLlamaCausalLMModel, + "bitnet": BitNetCausalLMModel, "doge": DogeCausalLMModel, "internlm2": InternLM2CausalLMModel, "llama4_text": Llama4CausalLMModel, @@ -471,6 +473,21 @@ def _create_default_registry() -> ModelRegistry: reg.register("nemotron_h", NemotronHCausalLMModel) + # --- Hybrid Conv+Attention (LFM2) --- + from mobius.models.lfm2 import Lfm2CausalLMModel + + reg.register("lfm2", Lfm2CausalLMModel) + + from mobius.models.lfm2_audio import Lfm2AudioModel + + reg.register("lfm2_audio", Lfm2AudioModel) + + # --- Moshi / PersonaPlex (audio-to-audio) --- + from mobius.models.moshi import MoshiModel + + reg.register("personaplex", MoshiModel) + reg.register("moshi", MoshiModel) + # --- Multimodal --- for name in ( "chameleon", @@ -728,6 +745,7 @@ def _create_default_registry() -> ModelRegistry: "apertus": "swiss-ai/Apertus-8B-Instruct-2509", "arcee": "arcee-ai/AFM-4.5B-Base", "diffllama": "kajuma/DiffLlama-0.3B-handcut", + "bitnet": "microsoft/bitnet-b1.58-2B-4T", "doge": "SmallDoge/Doge-20M", "dots1": "rednote-hilab/dots.llm1.inst", "exaone4": "LGAI-EXAONE/EXAONE-4.0-1.2B", @@ -820,6 +838,10 @@ def _create_default_registry() -> ModelRegistry: # --- Hybrid SSM+Attention --- "jamba": "ai21labs/Jamba-v0.1", "bamba": "ibm-fms/Bamba-9B", + "lfm2": "LiquidAI/LFM2-1.2B", + "lfm2_audio": "LiquidAI/LFM2-Audio-1.5B", + "personaplex": "nvidia/personaplex-7b-v1", + "moshi": "kyutai/moshiko-pytorch-bf16", # --- Multimodal --- "qwen2_vl": "Qwen/Qwen2-VL-2B-Instruct", @@ -1078,6 +1100,10 @@ def _create_default_registry() -> ModelRegistry: "falcon_mamba": "ssm", "jamba": "hybrid-ssm+attn", "bamba": "hybrid-mamba2+attn", + "lfm2": "hybrid-conv+attn", + "lfm2_audio": "audio-to-audio", + "personaplex": "audio-to-audio", + "moshi": "audio-to-audio", "qwen3_next": "moe+linear-attn", } diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 4b90bbac..277b0065 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -79,6 +79,7 @@ "Qwen3VLVisionRotaryEmbedding", "RMSNorm", "SelectiveScan", + "ShortConv", "SiLU", "SigmoidTopKGate", "SnakeBeta", @@ -225,6 +226,7 @@ apply_rms_norm, ) from mobius.components._rotary_embedding import initialize_rope +from mobius.components._short_conv import ShortConv from mobius.components._ssm import ( JambaSelectiveScan, Mamba2Scan, diff --git a/src/mobius/components/_short_conv.py b/src/mobius/components/_short_conv.py new file mode 100644 index 00000000..2f3f99b1 --- /dev/null +++ b/src/mobius/components/_short_conv.py @@ -0,0 +1,177 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""ShortConv: gated causal depthwise Conv1d for LFM2 conv layers. + +Implements the ``Lfm2ShortConv`` layer from HuggingFace transformers: + +1. ``in_proj(x)`` -> split into B, C, x chunks (3 x hidden_size) +2. ``B * x`` (element-wise gating) +3. Causal depthwise Conv1d on ``Bx`` +4. ``C * conv_out`` (output gating) +5. ``out_proj(y)`` + +The convolution is depthwise (groups=hidden_size) with causal padding +(left-pad by kernel_size-1). During inference, the ``conv_state`` +(B, hidden_size, kernel_size-1) is carried across steps. + +HuggingFace weight names:: + + conv.conv.weight → self.conv_weight (hidden_size, 1, kernel_size) + conv.in_proj.weight → self.in_proj.weight (3*hidden_size, hidden_size) + conv.out_proj.weight → self.out_proj.weight (hidden_size, hidden_size) + +HuggingFace reference: ``Lfm2ShortConv`` in ``modeling_lfm2.py``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from onnxscript import nn +from onnxscript._internal import builder + +from mobius.components._common import INT64_MAX, Linear + +if TYPE_CHECKING: + import onnx_ir as ir + + +class ShortConv(nn.Module): + """Gated causal depthwise Conv1d block (LFM2 ShortConv). + + Data flow:: + + x → in_proj → [B, C, x] + ↓ + B * x = Bx + ↓ + causal depthwise Conv1d(Bx) + ↓ + C * conv_out + ↓ + out_proj → y + + The convolution uses ``groups=hidden_size`` (depthwise) and left-pads + by ``kernel_size - 1`` for causal behavior during prefill. During + generation (single-step), the cached ``conv_state`` (B, hidden_size, + kernel_size - 1) replaces left-padding. + + Args: + hidden_size: Model hidden dimension. + kernel_size: Convolution kernel size (typically 3-4). + bias: Whether conv/proj layers include bias. + """ + + def __init__( + self, + hidden_size: int, + kernel_size: int, + bias: bool = False, + ): + super().__init__() + self._hidden_size = hidden_size + self._kernel_size = kernel_size + self._bias = bias + + # in_proj: hidden_size → 3 * hidden_size (B, C, x) + self.in_proj = Linear(hidden_size, 3 * hidden_size, bias=bias) + # out_proj: hidden_size → hidden_size + self.out_proj = Linear(hidden_size, hidden_size, bias=bias) + + # Depthwise conv weight: (hidden_size, 1, kernel_size) + # Stored as a plain parameter (not a sub-module) to match HF naming + self.conv_weight = nn.Parameter( + shape=[hidden_size, 1, kernel_size], + ) + if bias: + self.conv_bias = nn.Parameter(shape=[hidden_size]) + else: + self.conv_bias = None + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + conv_state: ir.Value | None = None, + ) -> tuple[ir.Value, ir.Value]: + """Forward pass. + + Args: + hidden_states: (B, S, hidden_size) input tensor. + conv_state: (B, hidden_size, kernel_size-1) past conv state, + or None for first step / prefill. + + Returns: + (output, new_conv_state): + output: (B, S, hidden_size) + new_conv_state: (B, hidden_size, kernel_size-1) + """ + # in_proj → (B, S, 3*hidden_size) → transpose to (B, 3*hidden_size, S) + projected = self.in_proj(op, hidden_states) + # Transpose to channels-first: (B, 3*H, S) + projected = op.Transpose(projected, perm=[0, 2, 1]) + + # Split into B, C, x along dim=1: each (B, hidden_size, S) + h = self._hidden_size + b_gate = op.Slice( + projected, + op.Constant(value_ints=[0]), + op.Constant(value_ints=[h]), + op.Constant(value_ints=[1]), + ) + c_gate = op.Slice( + projected, + op.Constant(value_ints=[h]), + op.Constant(value_ints=[2 * h]), + op.Constant(value_ints=[1]), + ) + x = op.Slice( + projected, + op.Constant(value_ints=[2 * h]), + op.Constant(value_ints=[3 * h]), + op.Constant(value_ints=[1]), + ) + + # Bx = B * x (element-wise gating) + bx = op.Mul(b_gate, x) # (B, hidden_size, S) + + # Causal depthwise Conv1d on Bx + # Left-pad by (kernel_size - 1) for causal convolution + pad_left = self._kernel_size - 1 + if conv_state is not None: + # Inference: prepend cached state (replaces left-padding) + # conv_state: (B, hidden_size, kernel_size-1) + bx_padded = op.Concat(conv_state, bx, axis=2) + else: + # Prefill or no cache: explicit left-padding + pads = op.Constant(value_ints=[0, 0, pad_left, 0, 0, 0]) + bx_padded = op.Pad(bx, pads, mode="constant") + + # Extract new conv_state: last (kernel_size - 1) timesteps of bx_padded + # bx_padded shape: (B, hidden_size, S + kernel_size - 1) + new_conv_state = op.Slice( + bx_padded, + op.Constant(value_ints=[-(self._kernel_size - 1)]), + op.Constant(value_ints=[INT64_MAX]), + op.Constant(value_ints=[2]), + ) + + # Depthwise Conv1d: groups=hidden_size + conv_inputs = [bx_padded, self.conv_weight] + if self.conv_bias is not None: + conv_inputs.append(self.conv_bias) + + conv_out = op.Conv( + *conv_inputs, + group=self._hidden_size, + ) + + # Output gating: y = C * conv_out + y = op.Mul(c_gate, conv_out) # (B, hidden_size, S) + + # Transpose back to (B, S, hidden_size) for out_proj + y = op.Transpose(y, perm=[0, 2, 1]) + y = self.out_proj(op, y) + + return y, new_conv_state diff --git a/src/mobius/components/_short_conv_test.py b/src/mobius/components/_short_conv_test.py new file mode 100644 index 00000000..8be9d94b --- /dev/null +++ b/src/mobius/components/_short_conv_test.py @@ -0,0 +1,217 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ShortConv: gated causal depthwise Conv1d (LFM2 conv layers).""" + +from __future__ import annotations + +from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius.components._short_conv import ShortConv + +_HIDDEN = 32 +_KERNEL = 3 +_BATCH = 2 +_SEQ = 5 + + +class TestShortConvParameters: + """Verify parameter shapes match HuggingFace Lfm2ShortConv.""" + + def test_in_proj_shape(self): + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + assert list(conv.in_proj.weight.shape) == [3 * _HIDDEN, _HIDDEN] + + def test_out_proj_shape(self): + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + assert list(conv.out_proj.weight.shape) == [_HIDDEN, _HIDDEN] + + def test_conv_weight_shape(self): + """Depthwise conv: (hidden_size, 1, kernel_size) matches Conv1d(groups=hidden).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + assert list(conv.conv_weight.shape) == [_HIDDEN, 1, _KERNEL] + + def test_no_bias_by_default(self): + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL, bias=False) + assert conv.conv_bias is None + assert conv.in_proj.bias is None + assert conv.out_proj.bias is None + + def test_bias_creates_parameters(self): + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL, bias=True) + assert conv.conv_bias is not None + assert list(conv.conv_bias.shape) == [_HIDDEN] + assert conv.in_proj.bias is not None + assert conv.out_proj.bias is not None + + +class TestShortConvPrefill: + """Prefill path: conv_state=None, uses left-padding.""" + + def test_output_shape(self): + """Output should be (B, S, hidden_size).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert out.shape is not None + # (B, S, hidden_size) + assert list(out.shape) == [_BATCH, _SEQ, _HIDDEN] + + def test_conv_state_shape(self): + """new_conv_state should be (B, hidden_size, kernel_size-1).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + # kernel_size - 1 = 2 for kernel=3 + assert new_state.shape is not None + assert list(new_state.shape) == [_BATCH, _HIDDEN, _KERNEL - 1] + + def test_pad_op_used_for_causal(self): + """Prefill path pads left by kernel_size-1 using ONNX Pad.""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert count_op_type(graph, "Pad") >= 1 + + def test_conv_op_present(self): + """Graph must contain ONNX Conv node (depthwise).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert count_op_type(graph, "Conv") >= 1 + + def test_single_step_prefill(self): + """Single-token prefill (S=1) still works: state shape (B, H, K-1).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, 1, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert list(out.shape) == [_BATCH, 1, _HIDDEN] + assert list(new_state.shape) == [_BATCH, _HIDDEN, _KERNEL - 1] + + +class TestShortConvIncremental: + """Incremental (decode) path: conv_state provided, uses Concat instead of Pad.""" + + def test_output_shape_decode_step(self): + """Single decode step: output (B, 1, hidden_size).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, 1, _HIDDEN]) + state = create_test_input(builder, "conv_state", [_BATCH, _HIDDEN, _KERNEL - 1]) + + out, new_state = conv(op, x, conv_state=state) + graph.outputs.extend([out, new_state]) + + assert list(out.shape) == [_BATCH, 1, _HIDDEN] + + def test_conv_state_shape_decode(self): + """new_conv_state stays (B, hidden_size, kernel_size-1).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, 1, _HIDDEN]) + state = create_test_input(builder, "conv_state", [_BATCH, _HIDDEN, _KERNEL - 1]) + + out, new_state = conv(op, x, conv_state=state) + graph.outputs.extend([out, new_state]) + + assert list(new_state.shape) == [_BATCH, _HIDDEN, _KERNEL - 1] + + def test_concat_used_not_pad(self): + """Incremental path uses Concat (not Pad) to prepend conv_state.""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, 1, _HIDDEN]) + state = create_test_input(builder, "conv_state", [_BATCH, _HIDDEN, _KERNEL - 1]) + + out, new_state = conv(op, x, conv_state=state) + graph.outputs.extend([out, new_state]) + + assert count_op_type(graph, "Concat") >= 1 + assert count_op_type(graph, "Pad") == 0 + + def test_multi_step_decode(self): + """Multi-step decode (S>1 with conv_state) works correctly.""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, 3, _HIDDEN]) + state = create_test_input(builder, "conv_state", [_BATCH, _HIDDEN, _KERNEL - 1]) + + out, new_state = conv(op, x, conv_state=state) + graph.outputs.extend([out, new_state]) + + assert list(out.shape) == [_BATCH, 3, _HIDDEN] + assert list(new_state.shape) == [_BATCH, _HIDDEN, _KERNEL - 1] + + +class TestShortConvGating: + """Verify the B/C/x gating structure is present in the graph.""" + + def test_mul_ops_for_gating(self): + """Graph must have at least 2 Mul nodes: B*x and C*conv_out.""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert count_op_type(graph, "Mul") >= 2 + + def test_in_proj_splits_into_three(self): + """in_proj expands H → 3H; verified by Slice ops (one per chunk B/C/x).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=_KERNEL) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + # Three Slice ops to extract B, C, x from in_proj output + assert count_op_type(graph, "Slice") >= 3 + + +class TestShortConvKernelSizes: + """Verify correctness with different kernel sizes.""" + + def test_kernel_4(self): + """kernel_size=4: conv_state shape (B, H, 3).""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=4) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert list(out.shape) == [_BATCH, _SEQ, _HIDDEN] + assert list(new_state.shape) == [_BATCH, _HIDDEN, 3] + + def test_kernel_2(self): + """kernel_size=2: conv_state shape (B, H, 1) — minimum rolling window.""" + conv = ShortConv(hidden_size=_HIDDEN, kernel_size=2) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [_BATCH, _SEQ, _HIDDEN]) + + out, new_state = conv(op, x, conv_state=None) + graph.outputs.extend([out, new_state]) + + assert list(out.shape) == [_BATCH, _SEQ, _HIDDEN] + assert list(new_state.shape) == [_BATCH, _HIDDEN, 1] diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 3f157185..b943adc7 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -12,6 +12,7 @@ "BartForConditionalGeneration", "BertModel", "Blip2Model", + "BitNetCausalLMModel", "BloomCausalLMModel", "CLIPVisionModel", "CTRLCausalLMModel", @@ -58,12 +59,15 @@ "Llama4CausalLMModel", "LLaVAModel", "LayerNormCausalLMModel", + "Lfm2AudioModel", + "Lfm2CausalLMModel", "LongcatFlashCausalLMModel", "MPTCausalLMModel", "Mamba2CausalLMModel", "MambaCausalLMModel", "MiniMaxCausalLMModel", "MoECausalLMModel", + "MoshiModel", "NanoChatCausalLMModel", "NemotronCausalLMModel", "NemotronHCausalLMModel", @@ -127,6 +131,7 @@ from mobius.models.bart import BartForConditionalGeneration from mobius.models.base import CausalLMModel, LayerNormCausalLMModel from mobius.models.bert import BertModel +from mobius.models.bitnet import BitNetCausalLMModel from mobius.models.blip2 import Blip2Model from mobius.models.chatglm import ChatGLMCausalLMModel from mobius.models.clip import CLIPVisionModel @@ -161,6 +166,8 @@ from mobius.models.internvl import InternVL2Model from mobius.models.jamba import JambaCausalLMModel from mobius.models.jetmoe import JetMoeCausalLMModel +from mobius.models.lfm2 import Lfm2CausalLMModel +from mobius.models.lfm2_audio import Lfm2AudioModel from mobius.models.llama4 import Llama4CausalLMModel from mobius.models.llava import LLaVAModel from mobius.models.longcat_flash import LongcatFlashCausalLMModel @@ -174,6 +181,7 @@ Phi3MoECausalLMModel, Qwen2MoECausalLMModel, ) +from mobius.models.moshi import MoshiModel from mobius.models.nanochat import NanoChatCausalLMModel from mobius.models.nemotron import NemotronCausalLMModel from mobius.models.nemotron_h import NemotronHCausalLMModel diff --git a/src/mobius/models/bitnet.py b/src/mobius/models/bitnet.py new file mode 100644 index 00000000..8f96ca17 --- /dev/null +++ b/src/mobius/models/bitnet.py @@ -0,0 +1,228 @@ +"""BitNet b1.58 — ternary-weight language model. + +BitNet uses {-1, 0, +1} weights (1.58 bits) with per-tensor ``weight_scale``. +The architecture is Llama-like with two key differences: + +1. **Sub-layer norms**: ``attn_sub_norm`` applied to attention output before + ``o_proj``, and ``ffn_sub_norm`` applied to gated activation before + ``down_proj``. +2. **Activation**: squared ReLU (``relu2``) instead of SiLU. + +HuggingFace stores weights as uint8-packed 2-bit ternary values. This module +unpacks them to float in ``preprocess_weights`` so the ONNX graph uses +standard ``Linear`` ops. A future optimisation could use ``MatMulNBits`` +with ``bits=2`` or a custom ternary kernel. + +Replicates HuggingFace's ``BitNetForCausalLM``. +""" + +from __future__ import annotations + +import onnx_ir as ir +import torch +from onnxscript._internal import builder + +from mobius._configs import ArchitectureConfig +from mobius.components._attention import ( + Attention, + StaticCacheState, + _apply_attention, + apply_rotary_pos_emb, +) +from mobius.components._mlp import MLP +from mobius.components._rms_norm import RMSNorm +from mobius.models.base import CausalLMModel + + +class BitNetAttention(Attention): + """Attention with a sub-layer RMSNorm before the output projection. + + HuggingFace name: ``BitNetAttention`` — identical to standard + multi-head attention except ``attn_sub_norm`` is applied to the + concatenated attention heads *before* ``o_proj``. + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + # Sub-layer norm applied to attention output before o_proj + self.attn_sub_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value | None, + position_embeddings: tuple | None = None, + past_key_value: tuple | None = None, + static_cache: StaticCacheState | None = None, + ): + # 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) + + # Optional Q/K normalization (inherited from base Attention) + if self.q_norm is not None and self.k_norm is not None: + if self._qk_norm_full: + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + else: + query_states = op.Reshape(query_states, [0, 0, -1, self.head_dim]) + key_states = op.Reshape(key_states, [0, 0, -1, self.head_dim]) + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + query_states = op.Reshape(query_states, [0, 0, -1]) + key_states = op.Reshape(key_states, [0, 0, -1]) + + # RoPE + if position_embeddings is not None: + query_states = apply_rotary_pos_emb( + op, + query_states, + position_embeddings=position_embeddings, + num_heads=self.num_attention_heads, + rotary_embedding_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + key_states = apply_rotary_pos_emb( + op, + key_states, + position_embeddings=position_embeddings, + num_heads=self.num_key_value_heads, + rotary_embedding_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + + # Grouped-query attention — (batch, seq_len, num_heads * head_dim) + attn_output, present_key, present_value = _apply_attention( + op, + query_states, + key_states, + value_states, + attention_bias, + past_key_value[0] if past_key_value is not None else None, + past_key_value[1] if past_key_value is not None else None, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + scale=self.scaling, + static_cache=static_cache, + ) + + # BitNet-specific: sub-layer norm before o_proj + attn_output = self.attn_sub_norm(op, attn_output) + attn_output = self.o_proj(op, attn_output) + return attn_output, (present_key, present_value) + + +class BitNetMLP(MLP): + """Gated MLP with a sub-layer RMSNorm before the down projection. + + HuggingFace name: ``BitNetMLP`` — same as standard gated MLP except + ``ffn_sub_norm`` (over ``intermediate_size``) is applied between the + gated activation and ``down_proj``. + """ + + def __init__( + self, + config: ArchitectureConfig, + linear_class: type | None = None, + ): + super().__init__(config, linear_class=linear_class) + # Sub-layer norm applied after gated activation, before down_proj + self.ffn_sub_norm = RMSNorm(config.intermediate_size, eps=config.rms_norm_eps) + + def forward(self, op: builder.OpBuilder, x: ir.Value): + # gate_proj → relu² → element-wise multiply with up_proj + gate = self.act_fn(op, self.gate_proj(op, x)) + up = self.up_proj(op, x) + hidden = op.Mul(gate, up) + # BitNet-specific: sub-layer norm before down_proj + hidden = self.ffn_sub_norm(op, hidden) + return self.down_proj(op, hidden) + + +class BitNetCausalLMModel(CausalLMModel): + """BitNet b1.58 causal language model. + + Builds on the standard CausalLMModel (Llama-like) and replaces each + decoder layer's Attention and MLP with BitNet variants that include + sub-layer RMSNorms. Weight preprocessing unpacks HuggingFace's + uint8-packed ternary weights to float32. + + Replicates HuggingFace's ``BitNetForCausalLM``. + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + # Replace standard Attention/MLP with BitNet variants + for layer in self.model.layers: + layer.self_attn = BitNetAttention(config) + layer.mlp = BitNetMLP(config) + + 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 + + # Delegate remaining processing (tie_word_embeddings, etc.) + return super().preprocess_weights(new_state_dict) + + +def _unpack_ternary_weights(packed: torch.Tensor) -> torch.Tensor: + """Unpack uint8-packed ternary weights to float tensor. + + Packing format (HuggingFace BitNet): + * Ternary values {-1, 0, +1} are mapped to unsigned {0, 1, 2} (add 1). + * Four 2-bit values are packed per byte: ``v0 | (v1<<2) | (v2<<4) | (v3<<6)``. + * Packed shape: ``[out_features // 4, in_features]``. + * Unpacked shape: ``[out_features, in_features]``. + + Args: + packed: uint8 tensor of shape ``[out_features // 4, in_features]``. + + Returns: + Float tensor of shape ``[out_features, in_features]`` with values + in {-1, 0, +1}. + """ + if packed.ndim < 2: + raise ValueError(f"Expected packed tensor with ≥2 dims, got shape {packed.shape}") + # Extract 4 two-bit values per byte + v0 = (packed >> 0) & 0x03 # bits 0-1 + v1 = (packed >> 2) & 0x03 # bits 2-3 + v2 = (packed >> 4) & 0x03 # bits 4-5 + v3 = (packed >> 6) & 0x03 # bits 6-7 + + # Interleave: stack along new dim then reshape to expand out_features + # packed shape: [out_features // 4, in_features] + # stacked shape: [out_features // 4, 4, in_features] + # result shape: [out_features, in_features] + stacked = torch.stack([v0, v1, v2, v3], dim=-2) + unpacked = stacked.reshape(-1, packed.shape[-1]) + + # Map unsigned {0, 1, 2} back to signed {-1, 0, +1} + return unpacked.float() - 1.0 diff --git a/src/mobius/models/bitnet_test.py b/src/mobius/models/bitnet_test.py new file mode 100644 index 00000000..86841818 --- /dev/null +++ b/src/mobius/models/bitnet_test.py @@ -0,0 +1,93 @@ +"""Unit tests for BitNet ternary weight unpacking.""" + +from __future__ import annotations + +import torch + +from mobius.models.bitnet import _unpack_ternary_weights + + +class TestUnpackTernaryWeights: + """Tests for the _unpack_ternary_weights helper.""" + + def test_output_shape(self): + """Packed [out//4, in] unpacks to [out, in].""" + packed = torch.zeros(8, 16, dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert result.shape == (32, 16) + + def test_values_in_ternary_set(self): + """All output values must be in {-1, 0, +1} for valid packed data.""" + # Build valid packed bytes: only 2-bit values 0, 1, 2 (not 3) + rng = torch.Generator().manual_seed(42) + values = torch.randint(0, 3, (4, 8, 4), generator=rng, dtype=torch.uint8) + packed = ( + values[..., 0] + | (values[..., 1] << 2) + | (values[..., 2] << 4) + | (values[..., 3] << 6) + ) + result = _unpack_ternary_weights(packed) + unique = set(result.unique().tolist()) + assert unique.issubset({-1.0, 0.0, 1.0}) + + def test_known_byte_round_trip(self): + """Verify bit extraction for a known byte. + + Packing: ternary {-1, 0, +1} → unsigned {0, 1, 2} + Byte = v0 | (v1 << 2) | (v2 << 4) | (v3 << 6) + + Example: values [-1, 0, +1, -1] → unsigned [0, 1, 2, 0] + Byte = 0b_00_10_01_00 = 0x24 = 36 + """ + packed = torch.tensor([[36]], dtype=torch.uint8) # [1, 1] + result = _unpack_ternary_weights(packed) + # Unpacks to [4, 1]: four values from the single byte + assert result.shape == (4, 1) + expected = torch.tensor([[-1.0], [0.0], [1.0], [-1.0]]) + assert torch.equal(result, expected) + + def test_all_zeros_byte(self): + """Byte 0x00: all four packed values are 0 → all -1.""" + packed = torch.tensor([[0]], dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert torch.all(result == -1.0) + + def test_all_ones_byte(self): + """Byte 0x55 = 0b_01_01_01_01: all four packed values are 1 → all 0.""" + packed = torch.tensor([[0x55]], dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert torch.all(result == 0.0) + + def test_all_twos_byte(self): + """Byte 0xAA = 0b_10_10_10_10: all four packed values are 2 → all +1.""" + packed = torch.tensor([[0xAA]], dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert torch.all(result == 1.0) + + def test_output_dtype_is_float(self): + """Result should be float32.""" + packed = torch.zeros(2, 4, dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert result.dtype == torch.float32 + + def test_invalid_ndim_raises(self): + """1-D input should raise ValueError.""" + packed = torch.zeros(4, dtype=torch.uint8) + try: + _unpack_ternary_weights(packed) + assert False, "Expected ValueError" + except ValueError: + pass + + def test_multi_column(self): + """Verify correct unpacking across multiple columns. + + Two bytes per row, 2 rows → packed [2, 2] → unpacked [8, 2]. + """ + # Column 0: all -1 (byte 0x00), Column 1: all +1 (byte 0xAA) + packed = torch.tensor([[0x00, 0xAA], [0x00, 0xAA]], dtype=torch.uint8) + result = _unpack_ternary_weights(packed) + assert result.shape == (8, 2) + assert torch.all(result[:, 0] == -1.0) + assert torch.all(result[:, 1] == 1.0) diff --git a/src/mobius/models/lfm2.py b/src/mobius/models/lfm2.py new file mode 100644 index 00000000..23df4421 --- /dev/null +++ b/src/mobius/models/lfm2.py @@ -0,0 +1,341 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""LFM2 hybrid ShortConv + Attention causal language model. + +LFM2 interleaves ShortConv (gated causal depthwise Conv1d) layers with +Transformer attention layers. Both layer types include a SiLU-gated MLP +(SwiGLU). + +Layer type selection via ``layer_types`` in config: + ``"conv"`` → ShortConv: in_proj → B*x gating → causal conv → C gating → out_proj + ``"full_attention"`` → standard GQA with QK norm and RoPE + +Architecture per layer: + conv layers: operator_norm → ShortConv → residual → ffn_norm → MLP → residual + attn layers: operator_norm → Attention → residual → ffn_norm → MLP → residual + +State per layer: + Conv: conv_state (batch, hidden_size, kernel_size-1) + Attention: standard KV cache (key + value) + +HuggingFace weight names (conv layer):: + + model.layers.N.conv.conv.weight → layers.N.conv.conv_weight + model.layers.N.conv.in_proj.weight → layers.N.conv.in_proj.weight + model.layers.N.conv.out_proj.weight → layers.N.conv.out_proj.weight + model.layers.N.operator_norm.weight → layers.N.operator_norm.weight + model.layers.N.ffn_norm.weight → layers.N.ffn_norm.weight + model.layers.N.feed_forward.w1.weight → layers.N.feed_forward.gate_proj.weight + model.layers.N.feed_forward.w3.weight → layers.N.feed_forward.up_proj.weight + model.layers.N.feed_forward.w2.weight → layers.N.feed_forward.down_proj.weight + +HuggingFace weight names (attention layer):: + + model.layers.N.self_attn.q_proj.weight → layers.N.self_attn.q_proj.weight + model.layers.N.self_attn.k_proj.weight → layers.N.self_attn.k_proj.weight + model.layers.N.self_attn.v_proj.weight → layers.N.self_attn.v_proj.weight + model.layers.N.self_attn.out_proj.weight → layers.N.self_attn.o_proj.weight + model.layers.N.self_attn.q_layernorm.weight → layers.N.self_attn.q_norm.weight + model.layers.N.self_attn.k_layernorm.weight → layers.N.self_attn.k_norm.weight + +HuggingFace reference: ``Lfm2ForCausalLM``. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +import torch +from onnxscript import nn +from onnxscript._internal import builder + +from mobius._configs import Lfm2Config +from mobius._weight_utils import tie_word_embeddings +from mobius.components import ( + MLP, + Attention, + Embedding, + Linear, + RMSNorm, + ShortConv, + create_attention_bias, + initialize_rope, +) + +if TYPE_CHECKING: + import onnx_ir as ir + +# --------------------------------------------------------------------------- +# Decoder layers +# --------------------------------------------------------------------------- + + +class Lfm2ConvDecoderLayer(nn.Module): + """LFM2 ShortConv layer: operator_norm → ShortConv → residual → ffn_norm → MLP. + + Args: + config: LFM2 architecture config. + """ + + def __init__(self, config: Lfm2Config): + super().__init__() + self.conv = ShortConv( + hidden_size=config.hidden_size, + kernel_size=config.short_conv_kernel, + bias=config.short_conv_bias, + ) + self.operator_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feed_forward = MLP(config) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None, + ): + """Forward pass. Returns (hidden_states, (conv_state,)). + + attention_bias and position_embeddings are unused by conv layers + but accepted for uniform interface with attention layers. + """ + del attention_bias, position_embeddings # unused + + # Pre-norm → ShortConv → residual + residual = hidden_states + hidden_states = self.operator_norm(op, hidden_states) + + conv_state = past_key_value[0] if past_key_value is not None else None + conv_out, new_conv_state = self.conv(op, hidden_states, conv_state) + hidden_states = op.Add(residual, conv_out) + + # MLP path with pre-norm + residual = hidden_states + hidden_states = self.ffn_norm(op, hidden_states) + hidden_states = self.feed_forward(op, hidden_states) + hidden_states = op.Add(residual, hidden_states) + + return hidden_states, (new_conv_state,) + + +class Lfm2AttentionDecoderLayer(nn.Module): + """LFM2 attention layer: operator_norm → Attention → residual → ffn_norm → MLP. + + Uses GQA with per-head QK norm and RoPE. + + Args: + config: LFM2 architecture config. + """ + + def __init__(self, config: Lfm2Config): + super().__init__() + self.self_attn = Attention(config) + self.operator_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.feed_forward = MLP(config) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None, + ): + """Forward pass. Returns (hidden_states, (key_cache, value_cache)).""" + # Pre-norm → Attention → residual + residual = hidden_states + hidden_states = self.operator_norm(op, hidden_states) + + hidden_states, present_kv = self.self_attn( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + ) + hidden_states = op.Add(residual, hidden_states) + + # MLP path with pre-norm + residual = hidden_states + hidden_states = self.ffn_norm(op, hidden_states) + hidden_states = self.feed_forward(op, hidden_states) + hidden_states = op.Add(residual, hidden_states) + + return hidden_states, present_kv + + +# --------------------------------------------------------------------------- +# Full model +# --------------------------------------------------------------------------- + + +class _Lfm2TextModel(nn.Module): + """LFM2 text backbone: embedding -> N x (ShortConv|Attention) -> norm. + + Layer types are selected based on ``config.layer_types``: + ``"conv"`` → Lfm2ConvDecoderLayer + ``"full_attention"`` → Lfm2AttentionDecoderLayer + """ + + def __init__(self, config: Lfm2Config): + super().__init__() + self._dtype = config.dtype + self.embed_tokens = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) + + layer_types = config.layer_types or [] + self.layers = nn.ModuleList([]) + for i in range(config.num_hidden_layers): + ltype = layer_types[i] if i < len(layer_types) else "full_attention" + if ltype == "conv": + self.layers.append(Lfm2ConvDecoderLayer(config)) + else: + self.layers.append(Lfm2AttentionDecoderLayer(config)) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = initialize_rope(config) + + def forward( + self, + op: builder.OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + hidden_states = self.embed_tokens(op, input_ids) + position_embeddings = self.rotary_emb(op, position_ids) + + attention_bias = create_attention_bias( + op, + input_ids=input_ids, + attention_mask=attention_mask, + dtype=self._dtype, + ) + + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + hidden_states = self.norm(op, hidden_states) + return hidden_states, present_key_values + + +class Lfm2CausalLMModel(nn.Module): + """LFM2 hybrid ShortConv+Attention causal language model. + + Uses ``HybridCausalLMTask`` with mixed ``"conv"`` and + ``"full_attention"`` layer types for the cache. + + HuggingFace reference: ``Lfm2ForCausalLM``. + """ + + default_task: str = "hybrid-text-generation" + category: str = "Hybrid Conv+Attention" + config_class: type = Lfm2Config + + def __init__(self, config: Lfm2Config): + super().__init__() + self.config = config + self.model = _Lfm2TextModel(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: builder.OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + hidden_states, present_key_values = self.model( + op, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + logits = self.lm_head(op, hidden_states) + return logits, present_key_values + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map HuggingFace Lfm2ForCausalLM weights to ONNX parameters. + + Handles: + 1. Weight tying (embed_tokens ↔ lm_head) + 2. MLP w1/w3/w2 → gate_proj/up_proj/down_proj rename + 3. Attention out_proj → o_proj rename + 4. QK norm: q_layernorm/k_layernorm → q_norm/k_norm + 5. Conv weight nesting: conv.conv.weight → conv.conv_weight + """ + if self.config.tie_word_embeddings: + tie_word_embeddings(state_dict) + + new_state_dict: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + new_key = _rename_lfm2_weight(key) + new_state_dict[new_key] = value + + return new_state_dict + + +# Regex for layer-level weight keys +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def _rename_lfm2_weight(key: str) -> str: + """Rename a single HF weight key to match ONNX module structure. + + Global renames: + model.embed_tokens.weight → model.embed_tokens.weight (no change) + model.norm.weight → model.norm.weight (no change) + + Per-layer renames (within model.layers.N): + conv.conv.weight → conv.conv_weight (nested module → parameter) + feed_forward.w1 → feed_forward.gate_proj + feed_forward.w3 → feed_forward.up_proj + feed_forward.w2 → feed_forward.down_proj + self_attn.out_proj → self_attn.o_proj + self_attn.q_layernorm → self_attn.q_norm + self_attn.k_layernorm → self_attn.k_norm + """ + m = _LAYER_RE.match(key) + if m is None: + return key + + idx = m.group(1) + rest = m.group(2) + + # Conv weight: conv.conv.weight → conv.conv_weight + rest = rest.replace("conv.conv.weight", "conv.conv_weight") + rest = rest.replace("conv.conv.bias", "conv.conv_bias") + + # MLP: w1 → gate_proj, w3 → up_proj, w2 → down_proj + rest = rest.replace("feed_forward.w1.", "feed_forward.gate_proj.") + rest = rest.replace("feed_forward.w3.", "feed_forward.up_proj.") + rest = rest.replace("feed_forward.w2.", "feed_forward.down_proj.") + + # Attention output projection: out_proj → o_proj + rest = rest.replace("self_attn.out_proj.", "self_attn.o_proj.") + + # QK norm: q_layernorm → q_norm, k_layernorm → k_norm + rest = rest.replace("self_attn.q_layernorm.", "self_attn.q_norm.") + rest = rest.replace("self_attn.k_layernorm.", "self_attn.k_norm.") + + return f"model.layers.{idx}.{rest}" diff --git a/src/mobius/models/lfm2_audio.py b/src/mobius/models/lfm2_audio.py new file mode 100644 index 00000000..02b2a759 --- /dev/null +++ b/src/mobius/models/lfm2_audio.py @@ -0,0 +1,626 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""LFM2-Audio: audio-to-audio model with hybrid conv+attention backbone. + +Architecture (4-model ONNX split): +1. **audio_encoder**: ConformerEncoder + adapter MLP + mel (B, n_mels, T) -> audio embeddings (B, T', hidden_size) +2. **embedding**: text token embed + audio codebook embed + text_ids + audio_features -> inputs_embeds +3. **decoder**: LFM2 backbone (takes inputs_embeds, not input_ids) + inputs_embeds -> text_logits + hybrid KV cache +4. **audio_decoder**: depthformer (per-codebook autoregressive transformer) + backbone_hidden -> codebook_logits (one codebook at a time) + +The decoder uses hybrid cache: "conv" layers carry conv_state, +"full_attention" layers carry standard KV cache. + +HuggingFace weight name prefixes:: + + lfm. -> decoder sub-model (LFM2 backbone) + conformer. -> audio_encoder.encoder (ConformerEncoder) + audio_adapter. -> audio_encoder.adapter (projection MLP) + audio_embedding. -> embedding.audio_embedding + depthformer. -> audio_decoder.depthformer + depth_linear. -> audio_decoder.depth_linear + depth_embeddings. -> audio_decoder.depth_embeddings + embedding_norm. -> audio_decoder.embedding_norm + +Reference: ``liquid_audio.model.lfm2_audio.LFM2AudioModel``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from onnxscript import nn +from onnxscript._internal import builder + +from mobius._configs import Lfm2AudioConfig +from mobius._weight_utils import tie_word_embeddings +from mobius.components import ( + FCMLP, + MLP, + Attention, + ConformerEncoder, + Embedding, + Linear, + RMSNorm, + create_attention_bias, + initialize_rope, +) +from mobius.models.lfm2 import Lfm2AttentionDecoderLayer, Lfm2ConvDecoderLayer + +if TYPE_CHECKING: + import onnx_ir as ir + + +# --------------------------------------------------------------------------- +# Audio Encoder sub-model +# --------------------------------------------------------------------------- + + +class _Lfm2AudioEncoder(nn.Module): + """ConformerEncoder + adapter MLP for mel -> LFM hidden_size projection. + + The adapter is a 2-layer MLP: + Linear(encoder_dim, hidden_size) -> GELU -> Linear(hidden_size, hidden_size) + + Weight names (HF):: + + conformer.* -> encoder.* + audio_adapter.model.0.{weight,bias} -> adapter.up_proj.{weight,bias} + audio_adapter.model.1.{weight,bias} -> (batch_norm, skipped) + audio_adapter.model.3.{weight,bias} -> adapter.down_proj.{weight,bias} + """ + + def __init__(self, config: Lfm2AudioConfig): + super().__init__() + audio = config.audio + assert audio is not None + + self.encoder = ConformerEncoder( + input_size=audio.num_mel_bins or 128, + attention_dim=audio.attention_dim or audio.d_model or 512, + attention_heads=audio.attention_heads or audio.encoder_attention_heads or 8, + num_blocks=audio.num_blocks or audio.encoder_layers or 17, + linear_units=(audio.linear_units or audio.encoder_ffn_dim or 2048), + kernel_size=audio.kernel_size or 9, + conv_channels=audio.conv_channels or 256, + t5_bias_max_distance=audio.t5_bias_max_distance or 500, + ) + # Adapter: encoder_dim -> hidden_size + encoder_dim = audio.attention_dim or audio.d_model or 512 + self.adapter = FCMLP( + hidden_size=encoder_dim, + intermediate_size=config.hidden_size, + activation="gelu", + bias=True, + ) + + def forward(self, op: builder.OpBuilder, input_features: ir.Value): + """Forward: mel (B, n_mels, T) -> (B, T', hidden_size).""" + # ConformerEncoder expects (B, T, n_mels); transpose from (B, n_mels, T) + input_features = op.Transpose(input_features, perm=[0, 2, 1]) + audio_features = self.encoder(op, input_features) + # Adapter MLP: (B, T', encoder_dim) -> (B, T', hidden_size) + return self.adapter(op, audio_features) + + +# --------------------------------------------------------------------------- +# Embedding sub-model +# --------------------------------------------------------------------------- + + +class _Lfm2AudioEmbedding(nn.Module): + """Embedding model for LFM2-Audio. + + Combines text token embeddings with audio feature embeddings. + In the actual model, a modality_flag tensor controls which positions + get text embeddings vs audio-in vs audio-out embeddings. + + For ONNX export, this takes pre-computed audio features and text_ids, + returning the combined inputs_embeds sequence. + + Weight names (HF):: + + lfm.embed_tokens.weight -> text_embed.weight + audio_embedding.embedding.weight -> audio_embed.weight + """ + + 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) + + +# --------------------------------------------------------------------------- +# Decoder sub-model (LFM2 backbone without embed_tokens) +# --------------------------------------------------------------------------- + + +# Reuse the decoder layers from the base LFM2 model — they are identical +# for the audio backbone. Lfm2AudioConfig inherits from Lfm2Config, so +# the constructors accept it directly. +_Lfm2AudioDecoderLayer = Lfm2AttentionDecoderLayer +_Lfm2AudioConvLayer = Lfm2ConvDecoderLayer + + +class _Lfm2AudioDecoder(nn.Module): + """LFM2 decoder backbone: takes inputs_embeds -> logits + cache. + + This is the LFM2 model minus the embedding layer. It takes + pre-assembled inputs_embeds (from the embedding model) and runs + the hybrid conv+attention backbone, then projects to vocab logits. + + The text LM head shares weights with embed_tokens (tied). + """ + + def __init__(self, config: Lfm2AudioConfig): + super().__init__() + self._dtype = config.dtype + + layer_types = config.layer_types or [] + self.layers = nn.ModuleList([]) + for i in range(config.num_hidden_layers): + ltype = layer_types[i] if i < len(layer_types) else "full_attention" + if ltype == "conv": + self.layers.append(_Lfm2AudioConvLayer(config)) + else: + self.layers.append(_Lfm2AudioDecoderLayer(config)) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = initialize_rope(config) + # LM head (tied with lfm.embed_tokens in preprocess_weights) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: builder.OpBuilder, + inputs_embeds: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(op, position_ids) + attention_bias = create_attention_bias( + op, + # Use a dummy input_ids shape from inputs_embeds + input_ids=position_ids, + attention_mask=attention_mask, + dtype=self._dtype, + ) + + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + hidden_states = self.norm(op, hidden_states) + logits = self.lm_head(op, hidden_states) + return logits, present_key_values + + +# --------------------------------------------------------------------------- +# Audio decoder sub-model (depthformer) +# --------------------------------------------------------------------------- + + +class _DepthformerLayer(nn.Module): + """Single depthformer layer with RMSNorm, Attention, and SwiGLU MLP. + + Architecture: RMSNorm -> Attention -> residual -> + RMSNorm -> SwiGLU MLP -> residual. + + The depthformer uses the same StandardBlock structure as the LFM2 + backbone: operator_norm + operator + ffn_norm + feed_forward. + """ + + def __init__(self, config: Lfm2AudioConfig): + super().__init__() + from mobius._configs import ArchitectureConfig + + # Create a mini-config for the depthformer attention + depthformer_dim = config.depthformer_dim + depthformer_heads = config.depthformer_heads + head_dim = depthformer_dim // depthformer_heads + + # Build attention config for the depthformer + attn_config = ArchitectureConfig( + hidden_size=depthformer_dim, + intermediate_size=depthformer_dim * 4, + num_attention_heads=depthformer_heads, + num_key_value_heads=depthformer_heads, + head_dim=head_dim, + hidden_act="silu", + attn_qk_norm=True, + rms_norm_eps=1e-5, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + ) + self.self_attn = Attention(attn_config) + self.operator_norm = RMSNorm(depthformer_dim, eps=1e-5) + self.ffn_norm = RMSNorm(depthformer_dim, eps=1e-5) + self.feed_forward = MLP(attn_config) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value | None, + position_embeddings: tuple, + past_key_value: tuple | None, + ): + residual = hidden_states + hidden_states = self.operator_norm(op, hidden_states) + hidden_states, present_kv = self.self_attn( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + ) + hidden_states = op.Add(residual, hidden_states) + residual = hidden_states + hidden_states = self.ffn_norm(op, hidden_states) + hidden_states = self.feed_forward(op, hidden_states) + hidden_states = op.Add(residual, hidden_states) + return hidden_states, present_kv + + +class _Lfm2AudioDecoderModule(nn.Module): + """Depthformer audio decoder for per-codebook token prediction. + + Takes backbone output + previous codebook embedding, runs through + depthformer layers, and produces logits for the current codebook. + + Architecture:: + + depth_linear(backbone_hidden) -> split by codebook_idx -> + + prev_embedding -> depthformer layers -> embedding_norm -> + codebook_head -> codebook_logits + + The codebook heads share weights with depth_embeddings (tied). + """ + + def __init__(self, config: Lfm2AudioConfig): + super().__init__() + depthformer_dim = config.depthformer_dim + + # Project backbone hidden -> codebook inputs + # depth_linear: (hidden_size) -> (codebooks * depthformer_dim) + self.depth_linear = Linear( + config.hidden_size, + config.num_codebooks * depthformer_dim, + bias=True, + ) + + # Depthformer layers + self.layers = nn.ModuleList([]) + for _ in range(config.depthformer_layers): + self.layers.append(_DepthformerLayer(config)) + + # Output norm + per-codebook heads + self.embedding_norm = RMSNorm(depthformer_dim, eps=1e-5) + + # 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] + ) + + self._depthformer_dim = depthformer_dim + self._num_codebooks = config.num_codebooks + + # Build a separate RoPE for depthformer + from mobius._configs import ArchitectureConfig + + rope_config = ArchitectureConfig( + hidden_size=depthformer_dim, + num_attention_heads=config.depthformer_heads, + head_dim=depthformer_dim // config.depthformer_heads, + rope_theta=config.rope_theta, + max_position_embeddings=config.max_position_embeddings, + ) + self.rotary_emb = initialize_rope(rope_config) + + def forward( + self, + op: builder.OpBuilder, + backbone_hidden: ir.Value, + prev_embedding: ir.Value, + codebook_idx: ir.Value, + past_key_values: list | None = None, + ): + """Forward pass for single-codebook prediction. + + Args: + backbone_hidden: (B, 1, hidden_size) from LFM2 decoder + prev_embedding: (B, 1, depthformer_dim) from previous codebook + codebook_idx: scalar int - which codebook to predict + past_key_values: depthformer KV cache + + Returns: + (codebook_logits, present_key_values) + """ + # Project backbone hidden to all codebook inputs + # (B, 1, hidden_size) -> (B, 1, codebooks * depthformer_dim) + projected = self.depth_linear(op, backbone_hidden) + + # Reshape to (B, codebooks, depthformer_dim) for gathering + # First squeeze the seq dim: (B, 1, C*D) -> (B, C*D) + projected_2d = op.Squeeze(projected, [1]) + projected_3d = op.Reshape( + projected_2d, + op.Constant( + value_ints=[ + -1, + self._num_codebooks, + self._depthformer_dim, + ] + ), + ) + # (B, codebooks, depthformer_dim) + + # Gather the codebook_idx slice: (B, 1, depthformer_dim) + # Reshape idx to (1, 1, 1) then expand to (B, 1, depthformer_dim) + idx_3d = op.Reshape(codebook_idx, op.Constant(value_ints=[1, 1, 1])) + # Build expand shape dynamically to match batch size at runtime + batch_dim = op.Shape(projected_3d, start=0, end=1) # (1,) containing B + expand_shape = op.Concat( + batch_dim, + op.Constant(value_ints=[1]), + op.Constant(value_ints=[self._depthformer_dim]), + axis=0, + ) # (3,) -> [B, 1, depthformer_dim] + idx_expanded = op.Expand(idx_3d, expand_shape) + depthformer_input = op.GatherElements(projected_3d, idx_expanded, axis=1) + # (B, 1, depthformer_dim) - unsqueeze back seq dim is already there + + # Add previous codebook embedding + hidden_states = op.Add(depthformer_input, prev_embedding) + + # Position IDs for depthformer (single step: just codebook_idx). + # Shape (B, 1) — derive batch dim from hidden_states at runtime. + batch_dim = op.Shape(hidden_states, start=0, end=1) # (1,) containing B + one_dim = op.Constant(value_ints=[1]) + position_shape = op.Concat(batch_dim, one_dim, axis=0) # (2,) -> [B, 1] + position_ids = op.Reshape(codebook_idx, position_shape) + position_embeddings = self.rotary_emb(op, position_ids) + + # Run depthformer layers + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=None, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + # Norm + logits for current codebook + hidden_states = self.embedding_norm(op, hidden_states) + + # Dynamic codebook head selection: + # stacked_head_weights: (num_codebooks, audio_vocab_size, dim) + # Gather by codebook_idx -> (audio_vocab_size, dim) + head_weight = op.Gather(self.stacked_head_weights, codebook_idx, axis=0) + # (audio_vocab_size, depthformer_dim) + head_weight_3d = op.Unsqueeze(head_weight, [0]) + # (1, audio_vocab_size, depthformer_dim) + # hidden_states: (B, 1, depthformer_dim) + # logits = hidden_states @ head_weight^T -> (B, 1, audio_vocab_size) + logits = op.MatMul(hidden_states, op.Transpose(head_weight_3d, perm=[0, 2, 1])) + + return logits, present_key_values + + +# --------------------------------------------------------------------------- +# Composite model +# --------------------------------------------------------------------------- + + +class Lfm2AudioModel(nn.Module): + """LFM2-Audio: audio-to-audio model. + + Exports as 4 ONNX models via AudioToAudioTask: + - audio_encoder: ConformerEncoder + adapter + - embedding: text + audio embedding fusion + - decoder: LFM2 hybrid backbone + - audio_decoder: depthformer per-codebook decoder + + HuggingFace reference: ``liquid_audio.model.lfm2_audio.LFM2AudioModel``. + """ + + default_task: str = "audio-to-audio" + category: str = "Audio-to-Audio" + config_class: type = Lfm2AudioConfig + + def __init__(self, config: Lfm2AudioConfig): + super().__init__() + self.config = config + + self.audio_encoder = _Lfm2AudioEncoder(config) + self.embedding = _Lfm2AudioEmbedding(config) + self.decoder = _Lfm2AudioDecoder(config) + self.audio_decoder = _Lfm2AudioDecoderModule(config) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map LFM2-Audio weights to ONNX sub-model parameters. + + Routes weights to sub-models by prefix: + lfm.embed_tokens.* -> embedding.text_embed.* + lfm.* -> decoder.* (backbone layers) + conformer.* -> audio_encoder.encoder.* + audio_adapter.* -> audio_encoder.adapter.* + audio_embedding.* -> embedding.audio_embed.* + depthformer.* -> audio_decoder.depthformer.* + depth_linear.* -> audio_decoder.depth_linear.* + depth_embeddings.* -> audio_decoder.depth_embeddings.* + embedding_norm.* -> audio_decoder.embedding_norm.* + """ + if self.config.tie_word_embeddings: + tie_word_embeddings( + state_dict, + embed_key="lfm.embed_tokens.weight", + head_key="lfm.lm_head.weight", + ) + + new_state_dict: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + new_key = _rename_lfm2_audio_weight(key) + if new_key is not None: + new_state_dict[new_key] = value + + # Stack per-codebook head weights into stacked_head_weights + # for dynamic codebook selection in the audio_decoder forward. + head_weights = [] + for i in range(self.config.num_codebooks): + wkey = f"audio_decoder.depth_embeddings.{i}.weight" + if wkey in new_state_dict: + head_weights.append(new_state_dict[wkey]) + if head_weights: + import torch + + new_state_dict["audio_decoder.stacked_head_weights"] = torch.stack( + head_weights, dim=0 + ) + + return new_state_dict + + +def _rename_lfm2_audio_weight(key: str) -> str | None: + """Rename a single HF weight key to ONNX module structure. + + Returns None if the weight should be skipped. + """ + import re + + # LFM backbone embed_tokens -> embedding.text_embed + if key.startswith("lfm.embed_tokens."): + return key.replace("lfm.embed_tokens.", "embedding.text_embed.") + + # LFM backbone layers -> decoder.layers + if key.startswith("lfm."): + rest = key[len("lfm.") :] + # model.layers.N patterns + m = re.match(r"^layers\.(\d+)\.(.+)$", rest) + if m: + idx = m.group(1) + layer_rest = m.group(2) + # Conv weight nesting + layer_rest = layer_rest.replace("conv.conv.weight", "conv.conv_weight") + layer_rest = layer_rest.replace("conv.conv.bias", "conv.conv_bias") + # MLP: w1->gate_proj, w3->up_proj, w2->down_proj + layer_rest = layer_rest.replace("feed_forward.w1.", "feed_forward.gate_proj.") + layer_rest = layer_rest.replace("feed_forward.w3.", "feed_forward.up_proj.") + layer_rest = layer_rest.replace("feed_forward.w2.", "feed_forward.down_proj.") + # Attention: out_proj->o_proj, layernorm->norm + layer_rest = layer_rest.replace("self_attn.out_proj.", "self_attn.o_proj.") + layer_rest = layer_rest.replace("self_attn.q_layernorm.", "self_attn.q_norm.") + layer_rest = layer_rest.replace("self_attn.k_layernorm.", "self_attn.k_norm.") + return f"decoder.layers.{idx}.{layer_rest}" + + # lfm.norm -> decoder.norm + return f"decoder.{rest}" + + # Conformer -> audio_encoder.encoder + if key.startswith("conformer."): + return key.replace("conformer.", "audio_encoder.encoder.") + + # Audio adapter -> audio_encoder.adapter + if key.startswith("audio_adapter."): + rest = key[len("audio_adapter.") :] + # audio_adapter.model.0.* -> adapter.up_proj.* + rest = rest.replace("model.0.", "up_proj.") + # audio_adapter.model.3.* -> adapter.down_proj.* + rest = rest.replace("model.3.", "down_proj.") + # Skip batch norm (model.1.*) + if "model.1." in key or "model.2." in key: + return None + return f"audio_encoder.adapter.{rest}" + + # Audio embedding + if key.startswith("audio_embedding."): + rest = key[len("audio_embedding.") :] + rest = rest.replace("embedding.", "audio_embed.") + return f"embedding.{rest}" + + # Depthformer layers + if key.startswith("depthformer.layers."): + rest = key[len("depthformer.layers.") :] + m = re.match(r"^(\d+)\.(.+)$", rest) + if m: + idx = m.group(1) + layer_rest = m.group(2) + # operator.qkv_proj -> self_attn.qkv_proj (if fused) + # operator.out_proj -> self_attn.o_proj + layer_rest = layer_rest.replace("operator.", "self_attn.") + layer_rest = layer_rest.replace("self_attn.out_proj.", "self_attn.o_proj.") + layer_rest = layer_rest.replace( + "self_attn.bounded_attention.q_layernorm.", + "self_attn.q_norm.", + ) + layer_rest = layer_rest.replace( + "self_attn.bounded_attention.k_layernorm.", + "self_attn.k_norm.", + ) + # MLP renames + layer_rest = layer_rest.replace("feed_forward.w1.", "feed_forward.gate_proj.") + layer_rest = layer_rest.replace("feed_forward.w3.", "feed_forward.up_proj.") + layer_rest = layer_rest.replace("feed_forward.w2.", "feed_forward.down_proj.") + return f"audio_decoder.layers.{idx}.{layer_rest}" + return None + + # Depth linear + if key.startswith("depth_linear."): + return key.replace("depth_linear.", "audio_decoder.depth_linear.") + + # Depth embeddings + if key.startswith("depth_embeddings."): + return key.replace("depth_embeddings.", "audio_decoder.depth_embeddings.") + + # Embedding norm (for depthformer output) + if key.startswith("embedding_norm."): + return key.replace("embedding_norm.", "audio_decoder.embedding_norm.") + + return key diff --git a/src/mobius/models/moshi.py b/src/mobius/models/moshi.py new file mode 100644 index 00000000..a1318af1 --- /dev/null +++ b/src/mobius/models/moshi.py @@ -0,0 +1,646 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Moshi / PersonaPlex: full-duplex speech-to-speech model. + +Architecture (3-model ONNX split): +1. **embedding**: text token + audio codec token fusion + text_ids (B, S) + audio_codes (B, S, num_codebooks) -> inputs_embeds (B, S, H) +2. **decoder**: causal transformer backbone + inputs_embeds -> text_logits + standard KV cache +3. **audio_decoder**: depthformer per-codebook autoregressive transformer + backbone_hidden (B, 1, H) + prev_embedding (B, 1, D) + codebook_idx -> + codebook_logits (B, 1, audio_logits_size) + depformer KV cache + +Unlike LFM2-Audio, Moshi has no mel-spectrogram audio encoder. Audio is +consumed and produced as RVQ codec token IDs, embedded directly. + +HuggingFace weight name prefixes:: + + text_emb. -> embedding.text_emb + emb.N. -> embedding.audio_emb.N (N=0..num_codebooks-1) + transformer.layers. -> decoder.layers + out_norm. -> decoder.out_norm + text_linear. -> decoder.lm_head + depformer_text_emb. -> audio_decoder.depth_text_emb + depformer_emb.N. -> audio_decoder.depth_emb.N (N=0..num_codebooks-2) + depformer_in. -> audio_decoder.stacked_depformer_in (stacked) + depformer.layers. -> audio_decoder.layers + linears. -> audio_decoder.stacked_output_heads (stacked) + +Reference: Défossez et al., "Moshi: a speech-text foundation model for +real-time dialogue" (2024). Base model: ``kyutai/moshiko-pytorch-bf16``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from onnxscript import nn +from onnxscript._internal import builder + +from mobius._configs import ArchitectureConfig, MoshiConfig +from mobius.components import ( + MLP, + Attention, + Embedding, + Linear, + RMSNorm, + create_attention_bias, + initialize_rope, +) + +if TYPE_CHECKING: + import onnx_ir as ir + + +# --------------------------------------------------------------------------- +# Embedding sub-model +# --------------------------------------------------------------------------- + + +class _MoshiEmbedding(nn.Module): + """Moshi embedding: text + audio codec tokens -> inputs_embeds. + + Combines the text token embedding with the sum of all per-codebook audio + token embeddings. Each timestep contributes one text token ID and + ``num_codebooks`` audio codec token IDs. + + Result:: + + inputs_embeds = text_emb[text_ids] + sum(audio_emb[i][audio_codes[...,i]]) + + Weight names (HF):: + + text_emb.weight -> text_emb.weight + emb.N.weight (N=0..K) -> audio_emb.N.weight + """ + + def __init__(self, config: MoshiConfig): + super().__init__() + self._num_codebooks = config.num_codebooks + # Moshi vocabulary has 32001 tokens (32000 text + 1 extra). + self.text_emb = Embedding(config.vocab_size + 1, config.hidden_size) + # Per-codebook audio token embeddings: audio_vocab_size x hidden_size each. + self.audio_emb = nn.ModuleList( + [ + Embedding(config.audio_vocab_size, config.hidden_size) + for _ in range(config.num_codebooks) + ] + ) + + def forward( + self, + op: builder.OpBuilder, + input_ids: ir.Value, + audio_codes: ir.Value, + ) -> ir.Value: + """Forward: (text_ids, audio_codes) -> inputs_embeds. + + Args: + input_ids: (batch, seq) int64 text token IDs + audio_codes: (batch, seq, num_codebooks) int64 audio codec codes + + Returns: + inputs_embeds: (batch, seq, hidden_size) + """ + # Text embedding: (batch, seq, hidden_size) + embeds = self.text_emb(op, input_ids) + + # Add per-codebook audio embeddings + for i, emb_table in enumerate(self.audio_emb): + # Gather codes for codebook i along axis=2: (batch, seq) + code_i = op.Gather(audio_codes, op.Constant(value_int=i), axis=2) + embeds = op.Add(embeds, emb_table(op, code_i)) + + return embeds + + +# --------------------------------------------------------------------------- +# Main transformer decoder layers +# --------------------------------------------------------------------------- + + +class _MoshiDecoderLayer(nn.Module): + """Single Moshi transformer layer. + + Architecture: PreNorm → Attention → residual → PreNorm → SwiGLU MLP → residual. + + Weight names (HF, under ``transformer.layers.N.``):: + + norm1.alpha [1,1,H] -> norm1.weight [H] (via reshape) + self_attn.in_proj_weight [3H,H] -> self_attn.{q,k,v}_proj.weight [H,H] + self_attn.out_proj.weight [H,H] -> self_attn.o_proj.weight [H,H] + norm2.alpha [1,1,H] -> norm2.weight [H] (via reshape) + gating.linear_in.weight [2I,H] -> gating.{gate,up}_proj.weight [I,H] + gating.linear_out.weight [H,I] -> gating.down_proj.weight [H,I] + """ + + def __init__(self, config: MoshiConfig): + super().__init__() + self.norm1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = Attention(config) + self.norm2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.gating = MLP(config) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None, + ): + # Pre-norm + causal self-attention + residual = hidden_states + hidden_states = self.norm1(op, hidden_states) + hidden_states, present_kv = self.self_attn( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + ) + hidden_states = op.Add(residual, hidden_states) + + # Pre-norm + SwiGLU MLP + residual = hidden_states + hidden_states = self.norm2(op, hidden_states) + hidden_states = self.gating(op, hidden_states) + hidden_states = op.Add(residual, hidden_states) + + return hidden_states, present_kv + + +class _MoshiDecoder(nn.Module): + """Moshi main causal transformer decoder. + + Weight names (HF):: + + transformer.layers.N.* -> layers.N.* (via preprocess_weights) + out_norm.alpha [1,1,H] -> out_norm.weight [H] + text_linear.weight -> lm_head.weight + """ + + def __init__(self, config: MoshiConfig): + super().__init__() + self.layers = nn.ModuleList( + [_MoshiDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.out_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + self.rotary_emb = initialize_rope(config) + + def forward( + self, + op: builder.OpBuilder, + inputs_embeds: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + hidden_states = inputs_embeds + # RoPE position embeddings + position_embeddings = self.rotary_emb(op, position_ids) + # Causal attention bias from the attention mask + attention_bias = create_attention_bias(op, attention_mask, hidden_states, position_ids) + + present_key_values = [] + for i, layer in enumerate(self.layers): + past_kv = past_key_values[i] if past_key_values is not None else None + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + # Final norm + LM head + hidden_states = self.out_norm(op, hidden_states) + logits = self.lm_head(op, hidden_states) + + return logits, present_key_values + + +# --------------------------------------------------------------------------- +# Depformer audio decoder layers +# --------------------------------------------------------------------------- + + +class _DepformerLayer(nn.Module): + """Single Moshi depformer layer. + + Shared causal attention over the codebook-position sequence (KV cache + accumulates one step per codebook), plus a per-codebook SwiGLU MLP + selected at runtime by ``codebook_idx``. + + Per-codebook gating weights are stacked as parameters so the correct + MLP can be selected with a single ``Gather`` on ``codebook_idx``. + + The attention uses ``num_heads = num_codebooks`` and + ``head_dim = depformer_dim`` (one full-dimensioned head per codebook). + + Weight names (HF, under ``depformer.layers.N.``):: + + norm1.alpha [1,1,D] -> norm1.weight [D] + self_attn.in_proj_weight [3*K*D,D] -> self_attn.{q,k,v}_proj.weight [K*D,D] + self_attn.out_proj.weight [K*D,D] -> self_attn.o_proj.weight [D,K*D] (transposed) + norm2.alpha [1,1,D] -> norm2.weight [D] + gating.M.linear_in.weight [2I,D] -> stacked_{gate,up}_proj (stacked over M) + gating.M.linear_out.weight [D,I] -> stacked_down_proj (stacked over M) + """ + + def __init__(self, config: MoshiConfig): + super().__init__() + depformer_dim = config.depformer_dim + num_codebooks = config.num_codebooks + interm = config.depformer_intermediate_size + + # Shared attention: num_heads=num_codebooks so each head covers one codebook. + attn_config = ArchitectureConfig( + hidden_size=depformer_dim, + num_attention_heads=num_codebooks, + num_key_value_heads=num_codebooks, + head_dim=depformer_dim, # head_dim = full depformer dim per head + hidden_act="silu", + rms_norm_eps=1e-5, + rope_theta=10000.0, + max_position_embeddings=num_codebooks, + ) + self.norm1 = RMSNorm(depformer_dim, eps=1e-5) + self.self_attn = Attention(attn_config) + self.norm2 = RMSNorm(depformer_dim, eps=1e-5) + + # Per-codebook stacked gating weights: + # (num_codebooks, intermediate_size, depformer_dim) for gate/up projections. + # (num_codebooks, depformer_dim, intermediate_size) for down projection. + self.stacked_gate_proj = nn.Parameter([num_codebooks, interm, depformer_dim]) + self.stacked_up_proj = nn.Parameter([num_codebooks, interm, depformer_dim]) + self.stacked_down_proj = nn.Parameter([num_codebooks, depformer_dim, interm]) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + codebook_idx: ir.Value, + position_embeddings: tuple, + past_key_value: tuple | None, + ): + # Pre-norm + shared causal attention (KV cache holds previous codebook steps) + residual = hidden_states + hidden_states = self.norm1(op, hidden_states) + hidden_states, present_kv = self.self_attn( + op, + hidden_states=hidden_states, + attention_bias=None, + position_embeddings=position_embeddings, + past_key_value=past_key_value, + ) + hidden_states = op.Add(residual, hidden_states) + + # Pre-norm + per-codebook SwiGLU MLP (selected by codebook_idx) + residual = hidden_states + hidden_states = self.norm2(op, hidden_states) + + # Select gating weights for the current codebook + # stacked_gate_proj: (num_codebooks, interm, dim) -> (interm, dim) + gate_w = op.Gather(self.stacked_gate_proj, codebook_idx, axis=0) + up_w = op.Gather(self.stacked_up_proj, codebook_idx, axis=0) + down_w = op.Gather(self.stacked_down_proj, codebook_idx, axis=0) + + # SwiGLU: silu(gate_proj(x)) * up_proj(x) -> down_proj + # gate_w: (interm, dim) -> op.Transpose -> (dim, interm) + gate = op.MatMul(hidden_states, op.Transpose(gate_w)) # (batch, 1, interm) + up = op.MatMul(hidden_states, op.Transpose(up_w)) # (batch, 1, interm) + # SiLU(gate) = gate * sigmoid(gate) + activated = op.Mul(gate, op.Sigmoid(gate)) + intermediate = op.Mul(activated, up) + # down_w: (dim, interm) -> op.Transpose -> (interm, dim) + hidden_states = op.MatMul(intermediate, op.Transpose(down_w)) # (batch, 1, dim) + + hidden_states = op.Add(residual, hidden_states) + return hidden_states, present_kv + + +class _MoshiAudioDecoder(nn.Module): + """Moshi depformer: per-codebook autoregressive depth transformer. + + At each inference step, takes the main transformer's hidden state plus + the previous codebook's embedding, runs it through the depformer layers, + and produces logits for the current codebook. + + The depformer KV cache accumulates codebook steps (not time steps). + Each call advances by one codebook (``codebook_idx`` = 0..num_codebooks-1). + + Weight names (HF):: + + depformer_text_emb.weight -> depth_text_emb.weight + depformer_emb.N.weight -> depth_emb.N.weight (N=0..num_codebooks-2) + depformer_in.N.weight -> stacked_depformer_in (stacked; N=0..num_codebooks-1) + depformer.layers.N.* -> layers.N.* (via preprocess_weights) + linears.N.weight -> stacked_output_heads (stacked; N=0..num_codebooks-1) + """ + + def __init__(self, config: MoshiConfig): + super().__init__() + depformer_dim = config.depformer_dim + num_codebooks = config.num_codebooks + audio_logits_size = config.audio_vocab_size - 1 # 2048: no padding token in output + + # Text depth embedding for codebook-0 input (from main text token) + self.depth_text_emb = Embedding(config.vocab_size + 1, depformer_dim) + # Audio depth embeddings for codebooks 1..num_codebooks-1 + self.depth_emb = nn.ModuleList( + [ + Embedding(config.audio_vocab_size, depformer_dim) + for _ in range(num_codebooks - 1) + ] + ) + + # Per-codebook input projections: (num_codebooks, depformer_dim, hidden_size) + # Gathered at runtime by codebook_idx to project backbone_hidden. + self.stacked_depformer_in = nn.Parameter( + [num_codebooks, depformer_dim, config.hidden_size] + ) + + # Depformer layers + self.layers = nn.ModuleList( + [_DepformerLayer(config) for _ in range(config.depformer_layers)] + ) + + # Per-codebook output heads: (num_codebooks, audio_logits_size, depformer_dim) + # Gathered at runtime by codebook_idx to produce logits. + self.stacked_output_heads = nn.Parameter( + [num_codebooks, audio_logits_size, depformer_dim] + ) + + # RoPE for depformer (codebook index as position) + rope_config = ArchitectureConfig( + hidden_size=depformer_dim, + num_attention_heads=num_codebooks, + head_dim=depformer_dim, + rope_theta=10000.0, + max_position_embeddings=num_codebooks, + ) + self.rotary_emb = initialize_rope(rope_config) + + self._num_codebooks = num_codebooks + self._depformer_dim = depformer_dim + + def forward( + self, + op: builder.OpBuilder, + backbone_hidden: ir.Value, + prev_embedding: ir.Value, + codebook_idx: ir.Value, + past_key_values: list | None = None, + ): + """Single-codebook forward pass. + + Args: + backbone_hidden: (batch, 1, hidden_size) from main transformer + prev_embedding: (batch, 1, depformer_dim) previous codebook embed + codebook_idx: scalar int64 — which codebook to predict + past_key_values: depformer KV cache (one entry per layer) + + Returns: + (codebook_logits, present_key_values) + """ + # Select input projection for current codebook + # stacked_depformer_in: (num_codebooks, depformer_dim, hidden_size) + # -> (depformer_dim, hidden_size) after Gather + in_proj = op.Gather(self.stacked_depformer_in, codebook_idx, axis=0) + + # Project backbone hidden: (batch, 1, hidden_size) @ (hidden_size, depformer_dim) + depformer_input = op.MatMul(backbone_hidden, op.Transpose(in_proj)) # (batch,1,D) + + # Add previous codebook embedding (provides depth autoregressive context) + hidden_states = op.Add(depformer_input, prev_embedding) + + # RoPE using codebook_idx as position (shape: [1, 1]) + codebook_pos = op.Reshape(codebook_idx, op.Constant(value_ints=[1, 1])) + position_embeddings = self.rotary_emb(op, codebook_pos) + + # Run depformer layers with per-codebook MLP selection + present_key_values = [] + past_kvs = ( + past_key_values if past_key_values is not None else [None] * len(self.layers) + ) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + codebook_idx=codebook_idx, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + # Select output head for current codebook + # stacked_output_heads: (num_codebooks, audio_logits_size, depformer_dim) + # -> (audio_logits_size, depformer_dim) after Gather + head_w = op.Gather(self.stacked_output_heads, codebook_idx, axis=0) + + # Logits: (batch, 1, depformer_dim) @ (depformer_dim, audio_logits_size) + codebook_logits = op.MatMul(hidden_states, op.Transpose(head_w)) # (batch,1,V) + + return codebook_logits, present_key_values + + +# --------------------------------------------------------------------------- +# Top-level model +# --------------------------------------------------------------------------- + + +class MoshiModel(nn.Module): + """Moshi/PersonaPlex audio-to-audio model (3-model ONNX split). + + Used with ``MoshiTask`` which builds: + - ``embedding``: text + audio codec token embeddings + - ``decoder``: 32-layer causal transformer, text logits + KV cache + - ``audio_decoder``: 6-layer depformer, per-codebook audio logits + + Weight names are mapped from the Moshi HuggingFace checkpoint format + via ``preprocess_weights``. + """ + + default_task: str = "moshi" + config_class: type = MoshiConfig + + def __init__(self, config: MoshiConfig): + super().__init__() + self.embedding = _MoshiEmbedding(config) + self.decoder = _MoshiDecoder(config) + self.audio_decoder = _MoshiAudioDecoder(config) + self._config = config + + @staticmethod + def preprocess_weights(state_dict: dict) -> dict: + """Map Moshi HF checkpoint weights to mobius module names. + + Transforms: + - ``norm*.alpha [1,1,H]`` → ``norm*.weight [H]`` (squeeze) + - ``self_attn.in_proj_weight [3H,H]`` → split into ``q/k/v_proj.weight`` + - ``self_attn.out_proj.weight`` → ``self_attn.o_proj.weight`` + - ``gating.linear_in.weight [2I,H]`` → split into ``gate/up_proj.weight`` + - ``gating.linear_out.weight`` → ``gating.down_proj.weight`` + - ``depformer_in.N.weight`` → stacked ``stacked_depformer_in`` + - ``linears.N.weight`` → stacked ``stacked_output_heads`` + - Per-codebook gating → stacked ``stacked_{gate,up,down}_proj`` + """ + new_sd: dict = {} + + # Count transformer and depformer layers + num_transformer_layers = sum( + 1 + for k in state_dict + if k.startswith("transformer.layers.") and k.endswith(".norm1.alpha") + ) + num_depformer_layers = sum( + 1 + for k in state_dict + if k.startswith("depformer.layers.") and k.endswith(".norm1.alpha") + ) + + # Count codebooks from emb.* keys + num_codebooks = sum( + 1 for k in state_dict if k.startswith("emb.") and k.endswith(".weight") + ) + + # ---------------------------------------------------------------- + # Embedding sub-model + # ---------------------------------------------------------------- + if "text_emb.weight" in state_dict: + new_sd["embedding.text_emb.weight"] = state_dict.pop("text_emb.weight") + for i in range(num_codebooks): + key = f"emb.{i}.weight" + if key in state_dict: + new_sd[f"embedding.audio_emb.{i}.weight"] = state_dict.pop(key) + + # ---------------------------------------------------------------- + # Main transformer decoder layers + # ---------------------------------------------------------------- + for i in range(num_transformer_layers): + prefix = f"transformer.layers.{i}" + dst = f"decoder.layers.{i}" + + # RMSNorm: alpha [1,1,H] → weight [H] + for norm in ("norm1", "norm2"): + alpha_key = f"{prefix}.{norm}.alpha" + if alpha_key in state_dict: + new_sd[f"{dst}.{norm}.weight"] = state_dict.pop(alpha_key).flatten() + + # Attention: split packed QKV in_proj_weight + in_proj_key = f"{prefix}.self_attn.in_proj_weight" + if in_proj_key in state_dict: + q, k, v = state_dict.pop(in_proj_key).chunk(3, dim=0) + new_sd[f"{dst}.self_attn.q_proj.weight"] = q + new_sd[f"{dst}.self_attn.k_proj.weight"] = k + new_sd[f"{dst}.self_attn.v_proj.weight"] = v + + # out_proj -> o_proj (square matrix, no transpose needed for hidden_size=head_dim*heads) + out_proj_key = f"{prefix}.self_attn.out_proj.weight" + if out_proj_key in state_dict: + new_sd[f"{dst}.self_attn.o_proj.weight"] = state_dict.pop(out_proj_key) + + # Gating MLP: split linear_in, rename linear_out + lin_in_key = f"{prefix}.gating.linear_in.weight" + if lin_in_key in state_dict: + gate, up = state_dict.pop(lin_in_key).chunk(2, dim=0) + new_sd[f"{dst}.gating.gate_proj.weight"] = gate + new_sd[f"{dst}.gating.up_proj.weight"] = up + + lin_out_key = f"{prefix}.gating.linear_out.weight" + if lin_out_key in state_dict: + new_sd[f"{dst}.gating.down_proj.weight"] = state_dict.pop(lin_out_key) + + # Decoder final norm and LM head + if "out_norm.alpha" in state_dict: + new_sd["decoder.out_norm.weight"] = state_dict.pop("out_norm.alpha").flatten() + if "text_linear.weight" in state_dict: + new_sd["decoder.lm_head.weight"] = state_dict.pop("text_linear.weight") + + # ---------------------------------------------------------------- + # Audio decoder (depformer) + # ---------------------------------------------------------------- + + # Text depth embedding + if "depformer_text_emb.weight" in state_dict: + new_sd["audio_decoder.depth_text_emb.weight"] = state_dict.pop( + "depformer_text_emb.weight" + ) + + # Per-codebook depth audio embeddings (codebooks 1..num_codebooks-1) + for i in range(num_codebooks - 1): + key = f"depformer_emb.{i}.weight" + if key in state_dict: + new_sd[f"audio_decoder.depth_emb.{i}.weight"] = state_dict.pop(key) + + # Stack per-codebook input projections: (num_codebooks, depformer_dim, hidden_size) + in_projs = [] + for i in range(num_codebooks): + key = f"depformer_in.{i}.weight" + if key in state_dict: + in_projs.append(state_dict.pop(key)) + if in_projs: + new_sd["audio_decoder.stacked_depformer_in"] = torch.stack(in_projs, dim=0) + + # Stack output heads: (num_codebooks, audio_logits_size, depformer_dim) + out_heads = [] + for i in range(num_codebooks): + key = f"linears.{i}.weight" + if key in state_dict: + out_heads.append(state_dict.pop(key)) + if out_heads: + new_sd["audio_decoder.stacked_output_heads"] = torch.stack(out_heads, dim=0) + + # Depformer layers + for i in range(num_depformer_layers): + dep_prefix = f"depformer.layers.{i}" + dst = f"audio_decoder.layers.{i}" + + # RMSNorm + for norm in ("norm1", "norm2"): + alpha_key = f"{dep_prefix}.{norm}.alpha" + if alpha_key in state_dict: + new_sd[f"{dst}.{norm}.weight"] = state_dict.pop(alpha_key).flatten() + + # Attention: split packed QKV (3 * num_codebooks * depformer_dim, depformer_dim) + in_proj_key = f"{dep_prefix}.self_attn.in_proj_weight" + if in_proj_key in state_dict: + q, k, v = state_dict.pop(in_proj_key).chunk(3, dim=0) + new_sd[f"{dst}.self_attn.q_proj.weight"] = q + new_sd[f"{dst}.self_attn.k_proj.weight"] = k + new_sd[f"{dst}.self_attn.v_proj.weight"] = v + + # out_proj is stored transposed: [K*D, D] in HF vs [D, K*D] in mobius. + out_proj_key = f"{dep_prefix}.self_attn.out_proj.weight" + if out_proj_key in state_dict: + # Transpose from [K*D, D] to [D, K*D] for Linear(K*D, D).weight = [D, K*D] + new_sd[f"{dst}.self_attn.o_proj.weight"] = state_dict.pop(out_proj_key).T + + # Per-codebook gating: stack into stacked_{gate,up,down}_proj + gate_projs, up_projs, down_projs = [], [], [] + for j in range(num_codebooks): + lin_in_key = f"{dep_prefix}.gating.{j}.linear_in.weight" + if lin_in_key in state_dict: + gate, up = state_dict.pop(lin_in_key).chunk(2, dim=0) + gate_projs.append(gate) + up_projs.append(up) + + lin_out_key = f"{dep_prefix}.gating.{j}.linear_out.weight" + if lin_out_key in state_dict: + down_projs.append(state_dict.pop(lin_out_key)) + + if gate_projs: + new_sd[f"{dst}.stacked_gate_proj"] = torch.stack(gate_projs, dim=0) + new_sd[f"{dst}.stacked_up_proj"] = torch.stack(up_projs, dim=0) + if down_projs: + new_sd[f"{dst}.stacked_down_proj"] = torch.stack(down_projs, dim=0) + + # Pass through any remaining weights unchanged + new_sd.update(state_dict) + return new_sd diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 6c91764e..d1f48e52 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -21,12 +21,14 @@ __all__ = [ "AdapterTask", "AudioFeatureExtractionTask", + "AudioToAudioTask", "CausalLMTask", "CodecTask", "ControlNetTask", "DenoisingTask", "FeatureExtractionTask", "HybridCausalLMTask", + "MoshiTask", "HybridQwenVLTask", "ImageClassificationTask", "ModelTask", @@ -54,6 +56,7 @@ from mobius._constants import OPSET_VERSION from mobius.tasks._adapter import AdapterTask from mobius.tasks._audio_feature_extraction import AudioFeatureExtractionTask +from mobius.tasks._audio_to_audio import AudioToAudioTask, MoshiTask from mobius.tasks._base import ModelTask from mobius.tasks._causal_lm import ( CausalLMTask, @@ -90,6 +93,8 @@ TASK_REGISTRY: dict[str, type[ModelTask]] = { "adapter": AdapterTask, "audio-feature-extraction": AudioFeatureExtractionTask, + "audio-to-audio": AudioToAudioTask, + "moshi": MoshiTask, "codec": CodecTask, "controlnet": ControlNetTask, "denoising": DenoisingTask, diff --git a/src/mobius/tasks/_audio_to_audio.py b/src/mobius/tasks/_audio_to_audio.py new file mode 100644 index 00000000..19339f82 --- /dev/null +++ b/src/mobius/tasks/_audio_to_audio.py @@ -0,0 +1,380 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Audio-to-audio task for end-to-end speech models. + +Builds separate ONNX models for audio-to-audio pipelines like +LFM2-Audio and Moshi/PersonaPlex. These models take audio in and +produce audio out, with an intermediate language model backbone. + +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) +""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import ArchitectureConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ( + ModelTask, + _make_graph, + _make_hybrid_cache_inputs, + _make_kv_cache_inputs, + _make_model, + _register_hybrid_cache_outputs, + _register_kv_cache_outputs, +) + + +class AudioToAudioTask(ModelTask): + """Multi-model split for audio-to-audio models. + + The module must provide sub-modules as attributes: + + - ``audio_encoder``: audio encoder taking mel/waveform input + - ``embedding``: embedding model fusing text + audio features + - ``decoder``: language model backbone with KV cache + - ``audio_decoder`` (optional): depthformer or codec decoder + + Each sub-module is wired into its own ONNX graph. + Supports both standard KV cache and hybrid (conv+attention) cache + for the decoder, selected automatically based on ``config.layer_types``. + """ + + def build( + self, + module: nn.Module, + config: ArchitectureConfig, + ) -> ModelPackage: + models: dict[str, ir.Model] = {} + + models["audio_encoder"] = self._build_audio_encoder(module.audio_encoder, config) + models["embedding"] = self._build_embedding(module.embedding, config) + models["decoder"] = self._build_decoder(module.decoder, config) + + if hasattr(module, "audio_decoder"): + models["audio_decoder"] = self._build_audio_decoder(module.audio_decoder, config) + + return ModelPackage(models, config=config) + + def _build_audio_encoder( + self, + audio_encoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build audio encoder: mel (batch, n_mels, time) -> audio features.""" + batch = ir.SymbolicDim("batch") + mel_seq = ir.SymbolicDim("mel_sequence_len") + n_mels = (config.audio.num_mel_bins or 128) if config.audio else 128 + + input_features = ir.Value( + name="input_features", + shape=ir.Shape([batch, n_mels, mel_seq]), + type=ir.TensorType(config.dtype), + ) + + graph, builder = _make_graph([input_features], name="audio_encoder") + audio_features = audio_encoder(builder.op, input_features) + + audio_features.name = "audio_features" + graph.outputs.append(audio_features) + return _make_model(graph) + + def _build_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build embedding: text_ids -> inputs_embeds.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + + input_ids = ir.Value( + name="input_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph, builder = _make_graph([input_ids], name="embedding") + inputs_embeds = embedding( + builder.op, + input_ids=input_ids, + ) + + inputs_embeds.name = "inputs_embeds" + graph.outputs.append(inputs_embeds) + return _make_model(graph) + + def _build_decoder( + self, + decoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build decoder: inputs_embeds -> logits + cache. + + Automatically uses hybrid cache (conv+attention) when + ``config.layer_types`` is set, otherwise standard KV cache. + """ + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + + inputs_embeds = ir.Value( + name="inputs_embeds", + shape=ir.Shape([batch, seq_len, config.hidden_size]), + type=ir.TensorType(config.dtype), + ) + attention_mask = ir.Value( + name="attention_mask", + shape=ir.Shape([batch, "past_seq_len + seq_len"]), + type=ir.TensorType(ir.DataType.INT64), + ) + position_ids = ir.Value( + name="position_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph_inputs = [inputs_embeds, attention_mask, position_ids] + + # Select cache type based on config + use_hybrid = config.layer_types is not None and any( + lt != "full_attention" for lt in config.layer_types + ) + if use_hybrid: + cache_inputs, past_key_values = _make_hybrid_cache_inputs( + config, config.dtype, batch, past_seq_len + ) + else: + cache_inputs, past_key_values = _make_kv_cache_inputs( + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + graph_inputs.extend(cache_inputs) + + graph, builder = _make_graph(graph_inputs, name="decoder") + logits, present_key_values = decoder( + builder.op, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + + logits.name = "logits" + graph.outputs.append(logits) + if use_hybrid: + _register_hybrid_cache_outputs( + graph, + present_key_values, + config.layer_types or [], + ) + else: + _register_kv_cache_outputs(graph, present_key_values) + return _make_model(graph) + + def _build_audio_decoder( + self, + audio_decoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build audio decoder: backbone hidden -> per-codebook logits. + + The audio decoder (depthformer) takes a backbone hidden state and + produces logits for one codebook at a time. Runtime handles the + autoregressive loop over codebooks. + + Inputs: + backbone_hidden: (batch, 1, hidden_size) - LM output embedding + prev_embedding: (batch, 1, depthformer_dim) - previous codebook + codebook_idx: scalar int64 - which codebook to predict + past KV cache for depthformer layers + + Outputs: + codebook_logits: (batch, 1, audio_vocab_size) + present KV cache + """ + depthformer_dim = getattr(config, "depthformer_dim", config.hidden_size) + depthformer_layers = getattr(config, "depthformer_layers", 6) + depthformer_heads = getattr(config, "depthformer_heads", 16) + depthformer_head_dim = depthformer_dim // depthformer_heads + + batch = ir.SymbolicDim("batch") + past_seq_len = ir.SymbolicDim("past_sequence_len") + + backbone_hidden = ir.Value( + name="backbone_hidden", + shape=ir.Shape([batch, 1, config.hidden_size]), + type=ir.TensorType(config.dtype), + ) + prev_embedding = ir.Value( + name="prev_embedding", + shape=ir.Shape([batch, 1, depthformer_dim]), + type=ir.TensorType(config.dtype), + ) + codebook_idx = ir.Value( + name="codebook_idx", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph_inputs = [backbone_hidden, prev_embedding, codebook_idx] + + # Depthformer KV cache (all attention layers) + kv_inputs, past_key_values = _make_kv_cache_inputs( + depthformer_layers, + depthformer_heads, + depthformer_head_dim, + config.dtype, + batch, + past_seq_len, + prefix="past_key_values", + ) + graph_inputs.extend(kv_inputs) + + graph, builder = _make_graph(graph_inputs, name="audio_decoder") + codebook_logits, present_kv = audio_decoder( + builder.op, + backbone_hidden=backbone_hidden, + prev_embedding=prev_embedding, + codebook_idx=codebook_idx, + past_key_values=past_key_values, + ) + + codebook_logits.name = "codebook_logits" + graph.outputs.append(codebook_logits) + _register_kv_cache_outputs(graph, present_kv) + return _make_model(graph) + + +class MoshiTask(AudioToAudioTask): + """Multi-model split for Moshi/PersonaPlex audio-to-audio models. + + Differs from :class:`AudioToAudioTask`: + + - No ``audio_encoder``: Moshi consumes audio as codec token IDs, + not mel spectrograms, so no waveform encoder is needed. + - ``embedding`` accepts both ``input_ids`` (text) and ``audio_codes`` + (shape: batch x seq x num_codebooks) and returns ``inputs_embeds``. + - ``audio_decoder`` KV cache is sized with ``head_dim = depformer_dim`` + (one full-dimensioned head per codebook) rather than + ``depformer_dim // depformer_num_heads``. + """ + + def build( + self, + module: nn.Module, + config: ArchitectureConfig, + ) -> ModelPackage: + models: dict[str, ir.Model] = {} + + # Moshi has no audio encoder — audio input is codec token IDs. + models["embedding"] = self._build_embedding(module.embedding, config) + models["decoder"] = self._build_decoder(module.decoder, config) + + if hasattr(module, "audio_decoder"): + models["audio_decoder"] = self._build_audio_decoder(module.audio_decoder, config) + + return ModelPackage(models, config=config) + + def _build_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build Moshi embedding: (text_ids, audio_codes) -> inputs_embeds.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + num_codebooks = getattr(config, "num_codebooks", 16) + + input_ids = ir.Value( + name="input_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + audio_codes = ir.Value( + name="audio_codes", + shape=ir.Shape([batch, seq_len, num_codebooks]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph, builder = _make_graph([input_ids, audio_codes], name="embedding") + inputs_embeds = embedding(builder.op, input_ids=input_ids, audio_codes=audio_codes) + + inputs_embeds.name = "inputs_embeds" + graph.outputs.append(inputs_embeds) + return _make_model(graph) + + def _build_audio_decoder( + self, + audio_decoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build Moshi depformer: backbone_hidden -> codebook_logits. + + Uses ``head_dim = depformer_dim`` (one full-size head per codebook), + matching PersonaPlex's packed-QKV depformer attention weights. + """ + depformer_dim = getattr(config, "depformer_dim", 1024) + depformer_layers = getattr(config, "depformer_layers", 6) + # Each head in the depformer covers one full depformer dimension. + depformer_heads = getattr( + config, "depformer_num_heads", getattr(config, "num_codebooks", 16) + ) + depformer_head_dim = depformer_dim # full head_dim = depformer_dim + + batch = ir.SymbolicDim("batch") + past_seq_len = ir.SymbolicDim("past_sequence_len") + + backbone_hidden = ir.Value( + name="backbone_hidden", + shape=ir.Shape([batch, 1, config.hidden_size]), + type=ir.TensorType(config.dtype), + ) + prev_embedding = ir.Value( + name="prev_embedding", + shape=ir.Shape([batch, 1, depformer_dim]), + type=ir.TensorType(config.dtype), + ) + codebook_idx = ir.Value( + name="codebook_idx", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph_inputs = [backbone_hidden, prev_embedding, codebook_idx] + + kv_inputs, past_key_values = _make_kv_cache_inputs( + depformer_layers, + depformer_heads, + depformer_head_dim, + config.dtype, + batch, + past_seq_len, + prefix="past_key_values", + ) + graph_inputs.extend(kv_inputs) + + graph, builder = _make_graph(graph_inputs, name="audio_decoder") + codebook_logits, present_kv = audio_decoder( + builder.op, + backbone_hidden=backbone_hidden, + prev_embedding=prev_embedding, + codebook_idx=codebook_idx, + past_key_values=past_key_values, + ) + + codebook_logits.name = "codebook_logits" + graph.outputs.append(codebook_logits) + _register_kv_cache_outputs(graph, present_kv) + return _make_model(graph) diff --git a/src/mobius/tasks/_base.py b/src/mobius/tasks/_base.py index 22d72723..653a3aac 100644 --- a/src/mobius/tasks/_base.py +++ b/src/mobius/tasks/_base.py @@ -20,8 +20,9 @@ _FUNCTIONS_DOMAIN = "pkg.mobius" # Cache state pair: (key, value) or (conv_state, ssm_state) for stateful -# layers. MLP-only layers are stateless and use (None, None). -StatePair = tuple[ir.Value, ir.Value] | tuple[None, None] +# layers. Single-state layers (conv/lightning) use a 1-tuple. +# MLP-only layers are stateless and use (None, None). +StatePair = tuple[ir.Value, ir.Value] | tuple[ir.Value] | tuple[None, None] class LinearAttentionDims(NamedTuple): @@ -194,6 +195,7 @@ def _make_hybrid_cache_inputs( Supported layer types: ``"full_attention"`` — standard KV cache (key + value). + ``"conv"`` — ShortConv conv_state only. ``"linear_attention"`` (DeltaNet) — conv_state + recurrent_state. ``"mamba"`` / ``"mamba2"`` — conv_state + ssm_state. ``"mlp"`` — stateless, produces ``(None, None)`` pair. @@ -262,6 +264,17 @@ def _make_hybrid_cache_inputs( ) flat.extend([conv_state, rec_state]) pairs.append((conv_state, rec_state)) + elif ltype == "conv": + # ShortConv layers: conv_state only (no SSM state) + # State: (batch, hidden_size, short_conv_kernel - 1) + short_conv_kernel = getattr(config, "short_conv_kernel", 3) + conv_state = ir.Value( + name=f"{prefix}.{i}.conv_state", + shape=ir.Shape([batch, config.hidden_size, short_conv_kernel - 1]), + type=ir.TensorType(dtype), + ) + flat.append(conv_state) + pairs.append((conv_state,)) # 1-tuple: conv has no second state elif ltype == "mlp": # MLP-only layers are stateless — no cache inputs needed pairs.append((None, None)) @@ -338,6 +351,11 @@ def _register_hybrid_cache_outputs( (state_a,) = states state_a.name = f"{prefix}.{i}.recurrent_state" graph.outputs.append(state_a) + elif ltype == "conv": + # ShortConv: single conv_state only + (state_a,) = states + state_a.name = f"{prefix}.{i}.conv_state" + graph.outputs.append(state_a) else: state_a, state_b = states if ltype == "linear_attention": diff --git a/testdata/cases/audio/lfm2-audio-1.5b.yaml b/testdata/cases/audio/lfm2-audio-1.5b.yaml new file mode 100644 index 00000000..0b79c012 --- /dev/null +++ b/testdata/cases/audio/lfm2-audio-1.5b.yaml @@ -0,0 +1,9 @@ +model_id: "LiquidAI/LFM2-Audio-1.5B" +revision: "main" +task_type: "audio-to-audio" +dtype: "float32" + +level: "L4" + +skip_reason: "LFM2-Audio 1.5B — requires custom audio pipeline and no standard golden generation." +notes: "LFM2-Audio 1.5B. Hybrid conv+attention backbone with ConformerEncoder for audio-to-audio." diff --git a/testdata/cases/causal-lm/bitnet-b1.58-2b.yaml b/testdata/cases/causal-lm/bitnet-b1.58-2b.yaml new file mode 100644 index 00000000..1b941a6d --- /dev/null +++ b/testdata/cases/causal-lm/bitnet-b1.58-2b.yaml @@ -0,0 +1,20 @@ +model_id: "microsoft/bitnet-b1.58-2B-4T" +revision: "04c3b9ad9361b824064a1f25ea60a8be9599b127" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +skip_reason: "Ternary weights require custom unpacking; golden JSON not yet generated." +notes: > + BitNet b1.58 2B-4T. Ternary {-1,0,+1} weight LLM with sub-layer + RMSNorms (attn_sub_norm, ffn_sub_norm) and squared ReLU activation. + Weights stored as 2-bit packed uint8 with per-tensor weight_scale. diff --git a/testdata/cases/causal-lm/lfm2-1.2b.yaml b/testdata/cases/causal-lm/lfm2-1.2b.yaml new file mode 100644 index 00000000..593ac60e --- /dev/null +++ b/testdata/cases/causal-lm/lfm2-1.2b.yaml @@ -0,0 +1,17 @@ +model_id: "LiquidAI/LFM2-1.2B" +revision: "main" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +skip_reason: "LFM2 1.2B — too large for CI golden data generation." +notes: "LFM2 1.2B. Hybrid ShortConv + attention architecture from Liquid AI." diff --git a/tests/_test_configs.py b/tests/_test_configs.py index d2fe8354..c1f44dc0 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -29,6 +29,7 @@ GraniteMoeHybridConfig, JambaConfig, JetMoeConfig, + Lfm2Config, LongcatFlashConfig, Mamba2Config, MambaConfig, @@ -100,6 +101,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ("llama", {}, True), ("mistral", {}, False), ("qwen2", {}, True), + ("bitnet", {"hidden_act": "relu2", "tie_word_embeddings": True}, True), ("cohere", {"tie_word_embeddings": True, "logit_scale": 0.0625}, True), ("cohere2", {"tie_word_embeddings": True, "logit_scale": 0.0625}, False), ("diffllama", {}, False), @@ -1078,6 +1080,36 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + # lfm2: hybrid ShortConv+Attention (requires Lfm2Config) + ( + "lfm2", + { + "_config_cls": Lfm2Config, + "num_hidden_layers": 4, + "layer_types": [ + "conv", + "conv", + "full_attention", + "conv", + ], + "attn_qk_norm": True, + "short_conv_kernel": 3, + "short_conv_bias": False, + }, + True, + ), + # lfm2: all attention layers (no conv) + ( + "lfm2", + { + "_config_cls": Lfm2Config, + "layer_types": ["full_attention"] * TINY_LAYERS, + "attn_qk_norm": True, + "short_conv_kernel": 3, + "short_conv_bias": False, + }, + False, + ), # gemma3n_text: all full attention (no sliding window) ( "gemma3n_text", @@ -2085,6 +2117,16 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ] +# --------------------------------------------------------------------------- +# Audio-to-audio model configs +# --------------------------------------------------------------------------- +AUDIO_TO_AUDIO_CONFIGS: list[tuple[str, dict, bool]] = [ + # Audio-to-audio models are tested via explicit test classes + # (TestBuildLfm2AudioGraph etc.) since they produce multi-model + # packages with different keys than standard single-model tasks. +] + + # --------------------------------------------------------------------------- # Aggregate lists # --------------------------------------------------------------------------- @@ -2097,6 +2139,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: + SSM_CONFIGS + VL_CONFIGS + SPEECH_CONFIGS + + AUDIO_TO_AUDIO_CONFIGS ) # Model types explicitly declared in configs above (may have duplicates — diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index c9c421e3..14cb692b 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -270,6 +270,11 @@ def test_graph_builds_without_weights(self, model_type: str, config_overrides: d assert f"present.{i}.ssm_state" in output_names, ( f"Missing present.{i}.ssm_state" ) + elif ltype == "conv": + # ShortConv: single conv_state only + assert f"present.{i}.conv_state" in output_names, ( + f"Missing present.{i}.conv_state" + ) else: assert f"present.{i}.key" in output_names, f"Missing present.{i}.key" assert f"present.{i}.value" in output_names, f"Missing present.{i}.value" @@ -3346,6 +3351,10 @@ def test_jamba_preprocess_weights_moe_renames(self): # Hybrid SSM+Attention dedicated tests "bamba", "jamba", + # Audio-to-audio dedicated tests + "lfm2_audio", + "moshi", + "personaplex", } # Registered model types that truly have no test coverage yet. @@ -3700,3 +3709,255 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: task = get_task(_default_task_for_model(model_type)) pkg = task.build(module, config) _assert_outputs_have_shapes_and_dtypes(pkg, model_type) + + +class TestBuildLfm2AudioGraph: + """Verify LFM2-Audio 4-model split builds correctly.""" + + def _lfm2_audio_config(self): + from mobius._configs import AudioConfig, Lfm2AudioConfig + + return Lfm2AudioConfig( + vocab_size=TINY_VOCAB, + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_hidden_layers=4, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_KV_HEADS, + head_dim=TINY_HEAD_DIM, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=1_000_000.0, + pad_token_id=0, + layer_types=["conv", "conv", "full_attention", "conv"], + attn_qk_norm=True, + short_conv_kernel=3, + short_conv_bias=False, + depthformer_layers=2, + depthformer_dim=TINY_HIDDEN, + depthformer_heads=TINY_HEADS, + num_codebooks=2, + audio_vocab_size=32, + audio=AudioConfig( + attention_dim=TINY_HIDDEN, + attention_heads=TINY_HEADS, + num_blocks=2, + linear_units=TINY_INTERMEDIATE, + kernel_size=3, + conv_channels=32, + t5_bias_max_distance=100, + num_mel_bins=16, + output_dim=TINY_HIDDEN, + ), + ) + + def test_4_model_package_structure(self): + """Build LFM2-Audio and verify 4-model split.""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + + assert "audio_encoder" in pkg, "Should have audio_encoder model" + assert "embedding" in pkg, "Should have embedding model" + assert "decoder" in pkg, "Should have decoder model" + assert "audio_decoder" in pkg, "Should have audio_decoder model" + + def test_audio_encoder_io(self): + """Verify audio_encoder: input_features -> audio_features.""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + + enc = pkg["audio_encoder"] + enc_inputs = {inp.name for inp in enc.graph.inputs} + enc_outputs = {out.name for out in enc.graph.outputs} + assert "input_features" in enc_inputs + assert "audio_features" in enc_outputs + + def test_decoder_hybrid_cache(self): + """Verify decoder has hybrid cache (conv_state + KV).""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + + dec = pkg["decoder"] + dec_inputs = {inp.name for inp in dec.graph.inputs} + dec_outputs = {out.name for out in dec.graph.outputs} + + assert "inputs_embeds" in dec_inputs + assert "logits" in dec_outputs + # Hybrid cache: conv layers have conv_state, attn has key/value + assert "past_key_values.0.conv_state" in dec_inputs + assert "present.2.key" in dec_outputs + + def test_audio_decoder_io(self): + """Verify audio_decoder: backbone_hidden -> codebook_logits.""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + + adec = pkg["audio_decoder"] + adec_inputs = {inp.name for inp in adec.graph.inputs} + adec_outputs = {out.name for out in adec.graph.outputs} + + assert "backbone_hidden" in adec_inputs + assert "prev_embedding" in adec_inputs + assert "codebook_idx" in adec_inputs + assert "codebook_logits" in adec_outputs + + def test_onnx_checker_passes(self): + """Run ONNX checker on all 4 sub-models.""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + _run_onnx_checker(pkg, "lfm2_audio") + + def test_outputs_have_shapes_and_dtypes(self): + """Verify all sub-model outputs have shape/dtype info.""" + from mobius.models.lfm2_audio import Lfm2AudioModel + from mobius.tasks._audio_to_audio import AudioToAudioTask + + config = self._lfm2_audio_config() + module = Lfm2AudioModel(config) + task = AudioToAudioTask() + pkg = task.build(module, config) + _assert_outputs_have_shapes_and_dtypes(pkg, "lfm2_audio") + + +class TestBuildMoshiGraph: + """Verify Moshi/PersonaPlex 3-model split builds correctly.""" + + def _moshi_config(self): + from mobius._configs import MoshiConfig + + return MoshiConfig( + vocab_size=TINY_VOCAB, + hidden_size=TINY_HIDDEN, + intermediate_size=TINY_INTERMEDIATE, + num_hidden_layers=TINY_LAYERS, + num_attention_heads=TINY_HEADS, + num_key_value_heads=TINY_HEADS, + head_dim=TINY_HEAD_DIM, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-5, + rope_type="default", + rope_theta=10000.0, + depformer_dim=32, + depformer_layers=2, + depformer_num_heads=2, + depformer_intermediate_size=32, + num_codebooks=2, + audio_vocab_size=10, + ) + + def test_3_model_package_structure(self): + """Build Moshi and verify 3-model split (no audio_encoder).""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + + assert "embedding" in pkg, "Should have embedding model" + assert "decoder" in pkg, "Should have decoder model" + assert "audio_decoder" in pkg, "Should have audio_decoder model" + assert "audio_encoder" not in pkg, "Moshi has no audio_encoder" + + def test_embedding_io(self): + """Verify embedding: (input_ids, audio_codes) -> inputs_embeds.""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + + emb = pkg["embedding"] + emb_inputs = {inp.name for inp in emb.graph.inputs} + emb_outputs = {out.name for out in emb.graph.outputs} + assert "input_ids" in emb_inputs + assert "audio_codes" in emb_inputs + assert "inputs_embeds" in emb_outputs + + def test_decoder_kv_cache(self): + """Verify decoder has standard KV cache.""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + + dec = pkg["decoder"] + dec_inputs = {inp.name for inp in dec.graph.inputs} + dec_outputs = {out.name for out in dec.graph.outputs} + assert "inputs_embeds" in dec_inputs + assert "logits" in dec_outputs + assert any("past_key_values" in n for n in dec_inputs) + + def test_audio_decoder_io(self): + """Verify audio_decoder: backbone_hidden -> codebook_logits.""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + + adec = pkg["audio_decoder"] + adec_inputs = {inp.name for inp in adec.graph.inputs} + adec_outputs = {out.name for out in adec.graph.outputs} + assert "backbone_hidden" in adec_inputs + assert "prev_embedding" in adec_inputs + assert "codebook_idx" in adec_inputs + assert "codebook_logits" in adec_outputs + + def test_onnx_checker_passes(self): + """Run ONNX checker on all 3 sub-models.""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + _run_onnx_checker(pkg, "moshi") + + def test_outputs_have_shapes_and_dtypes(self): + """Verify all sub-model outputs have shape/dtype info.""" + from mobius.models.moshi import MoshiModel + from mobius.tasks._audio_to_audio import MoshiTask + + config = self._moshi_config() + module = MoshiModel(config) + task = MoshiTask() + pkg = task.build(module, config) + _assert_outputs_have_shapes_and_dtypes(pkg, "moshi") diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 559d5ab5..5ddd8d44 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -181,6 +181,7 @@ def _all_registered_with_test_id() -> dict[str, str]: "qwen3_vl": "VL model — requires image inputs", # --- Audio / speech models (require audio inputs) --- "data2vec-audio": "Audio model — requires audio inputs", + "lfm2_audio": "Multi-model audio architecture — tested via TestBuildLfm2AudioGraph", "hubert": "Audio model — requires audio inputs", "musicgen": "Audio model — requires audio inputs", "seamless_m4t": "Audio model — requires audio inputs", diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 607c0739..3593fa2f 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -87,6 +87,10 @@ # Zamba weight-tying references layers.2.shared_transf (the third layer) but # the tiny config only has 2 layers — HF tie_weights validation crashes. "zamba": "Zamba weight-tying requires num_layers > 2; tiny 2-layer config causes HF tie_weights error", + # Moshi/PersonaPlex: audio-to-audio models using the 'moshi' library, not transformers. + # config.json is minimal (only model_type + version), HF AutoModelForCausalLM cannot load. + "personaplex": "PersonaPlex uses moshi library; config.json not compatible with HF AutoConfig", + "moshi": "Moshi uses moshi library; config.json not compatible with HF AutoConfig", } # Per-model atol overrides for L3 synthetic parity.