WWDC 24：用 Core ML 在 Mac 上运行 Mistral 7B

WWDC' 24 是 Apple 正式发布 Apple Intelligence 并重申其对高效、私密和设备端 AI 承诺的时刻。在主题演讲和随后的会议中，他们展示了 Apple Intelligence，它支持大量 AI 增强功能，展示了日常任务的实用用途。这些不是 *为了 AI 而 AI* 的花哨演示。这些是节省时间、恰当（且有趣！）的助手，与应用程序和操作系统深度集成，同时也为开发者提供了多种方式将这些功能纳入他们自己的应用程序中。

Apple Intelligence 功能之所以能如此出色地工作，是因为垂直集成的软件栈充分利用了 Apple Silicon 的能力。Apple 还为开发者提供了一个在设备端运行模型的平台，称为 Core ML。这个软件栈允许你在 Apple Silicon 硬件上所有 3 个计算单元（CPU、GPU 和神经引擎）上运行 ML 模型。

在这篇博客文章中，我们将探索一些最佳的新 Core ML 功能，以复现 Apple 在 WWDC'24 使用 Core ML 在设备端部署机器学习和 AI 模型会议中展示的 Mistral 7B 示例，他们使用 swift-transformers 的一个分支在 Mac 上运行最先进的 LLM。这是一个高质量模型，拥有超过 70 亿参数，推动了当今消费级硬件的极限。你也可以查看 WWDC'24 将你的机器学习和 AI 模型带到 Apple silicon 会议，其中展示了 Mistral 7B 转换过程的一部分。

让我们看看如何尽可能高效地运行它，并了解 iOS 18 和 macOS Sequoia 中可用的新工具。

这就是我们今天要构建的内容：

TL;DR

通过这篇博客文章，你将了解伴随最新 macOS 版本的所有新功能，并且你将成功在 Mac 上使用不到 4GB 内存运行一个 70 亿参数的模型。

步骤 1：克隆 swift-transformers 仓库的 preview 分支：git clone -b preview https://github.com/huggingface/swift-transformers 步骤 2：从 this HF Mirror repo 下载转换后的 Core ML 模型步骤 3：使用 Swift 运行推理：swift run transformers "Best recommendations for a place to visit in Paris in August 2024:" --max-length 200 Mistral7B-CoreML/StatefulMistralInstructInt4.mlpackage

WWDC' 24 中最佳的新 Core ML 功能

以下是我们将用于在 Mac 上运行 Mistral 7B 的一些最具影响力的 WWDC' 24 Core ML 功能。

Swift Tensor

我们想强调的第一个功能是一个全新的 Swift 类型，用于处理 ML 张量。这些是每个 ML 框架都使用的多维数据结构。从事 ML 的 Python 开发者熟悉 numpy 数组或 torch 张量，它们提供了方便、高级的接口来轻松操作这些大型多维矩阵。新的 MLTensor 类型提供了一个高级抽象，模仿 Python 框架中可用的抽象，极大地简化了在 Swift 中处理张量数据。

Core ML 已经有多维数据类型，形式为 MLMultiArray 和 MLShapedArray。然而，它们仅用于数据存储和简单操作，比如包装你的数据并将其作为输入发送到 Core ML 模型，或从 Core ML 模型解包结果。然而，使用这些 API 操作张量数据很困难。只提供了少数基本操作，你可能不得不通过将底层存储作为数字数据的不透明指针来访问，自己编写操作。这既耗时又容易出错。

新的 Tensor 类型提供了一种高级抽象，模仿了 Python 框架中可用的抽象，极大地简化了 Swift 中张量数据的处理。考虑一个我们想要移植到 Core ML 的语言模型，就像我们想要的那样。语言模型接受一个输入令牌序列，并输出词汇表中所有令牌的概率估计，这意味着具有高概率的令牌有很大机会是输入的合理延续。应用程序的工作是根据这些概率选择最佳的下一个令牌附加到序列中。Tensor 类型使得无需自定义代码即可轻松处理这些操作。

