发布 Swift Transformers：在苹果设备上运行设备端 LLM

发布于 2023 年 8 月 8 日

我非常尊敬 iOS/Mac 开发者。我从 2007 年开始为 iPhone 编写应用程序，当时甚至还没有 API 或文档。新设备在约束空间上采纳了一些不熟悉的决策，结合了功耗、屏幕空间、UI 习惯、网络访问、持久性和延迟等因素，这与我们之前习惯的有所不同。然而，这个社区很快就设法创造出与新范式融为一体的顶级应用程序。

我相信机器学习（ML）是一种构建软件的新方式，而且我知道许多 Swift 开发者希望将 AI 功能集成到他们的应用中。ML 生态系统已经非常成熟，有数千个模型可以解决各种各样的问题。此外，大型语言模型（LLM）最近已成为近乎通用的工具——只要我们能将任务建模为处理文本或类文本数据，它们就可以适应新的领域。我们正在见证计算历史上的一个决定性时刻，LLM 正走出研究实验室，成为每个人的计算工具。

然而，在应用程序中使用像 Llama 这样的 LLM 模型涉及多个任务，许多人需要独自面对和解决。我们一直在探索这个领域，并希望与社区继续合作。我们的目标是创建一套工具和构建模块，帮助开发者更快地进行构建。

今天，我们发布这篇指南，介绍在 Mac 上使用 Core ML 运行像 Llama 2 这样的模型所需的步骤。我们还发布了 alpha 版的库和工具，以支持开发者的这一旅程。我们呼吁所有对 ML 感兴趣的 Swift 开发者——这是否意味着*所有* Swift 开发者？——通过提交 PR、报告 bug 或提出意见来共同改进这一切。

冲呀！

今日发布

swift-transformers，一个正在开发中的 Swift 包，旨在 Swift 中实现类似 transformers 的 API，专注于文本生成。它是 swift-coreml-transformers 的演进版本，目标更广泛：Hub 集成、支持任意分词器以及可插拔模型。
swift-chat，一个演示如何使用该包的简单应用程序。
exporters 的更新版本，一个用于 transformers 模型的 Core ML 转换包。
transformers-to-coreml 的更新版本，一个基于 exporters 构建的无代码 Core ML 转换工具。
一些转换好的模型，例如 Llama 2 7B 或 Falcon 7B，可直接用于这些文本生成工具。

任务概览

当我发布推文展示 Falcon 或 Llama 2 在我的 Mac 上运行时，我收到了许多其他开发者的提问，他们想知道如何将这些模型转换为 Core ML，因为他们也想在自己的应用中使用它们。转换是一个关键步骤，但这只是拼图的第一块。我写这些应用的真正原因是为了面对任何其他开发者都会遇到的问题，并找出我们可以在哪些方面提供帮助。在这篇文章的其余部分，我们将介绍其中一些任务，并解释我们在哪些方面有（或没有）工具可以提供帮助。

转换为 Core ML。我们将以 Llama 2 为例进行实际操作。
优化技术，使您的模型（和应用）运行更快并消耗尽可能少的内存。这是一个贯穿整个项目且没有一蹴而就的解决方案的领域。
swift-transformers，我们用于帮助处理一些常见任务的新库。
- 分词器。分词是将文本输入转换为模型处理的实际数字集（以及从生成的预测结果转回文本）的方法。这比听起来要复杂得多，因为有许多不同的选项和策略。
- 模型和 Hub 包装器。如果我们想支持 Hub 上各种各样的模型，就不能硬编码模型设置。我们创建了一个简单的 LanguageModel 抽象以及各种工具，用于从 Hub 下载模型和分词器配置文件。
- 生成算法。语言模型被训练来预测在一段文本序列后可能出现的下一个词元的概率分布。我们需要多次调用模型来生成文本输出，并在每一步选择一个词元。决定下一步应该选择哪个词元有很多方法。
- 支持的模型。并非所有模型系列都（尚）受支持。
swift-chat。这是一个小型应用，简单展示了如何在项目中使用 swift-transformers。
缺失部分/未来计划。一些重要但尚未实现的功能，作为未来工作的方向。
资源。所有项目和工具的链接。

转换为 Core ML

Core ML 是苹果的原生机器学习框架，也是其使用的文件格式的名称。将模型从（例如）PyTorch 转换为 Core ML 后，您就可以在 Swift 应用中使用它。Core ML 框架会自动选择最适合运行模型的硬件：CPU、GPU 或称为神经网络引擎（Neural Engine）的专用张量单元。根据您的系统特性和模型细节，也可能组合使用这些计算单元。

