文章

WhisperKit [译]

我们于 2023 年 11 月成立 Argmax,旨在赋能全球开发者和企业在用户设备上部署商业规模的推理工作负载。对生产环境中 Whisper 推理需求的快速增长促使我们将其作为我们的首个项目。

今天,我们非常激动地宣布,WhisperKit 项目测试版已在 MIT 许可下开源!

  • Swift 包:只需两行代码,即可在应用中实现 Whisper 推理。

  • 示例应用:适用于 iOS 和 macOS 的 TestFlight 版本。

  • Python 工具:用于在 Mac 上优化和评估 Whisper。

我们发布这个测试版的目标是收集开发者反馈,并在几周内实现稳定的发布候选版本,以便 WhisperKit 加快设备上推理的 生产化 进程!

行动号召

如果您对以下内容感兴趣…

  • 技术概览:请继续阅读本文。如有任何反馈,请与我们联系

  • 扩展 WhisperKit:请分支并向我们提交拉取请求!这里是贡献指南

  • 在您的应用中测试 WhisperKit:只需在 Xcode 中添加它作为包依赖,即可准备好部署:

  • 可扩展性。将其模块化为 Swift 协议,自定义行为的实现简化为单个协议扩展。我们还优化了代码的可读性,让扩展 WhisperKit 变得更加容易上手。贡献指南和路线图资源旨在为开发者提供达到功能完备和稳定发布候选版本的时间表。

  • 可预测性。我们采用了以准确性为导向的开发方式,内部测试基础设施会验证代码和模型提交在 Whisper 准确性评估基准测试中的有效性,这些测试包括 librispeech(约 2.6k 短音频剪辑,总计约 5 小时)和 earnings22(约 120 长音频剪辑,总计约 120 小时)数据集。定期测试的结果可在此处查看。这种方法使我们能够检测并减轻由 WhisperKit 代码更改以及软件栈更低层的性能和功能退化所导致的 推理质量 问题(详见下文)。这有助于我们更快地发现和修复大部分问题。此外,我们提供了客户级 SLA,以确保特定模型和设备版本的开发者或企业在最短时间内检测和修复所有问题。

  • 自动部署。我们在此处发布并提供了多种兼容 WhisperKit 的模型。WhisperKit 包括 API,用于列出并从此服务器下载特定版本。我们的示例应用基于这套空中部署基础设施构建。此外,开发者可以利用 whisperkittools(Python 工具包)发布定制 Whisper 版本,在自己的数据集上进行评估,并在生产环境中提供服务。

性能聚焦:实时性

WhisperKit 的目标是在苹果硅芯片上实现 最低延迟最高吞吐量。对于今天的测试版发布,我们首先专注于 最低延迟,以便在 iPhone、iPad 和 Mac 上实现流式转录应用。

实现实时速度的流式转录极具挑战性,因为 Whisper 并非设计用于低延迟音频处理。以下是在转录第一个文本标记之前需要发生的步骤:

第 1 步: 定期积累足够的音频,以包含可理解的语音块,例如 1-2 秒。这一步无论如何都会产生延迟。

第 2 步: 根据可用的历史音频量,将音频填充或剪裁至 Whisper 预期的剪辑长度,即 30 秒。

第 3 步: 对这 30 秒音频剪辑运行一个 0.6b 变压器编码器,或等效地处理 1500 音频标记。

第 4 步: 在第 3 步输出的基础上,通过 1b 变压器解码器运行 4 个特殊标记,以配置语言和任务设置。

第 5 步: 最终生成第一个文本标记!然后,继续生成标记,直到观察到特殊的结束标记。

第 1 步至第 4 步的总延迟决定了 首个标记时间 的延迟,这是衡量交互式转录用户体验是否迅速的一个指标。

第 5 步用于测量 每秒标记数,这是用户语音被转录为文本速度的一个指标。

在录制实时音频块时,重复第 1 步至第 5 步,直至用户结束转录会话。

如果第 1 步至第 5 步的总延迟小于或等于音频块的长度,则 Whisper 推理被认为是实时的。例如,如果每秒调用一次转录管道(第 1-5 步),而管道只用 500 毫秒完成,则推理速度为 2 倍实时(即 0.5 的 实时因子)。

如果推理速度哪怕稍微慢于实时,音频块的入队速度将超过推理管道的处理能力。在这种情况下,队列将不断增长,直到用户因不可接受的延迟停止说话或者陈旧的音频被丢弃。

我们对 Whisper 管道进行了端到端的优化,以最小化这些不理想的情况。由于实时性标准取决于说话速度,我们无法保证所有可能的语音都能实现实时性能。相反,我们优化了管道的性能,以 最大化每个音频块可解码的文本标记数

特别关注的是 openai/whisper-large-v3 变体,因为这是目前最难部署的版本。我们以 1 秒转录 1 秒语音 作为起始预算。

音频编码器的加速