当我们发布 swift-transformers 时，我们编写了大量代码（后来由社区扩展，谢谢！❤️）来帮助进行输入准备（将单词转换为令牌）和输出后处理。例如，查看我们使用 Accelerate 的 softmax 操作。使用 MLTensor 时，所有这些都可以移除，因为 softmax 是开箱即用的！

有状态缓冲区

在 WWDC' 24 之前，Core ML 模型本质上是一个纯无状态函数，提供输入并返回一些输出。然而，有时您需要保持依赖于先前计算的状态。函数式编程维护状态的方法是添加一个额外的输入/输出对。因此，根据您的输入和状态，模型计算输出和新状态。这种方法没有问题，事实上，像 JAX 这样的高性能框架就是这样工作的。

然而，存在实际限制：有状态数据需要每次调用模型时作为输入发送到模型并作为输出检索。如果有状态数据很大，那么所有这些来回都会增加开销并减慢速度。这对于 LLMs 尤为重要，因为您必须运行许多迭代才能生成序列。性能瓶颈通常是您计算机的内存带宽（即，您可以将内容移动到 GPU 并返回的速度）。有状态模型通过为状态数据保留一块内存并将其保存在 GPU 上来解决这个问题，这样您就不必每次使用模型时都发送和接收它。

有状态缓冲区是在这个 WWDC' 24 会议中引入的，使用了一个简单的玩具示例，易于理解但不代表像 LLMs 这样的大模型的实际用途。基于 transformers 的模型的 LLM 性能技巧是键值缓存（称为 kv-caching）。如下图所示，它通过缓存先前步骤中执行的先前操作的结果，避免了关键注意力块中昂贵的矩阵乘法。我们不会深入细节，但要点是：kv-cache 显著提高了性能，并且它需要一大块内存，是使用有状态缓冲区的完美候选者。这里有一个关于有状态模型的 coremltools 用户指南更新。

新的量化技术

在 WWDC 23 中，我们探索了一种非常酷的技术，称为 palletization，我们展示了它如何帮助将文本到图像模型，如 Stable Diffusion，带到 Mac 和 iPhone 上。

虽然这些技术可以让你大幅减小模型体积，但如果过度使用，对质量的影响将是巨大的。较大的模型受此影响更大，因为其权重数据具有广泛的动态范围。试图创建一个能涵盖所有可能值的小型查找表（LUT）变得越来越困难。WWDC 24 引入的解决方案是每次专注于数据的一小部分，并为同一张量的不同区域创建多个查找表。

这些方法（分块量化）使我们能够将模型压缩到低至 4 位精度。我们不再使用 4 个字节（一个 float32 数字的大小）来表示每个模型参数，而是每个参数只需半个字节（一个半字节）。这使得模型大小缩减为原来的八分之一（减去分块量化表带来的一些额外开销），或者与 float16 精度相比缩小为四分之一。

多功能支持

本示例中我们不会使用这一功能，但我们想在此提及，因为它在 WWDC 24 上被引入，我们将在后续的一些工作中展示它。多功能支持本质上允许你将 LoRA 适配器打包进生成式模型中，以便将同一个模型（附带一小组称为适配器的额外参数）用于不同的任务。LoRA 是社区首选的大模型微调技术。例如，在扩散模型中，你可以使用 LoRA 生成不同风格的图像，如照片级写实或卡通风格。我们认为 LoRA 是支撑 Apple Genmoji 实现的部分方案。对于语言模型，LoRA 适配器可用于将通用 LLM 适配到特定任务或领域。

要了解更多关于 LoRA 的信息，你可以查看这篇文章。

要了解更多关于多功能的信息，你可以查看 Apple coremltools 用户指南此处。

将 Mistral 7B 转换为 Core ML

高效运行大型语言模型最重要的组件是 kv-cache。如上所述，它是 WWDC 24 发布的新有状态模型功能的绝佳候选。transformers 库中的模型已经使用了高效依赖 kv-caching 的注意力实现。然而，默认实现是针对 Nvidia GPU 优化的，而该硬件与 Apple Silicon 有着不同的约束条件。在 Core ML 的情况下，我们需要预先分配完整的缓存缓冲区，并确保每次调用模型时都在原地更新缓冲区。这避免了低效的内存分配和张量拼接，同时也是 Core ML 有状态缓冲区的一项要求。

