Parakeet TDT 0.6B V2 on Ascend NPU (torch_npu + NeMo)

1. 简介

本文档记录 nvidia/parakeet-tdt-0.6b-v2 英文 ASR 模型在 Ascend 910B4 NPU（Atlas 800T A2）上的适配、部署与验证结果。

该模型是 NVIDIA 发布的 6 亿参数自动语音识别模型，采用 FastConformer-TDT 架构，支持英文语音转文本、标点恢复、大小写保留以及词级/字符级时间戳预测。

适配方式：使用 torch_npu 将 NeMo ASR 模型加载到 NPU 设备，通过 .to('npu:0') 迁移后直接调用 transcribe()，前向传播无需修改代码。

相关获取地址：

• 模型权重（HuggingFace）：https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2
• NeMo 工具包：https://github.com/NVIDIA/NeMo

2. 验证环境

• NPU：1 逻辑卡（npu:0）
• 模型路径：./parakeet-tdt-0.6b-v2.nemo
• 测试音频：sample.wav（7.43 秒，16kHz 单声道）

3. 部署方式

3.1 依赖安装

pip install nemo_toolkit[asr] torch_npu librosa

3.2 模型加载（Python）

import torch
import torch_npu
import nemo.collections.asr as nemo_asr

# 设置 NPU 设备
torch.npu.set_device("npu:0")

# 加载模型并迁移至 NPU
model = nemo_asr.models.ASRModel.restore_from("./parakeet-tdt-0.6b-v2.nemo")
model = model.to("npu:0")
model.eval()
model.freeze()

3.3 推理

outputs = model.transcribe(["sample.wav"])
print(outputs[0].text)

带时间戳：

outputs = model.transcribe(["sample.wav"], timestamps=True)
print(outputs[0].timestamp["word"])

4. Smoke 验证

基础检查项：

# 环境检查
python check_accuracy_run_perf.py --dry-run

# 一键烟测试（精度 + 性能）
python check_accuracy_run_perf.py --audio sample.wav

已验证通过项：

5. 性能参考

测试条件：单文件 sample.wav（7.43 秒），1 轮 warmup + 1 轮实测。

注意： TDT 解码器使用了 CUDA 图优化的循环，在 NPU 上会回退到非图模式，因此当前 NPU 推理延迟高于 CPU 参考。这是 NeMo 层面 CUDA 图优化的限制，非 NPU 本身的计算瓶颈。

多轮性能基准测试可通过以下命令执行：

# 5 轮 benchmark（含 1 轮 warmup）
python accuracy_run_perf.py --audio sample.wav --warmup 1 --runs 5

6. 精度评测

使用 accuracy_run.py 对单文件进行 NPU vs CPU 逐字精度比对。

验证命令：

# 精度验证 + 结果保存
python accuracy_run.py --audio sample.wav --output results.json

7. 注意事项

1. CUDA 图回退：TDT 解码器内部使用了 CUDA 图优化（cudagraph）对循环进行加速，在 NPU 上该优化不可用，会回退到 eager 模式执行。这是 NeMo 框架层面的限制，不影响精度。若需 NPU 推理加速，建议关注 NeMo 后续版本的昇腾适配进展。

2. 模型格式：本模型使用 .nemo 格式（NeMo 原生打包格式），通过 ASRModel.restore_from() 加载。不支持直接导入 PyTorch .pth 权重。

3. 内存占用：模型加载至少需要 2 GB 内存。NPU 显存方面，6 亿参数模型在 910B4 上推理显存占用约 2–3 GB，单卡即可满足。

4. 音频格式：输入音频需为 16kHz 单声道 WAV/FLAC 格式。非标准格式可通过 librosa 等工具预处理。

5. 环境脚本：适配仓库提供以下工具脚本：

   脚本	用途
   inference.py	单文件 NPU 推理（支持可选时间戳）
   accuracy_run.py	精度验证（NPU vs CPU 参考）+ 性能计时
   accuracy_run_perf.py	多轮性能基准测试（延迟、RTFx）
   check_accuracy_run_perf.py	环境检查 + 烟测试集成入口