Add NeMo .nemo integration + FastConformer-RNNT streaming ASR model (#359)

## Summary

Adds support for **`nvidia/nemotron-speech-streaming-en-0.6b`** (a
FastConformer-CacheAware-RNNT streaming ASR model distributed as a NeMo
`.nemo` archive) in two layers:

1. **New `.nemo` integration** (`src/mobius/integrations/nemo/`) —
analogous to the existing GGUF integration. An archive reader
(`NeMoArchive`), a config mapping (`model_config.yaml` →
`ArchitectureConfig`), and a `build_from_nemo` pipeline. Resolves local
paths and HuggingFace Hub references.

2. **FastConformer-RNNT model** (`models/nemo_rnnt.py` +
`tasks/_rnnt.py`) — emitted as three ONNX sub-models (encoder /
prediction decoder / joint) via a new `RNNTTask`.

## Architecture

- **Encoder:** 8× causal `dw_striding` conv subsampling → 24 Conformer
layers (relative-position MHA with `rel_shift`, FF/attn/conv/FF macaron
structure, layer-norm conv module). Positional encodings built
dynamically in-graph.
- **Prediction:** embedding + 2-layer LSTM with PyTorch→ONNX gate
reordering; zero start-of-sequence embedding (NeMo `add_sos=True`).
- **Joint:** enc/pred projections → ReLU → vocab logits → log-softmax.

## Validation

| Level | Coverage | Result |
|-------|----------|--------|
| L1 | Graph build + I/O shapes (random weights) | 7 pass |
| L4 | Parity vs committed NeMo golden | encoder ~5e-7, decoder ~5e-7,
joint ~2.7e-5 |
| L5 | Greedy decode + incremental-vs-one-shot decoder state | pass |

Parity validated against `nemo_toolkit` 2.7.3 reference.

## Contract / limitations

- **Offline full-context only** — the encoder consumes the full feature
sequence (causal convs, no cache-aware streaming state); chunked
streaming is out of scope for this export.
- **batch=1 / equal-length** — no length input/output yet.
- **fp32 only** — the builder rejects non-fp32 (graph emits fp32
constants); unsupported NeMo encoder variants
(non-`dw_striding`/`rel_pos`/`layer_norm`, `xscaling`) are rejected with
a clear error.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

.github/copilot-instructions.md

@@ -98,8 +98,9 @@ and wrapper modules for nesting. See the `weight-name-alignment` skill.
 - Ruff for linting/formatting (line-length=95, Python 3.10+)
 - MyPy strict mode (excludes `*_test.py`)
 - Use `op.Shape(x, start=i, end=i+1)` for single dimension extraction
-- Use ONNX opset 23 `op.Attention` with `q_num_heads`/`kv_num_heads`
-  attributes (not `num_heads`)
+- Use the ONNX opset-24 `op.Attention` (introduced in opset 23) with
+  `q_num_heads`/`kv_num_heads` attributes (not `num_heads`). The codebase
+  emits opset 24 graphs (`OPSET_VERSION` in `src/mobius/_constants.py`).
 
 ### Protobuf prohibition
 

docs/design/gguf-support-proposal.md

@@ -779,7 +779,7 @@ the fp16/fp32 activation, `B` is the packed uint8 weight blob.
    to fuse (e.g., unusual shape, unsupported config), performance
    degrades to "dequantize then fp matmul" — 2-4× slower.
 
-4. **INT4 requires opset 21+**: Our codebase uses opset 23, so this
+4. **INT4 requires opset 21+**: Our codebase uses opset 24, so this
    isn't a blocker, but older runtimes (pre-2024) can't load int4 QDQ
    models.
 

examples/nemotron_fastconformer_rnnt.py

@@ -0,0 +1,107 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Regenerate the FastConformer-RNNT golden reference used by the L4 parity test.
+
+This produces ``testdata/golden/speech/nemotron_fastconformer_rnnt.npz`` by
+running the *real* NeMo model through the NeMo toolkit (the ground-truth
+reference implementation). It must be run inside an environment that has
+``nemo_toolkit`` installed (it is **not** a mobius runtime dependency)::
+
+    python -m venv /tmp/nemo_ref_venv
+    source /tmp/nemo_ref_venv/bin/activate
+    pip install "nemo_toolkit[asr]==2.7.3"
+    python scripts/generate_nemo_rnnt_golden.py \
+        --model nvidia/nemotron-speech-streaming-en-0.6b \
+        --revision 7a9b763e6c5fb103da690219c049fac917aa50b1 \
+        --out testdata/golden/speech/nemotron_fastconformer_rnnt.npz
+
+The committed ``.npz`` stores only the arrays needed by the parity test plus a
+``meta`` JSON blob (model id, revision, NeMo version, dtype, seed, ids) so the
+reference is self-describing and auditable.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import numpy as np
+import torch
+
+# Deterministic input feature / token fixtures (also recorded in metadata).
+_SEED = 0
+_T = 131
+_FEAT_DIM = 128
+_TOKENS = [3, 5, 7, 9]
+_SOS_ID = 1025  # rnnt_num_classes (1024) + 1; zero start-of-sequence embedding
+_BLANK_ID = 1024
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="nvidia/nemotron-speech-streaming-en-0.6b")
+    parser.add_argument(
+        "--revision",
+        default="7a9b763e6c5fb103da690219c049fac917aa50b1",
+        help="HuggingFace Hub commit SHA to pin the reference model.",
+    )
+    parser.add_argument(
+        "--out",
+        default="testdata/golden/speech/nemotron_fastconformer_rnnt.npz",
+    )
+    args = parser.parse_args()
+
+    import nemo  # type: ignore[import-not-found]
+    import nemo.collections.asr as nemo_asr  # type: ignore[import-not-found]
+    from huggingface_hub import hf_hub_download
+
+    torch.manual_seed(_SEED)
+
+    nemo_path = hf_hub_download(
+        repo_id=args.model,
+        filename="nemotron-speech-streaming-en-0.6b.nemo",
+        revision=args.revision,
+    )
+    model = nemo_asr.models.ASRModel.restore_from(nemo_path, map_location="cpu")
+    model.eval()
+
+    feats = torch.randn(1, _FEAT_DIM, _T)
+    length = torch.tensor([_T])
+    tokens = torch.tensor([_TOKENS])
+    tlen = torch.tensor([len(_TOKENS)])
+
+    with torch.no_grad():
+        enc_out, _ = model.encoder(audio_signal=feats, length=length)
+        # decoder.predict(add_sos=True) prepends a zero SOS vector -> (B, H, U+1)
+        pred_out = model.decoder(targets=tokens, target_length=tlen)[0]
+        # Bypass fuse_loss_wer by calling the inner joint on (B, T, D) layout.
+        joint_out = model.joint.joint(enc_out.transpose(1, 2), pred_out.transpose(1, 2))
+
+    meta = {
+        "model_id": args.model,
+        "revision": args.revision,
+        "nemo_version": nemo.__version__,
+        "dtype": "float32",
+        "seed": _SEED,
+        "feat_dim": _FEAT_DIM,
+        "input_frames": _T,
+        "tokens": _TOKENS,
+        "sos_id": _SOS_ID,
+        "blank_id": _BLANK_ID,
+    }
+
+    np.savez_compressed(
+        args.out,
+        feats=feats.numpy().astype(np.float32),
+        enc_out=enc_out.numpy().astype(np.float32),
+        tokens=tokens.numpy().astype(np.int64),
+        pred_out=pred_out.numpy().astype(np.float32),
+        joint_out=joint_out.numpy().astype(np.float32),
+        meta=np.array(json.dumps(meta)),
+    )
+    print(f"saved {args.out}\n{json.dumps(meta, indent=2)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_nemo_rnnt_golden.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Regenerate the cache-aware streaming golden for the FastConformer-RNNT test.
+
+This produces ``testdata/golden/speech/nemotron_fastconformer_rnnt_streaming.npz``
+by driving the *real* NeMo encoder's streaming ``forward`` (with explicit
+``cache_last_channel`` / ``cache_last_time`` / ``cache_last_channel_len`` state)
+over two consecutive feature chunks. It must run inside an environment that has
+``nemo_toolkit`` installed (not a mobius runtime dependency)::
+
+    python -m venv /tmp/nemo_ref_venv
+    source /tmp/nemo_ref_venv/bin/activate
+    pip install "nemo_toolkit[asr]==2.7.3"
+    python scripts/generate_nemo_rnnt_streaming_golden.py \
+        --model nvidia/nemotron-speech-streaming-en-0.6b \
+        --revision 7a9b763e6c5fb103da690219c049fac917aa50b1 \
+        --out testdata/golden/speech/nemotron_fastconformer_rnnt_streaming.npz
+
+To keep the committed reference small, only the per-chunk feature inputs and the
+encoder outputs / lengths / cache-length scalars are stored (not the full
+multi-megabyte cache tensors). The streaming parity test validates cache
+correctness implicitly by chaining chunk-0's ONNX output caches into chunk-1 and
+matching chunk-1's encoder output against this NeMo reference.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import numpy as np
+import torch
+
+_SEED = 0
+_FEAT_DIM = 128
+_CHUNK = 120  # feature frames per streaming chunk
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="nvidia/nemotron-speech-streaming-en-0.6b")
+    parser.add_argument(
+        "--revision",
+        default="7a9b763e6c5fb103da690219c049fac917aa50b1",
+        help="HuggingFace Hub commit SHA to pin the reference model.",
+    )
+    parser.add_argument(
+        "--out",
+        default="testdata/golden/speech/nemotron_fastconformer_rnnt_streaming.npz",
+    )
+    args = parser.parse_args()
+
+    import nemo  # type: ignore[import-not-found]
+    import nemo.collections.asr as nemo_asr  # type: ignore[import-not-found]
+    from huggingface_hub import hf_hub_download
+
+    torch.manual_seed(_SEED)
+
+    nemo_path = hf_hub_download(
+        repo_id=args.model,
+        filename="nemotron-speech-streaming-en-0.6b.nemo",
+        revision=args.revision,
+    )
+    model = nemo_asr.models.ASRModel.restore_from(nemo_path, map_location="cpu")
+    model.eval()
+    enc = model.encoder
+    enc.setup_streaming_params()
+
+    def step(feats, ch, ct, cl):
+        with torch.no_grad():
+            return enc(
+                audio_signal=feats,
+                length=torch.tensor([_CHUNK]),
+                cache_last_channel=ch,
+                cache_last_time=ct,
+                cache_last_channel_len=cl,
+            )
+
+    ch0, ct0, cl0 = enc.get_initial_cache_state(batch_size=1)
+    f0 = torch.randn(1, _FEAT_DIM, _CHUNK)
+    out0, len0, ch1, ct1, cl1 = step(f0, ch0, ct0, cl0)
+    f1 = torch.randn(1, _FEAT_DIM, _CHUNK)
+    out1, len1, ch2, ct2, cl2 = step(f1, ch1, ct1, cl1)
+    del ch2, ct2  # full out-caches of the second chunk are not part of the golden
+
+    meta = {
+        "model_id": args.model,
+        "revision": args.revision,
+        "nemo_version": nemo.__version__,
+        "dtype": "float32",
+        "seed": _SEED,
+        "feat_dim": _FEAT_DIM,
+        "chunk_frames": _CHUNK,
+        "last_channel_cache_size": int(enc.streaming_cfg.last_channel_cache_size),
+        "drop_extra_pre_encoded": int(enc.streaming_cfg.drop_extra_pre_encoded),
+    }
+
+    np.savez_compressed(
+        args.out,
+        f0=f0.numpy().astype(np.float32),
+        f1=f1.numpy().astype(np.float32),
+        out0=out0.numpy().astype(np.float32),
+        out1=out1.numpy().astype(np.float32),
+        len0=len0.numpy().astype(np.int64),
+        len1=len1.numpy().astype(np.int64),
+        cl1=cl1.numpy().astype(np.int64),
+        cl2=cl2.numpy().astype(np.int64),
+        meta=np.array(json.dumps(meta)),
+    )
+    print(f"saved {args.out}\n{json.dumps(meta, indent=2)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_nemo_rnnt_streaming_golden.py

@@ -36,6 +36,7 @@
     "build_diffusers_pipeline",
     "build_from_gguf",
     "build_from_module",
+    "build_from_nemo",
     "components",
     "ep_capabilities",
     "ep_registry",
@@ -89,4 +90,5 @@
 )
 from mobius._weight_loading import apply_weights
 from mobius.integrations.gguf import build_from_gguf