为了实现这一目标，我们必须使用一种考虑这些因素的不同注意力实现。这需要修改 Mistral 架构的 transformers 建模代码，具体在这段代码片段中完成。

注意：如果你想跟随并复现整个转换过程（或转换另一个基于 Mistral 的模型，比如不同的微调版本），你可以使用这个脚本来运行所有转换步骤。

追踪与转换

第一步是加载模型。我们将使用带有原地缓存方法的修补实现。

MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
torch_model = StatefulMistralForCausalLM(MODEL_ID)
torch_model.eval()

在运行 Core ML 转换之前，我们需要用示例输入来追踪模型。这个过程会记录在这些输入上执行的张量操作，追踪得到的图将在转换期间被翻译为 Core ML 操作。我们使用示例输入来追踪模型；不需要真实数据。

input_ids = torch.zeros((1, 2), dtype=torch.int32)
causal_mask = torch.zeros((1, 1, 2, 5), dtype=torch.float32)

traced_model = torch.jit.trace(torch_model, [input_ids, causal_mask])

语言模型的输入是长度可变的 token 序列。我们允许输入从单个 token 增长到最大上下文长度 2048。我们可以使用 coremltools 的范围维度来指定这些边界。

query_length = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)
end_step_dim = ct.RangeDim(lower_bound=1, upper_bound=2048, default=1)

inputs = [
    ct.TensorType(shape=(1, query_length), dtype=np.int32, name="inputIds"),
    ct.TensorType(shape=(1, 1, query_length, end_step_dim), dtype=np.float16, name="causalMask"),
]

outputs = [ct.TensorType(dtype=np.float16, name="logits")]

除了序列 token（在上面的示例中称为 inputIds）之外，还有另一个输入称为 causalMask，它指定模型需要关注的 token。这主要在使用批处理同时生成多个序列时使用。请查看此处的示例运行器中如何使用这些输入。

在这种情况下，一个批次内的所有输入序列必须具有相同的长度，因此我们使用填充 token 和因果掩码来告诉模型，填充 token 不应被视为输入。

状态准备

PyTorch 建模代码使用 keyCache 和 valueCache 作为缓存缓冲区的名称来保存 kv-cache。这些块按最大上下文长度（2048）分配。我们使用 coremltools 新的 StateType 来指定这些块在转换期间必须转换为有状态的 Core ML 缓冲区。

# Specify kv-cache states by using `StateType`.

states = [
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="keyCache",
    ),
    ct.StateType(
        wrapped_type=ct.TensorType(shape=torch_model.kv_cache_shape, dtype=np.float16),
        name="valueCache",
    ),
]

Core ML 转换

要将模型转换为 Core ML，我们需要指定输入和输出类型以及状态。转换后的模型将使用 float16 精度，因为这是我们对输入数据指定的精度。我们还需要将最低部署目标指定为 iOS18，因为这是这些功能可用的地方。（我们也可以使用 macOS15，它指的是相同的转换目标。）

mlmodel_fp16 = ct.convert(
    traced_model,
    inputs=inputs,
    states=states,
    outputs=outputs,
    minimum_deployment_target=ct.target.iOS18,
    skip_model_load=True,
)

模型压缩

使用上述新的块大小量化策略，我们使用块大小为 32 的 4 位线性量化。这将大大减小模型大小并使模型运行得更快。尽管计算仍将以 float16 执行，但权重以 4 位模式传输并即时解压缩，这比传输大量 16 位权重更高效。

量化参数配置如下：

op_config = ct.optimize.coreml.OpLinearQuantizerConfig(
    mode="linear_symmetric",
    dtype="int4",
    granularity="per_block",
    block_size=32,
)
config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)

让我们使用该配置来量化模型。下面这行代码将需要几分钟才能运行：