为了了解在实际中转换模型是什么样的，我们将看看如何转换最近发布的 Llama 2 模型。这个过程有时可能很复杂，但我们提供了一些工具来提供帮助。这些工具并不总是有效，因为新模型层出不穷，我们需要不断进行调整和修改。

我们推荐的方法是

使用 transformers-to-coreml 转换 Space

这是一个基于 exporters（见下文）构建的自动化工具，它对您的模型要么有效，要么无效。它不需要编码：输入 Hub 模型标识符，选择您计划使用模型的任务，然后点击应用。如果转换成功，您可以将转换后的 Core ML 权重推送到 Hub，就大功告成了！

您可以访问该 Space 或直接在这里使用

使用 exporters，一个基于 Apple 的 coremltools（见下文）构建的 Python 转换包。

这个库为您提供了更多配置转换任务的选项。此外，它还允许您创建自己的转换配置类，您可以用它来进行额外的控制或解决转换问题。

使用 coremltools，Apple 的转换包。

这是最底层的方法，因此提供了最大的控制权。对于某些模型（尤其是新模型），它仍然可能失败，但您总是有机会深入源代码，尝试找出原因。

关于 Llama 2 的好消息是，我们已经完成了前期工作，使用上述任何一种方法都可以成功转换。坏消息是，它在发布时*转换失败*，我们不得不进行一些修复才能支持它。我们在附录中简要回顾了当时发生的情况，这样您就可以了解当事情出错时该怎么做。

重要的经验教训

我跟踪了一些最近模型（Llama 2、Falcon、StarCoder）的转换过程，并将我学到的知识应用到了 exporters 和 transformers-to-coreml Space 中。以下是一些要点的总结

如果您必须使用 coremltools，请使用最新版本：7.0b1。尽管技术上是 beta 版，但我已经使用它几周了，它真的很好：稳定，包含大量修复，支持 PyTorch 2，并有高级量化工具等新功能。
exporters 在转换文本生成任务时不再对输出应用 softmax。我们意识到这对于某些生成算法是必要的。
exporters 现在默认对文本模型使用固定序列长度。Core ML 有一种方法可以指定“灵活形状”，这样您的输入序列可以有任何长度，例如在 1 到 4096 个词元之间。我们发现灵活输入只能在 CPU 上运行，而不能在 GPU 或神经网络引擎上运行。更多调查即将展开！

我们将继续向我们的工具中添加最佳实践，这样您就不必再次发现同样的问题。

优化

如果转换后的模型不能在您的目标硬件上快速运行并尊重系统资源，那么转换模型就毫无意义。本文提到的模型对于本地使用来说相当大，我们有意使用它们来挑战当前技术的极限，并了解瓶颈所在。

我们确定了几个关键的优化领域。它们对我们来说是一个非常重要的课题，也是当前和未来工作的主题。其中一些包括

缓存先前生成的注意力键和值，就像 transformers 模型在 PyTorch 实现中所做的那样。注意力分数的计算需要在迄今为止生成的整个序列上运行，但所有过去的键值对都已在之前的运行中计算过了。我们目前*没有*为 Core ML 模型使用任何缓存机制，但计划这样做！
使用离散形状而不是小的固定序列长度。不使用灵活形状的主要原因是它们与 GPU 或神经网络引擎不兼容。次要原因是没有缓存，如上所述，生成会随着序列长度的增长而变慢。使用一组离散的固定形状，再加上缓存键值对，应该可以实现更大的上下文大小和更自然的聊天体验。
量化技术。我们已经在 Stable Diffusion 模型的背景下探索过它们，并对它们可能带来的选择感到非常兴奋。例如，6 位调色板化可以减小模型大小并且资源效率高。混合位量化是一种新技术，可以在对模型质量影响很小的情况下实现 4 位量化（平均）。我们也计划在语言模型上研究这些主题！

对于生产应用，请考虑使用较小的模型进行迭代，尤其是在开发期间，然后应用优化技术来选择您用例所能承受的最小模型。

`swift-transformers`

swift-transformers 是一个正在开发中的 Swift 包，旨在为 Swift 开发者提供类似 transformers 的 API。让我们看看它有什么以及缺少什么。

分词器

分词解决了两个互补的任务：将文本输入适配为模型使用的张量格式，并将模型的结果转换回文本。这个过程很微妙，例如

