Fork me on GitHub

一文搞懂 NVIDIA 在 GPU 上高效部署语音 AI 模型的最新应用

以下文章来源于 https://zhuanlan.zhihu.com/p/680427251

导读 本文将介绍如何在 GPU 上高效地部署语音 AI 模型。

主要内容如下:

  1. 语音 AI 部署背景介绍

  2. 基于 GPU 的 ASR 解决方案介绍

  3. 基于 GPU 的 TTS 解决方案介绍

  4. 未来规划

  5. 问答环节

分享嘉宾|刘川 NVIDIA 解决方案架构经理

编辑整理|吕晶

内容校对|李瑶

出品社区|DataFun

01语音 AI 部署背景介绍

首先介绍下搭建语音识别和语音生成类工作管线的痛点与挑战。



首先,AI 模型的部署,有端上和云上两种不同的方式。在云上部署时,常常面对服务延时高、并发路数低、部署成本高等问题。我们希望通过更有效地利用 GPU 资源,服务更多的用户,同时降低部署成本。

第二,语音 AI 与传统的 CV 算法不同,其工作管线更为复杂,通常包含多个模块,并且需要处理流之间的状态维护、管理以及状态切换。这使得开发难度大,有时简单的 Python 脚本调度并不高效。

第三,当前许多从事语音 AI 服务的实践者开始探索使用大型模型,如 Whisper,来完成语音识别和语音模型的任务。然而,使用大型模型带来了更大的计算需求,因此迫切需要提升大语言模型在 ASR、TTS 领域中的推理效率。

02基于 GPU 的 ASR 解决方案介绍

1. 基于 Triton 的 ASR 部署方案



过去一年中,NVIDIA 与 WeNet 社区紧密合作,利用 Triton Inference Server 这一推理服务的部署框架,构建了在 GPU 上部署 WeNet 模型的工作管线。在 WeNet 的架构下,进行了基本的音频特征提取(Fbank 特征),并采用了 U2++ 或 Conformer 的编码器来处理音频信号。解码的方式有很多选择,我们通过 CTC 的 prefix Beam search 和 Transformer 或者 Conformer 的 decoder 一起完成(类似于 U2++ 解码)。

整体来讲,通过 TensorRT 加速模型推理,同时通过 Triton 部署整个管线。

在利用 Triton 部署整个 WeNet 管线的过程中,解决了如下一些问题:

  • Challenge1:如何有序组合各个模块, 使 tensor 在 pipeline 中高效传递?如何让多个模块能够流水线并行起来?

在使用 Triton 时,它不仅能够有序组合多个模块,还允许用户外接更多所需的模块。例如,通过 Triton 的 business logic scripting 功能,可以添加 VAD 模块或者进行音频切分,并将其与 ASR 模块连接起来。这样的流程串接都可以在 Triton 中实现,并且各个部分可以并行进行,实现高效的流水线运行。

  • Challenge2:如何将自定义的前处理(解码、特征提取等)操作与其他模型一起部署在 Triton 上?

Triton 不仅能够部署深度学习模型,还支持许多自定义操作,包括 Python 和 C++ 的自定义前后处理。这些操作可以与深度学习模型对接,这也是 Triton 提供的 custom backend 功能。因此,NVIDIA 的 custom backend 可以与 deep learning、framework backend 共同部署,并实现它们之间的桥接。

  • Challenge3:CTC 对 rescoring 模块的调用涉及条件判断,如何串联带有逻辑分支结构的 pipeline?

在部署前后处理中,尤其是在解码模块,我们经常需要进行条件判断等逻辑操作。这些逻辑操作可以用 Python 的 business logic scripting 去实现,从而建立带有逻辑分支结构的工作关系。

  • Challenge4:对于非流式应用,如何提升推理时的 GPU 利用率?

