[论文解读] A GPU-based WFST Decoder with Exact Lattice Generation

Zhehuai Chen, Justin Luitjens+4 more|arXiv (Cornell University)|Apr 9, 2018

Speech Recognition and Synthesis参考文献 28被引用 3

一句话总结

本论文提出了一种基于GPU加速的Kaldi WFST解码器，支持精确的解码轨迹生成，并通过利用原子操作进行token重组、动态负载均衡以及优化的轨迹处理，实现了相较于CPU解码3–15倍的加速。该系统支持通用的声学模型与语言模型，可在多种GPU架构上运行，并通过序列并行与多进程服务（MPS）实现高达46倍的加速，同时保持与CPU解码相同的识别结果与置信度度量。

ABSTRACT

We describe initial work on an extension of the Kaldi toolkit that supports weighted finite-state transducer (WFST) decoding on Graphics Processing Units (GPUs). We implement token recombination as an atomic GPU operation in order to fully parallelize the Viterbi beam search, and propose a dynamic load balancing strategy for more efficient token passing scheduling among GPU threads. We also redesign the exact lattice generation and lattice pruning algorithms for better utilization of the GPUs. Experiments on the Switchboard corpus show that the proposed method achieves identical 1-best results and lattice quality in recognition and confidence measure tasks, while running 3 to 15 times faster than the single process Kaldi decoder. The above results are reported on different GPU architectures. Additionally we obtain a 46-fold speedup with sequence parallelism and multi-process service (MPS) in GPU.

研究动机与目标

• 通过GPU并行计算加速自动语音识别（ASR）中的WFST解码，而无需依赖特定的声学或语言模型。
• 通过支持精确轨迹生成的两遍解码机制，解决GPU解码中大型语言模型（LMs）内存占用高且可扩展性差的问题。
• 设计一种通用、可移植的GPU解码器，支持旧型GPU架构（如Kepler），并可无缝集成至现有的Kaldi工作流。
• 通过优化的token重组、动态负载均衡以及减少同步与内存传输开销的并行轨迹处理，提升解码效率。

提出的方法

• 利用CUDA的atomicMin原子操作将token重组实现为GPU原子操作，消除串行化瓶颈，实现在Viterbi束搜索中的完全并行化。
• 设计一种动态负载均衡策略，将token传递任务在GPU线程间合理分配，最小化空闲时间并提升线程利用率。
• 重新设计适用于GPU执行的精确轨迹生成与剪枝算法，减少数据移动与内核启动开销。
• 采用两遍解码框架：第一遍生成用于后处理的精确轨迹，第二遍使用大型语言模型进行重打分。
• 利用序列并行与多进程服务（MPS）减少上下文切换，在现代GPU上实现高达46倍的加速。
• 将系统作为开源扩展集成至Kaldi工具包中，确保与所有已发布Kaldi语音识别配方的向后兼容性。

实验结果

研究问题

• RQ1 基于GPU的WFST解码能否在保持1-best识别结果与轨迹质量完全一致的前提下，实现相较于CPU解码的显著加速？
• RQ2 在GPU上如何高效并行化Viterbi束搜索中的token重组，同时避免精度损失与同步瓶颈？
• RQ3 在WFST解码的token传递过程中，何种动态负载均衡策略能有效平衡GPU线程间的负载？
• RQ4 精确轨迹生成与剪枝能否在GPU上高效并行化，以减少内存传输与内核启动开销？
• RQ5 在真实ASR系统中，序列并行与多进程服务（MPS）能在多大程度上进一步加速基于GPU的WFST解码？

主要发现

• 该GPU解码器在多种GPU架构上相较单进程CPU解码器实现了3–15倍的加速，且1-best识别结果与轨迹质量完全一致。
• 当使用序列并行与多进程服务（MPS）时，系统实现46倍加速，而相同条件下CPU仅实现1.8倍加速。
• 使用原子操作进行token重组消除了临界区，实现了对Kepler等旧型GPU架构的兼容。
• 动态负载均衡相较于静态负载均衡与图划分方法性能更优，1-best解码实现15倍加速，轨迹解码实现9.7倍加速。
• 轨迹剪枝显著提升了性能，通过减少数据传输与内核启动成本，GPU版本在速度与可扩展性方面均优于CPU。
• 解码器在不同语言模型大小（13MB至258MB）下保持一致性能，12GB GPU上的内存占用上限为11GB，表明内存优化效果显著。