第 3 步中提到的音频编码器之前是通过 whisper.cpp 项目使用 Core ML 优化的,遵循了在苹果神经引擎上部署变压器的原则。简单来说,这项工作引入了一组 神经引擎编译器提示作为 PyTorch 代码,当转换为 Core ML 时,它们转化为高性能模型。由于神经引擎编译器是苹果的私有框架,这相当于是对编译器的黑箱优化,以提高硬件利用率。

这为加速 Whisper 的音频编码器提供了良好的起点,我们很高兴看到我们以前的工作被各种开源项目采用。

WhisperKit 中,我们进一步改进了这些编译器提示,为 iPhone 12 至 15 相比以往实现带来了 1.86x-2.85x 的速度提升:

Xcode Core ML 性能报告 显示了在应用我们的改进前后的 预测 延迟情况,应用于 openai/whisper-base。不幸的是,尽管多次尝试,Xcode 未能为 large-v3 变体在 iPhone 上生成相同的报告。根据我们的应用内测量,我们发现 openai/whisper-large-v3 在 iPhone 上的速度提升了 1.85 倍(与上述结果一致),延迟降至约 350 毫秒(从约 650 毫秒下降)。这意味着我们从 1 秒的预算中 节省了约 300 毫秒 用于解码更多语音标记!我们还有 650 毫秒的剩余预算,而且还没开始解码。

请注意,我们为带有此优化的模型添加了后缀 turbo,以便与我们针对 GPU 的另一种实现区分开(更多详情,请与我们联系)。

文本解码器的加速

除了音频编码器,我们还将文本解码器映射到了神经引擎,使 iPhone 部署可以独立利用神经引擎,达到高能效和最低延迟。

此外,我们显著提高了先前工作中混合位调色板化 (MBP) 的准确性保持,并生成了量化变体以及原始的 float16 版本的两种模型,以减少内存使用。这是一个可选的优化,所有 Mac 都能够为所有 Whisper 变体使用 float16 精度。只有在运行 large-v3 变体的旧 iPhone 上才需要压缩,因为 RAM 可用性减少。由于这不是无损优化,我们引入了 推理质量(QoI) 的概念,并为所有压缩变体计算它。您可以在此处了解更多关于 QoI 的结果和背景。

除了内存节省外,这种压缩技术在神经引擎上得到了本地加速,但延迟改进主要对文本解码器显著,因为压缩权重导致主要是带宽受限的工作负载加速。

通过所有这些优化,文本解码器的延迟在 50 到 100 毫秒之间,取决于设备,使我们能够在剩余的 650 毫秒预算内生成 6-12 个文本标记

特殊标记的 KV 缓存预计算

最后,请记住,前几个标记是与语言和任务设置相关的特殊标记,与语音内容无关。人们可以考虑在编译时预先计算所需的前向传递,并预先填充 KV 缓存,而不是为每个音频块窗口重新计算它们。这将使我们能够计算 3 个实际文本标记,而不是每次都重新计算特殊标记。这也将通过最多 3 倍文本解码器前向传递延迟的时间 减少首个标记时间

不幸的是,与这些特殊标记相关的 KV 缓存值也是音频输入的函数,而音频输入只能在运行时获得!尽管如此,我们有一个强烈的直觉,即 Whisper 的训练目标并没有因为因果掩蔽和内容独立的梯度平均而在音频嵌入和与特殊标记相关的 KV 缓存值之间引入强相关性。我们通过对一批随机输入数据(而不是实际音频编码)预计算所有任务和语言设置的 KV 缓存值来测试这种直觉。然后,我们创建了一个查找表,在推理期间获取 KV 缓存值,而不是运行文本解码器。

我们通过重新运行所有评估并将此优化纳入 large-v3turbo 版本来验证这种方法。与之前的 QoI 和 WER 指标相比,结果没有变化!

我们还测试了一些无效假设,例如使用所有零和简单的高斯噪声预填充 KV 缓存,但这些都未能产生任何有意义的结果,因此我们没有进一步测试它们。通过这种优化,我们现在能够在流式模式下每 1 秒音频块约 生成 9-15 个标记,这对于大多数日常语音来说足够快。

稳定发布前的剩余工作

我们期待您的反馈和贡献,以实现稳定发布!

与此同时,我们计划在 WhisperKit 第一次稳定发布之前至少包括以下功能和改进:

  • 性能报告:在我们的示例应用中创建,以轻松重现所有设备和模型版本上的性能结果,并在 GitHub 上创建一个社区基准页面来众包结果。

  • 异步批量预测,用于 单进程 高吞吐量批处理。今天的测试版发布需要多个并发进程才能实现最大吞吐量。

  • watchOS 示例应用,并包括对 Apple Watch 目标的测试覆盖。此源代码已在 WhisperKit 中。

  • 我们的 Metal基础推理引擎 作为备选后端,以更灵活和更高性能地利用 GPU。



原文作者:Argmax
原文链接:https://www.takeargmax.com/blog/whisperkit

本文由作者按照 CC BY 4.0 进行授权