对于非流式应用,Triton 提供了动态批处理和多模型实例并行执行等功能。这些功能可以在将 ASR 服务部署在 GPU 上时,最大化 GPU 的利用率,提升整个服务的吞吐。

  • Challenge5:对于流式应用,如何保持会话,使每个 chunk 清晰自己属于哪个流?

对于流式应用,Triton 内置了流式状态管理器,能够有序地管理整个流式推理过程。每个流中的每个 chunk 都能清晰地知道自己属于哪个流,状态管理器在状态发生切换时能够自动地管理这个过程。最终,在推理完成后,客户端也能明确地了解每个 chunk 的识别结果属于最初输入的哪个流。这种有序的管理确保了流式应用中音频输入输出流程的清晰性。

  • Challenge6:对于流式应用,如何提升推理时 GPU 利用率?

针对流式应用,我们致力于提升 GPU 利用率。通过动态批处理,能够有效推理多个流中的多个 chunk,使 GPU 以批处理方式高效运行,从而最大程度地提升 GPU 资源利用率。

  • Challenge7:对于流式应用,如何便捷地管理和维护多个流的状态?

在流式应用方面,NVDIA 提供了便捷的流状态管理功能。这包括对多个流的状态进行有效管理和维护,确保整个流式应用的状态处理更加高效和便捷。

  • Challenge8:Conformer 推理速度是否能进一步提升?

对于类似 Conformer 和 Transformer 这样的模型,我们不断探索加速方法,主要集中在 TensorRT 这一加速框架。通过 TensorRT 对模型本身的推理速度进行优化,有效提高整个服务的效率。

基于 Triton 的 ASR 部署方案实际上可以扩展到更通用的 ASR 工作管线,并不局限于 WeNet 的模型和管线。我们利用 Triton,可以在很多其他的开源 ASR 方案架构下尽可能地提升 ASR 服务的吞吐。例如,以下是 NVIDIA 对除 WeNet 之外其他一些开源 ASR 项目贡献的 Triton 部署方案,如果感兴趣可以通过下方链接了解。



2. 基于 GPU 的 ASR 的 decoding



接下来介绍在解码方面的加速优化工作,主要包括两个方面。

  • 对于 CTC prefix Beam search 的 CUDA 加速实现

之前,CTC prefix beam search 主要在 CPU 上运行,而深度学习模型(如Conformer、Transformer、Zip Former)通常在 GPU 上运行。这导致深度学习模型和 CTC 解码的 prefix beam search 在 GPU 和 CPU 上分别运行,可能产生一些 CPU 到 GPU 之间的拷贝开销。为解决这个问题,我们通过 CUDA 实现了 CTC prefix beam search,将解码循环包含在一个 GPU 的 kernel 中,并实现了 beam search 的查找。

这个方案速度非常快,相较于 Touch Audio 中的 Flashlight decoder 在 CPU 上的版本,我们的 CUDA decoder 的速度提升了十倍以上,而且精度相比于 Flashlight decoder 也更好。如果感兴趣,可以通过 Torch Audio 直接尝试。但值得注意的是,CTC 目前还不支持热词和语言模型,它只是一个单纯的 prefix beam search,这一局限性可能会影响在生产环境中的应用。

  • 基于 CUDA 的 TLG 解码

因为以上问题,我们还提供了另一种解码实现,基于 CUDA 的 TLG 解码。TLG 解码在包括 Kaldi 在内的很多 ASR 系统上应用广泛。TLG 解码的速度相较于 CTC 加上一些 blank skip 的手段后,可能没有太明显的优势。但 CUDA 的 TLG 使得深度学习模型和解码全链路都在 GPU 上运行,避免了 CPU 到 GPU 之间的内存复制开销。有时 CPU 的压力也很大,因为 CPU 可能会运行其他前后处理操作,因此使用 CUDA 基础的 TLG 解码可以缓解 CPU 上的压力。TLG 解码的优势在于精度相比 CTC 更好,虽然速度可能稍慢一些,因此可以根据需要选择是否使用基于 CUDA 的 TLG 解码。如果想尝试,可以在 WeNet 的开源方案中找到相关信息。