mlmodel_int4 = ct.optimize.coreml.linear_quantize_weights(mlmodel_fp16, config=config)

mlmodel_int4.save("StatefulMistral7BInstructInt4.mlpackage")

转换和量化完成后还有最后一步。我们需要包含一段额外的元数据，指明我们使用的模型标识符（mistralai/Mistral-7B-Instruct-v0.3）。Swift 代码将使用它从 Hub 下载 tokenizer 文件。分词是将文本数据转换为模型使用的数值表示，每个模型都不同。

mlmodel_int4._spec.description.metadata.userDefined.update({
    "co.huggingface.exporters.name": MODEL_ID
})

生成的模型是一个约 3.8G 的 mlpackage，相比之下，float16 转换会产生 14G。你可以在 Hub 上找到它。

使用 Swift 运行 Mistral 7B

如果你按照上述步骤操作或从 Hub 下载了模型，你可以使用 swift-transformers 的 preview 分支在本地运行它。Apple 工程师将其贡献给了该项目，包括以下重要功能：

完全支持 Tensor，这极大地简化了预处理和后处理任务，并允许我们删除许多行低级、令人困惑且脆弱的代码。
支持 Stateful API 的 Swift 对应版本。

由于采用这些功能是一项破坏性变更，并且需要 iOS 18 或 macOS 15，我们暂时会将其保留在 preview 分支中。

要从命令行运行该模型，请先从 GitHub 仓库克隆 preview 分支：

    git clone -b preview https://github.com/huggingface/swift-transformers

然后运行 CLI 来测试该模型：

#to run in release mode, pass -c release
swift run transformers "Best recommendations for a place to visit in Paris in August 2024:" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

为了更方便测试，你还可以使用 swift-chat，这是我们编写的一个简单应用，用来展示如何在其中集成 swift-transformers 包。你也必须使用 preview 分支。本文开头展示了 swift-chat 运行转换后的 Mistral 模型的示例。

用 Python 运行 Mistral 7B

对于更熟悉 Python 的读者来说，这同样很简单！

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Best recommendations for a place to visit in Paris in August 2024:"

coremltools 让用 Python 运行 Core ML 模型也同样简单。

下一步是什么？

我们对今年 Core ML 和 coremltools 的进展感到非常兴奋，也期待看到大量第三方应用利用 ML 模型来解决人们需要的真实任务。就我们而言，我们致力于让这件事尽可能简单，以便开发者能够专注于创建出色的应用。我们的规划中有几件事：

这里介绍的模型更新非常适合 Mac 电脑上的 GPU。Core ML 可以使用 Neural Engine，而它在 iPhone 上尤其高效。要从 Neural Engine 获得最佳性能，还需要做一些额外适配，我们计划在几个示例模型上开展这项工作。这项工作将基于 Apple 这篇 2022 年（且仍然非常相关）文章中讨论的经验。我们不会在 iPhone 上运行 Mistral 7B，但有几个更小的模型，比如 Apple 的 OpenELM 或 DCLM，非常适合作为探索对象！
这里展示的代码具有高度实验性。随着夏季推进，我们计划采用这些方法，并将它们整合到 exporters 中，这是一个旨在将 transformers 模型转换为 Core ML 的 Python 工具。希望很快你就能非常轻松地转换许多有趣的模型架构。
我们会继续维护 swift-transformers 的 preview 分支，以便在新功能或 API 变更发布时将其纳入。如果你感兴趣，请持续关注它！

你可以如何提供帮助？

Apple 在 WWDC 上发布的工具有助于我们实现长期目标：让 AI 变得简单，并让所有人都能使用。我们很想知道你能把它们带到哪里。我们展示的示例是实验性的，但你可以用它把任何 Mistral 微调模型转换为 Core ML——如果你这样做了，请告诉我们！如果你想尝试其他模型架构，欢迎向 swift-transformers 的 preview 分支提交 issue 或 PR——我们会尽量帮你上手！

再也没有比今天更好的时机，把你的创造力用于解决你感兴趣的问题了！去尝试吧，玩得开心，并告诉我们怎样才能帮到你。