NeMo ASR 语言模型定制实战:NGPU-LM GPU 加速 N-gram 浅层融合解码
2026/9/13 17:18:24 网站建设 项目流程

NeMo ASR 语言模型定制实战:NGPU-LM GPU 加速 N-gram 浅层融合解码

【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech

在 NeMo(NVIDIA NeMo Speech 仓库)中,ASR 系统无需重新训练声学模型,只需在解码阶段引入外部语言模型(LM)做浅层融合(Shallow Fusion),就能显著提升识别准确率。本文围绕仓库中的 NGPU-LM 语言模型融合文档 展开,完整覆盖浅层融合的原理、NGPU-LM 的训练流程、CTC / RNN-T / TDT / AED 四类模型的融合解码命令与全部可调参数,并结合源码剖析 NeMo 中 GPU 批处理 N-gram 查询的底层实现。读完后,你可以独立完成:用 KenLM 训练 n-gram 语言模型、将其转换为.nemo格式加载进解码器、并在各类 NeMo ASR 模型上配置 greedy / beam 解码的 NGPU-LM 融合。

浅层融合:不重训 ASR 模型提升准确率

ASR 解码时,声学模型输出的概率会与一个独立外部语言模型的得分结合,最终选择合并得分最高的词序列作为转写结果。外部 LM 通常在海量文本上训练,能够捕捉语言的统计规律、句法结构和上下文依赖,从而预测更合理的词序列,纠正 ASR 模型可能的错误。

这一机制对领域适配尤其有价值:把外部 LM 训练在医疗、法律或技术文档等特定领域文本上,它便习得了该领域的词汇,进而引导 ASR 解码产生更准确、更符合上下文语境的结果。传统上,浅层融合是在beam search 解码中完成的——beam search 会探索多个有希望的假设路径以寻找最可能的转写。

NGPU-LM:GPU 加速的 N-gram 语言模型

训练传统 n-gram 语言模型最常用的库是 KenLM,它以高效的 CPU 实现著称,但在高吞吐场景(尤其是大规模数据)下,CPU 依赖会成为瓶颈。

NGPU-LM是统计型 n-gram 语言模型的 GPU 加速实现。它采用通用 trie(前缀树)数据结构,支持快速、批处理的查询(详见论文 "NGPU-LM: GPU-Accelerated N-Gram Language Model for Context-Biasing in Greedy ASR Decoding", arXiv:2505.22857)。它带来两个关键能力:

  1. greedy 解码中的浅层融合:在标准 greedy 解码与"beam search + LM"之间提供了中间档位——保留 greedy 的速度与简洁,同时夺回相当一部分原本只有 beam search + 外部 LM 才能获得的准确率。
  2. 完全 GPU 化的 beam search:NeMo 为所有主流 ASR 模型类型提供了高效的、全 GPU 化的 beam search 实现,使得 beam 解码的实时率(RTFx)接近 greedy 解码——在 batch size 为 32 时,beam 与 greedy 解码的 RTFx 差距仅约 20%。这些实现内嵌了 NGPU-LM,使用户在 beam 模式下也能兼顾速度与解码定制(详见论文 "Pushing the Limits of Beam Search Decoding for Transducer-based ASR Models", arXiv:2506.00185)。

NGPU-LM 融合目前支持 **BPE 类 ASR 模型(CTC、RNNT、TDT、AED)**的 greedy 和 beam 解码。

从源码结构看,NGPU-LM 的核心实现位于 ngram_lm 子模块,其中NGramGPULanguageModel提供from_file/from_arpa等类方法,内部使用 NumPy 存储的 suffix tree(加权接受器)加 Triton kernel(ngram_lm_triton.py)在 GPU 上完成批处理推进;同目录还保留了KenLMBatchedWrapper作为基于 KenLM Python 包的逐样本参考实现,用于对照测试。

训练 NGPU-LM

NGPU-LM 基于 KenLM 生成的.ARPA文件构建。训练脚本为 train_kenlm.py,基本用法:

python train_kenlm.py nemo_model_file=<path to the .nemo file of the model> \ train_paths=<list of paths to the training text or JSON manifest files> \ kenlm_bin_path=<path to the bin folder of KenLM library> \ kenlm_model_file=<path to store the binary KenLM model> \ ngram_length=<order of N-gram model> \ preserve_arpa=true \ save_nemo=True

生成的.ARPA文件可直接用于 GPU 解码;但为获得更快性能,推荐把save_nemo设为true,将模型转换为.nemo格式。