3. NST 训练方法优化

在算法方面,我们聚焦于半监督的 ASR 训练,采用了 NST 方法。



首先介绍一下 NST 的训练过程:先使用有监督数据训练 teacher 模型,然后使用该模型在无监督数据上运行,生成伪标签。通过筛选,选择有用或合理的带有伪标签的数据,与有监督数据一起训练 student 模型。训练完成后,student 模型成为新的 teacher 模型,进行反复迭代。这种方式充分利用了有监督和无监督数据,得到了我们需要的 ASR 模型。

然而,NST 方法面临的一个主要挑战就是噪声。如果有监督和无监督数据集的 Domain 相似,该方法产生的模型精度是可靠的。但如果 Domain 之间存在较大差异,模型在某些领域上表现较差。为了解决领域差异问题,我们提出了一种语言模型筛选的方法。

该方法的核心思想是我们通过大量实验发现 CER Hypo (带目标 domain 语言模型的解码识别结果与不带语言模型解码方式识别的结果之间的 CER) 与 CER Label (带目标 domain 语言模型的解码识别结果与无标注数据的真实标签之间的 CER)总是呈正相关性的,一般前者高后者也会较高,前者低后者也会较低。由于我们实际上并没有无标注数据的真实标签,因此我们可以用前者的大小来估计后者,从而判断 teacher 模型生成的伪标签是否可靠。

这样一来,我们可以为 CER Hypo 设定一个阈值,据此来筛选无监督数据,剔除那些 CER Hypo 较高(也意味着 CER Label 较高)的数据。经过这一筛选过程,使用剩下的数据进行学生模型的训练,以提高在无监督数据集上的模型精度。

在实验中,我们观察到使用 LM filter 后,CER 在不同数据集上都有显著下降,分别为 11.13% (AIshell1 有标注数据 + WenetSpeech 无标注数据)和 15.86% (AIShell2 有标注数据 + WenetSpeech 无标注数据)。这表明 LM filter 对于降低无监督数据中的噪声、提高模型性能起到了积极的作用。

4. 基于 Tensor-LLM 的 Whisper 模型加速推理



第四方面的工作是关于 ASR 大模型的推理加速,其中 Whisper 模型在多语种 ASR 领域取得了显著进展。尽管 Encoder-Decoder 算法在很久以前就存在,但是 Whisper 模型证明了大语言模型在 ASR 领域的强大能力,能够使用单一模型准确预测多种不同语言的 ASR 任务。然而,使用 Whisper 进行 ASR 推理面临一些挑战,主要包括模型的大小带来的推理速度下降和对 GPU 显存的大量消耗。



为了尽可能加速大语言模型的推理,解决当前使用现有框架进行 Whisper 推理可能会遇到的问题,在今年十月,NVIDIA 推出了大语言模型推理加速框架,名为 TensorRT-LLM(TensorRT Large Language Model)。TensorRT-LLM 是建立在 TensorRT 基础上的,专门用于加速大语言模型推理的开源框架。使用 TensorRT-LLM 进行推理,可以利用其中提供的各种 layer 或 OP 的 API 来构建自己的网络结构。

在 TensorRT-LLM 中,我们提供了许多预先构建的模型示例,包括 GPT、LLAMA、百川、chatGLM 等主流模型。当然,TensorRT-LLM 也天然地支持 Whisper。TensorRT-LLM 用于大型语言模型的推理,具有多项优势,有助于提高效率。如果需要进行多机多卡推理,TensorRT-LLM 具有内置的通信组件,支持 Tensor 并行和流水线并行。

此外,我们还整合了社区中几乎所有的大型语言模型推理加速技术,如 Flash Attention、MQH、GQA、KV Cache 以及用于服务端优化的 Inflight Batching (Continuous Batching)和 Paged Attention 等。还采用低精度量化技术,如 INT4、INT8、FP8,以提高推理速度。