我们应该使用单词、字符、字符组还是字节？
我们应该如何处理小写和大写字母？我们是否应该处理这种差异？
我们应该移除重复的字符，比如空格吗？或者它们是重要的？
我们如何处理不在模型词汇表中的单词？

有几种通用的分词算法，以及许多不同的规范化和预处理步骤，这些对于有效使用模型至关重要。transformers 库决定将所有这些操作抽象到同一个库（tokenizers）中，并将这些决策表示为与模型一起存储在 Hub 上的配置文件。例如，这是 Llama 2 分词器配置中描述*仅规范化步骤*的摘录

  "normalizer": {
    "type": "Sequence",
    "normalizers": [
      {
        "type": "Prepend",
        "prepend": "▁"
      },
      {
        "type": "Replace",
        "pattern": {
          "String": " "
        },
        "content": "▁"
      }
    ]
  },

它的解读是这样的：规范化是按顺序应用的一系列操作。首先，我们在输入字符串前 `Prepend`（前置）字符 `_`。然后我们将所有空格替换为 `_`。潜在的操作列表非常庞大，它们可以应用于正则表达式匹配，并且必须以非常特定的顺序执行。tokenizers 库中的代码处理了 Hub 中所有模型的这些细节。

相比之下，在其他领域（如 Swift 应用）中使用语言模型的项目，通常会采取将这些决策硬编码到应用源代码中的方式。对于少数几个模型来说这没问题，但这样一来就很难用不同的模型替换现有模型，而且很容易出错。

我们在 swift-transformers 中所做的，是在 Swift 中复制这些抽象，这样我们只需要编写一次，每个人都可以在他们的应用中使用它们。我们才刚刚开始，所以覆盖范围还很小。欢迎在仓库中提交 issue 或贡献您自己的代码！

具体来说，我们目前支持 BPE（字节对编码）分词器，这是当今使用的三种主要系列之一。GPT 模型、Falcon 和 Llama 都使用这种方法。对 Unigram 和 WordPiece 分词器的支持将在稍后推出。我们还没有移植所有可能的规范化器、预分词器和后处理器——只移植了我们在转换 Llama 2、Falcon 和 GPT 模型时遇到的那些。

这是在 Swift 中使用 Tokenizers 模块的方法

import Tokenizers

func testTokenizer() async throws {
    let tokenizer = try await AutoTokenizer.from(pretrained: "pcuenq/Llama-2-7b-chat-coreml")
    let inputIds = tokenizer("Today she took a train to the West")
    assert(inputIds == [1, 20628, 1183, 3614, 263, 7945, 304, 278, 3122])
}

然而，您通常不需要自己对输入文本进行分词——Generation 代码会处理好这件事。

模型和 Hub 包装器

如上所述，transformers 大量使用存储在 Hub 上的配置文件。我们准备了一个简单的 Hub 模块来从 Hub 下载配置文件，用于实例化分词器和检索有关模型的元数据。

关于模型，我们创建了一个简单的 LanguageModel 类型作为 Core ML 模型的包装器，专注于文本生成任务。通过使用协议，我们可以用相同的 API 查询任何模型。

为了检索您使用的模型的适当元数据，swift-transformers 依赖于一些自定义元数据字段，这些字段必须在转换 Core ML 文件时添加。swift-transformers 将使用这些信息从 Hub 下载所有必要的配置文件。以下是我们在 Xcode 的模型预览中使用的字段

exporters 和 transformers-to-coreml 会自动为您添加这些字段。如果您手动使用 coremltools，请确保自己添加它们。

生成算法

语言模型被训练来预测在输入序列后可能出现的下一个词元的概率分布。为了构成一个响应，我们需要多次调用模型，直到它产生一个特殊的*终止*词元，或者我们达到我们期望的长度。决定下一个最佳词元有很多方法。我们目前支持其中两种

贪婪解码（Greedy decoding）。这是最直接的算法：选择概率最高的词元，将其附加到序列中，然后重复。对于相同的输入序列，这将总是产生相同的结果。
top-k 采样。选择 `top-k`（其中 `k` 是一个参数）个最可能的词元，然后使用像 `temperature` 这样的参数从中随机*采样*，这将增加变异性，但代价是可能导致模型跑题并失去对内容的跟踪。

其他方法如“核采样”将在稍后推出。我们推荐阅读这篇博文（最近更新），它对生成方法及其工作原理进行了出色的概述。像辅助生成这样的复杂方法对于优化也可能非常有用！