+from mobius.integrations.nemo import build_from_nemo
 from mobius.tasks import CausalLMTask, ModelTask

src/mobius/__init__.py

@@ -161,6 +161,21 @@ def _cmd_build(args: argparse.Namespace) -> None:
             _save_package(pkg, output_dir, args, optimize, component_filter)
             return
 
+    # Auto-detect NeMo .nemo archives (local file or HF ref like
+    # 'owner/repo:model.nemo'). Routes to the NeMo import path; reuses the
+    # standard build args (--dtype, --ep, --external-data) and save logic.
+    if args.model and args.model.endswith(".nemo"):
+        from mobius.integrations.nemo import build_from_nemo
+
+        print(f"Detected NeMo archive: {args.model}")
+        pkg = build_from_nemo(
+            args.model,
+            dtype=dtype_override,
+            execution_provider=execution_provider,
+        )
+        _save_package(pkg, output_dir, args, optimize, component_filter)
+        return
+
     # Build from HuggingFace model ID or local config
     if args.config:
         import transformers

src/mobius/__main__.py

@@ -432,6 +432,30 @@ class ArchitectureConfig(BaseModelConfig):
     audio_t5_bias_max_distance: int | None = None
     audio_token_id: int | None = None
 
+    # FastConformer-RNNT (NeMo) config — see models/nemo_rnnt.py.
+    # Encoder
+    fastconformer_subsampling_factor: int = 8
+    fastconformer_subsampling_conv_channels: int = 256
+    fastconformer_conv_kernel_size: int = 9
+    fastconformer_pos_emb_max_len: int = 5000
+    fastconformer_xscaling: bool = False
+    # Number of input mel features to the encoder.
+    fastconformer_feat_in: int = 128
+    # Chunked-limited attention context: [left_context, right_context] in frames.
+    fastconformer_att_context_size: tuple[int, int] = (70, 13)
+    # Cache-aware streaming: per-layer last-channel (attention) cache length in
+    # subsampled frames (NeMo ``last_channel_cache_size`` = att_context left) and
+    # the number of leading subsampled frames dropped per chunk (NeMo
+    # ``drop_extra_pre_encoded``).
+    fastconformer_streaming_cache_size: int = 70
+    fastconformer_streaming_drop_extra: int = 2
+    # RNN-T prediction network + joint
+    rnnt_pred_hidden: int | None = None
+    rnnt_pred_rnn_layers: int = 1
+    rnnt_joint_hidden: int | None = None
+    # Number of acoustic classes excluding the blank symbol (vocab without blank).
+    rnnt_num_classes: int | None = None
+
     # LoRA config (for multimodal models like Phi4-MM)
     speech_lora: dict | None = None
 