最后,TensorRT-LLM 是原生支持与 Triton Inference Server 结合使用的,通过 TensorRT-LLM 加速的大型语言模型,可以轻松在 Triton 上进行服务部署。这是 NVIDIA 提出的一种非常实用且目前性能最先进的大型模型推理框架。





在 TensorRT-LLM 中,我们利用了一些特性来提升 Whisper 模型的推理速度。

首先是对 Whisper 的 encoder & decoder 的 GPU kernel 进行了优化。利用 TensorRT-LLM 中的 Fused MHA kernel,应用 Flash Attention 的优化技术,加速处理输入一个长序列的多头注意力推理的速度。在解码器生成 token 的过程中,利用 Masked Multi-Head Attention(MHA)的高性能 GPU kernel,用于处理解码过程中的 Q 等于 1,而 KV 相对较长的情况。除了多头注意力之外,还针对推理过程中的一些矩阵运算、Layernorm 等操作,专门实现了高性能的 GPU kernel,以替代 PyTorch 中相对低效的原生 kernel。此外,我们对 op 与 op 之间,layer 和 layer 之间进行了一些融合操作,以提高运行效率。

其次,是对大语言模型推理中 KV Cache 的优化。在 Whisper 中,有两种类型的 Key / Value 缓存由 TensorRT-LLM 管理。第一种是 self attention 的缓存,它在解码器迭代生成 token 的过程中持续生成。第二种是 cross attention 的 key 和 value,在编码器完成后,通过几个线性层获取 cross attention的key 和 value 的缓存,这也被缓存以用于后续的每个 token 生成。

此外,TensorRT-LLM 还支持对 Whisper 进行 INT 8 的量化。它将权重量化到 INT 8 的精度,但计算时仍然使用 FP 16。这种做法的优势包括节省显存和减少 weight 从显存复制到寄存器的 IO 开销。

这些优化使得 Whisper 在 V100 上的推理速度比 Faster Whisper 提高了近 40%,同时保持更低的识别错误率。TensorRT-LLM 还支持批量推断,而 Faster Whisper 不支持。

这些优势使得 TensorRT-LLM 成为 Whisper 推理的高效框架,而且我们还在逐步实现其他一些功能,包括 Inflight batching、Paged attention 和 Speculative decoding 等等。



使用 TensorRT-LLM 加速 Whisper,步骤如下:

首先,需要一个经过训练的 Whisper 模型,然后使用 TensorRT 进行模型初始化。由于 Whisper 是在 TensorRT 中原生支持的,所以只需运行初始化步骤,读取权重并构建网络即可。对于不是原生支持的模型,需要使用 Python API 重新构建模型,但 Whisper 不需要。

完成初始化后,进入第二步,即 engine building 过程,这是一个编译的过程,根据 Whisper 网络结构为每个操作和每个层找到最高效的 kernel 实现,并进行一些 kernel fusion 或 layer fusion 的优化操作。

最终得到 TensorRT-LLM 的 Whisper engine 后,可以使用 TRT-LLM 提供的 C++ 或 Python 运行,或直接部署到 Triton Inference Server 上运行。这个流程相对简单。如果在生产中使用 Whisper 模型,建议使用 TensorRT-LLM 进行高性能加速,可显著提升 Whisper 推理速度并节省成本。

03基于 GPU 的 TTS 解决方案介绍

1. 基于 FastPitch+Hifi-GAN 的 Streaming TTS 效果优化

NVIDIA 在 TTS 领域也做了一些供大家参考的工作,例如提供了高效的流式 TTS 部署方案,利用 TensorRT 加速模型推理速度,并通过 Triton Inference Server 实现了高效的流水线。今年,我们对流式 TTS 的效果进行了提升,主要集中在两个方面。



