MOSS-TTS-Nano-100M-ONNX on Ascend NPU

1. 简介

本文档记录 MOSS-TTS-Nano-100M-ONNX 在昇腾 NPU（Ascend 910B4）环境的适配与验证结果。

MOSS-TTS-Nano 是 OpenMOSS 开源的多语言轻量级语音合成（TTS）模型，采用 Audio Tokenizer + GPT2 解码器 的纯自回归架构，特点包括：

• 约 0.1B 参数，模型体积小巧
• 16 VQ (Vector Quantization) 通道，每帧 16 个 audio token
• 12 层全局 Transformer + 1 层局部 Transformer
• RoPE 位置编码、GELU 激活、tie_word_embeddings
• 支持中英文等多语言文本输入
• 输出音频格式：48 kHz，双声道

原仓库提供 ONNX 导出版本（MOSS-TTS-Nano-100M-ONNX），专为浏览器和 CPU 无依赖部署设计。本次昇腾 NPU 适配采用 PyTorch 原生版本作为运行基座，原因如下：

1. ONNX Runtime CANN EP：标准 onnxruntime PyPI 包不包含昇腾 CANN Execution Provider；尝试安装 onnxruntime-cann 包但当前环境无法顺利完成。
2. ATC 转换：使用 CANN 的 atc 工具将 ONNX 转为 .om 格式时，报错缺少 LayerNorm、Gelu、MatMulV2 等算子的支持（ONNX opset 17 与当前 CANN 8.5.1 OPP 算子库不完全兼容）。
3. PyTorch 原生适配：torch-npu 已完整支持模型所需全部 PyTorch 算子（Linear、Embedding、Softmax、LayerNorm、GELU、RoPE、Eager Attention 等），可直接在 NPU 上运行，无需额外算子开发。

结论：ONNX 与 PyTorch 版本共享同一套权重和架构，PyTorch 路径是当前昇腾 NPU 上最稳定、最快速的跑通方案。

相关获取地址：

• ONNX 权重与源码（AtomGit）：https://atomgit.com/OpenMOSS/MOSS-TTS-Nano-100M-ONNX
• PyTorch 权重（HuggingFace）：https://huggingface.co/OpenMOSS-Team/MOSS-TTS-Nano
• 音频编解码器（HuggingFace）：https://huggingface.co/OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano

参考文档：

• https://github.com/OpenMOSS/MOSS-TTS-Nano

2. 验证环境

• NPU：1 逻辑卡（Ascend 910B4）
• TTS 模型路径：/opt/atomgit/models/moss-tts-nano
• 音频编解码器路径：/opt/atomgit/models/moss-audio-tokenizer-nano
• 模型精度：float32
• 模型参数量：~117M

3. 推理启动

3.1 环境变量

export ASCEND_RT_VISIBLE_DEVICES=0
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export TASK_QUEUE_ENABLE=1

3.2 离线推理

python inference.py \
  --text "Hello world." \
  --device npu \
  --max-new-frames 100 \
  --output-tokens output_tokens.json

参数说明：

3.3 推理示例输出

[NPU] Loading model from /opt/atomgit/models/moss-tts-nano...
[NPU] Model loaded. Params: 117.3M
[NPU] Generating audio tokens for: 'Hello world.'
[NPU] Generated 100 frames, shape=(1, 100, 16)
[NPU] Generation time: 9.265s
[NPU] Throughput: 10.79 frames/s
[NPU] Tokens saved to output_tokens.json

4. Smoke 验证

已验证通过的推理命令：

# 英文短文本
python inference.py --text "Hello world." --device npu --max-new-frames 50

# 中文短文本
python inference.py --text "你好，世界。" --device npu --max-new-frames 50

# 贪婪解码（默认）
python inference.py --text "The quick brown fox jumps over the lazy dog." --device npu --max-new-frames 100

# CPU 对比
python inference.py --text "Hello world." --device cpu --max-new-frames 50

验证结果：

• 所有测试用例均正常生成音频 token 序列
• NPU 推理无报错，无内存溢出
• 输出 token shape 符合预期 [batch, frames, 16]

5. 性能参考

测试条件：NPU 预热 2 轮后连续 5 轮取平均。

压测脚本：python benchmark.py

6. 精度评测

使用 accuracy_eval.py 对 CPU (float32) 与 NPU (float32) 做了贪婪解码一致性评测。

判定说明：TTS 模型在 NPU 上与 CPU 的贪婪解码输出完全一致，说明 torch-npu 对模型所用算子（LayerNorm、GELU、RoPE、Eager Attention 等）的实现与 CPU 端数值对齐良好。

7. 注意事项

1. Attention 实现选择：模型默认使用 flash_attention_2，但当前昇腾 NPU 环境建议显式设置为 eager（或 sdpa），以避免 flash_attn 库缺失导致的加载失败。

2. 模型加载路径：inference.py 和 accuracy_eval.py 通过 sys.path.insert 将模型目录加入 Python 路径，以便加载自定义的 modeling_moss_tts_nano.py。请确保模型路径下包含 modeling_moss_tts_nano.py、configuration_moss_tts_nano.py 等自定义代码文件。

3. 音频波形输出：本交付件专注于验证 TTS 核心模型（LLM 部分）在 NPU 上的推理能力，输出为音频 token 序列。如需生成最终 .wav 音频文件，需要额外加载 MOSS-Audio-Tokenizer-Nano 进行编解码，该组件同样可以在 NPU 上运行。

4. ONNX 路径说明：本次验证基于 PyTorch 原生模型。若后续 onnxruntime-cann 或 CANN ATC 工具链完善了对 opset 17 的 LayerNorm/GELU 等算子支持，可将 ONNX 模型直接通过 CANN EP 或 .om 格式在 NPU 上运行。

8. 交付件清单