src/mobius/_configs/_base.py

@@ -37,6 +37,7 @@
     DeepSeekV3CausalLMModel,
     DiffLlamaCausalLMModel,
     DogeCausalLMModel,
+    EncDecRNNTModel,
     Ernie45MoECausalLMModel,
     ErnieCausalLMModel,
     ExaOne4CausalLMModel,
@@ -669,6 +670,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "wav2vec2-conformer": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"),
     "wavlm": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"),
     "mms": ModelRegistration(Wav2Vec2ForCTCModel, task="ctc-asr", config_class=MMSConfig),
+    "fastconformer_rnnt": ModelRegistration(EncDecRNNTModel, task="fastconformer-rnnt"),
 }
 
 

src/mobius/_registry.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""NeMo ``.nemo`` model import support for mobius.
+
+This package loads NVIDIA NeMo ``.nemo`` archives and converts them to ONNX
+models using the standard graph construction pipeline.
+
+Usage::
+
+    from mobius.integrations.nemo import build_from_nemo
+
+    pkg = build_from_nemo("nvidia/nemotron-speech-streaming-en-0.6b")
+"""
+
+from __future__ import annotations
+
+from mobius.integrations.nemo._builder import build_from_nemo
+from mobius.integrations.nemo._genai_config import write_genai_bundle
+
+__all__ = ["build_from_nemo", "write_genai_bundle"]