#!/usr/bin/env python3 """Export nvidia/canary-qwen-2.5b (NeMo SALM) to the onnx-asr "speech-llm" contract. Graphs ------ encoder.onnx input_features (1,128,T) + input_features_lens (1,) -> audio_embeds (1,L,2048) + audio_embeds_lens (1,) embed_tokens.onnx input_ids (1,S) -> inputs_embeds (1,S,2048) decoder.onnx inputs_embeds (1,S,2048) + attn_bias (1,1,S,P+S) + position_ids (1,S) + past_key_values.{i}.{key,value} (1,8,P,128) -> logits (1,S,V) + present.{i}.{key,value} (1,8,P+S,128) """ from __future__ import annotations import argparse import json import shutil from pathlib import Path import torch from torch import nn REPO = "nvidia/canary-qwen-2.5b" USER_PROMPT = "Transcribe the following: " class EncoderExport(nn.Module): """FastConformer encoder + modality adapter + linear projection to LLM space.""" def __init__(self, perception: nn.Module): super().__init__() self.encoder = perception.encoder self.modality_adapter = perception.modality_adapter self.proj = perception.proj def forward(self, input_features: torch.Tensor, input_features_lens: torch.Tensor): encoded, encoded_len = self.encoder(audio_signal=input_features, length=input_features_lens) encoded, encoded_len = self.modality_adapter(audio_signal=encoded, length=encoded_len) return self.proj(encoded.transpose(1, 2)), encoded_len class EmbedExport(nn.Module): def __init__(self, embed_tokens: nn.Module): super().__init__() self.embed_tokens = embed_tokens def forward(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) class DecoderExport(nn.Module): """Qwen3 causal LM over inputs_embeds with a flat past/present KV cache.""" def __init__(self, language_model: nn.Module, lm_head: nn.Module, num_layers: int): super().__init__() self.language_model = language_model self.lm_head = lm_head self.num_layers = num_layers def forward(self, inputs_embeds, attn_bias, position_ids, *past): from transformers.cache_utils import DynamicCache cache = DynamicCache() for i in range(self.num_layers): cache.update(past[2 * i], past[2 * i + 1], i) out = self.language_model( inputs_embeds=inputs_embeds, attention_mask=attn_bias, position_ids=position_ids, past_key_values=cache, use_cache=True, ) logits = self.lm_head(out.last_hidden_state) present: list[torch.Tensor] = [] for i in range(self.num_layers): if hasattr(cache, "layers"): present.append(cache.layers[i].keys) present.append(cache.layers[i].values) else: present.append(cache.key_cache[i]) present.append(cache.value_cache[i]) return (logits, *present) MAX_PROTOBUF = 1_900_000_000 def consolidate(src: Path, dst: Path) -> None: import onnx raw_size = sum(p.stat().st_size for p in src.parent.rglob("*") if p.is_file()) model = onnx.load(str(src), load_external_data=True) dst.parent.mkdir(parents=True, exist_ok=True) data_path = dst.with_suffix(".onnx_data") if data_path.exists(): data_path.unlink() if raw_size > MAX_PROTOBUF: onnx.save( model, str(dst), save_as_external_data=True, all_tensors_to_one_file=True, location=data_path.name, size_threshold=1024, convert_attribute=False, ) else: onnx.save(model, str(dst)) def export_graph(module, inputs, name: str, out: Path, tmp: Path, **kwargs) -> None: stage = tmp / name stage.mkdir(parents=True, exist_ok=True) torch.onnx.export(module, inputs, str(stage / f"{name}.onnx"), do_constant_folding=True, dynamo=False, **kwargs) consolidate(stage / f"{name}.onnx", out / f"{name}.onnx") shutil.rmtree(stage, ignore_errors=True) def quantize(name: str, out: Path, tmp: Path) -> None: from onnxruntime.quantization import QuantType, quantize_dynamic stage = tmp / f"{name}_int8" stage.mkdir(parents=True, exist_ok=True) big = (out / f"{name}.onnx").stat().st_size + ( (out / f"{name}.onnx_data").stat().st_size if (out / f"{name}.onnx_data").exists() else 0 ) > 1_500_000_000 quantize_dynamic( out / f"{name}.onnx", stage / f"{name}_int8.onnx", weight_type=QuantType.QInt8, use_external_data_format=big, extra_options={"MatMulConstBOnly": True}, ) consolidate(stage / f"{name}_int8.onnx", out / f"{name}_int8.onnx") shutil.rmtree(stage, ignore_errors=True) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--out", type=Path, default=Path(__file__).parent / "onnx") parser.add_argument("--opset", type=int, default=17) parser.add_argument("--only", default="") parser.add_argument("--skip-int8", action="store_true") args = parser.parse_args() only = set(args.only.split(",")) if args.only else {"encoder", "embed_tokens", "decoder", "meta"} from nemo.collections.common.prompts import PromptFormatter from nemo.collections.speechlm2.models import SALM out = args.out out.mkdir(parents=True, exist_ok=True) tmp = out / "_staging" print("loading SALM...", flush=True) model = SALM.from_pretrained(REPO, map_location="cpu") model = model.to(torch.float32).eval() # Fold the LoRA adapters into the frozen Qwen3 weights so the export is a plain causal LM. from peft import PeftModel if isinstance(model.llm, PeftModel): print("merging LoRA...", flush=True) model.llm = model.llm.merge_and_unload() try: model.llm.set_attn_implementation("sdpa") except Exception as exc: # noqa: BLE001 print("attn_implementation:", exc) text_cfg = model.llm.config num_layers = text_cfg.num_hidden_layers num_kv = text_cfg.num_key_value_heads head_dim = getattr(text_cfg, "head_dim", text_cfg.hidden_size // text_cfg.num_attention_heads) hidden = text_cfg.hidden_size features_size = model.perception.preprocessor.featurizer.nfilt print(f"layers={num_layers} kv={num_kv} head_dim={head_dim} hidden={hidden} mel={features_size}", flush=True) with torch.inference_mode(): if "encoder" in only: print("exporting encoder...", flush=True) enc = EncoderExport(model.perception).eval() dummy = torch.randn(1, features_size, 400) dummy_len = torch.tensor([400], dtype=torch.int64) export_graph( enc, (dummy, dummy_len), "encoder", out, tmp, input_names=["input_features", "input_features_lens"], output_names=["audio_embeds", "audio_embeds_lens"], dynamic_axes={ "input_features": {2: "frames"}, "audio_embeds": {1: "audio_seq"}, }, opset_version=args.opset, ) if "embed_tokens" in only: print("exporting embed_tokens...", flush=True) export_graph( EmbedExport(model.embed_tokens).eval(), (torch.zeros(1, 5, dtype=torch.int64),), "embed_tokens", out, tmp, input_names=["input_ids"], output_names=["inputs_embeds"], dynamic_axes={"input_ids": {1: "seq"}, "inputs_embeds": {1: "seq"}}, opset_version=args.opset, ) if "decoder" in only: print("exporting decoder...", flush=True) dec = DecoderExport(model.llm.model, model.llm.lm_head, num_layers).eval() s, p = 3, 2 dummy_past = [] for _ in range(num_layers): dummy_past.append(torch.randn(1, num_kv, p, head_dim)) dummy_past.append(torch.randn(1, num_kv, p, head_dim)) past_names, present_names = [], [] for i in range(num_layers): past_names += [f"past_key_values.{i}.key", f"past_key_values.{i}.value"] present_names += [f"present.{i}.key", f"present.{i}.value"] dyn = { "inputs_embeds": {1: "seq"}, "attn_bias": {2: "seq", 3: "total"}, "position_ids": {1: "seq"}, "logits": {1: "seq"}, } for name in past_names: dyn[name] = {2: "past"} for name in present_names: dyn[name] = {2: "total"} export_graph( dec, ( torch.randn(1, s, hidden), torch.zeros(1, 1, s, p + s), torch.arange(p, p + s, dtype=torch.int64)[None], *dummy_past, ), "decoder", out, tmp, input_names=["inputs_embeds", "attn_bias", "position_ids", *past_names], output_names=["logits", *present_names], dynamic_axes=dyn, opset_version=args.opset, ) if "meta" in only: print("writing vocab and config...", flush=True) tokenizer = model.tokenizer vocab = dict(tokenizer.tokenizer.get_vocab()) with (out / "vocab.json").open("wt", encoding="utf-8") as f: json.dump(vocab, f, ensure_ascii=False) locator_id = model.audio_locator_tag_id formatter = PromptFormatter.resolve(model.cfg.prompt_format)(tokenizer) ids = formatter.encode_dialog( turns=[{"role": "user", "content": f"{USER_PROMPT}{model.audio_locator_tag}"}] )["input_ids"].tolist() cut = ids.index(locator_id) asr_config = { "model_type": "speech-llm", "features_size": features_size, "preprocessor": f"nemo{features_size}", "hidden_size": hidden, "num_layers": num_layers, "num_key_value_heads": num_kv, "head_dim": head_dim, "eos_token_ids": sorted({tokenizer.eos_id, 151643}), "max_sequence_length": 448, "max_frames": 100_000, "hop_length": 160, "prompt_prefix_ids": ids[:cut], "prompt_suffix_ids": ids[cut + 1 :], "language_prompt_ids": {}, "user_prompt": USER_PROMPT, "source_model": REPO, } with (out / "config.json").open("wt", encoding="utf-8") as f: json.dump(asr_config, f, ensure_ascii=False, indent=2) if not args.skip_int8: for name in ("encoder", "embed_tokens", "decoder"): if name in only: print(f"quantizing {name}...", flush=True) quantize(name, out, tmp) shutil.rmtree(tmp, ignore_errors=True) for path in sorted(out.glob("*")): print(path.name, path.stat().st_size) if __name__ == "__main__": main()