支持的模型

到目前为止，我们已经用少数几个模型测试了 swift-transformers，以验证主要的设计决策。我们期待尝试更多模型！

Llama 2。
Falcon。
StarCoder 模型，基于 GPT 架构的一个变体。
GPT 系列，包括 GPT2、distilgpt、GPT-NeoX、GPT-J。

`swift-chat`

swift-chat 是一个基于 swift-transformers 构建的简单演示应用。其主要目的是展示如何在您的代码中使用 swift-transformers，但它也可以用作模型测试工具。

要使用它，请从 Hub 下载一个 Core ML 模型或创建您自己的模型，然后从 UI 中选择它。所有相关的模型配置文件将从 Hub 下载，使用元数据信息来识别这是哪种模型类型。

首次加载新模型时，需要一些时间来准备。在此阶段，CoreML 框架将编译模型并根据您的机器规格和模型结构决定在哪些计算设备上运行它。这些信息会被缓存并在未来的运行中重复使用。

该应用有意设计得简单，以使其易于阅读和简洁。它也缺少一些功能，主要是由于当前模型上下文大小的限制。例如，它没有任何“系统提示”的设置，这些提示对于指定您的语言模型的行为甚至其个性非常有用。

缺失部分/未来计划

如前所述，我们才刚刚开始！我们接下来的优先事项包括

编码器-解码器模型，如 T5 和 Flan。
更多分词器：支持 Unigram 和 WordPiece。
额外的生成算法。
支持键值缓存以进行优化。
使用离散序列形状进行转换。与键值缓存相结合，这将允许更大的上下文。

让我们知道您认为我们接下来应该做什么，或者前往仓库查看 “Good First Issues” 来一试身手！

结论

我们介绍了一套工具，帮助 Swift 开发者将语言模型集成到他们的应用中。我迫不及待地想看看您用它们创造出什么，并期待在社区的帮助下改进它们！请随时与我联系 :)

附录：以硬核方式转换 Llama 2

除非您遇到过 Core ML 转换问题并准备好与之战斗，否则您可以安全地忽略此部分 :)

根据我的经验，PyTorch 模型使用 coremltools 转换到 Core ML 失败的两个常见原因是

不支持的 PyTorch 操作或操作变体

PyTorch 有*大量*操作，所有这些操作都必须映射到一个中间表示（MIL，即*模型中间语言*），然后该中间表示又被转换为原生的 Core ML 指令。PyTorch 的操作集不是静态的，所以新的操作也必须添加到 coremltools 中。此外，一些操作非常复杂，可以在其参数的各种奇特组合上工作。一个最近添加的、非常复杂的操作示例是*缩放点积注意力*，在 PyTorch 2 中引入。一个部分支持的操作示例是 einsum：并非所有可能的方程都能翻译成 MIL。

边缘情况和类型不匹配

即使对于支持的 PyTorch 操作，也很难确保转换过程在所有可能的输入和所有不同的输入类型上都能正常工作。请记住，单个 PyTorch 操作可以有多个针对不同设备（cpu、CUDA）、输入类型（整数、浮点数）或精度（float16、float32）的后端实现。所有组合的乘积是惊人的，有时模型使用 PyTorch 代码的方式会触发一个可能没有被考虑或测试过的转换路径。

这就是我第一次尝试使用 coremltools 转换 Llama 2 时发生的情况

通过比较不同版本的 transformers，我发现问题是在引入这行代码时开始出现的。这是最近一次 `transformers` 重构的一部分，旨在更好地处理*所有*使用因果掩码的模型中的因果掩码，所以这对其他模型来说也是一个大问题，不仅仅是 Llama。

错误截图告诉我们，在填充掩码张量时存在类型不匹配。问题来自于行中的 `0`：它被解释为一个 `int`，但要填充的张量包含 `float`，而使用不同类型被转换过程拒绝了。在这种特殊情况下，我为 coremltools 提出了一个补丁，但幸运的是这种情况很少需要。在许多情况下，您可以修补您的代码（在 `transformers` 的本地副本中使用 `0.0` 就可以了），或者创建一个“特殊操作”来处理异常情况。我们的 exporters 库对自定义的特殊操作有很好的支持。请参阅这个例子，了解一个缺失的 `einsum` 方程，或者这个例子，了解一个使 `StarCoder` 模型在新版 `coremltools` 发布前能正常工作的变通方法。

幸运的是，coremltools 对新操作的覆盖范围很好，团队反应也非常快。