首先,我们发现许多 TTS 模型采用并行模型,其结构通常是非自回归的,并使用卷积等网络层一次性生成所有音频帧。然而,这种并行模型并不适合流式 TTS 合成。尤其在 chunk 之间的接缝处可能存在抖动瑕疵。因此,我们引入了一种 Incremental FastPitch 的方法,将完全并行的 FastPitch 转换为基于 chunk 的 FastPitch。通过使用 casual 卷积替代常规普通卷积,并采用基于 chunk 的 mask MultiHeadAttention,可以控制 chunk size 和 history size。

这种 Mask 使得每个 chunk 内的帧不仅可以看到 chunk 内的其他帧,还能够看到之前 chunk 的帧,通过这种方式实现了基于 chunk 的 FastPitch 并且使得 chunk 之间的信息可以互相关联提升流式 TTS 的质量。Incremental FastPitch 的训练过程仍然可以利用带 mask 的注意力机制来实现并行运算。在推理过程中,可以逐个 chunk 地生成,实现类似迭代的自回归生成过程,从而在流式生成中考虑到历史信息,提升生成效果。



其次,我们采用了 stream GAN 逻辑,即在 Hifi GAN 的训练中,利用 discriminator 强制学习如何让 Generator 生成两个能够良好拼接的连续音频 chunk。



基于先前提到的两种流式 TTS 优化方案,我们开发了相应的推理服务框架,同样基于 Triton Inference Server 加上 TensorRT。

在这个框架中,使用 Triton 的 C++ custom backend 实现了高性能的 TTS 调度器。该调度器负责组织整个 TTS 管线的各个模块,并在这个过程中实现了"Inflight batching",即连续批处理。新进来的请求可以随时加入到正在执行的 batch 中。新合成的音频 chunk 会以流式方式返回给 Triton Client。已完成的请求会立即终止,为新到达的请求腾出 slot。

此外,我们使用 Triton 的 Ensemble Model 功能以零代码的方式组合了需要同时运行的多个模块,如 Front End、声学模型的编码器总是要对输入的文本共同做一次处理。我们使用 Triton Ensemble 将它们无缝组合在一起,而无需编写任何串联代码,实现了零代码的模型串联功能。对于声学模型的解码器、vocoder 以及最后的 Chunk 拼接的 Blender,同样使用 Triton ensemble 功能以零代码的方式将它们组合在一起。最后,对于每个模型,包括声学模型的 encoder 和 vocoder,都使用 TensorRT 来加速推理。

2. 关于声音克隆的参考工作

关于声音克隆的工作



首先,我们开发了一个 Multi-speaker FastPitch 的训练方案,并在开源项目中提供了这个解决方案。在我们的实验中,混合了三个开源数据集(AIShell 3、CSS 10 Chinese、LJSpeech),共计 220 个 Speaker,进行了训练。

获得了多说话人的 FastPitch 模型后,则可进入声音克隆的 Finetuning 阶段。用户首先录制 20 句话。接着,根据用户的声音在训练集中找到一个与之最相似的 Speaker,用其声音的 embedding 初始化用户的 Speaker Embedding。然后,使用用户上传的 20 句话进 finetuning,保持 FastPitch 不变,对 finetuning Speaker 的 Embedding 进行微调。最终,该模型使用 Multi-speaker FastPitch 生成与用户声音相似的音频效果。



为了实现这一思路的工程化,我们同样采用 Triton Inference Server 进行部署。在此我们配置了两个 Triton Server,一个用于声音克隆,另一个用于 TTS 生成。用户录制的 20 句话可以通过 Triton 的客户端上传到声音克隆 Triton Server。在声音克隆 Triton Server 上,运行声音克隆的 finetuning 过程,使用 Triton 的 Python backend,在其中引用 PyTorch 包来实现 finetuning 流程。在这个管线中,能够同时处理多个用户的请求,以 batch 的方式进行 finetuning 提高 GPU 利用率和并发效率。