关于 n-gram 阶数的重要提示:对 BPE 类模型,推荐 n-gram 阶数取6;更高阶可能需要重新编译 KenLM 才能支持。

训练脚本的完整参数与内部流程

完整参数列表(同时参见 n-gram LM 工具文档):

参数类型默认值说明
nemo_model_filestr必填ASR 模型.nemo文件路径,或预训练 NeMo 模型名,用于提取 tokenizer
train_pathsList[str]必填训练文件或目录列表,支持纯文本、.jsonmanifest、.json.gz
kenlm_model_filestr必填KenLM 二进制模型文件的保存路径
kenlm_bin_pathstr必填KenLM 的bin目录路径
ngram_lengthint必填N-gram 模型的阶数
ngram_pruneList[int][0]N-gram 剪枝阈值列表,如 [0,0,1]
cache_pathstr""保存分词后文件的缓存路径
preserve_arpaboolFalse是否保留中间 ARPA 文件
verboseint1详细级别
save_nemoboolFalse是否以.nemo格式保存 LM(供 NGramGPULanguageModel 使用)
normalize_unk_nemoboolTrue是否在 NGramGPULanguageModel 中归一化 UNK token

结合 train_kenlm.py 的源码,训练流程分三步:

  1. 自动检测分词层级kenlm_utils.setup_tokenizer(nemo_model_file)从 ASR 模型中提取 tokenizer,并自动判断是字符级还是 BPE(subword)级。若为 subword 级,会追加--discount_fallback参数——这是训练 BPE 类模型所必需的。脚本用一个技巧把 BPE 子词 token 编码为 unicode 字符,从而训练一个"字符级"的 KenLM 模型,使其兼容下游解码器。
  2. 调用 KenLMlmplz:以lmplz -o <ngram_length> --arpa ... --prune ...生成 ARPA 文件(中间产物为<kenlm_model_file>.tmp.arpa);再调用build_binary trie生成二进制模型。若指定cache_path,分词后的文本会先落盘缓存,便于重复训练时复用。
  3. 可选的.nemo转换:若save_nemo=True,脚本通过NGramGPULanguageModel.from_arpa(lm_path=..., vocab_size=full_vocab_size, normalize_unk=...)将 ARPA 转换为.nemo检查点(保存为<kenlm_model_file>.nemo),供解码器直接加载。

依赖安装:KenLM 并非 NeMo 的默认依赖。可从源码结构看,需按 install_beamsearch_decoders.sh 的说明安装,或使用 Dockerfile.ngramtools 构建包含全部依赖的 Docker 镜像。

解码前的配置规则

运行 NGPU-LM 融合推理时,必须在解码配置中指定ngram_lm_modelngram_lm_alpha两个字段:

  • CTC、RNNT、TDT模型:这两个字段应设置在各自的greedybeam子配置中;
  • greedy 模式的 AED 模型:把 beam size 设为 1,并在beam子配置中指定这两个字段(因为 AED 的 greedy 解码走 beam 路径)。

CTC 解码与 NGPU-LM

Greedy 搜索

python examples/asr/speech_to_text_eval.py \ pretrained_name=nvidia/parakeet-ctc-1.1b \ amp=false \ amp_dtype=bfloat16 \ matmul_precision=high \ compute_dtype=bfloat16 \ presort_manifest=true \ cuda=0 \ batch_size=32 \ dataset_manifest=<path to the evaluation JSON manifest file> \ ctc_decoding.greedy.ngram_lm_model=<path to the .nemo/.ARPA file of the NGPU-LM model> \ ctc_decoding.greedy.ngram_lm_alpha=0.2 \ ctc_decoding.greedy.allow_cuda_graphs=True \ ctc_decoding.strategy="greedy_batch"

Beam 搜索:CTC beam 搜索中每条假设按如下公式打分:

final_score = acoustic_score + ngram_lm_alpha * lm_score + beam_beta * seq_length
  • acoustic_score:ASR 声学模型预测的分数;
  • lm_score:NGPU-LM 预测的分数;
  • ngram_lm_alpha:语言模型的权重;
  • beam_beta:计入序列长度的惩罚项。

批量 GPU CTC 解码策略beam_batch的可调参数:

参数类型默认值说明
ngram_lm_alphafloat必填语言模型得分的权重因子
beam_sizeint4Beam 宽度
beam_betafloat1对插入的惩罚,用于控制 beam 搜索中插入/删除错误的权衡
beam_thresholdfloat20与最佳假设比较以剪枝候选假设的阈值
python examples/asr/speech_to_text_eval.py \ pretrained_name=nvidia/parakeet-ctc-1.1b \ amp=false \ amp_dtype=bfloat16 \ matmul_precision=high \ compute_dtype=bfloat16 \ presort_manifest=true \ cuda=0 \ batch_size=32 \ dataset_manifest=<path to the evaluation JSON manifest file> \ ctc_decoding.beam.ngram_lm_model=<path to the .nemo/.ARPA file of the NGPU-LM model> \ ctc_decoding.beam.ngram_lm_alpha=0.2 \ ctc_decoding.beam.beam_size=12 \ ctc_decoding.beam.beam_beta=1.0 \ ctc_decoding.strategy="beam_batch" \ ctc_decoding.beam.allow_cuda_graphs=True

RNN-T / TDT 解码与 NGPU-LM

Greedy 搜索

python examples/asr/speech_to_text_eval.py \ pretrained_name=nvidia/parakeet-rnnt-1.1b \ amp=false \ amp_dtype=bfloat16 \ matmul_precision=high \ compute_dtype=bfloat16 \ presort_manifest=true \ cuda=0 \ batch_size=32 \ dataset_manifest=<path to the evaluation JSON manifest file> \ rnnt_decoding.greedy.ngram_lm_model=<path to the .nemo/.ARPA file of the NGPU-LM model> \ rnnt_decoding.greedy.ngram_lm_alpha=0.2 \ rnnt_decoding.greedy.allow_cuda_graphs=True \ rnnt_decoding.strategy="greedy_batch"

提示:使用 TDT 模型推理时,需在pretrained_name中提供预训练 TDT 模型(例如nvidia/parakeet-tdt_ctc-1.1b)。

Beam 搜索:RNN-T / TDT beam 搜索的打分公式为:

final_score = acoustic_score + ngram_lm_alpha * lm_score

最终假设基于归一化得分final_score / seq_length选出。

Transducer 模型的 Blank 打分:Transducer 模型含 blank 符号()用于帧间转移,而 LM 并不对 blank 建模。浅层融合中,LM 通常只作用于非 blank token:

ln p_tot[k] = ln p[k] + λ · ln p_LM[k] (k ∈ V) ln p_tot[k] = ln p[∅] (k = ∅)

在高 LM 权重下,这会导致 blank 预测过多、删除错误上升。NeMo 支持一种 blank 感知的打分方式,通过调整 LM 贡献来平衡预测:

ln p_tot[k] = ln p[k] + λ · ln((1 - p[∅]) · p_LM[k]) (k ∈ V) ln p_tot[k] = (1 + λ) · ln p[∅] (k = ∅)

早期剪枝 vs 晚期剪枝:LM 与 ASR 得分可以在不同阶段合并:

  • Early pruning(早期剪枝):ASR 先选出 top 假设,再施加 LM 重打分。对小 beam 更高效;
  • Late pruning(晚期剪枝):ASR 与 LM 得分先合并再剪枝,更精确,但需要对全词表做 LM 查询。

对 Transducer 模型,采用 blank 感知打分法的晚期剪枝通常优于标准方式。

Beam 搜索策略:NeMo 的全批处理(fully batched)实现支持以下策略:

  • malsd_batch:修改版 Alignment-Length Synchronous Decoding(Saon et al., ICASSP 2020)的全批处理实现,支持 RNNT 和 TDT 模型;
  • maes_batch:修改版 Adaptive Expansion Search(Kim et al., IEEE SPL 2020)的全批处理实现,仅支持 RNNT 模型,且不支持 CudaGraphs。

对应参数(malsd_batch/maes_batch策略):

参数类型适用策略默认值说明
ngram_lm_alphafloatmalsd_batch, maes_batch必填语言模型得分的权重因子
beam_sizeintmalsd_batch, maes_batch4Beam 宽度
pruning_modestrmalsd_batch, maes_batchlate剪枝模式,可选earlylate
blank_lm_score_modestrmalsd_batch, maes_batchlm_weighted_fullblank 符号打分模式,可选no_scorelm_weighted_full
max_symbols_per_stepintmalsd_batch10每步最多发出的符号数,避免无限循环
maes_num_stepintmaes_batch2自适应扩展的步数
maes_expansion_betafloatmaes_batch1.0除 beam size 外允许的最大前缀扩展数
maes_expansion_gammafloatmaes_batch2.3与最佳假设比较以剪枝候选假设的阈值
python examples/asr/speech_to_text_eval.py \ pretrained_name=nvidia/parakeet-rnnt-1.1b \ amp=false \ amp_dtype=bfloat16 \ matmul_precision=high \ compute_dtype=bfloat16 \ presort_manifest=true \ cuda=0 \ batch_size=32 \ dataset_manifest=<path to the evaluation JSON manifest file> \ rnnt_decoding.beam.ngram_lm_model=<path to the .nemo/.ARPA file of the NGPU-LM model> \ rnnt_decoding.beam.ngram_lm_alpha=0.2 \ rnnt_decoding.beam.beam_size=12 \ rnnt_decoding.beam.pruning_mode="late" \ rnnt_decoding.beam.blank_lm_score_mode="lm_weighted_full" \ rnnt_decoding.beam.allow_cuda_graphs=True \ rnnt_decoding.strategy="malsd_batch"

AED 解码与 NGPU-LM

python examples/asr/speech_to_text_eval.py \ pretrained_name="nvidia/canary-1b" \ amp=false \ amp_dtype=bfloat16 \ matmul_precision=high \ compute_dtype=bfloat16 \ presort_manifest=true \ cuda=0 \ batch_size=32 \ dataset_manifest=<dataset_manifest> \ multitask_decoding.beam.beam_size=4 \ multitask_decoding.beam.ngram_lm_model=<path to the .nemo/.ARPA file of the NGPU-LM model> \ multitask_decoding.beam.ngram_lm_alpha=0.2 \ multitask_decoding.strategy="beam"

提示:AED 模型的 NGPU-LM greedy 解码,请使用beam_size=1的 beam 搜索。

源码视角:融合模型如何进入解码器

从源码结构看,ngram_lm_modelngram_lm_alpha参数最终会传入各解码器的构造函数。以 ctc_greedy_decoding.py 中的GreedyBatchedCTCInfer为例:

  • 构造函数接收ngram_lm_model: Optional[str | Path]ngram_lm_alpha: float;当提供 LM 路径时,通过NGramGPULanguageModel.from_file(lm_path=..., vocab_size=blank_id)加载为融合模型并加入fusion_models列表;
  • forward阶段:若fusion_models非空,解码走_greedy_decode_logprobs_batched_fusion_models_*路径;在 CUDA 设备上且允许allow_cuda_graphs=True时,优先使用 CUDA Graphs 实现(对应文档命令中的allow_cuda_graphs=True参数),否则回退到 PyTorch 循环路径。CUDA Graphs 状态由CTCDecoderCudaGraphsState管理,其中预分配了batch_lm_states(每个 batch 样本的 LM trie 状态)、lm_scores等张量,说明 LM 的 trie 状态推进与 CTC 逐帧解码被融合在同一计算图中,以避免 CPU 开销;
  • 同一文件还支持与boosting_tree(关键词增强)模型叠加,二者共享 fusion 框架;
  • 其余策略的实现分别位于 ctc_beam_decoding.py、rnnt_malsd_batched_computer.py、rnnt_maes_batched_computer.py、tdt_malsd_batched_computer.py、multitask_beam_decoding.py 等文件中,与文档中各模型的strategy取值一一对应。

小结与参考资料

  • NGPU-LM 让 NeMo 在 greedy 与 beam 两种解码下都能做 GPU 端 n-gram 浅层融合,且 beam 解码速度可接近 greedy;
  • 训练链路上:KenLMlmplz生成 ARPA(BPE 模型需--discount_fallback,建议 6 阶),save_nemo=True产出.nemo加速格式;
  • 推理链路上:CTC/RNNT/TDT 把ngram_lm_modelngram_lm_alpha放在greedy/beam子配置中,AED 则通过multitask_decoding.beam配置并可用beam_size=1模拟 greedy;
  • Transducer 模型推荐使用pruning_mode="late"+blank_lm_score_mode="lm_weighted_full"的组合。

主要参考文献:

  1. V. Bataev, A. Andrusenko, L. Grigoryan, A. Laptev, V. Lavrukhin, B. Ginsburg.NGPU-LM: GPU-Accelerated N-Gram Language Model for Context-Biasing in Greedy ASR Decoding. arXiv:2505.22857, 2025.
  2. L. Grigoryan, V. Bataev, A. Andrusenko, H. Xu, V. Lavrukhin, B. Ginsburg.Pushing the Limits of Beam Search Decoding for Transducer-based ASR Models. arXiv:2506.00185, 2025.
  3. G. Saon, Z. Tüske, K. Audhkhasi.Alignment-Length Synchronous Decoding for RNN Transducer. ICASSP 2020, pp. 7804–7808.
  4. J. Kim, Y. Lee, E. Kim.Accelerating RNN Transducer Inference via Adaptive Expansion Search. IEEE Signal Processing Letters, vol. 27, pp. 2019–2023, 2020.

【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询