最后,将每个 Speaker 的 Embedding 返回给客户端,并将 finetuning 完成的 Speaker Embedding 存储在 Embedding Pool 中。用户想要生成自己的声音只需获取返回的 Speaker Embedding,并访问我们的 Triton TTS Server,即我们部署的 Multi-speaker TTS 模型。该模型使用 fine-tuning 完成的 Speaker Embedding 生成与用户声音相似的音频。这就是我们的声音克隆工作的流程。

04未来规划



最后分享一下我们对未来工作的规划。

  • ASR 方面

首先,不断优化 Whisper 的推理,利用新的一些技术实现加速;使用 TensorRT-LLM + Triton 作为 Whisper 部署方案;提供更多的 Whisper finetuning 方案。

另外,我们正在做的一项工作是 code-switched AST 训练方案。之前在 NeMo 中已经实现了基于字符的方法,现在正在尝试基于拼音的东亚语言的方案。

  • TTS 方面

首先是提供针对其他 SOTA 模型的加速方案;

第二是提供多语种的 TTS 方案;

第三是基于 LLM 的 TTS 方案。

以上就是本次分享的内容,谢谢大家。

05问答环节

Q1:CER-Hypo 是什么意思?

A1:在 NST 训练过程中,因为我们没有无监督数据的真实标签,计算 CER-Hypo 的目的是为了预估 teacher 模型在无监督数据上得到的伪标签和真实标签之间的差距,进而进行数据筛选。CER-Hypo 的计算是使用带大语言模型解码方式和不带语言模型解码方式解码出来的 token 的 CER 误差。其中大语言模型是利用目标 domain 有监督的语料训练出来的。

Q2:TensorRT 支持消费级显卡应用吗?比如 4090。

A2:理论上是支持的,但是我们从来没有在消费级显卡上去做过任何的测试优化,所以不能保证消费及显卡上是否会遇到什么问题,比如性能问题或者 bug。

Q3:如何基于 TRT-LLM 开发自己的量化算法,现在感觉写的与 TRT 有点强耦合,缺乏开发文档指导。

A3:NVIDIA 后期会推出一个专门做量化算法的库,这个库暂时没有公开的 release,在此先预告一下,这个库里面包含了很多做大语言模型量化的算法,包括 AWQ,GPTq, FP 8,Smooth Quantization 等等,同时它也支持去扩展自己的量化算法。具体情况可以联系 NVIDIA 进行咨询。

Q4:动态 batch 需要依赖 Triton 吗?

A4:分情况,如果是普通的模型,例如 Conformer 这种的,是依赖 Triton 的,因为 Triton 里面实现了动态 batch 的功能。但如果是大语言模型,做 inflight batching 的话,可以不依赖 Triton,在 TensorRT-LLM 里面 natively 支持了 inflight batch 这样一种动态 batch 的技术,所以不需要依赖 Triton 。

Q5:Whisper 可以流式识别么?

A5:理论上是可以的,从工程化的角度上讲,理论上 TensorRTLLM 后续也会去支持,但从生产实践的角度上讲,很少有人这么做。主要原因还是推理效率的问题。目前看很少会有人去用 Whisper 来做流式的识别。

Q6:TensorRTLLM 是完全开源的么?

A6:目前不是完全开源的。有一些核心的算子,比如 Flash attention 的算子,它的具体实现是不开源的。还有 inflight batch 这套调度的机制是没有开源的,其他部分都是开源的。

Q7:TRT-LLM 相比其他加速框架的优劣?

A7:优势很明显,就是速度最快,劣势是可能没那么好用,比如对于一个新的模型,还是要自己用 API 去重新搭,并不能做自动的转换。但是对于已经支持的模型是不需要去自己搭的,目前对绝大部分常用的大语言模型都已经支持了。

以上就是本次分享的内容,谢谢大家。




本文地址:https://www.6aiq.com/article/1706523622168
本文版权归作者和AIQ共有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出