WWDC 24: 使用 Core ML 运行 Mistral 7B
WWDC’ 24 是苹果正式发布 Apple Intelligence 并重申其对高效、私密、设备端 AI 承诺的时刻。在主题演讲及随后的分会场中，他们展示了 Apple Intelligence，它驱动了大量 AI 增强功能，展示了 AI 在日常任务中的实际应用。这些不是*为了 AI 而 AI* 的华丽演示，而是能节省时间、恰当（且有趣！）的助手，它们深度集成于应用程序和操作系统中，并为开发者提供了多种方式将这些功能融入他们自己的应用中。
Apple Intelligence 的各项功能之所以能如此出色，得益于其垂直整合的软件栈，它能够最大限度地发挥 Apple Silicon 的能力。苹果还为开发者提供了一个在设备上运行模型的平台，即 Core ML。这个软件栈允许你在 Apple Silicon 硬件上所有三种计算单元（CPU、GPU 和神经网络引擎）上运行机器学习模型。
在本篇博文中，我们将探讨一些 Core ML 最出色的新功能，以复现苹果在 WWDC'24 使用 Core ML 在设备端部署机器学习和 AI 模型分会场中展示的 Mistral 7B 示例。在该示例中，他们使用了一个 swift-transformers 的 fork 版本，在 Mac 上运行了一个顶尖的 LLM。这是一个拥有超过 70 亿参数的高质量模型，它挑战了当今消费级硬件的极限。你也可以查看 WWDC'24 的将你的机器学习和 AI 模型引入 Apple silicon 分会场，其中展示了 Mistral 7B 转换过程的一部分。
让我们看看需要采取哪些步骤来尽可能高效地运行它，并学习 iOS 18 和 macOS Sequoia 中可用的新工具。
这是我们今天要构建的内容
太长不看版
读完这篇博文，你将了解最新 macOS 版本带来的所有新特性，并成功地在你的 Mac 上用不到 4GB 的内存运行一个 70 亿参数的模型。
第一步：克隆 swift-transformers 仓库的 preview 分支：git clone -b preview https://github.com/huggingface/swift-transformers 第二步：从 此 Hugging Face 仓库 下载转换后的 Core ML 模型 第三步：使用 Swift 运行推理：swift run transformers "2024年8月在巴黎游览的最佳推荐地点是：" --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 的语言模型。语言模型接收一个 token 输入序列，并输出对词汇表中所有 token 概率的估计，这意味着概率高的 token 很有可能是输入的合理延续。应用程序的工作是根据这些概率选择最佳的下一个 token 来追加到序列中。Tensor 类型使得处理这些操作变得容易，无需自定义代码。
当我们发布 swift-transformers 时，我们编写了大量代码（后来由社区扩展，谢谢！❤️）来帮助进行输入准备（将单词转换为 token）和输出后处理。例如，可以查看我们使用 Accelerate 实现的 softmax 操作。所有这些在使用 MLTensor 时都可以移除，因为 softmax 是开箱即用的！
有状态缓冲区 (Stateful Buffers)
在 WWDC’ 24 之前，Core ML 模型本质上是一个纯粹的无状态函数，你提供输入并返回一些输出。然而，有时你需要保持一个依赖于先前计算的状态。维护状态的函数式编程方法是增加一个额外的输入/输出对。因此，模型根据你的输入和状态计算输出和新状态。这种方法没有错，事实上，像 JAX 这样的高性能框架就是这样工作的。
然而，这种方法存在实际限制：每次调用模型时，都需要将有状态的数据作为输入发送给模型，并作为输出取回。如果有状态的数据很大，这种来回传输会增加开销并降低速度。这对于 LLM 尤其重要，因为你需要运行多次迭代来生成一个序列。性能瓶颈通常是你的计算机内存带宽（即，你将数据移至 GPU 及返回的速度有多快）。有状态模型通过为状态数据预留一块内存并将其保留在 GPU 上来解决这个问题，这样你就不必每次使用模型时都发送和接收它了。
有状态缓冲区是在 这次 WWDC' 24 分会场 中通过一个易于理解但不代表像 LLM 这样大型模型的实际应用的玩具示例来介绍的。对于基于 Transformer 的模型，一个 LLM 性能技巧是键值缓存（称为 kv-caching）。如下图所示，它通过缓存先前步骤中执行的操作结果来避免关键的注意力模块中昂贵的矩阵乘法。我们不会深入细节，但关键点是：kv-cache 显著提高了性能，并且它需要一大块内存，这是使用有状态缓冲区的完美候选者。这里有一份关于有状态模型的 coremltools 用户指南更新。
新的量化技术
在 WWDC 23 中，我们探讨了一种非常酷的技术，称为调色板化 (palletization)，并展示了它如何帮助将文生图模型，例如 Stable Diffusion，带到 Mac 和 iPhone 上。
虽然这些技术可以让你大幅减小模型尺寸，但如果压缩得太过，对质量的影响是巨大的。更大的模型受此影响更严重，因为权重数据的动态范围很广。创建一个能够捕捉所有可能值的小型查找表（LUT）变得越来越困难。WWDC 24 中引入的解决方案是每次关注数据的一小部分，并为同一张量的不同区域创建多个查找表。
这些方法（分块量化）使我们能够将模型压缩到低至 4 位精度。我们不再使用 4 字节（一个 float32 数字的大小）来表示每个模型参数，而是可以只用半个字节（一个 nibble）。这使得模型大小减少了 8 倍（减去一些用于分块量化表的开销），或者与 float16 精度相比小了 4 倍。
多函数支持
虽然这个例子我们不会用到这个功能，但我们想在这里提一下，因为它是在 WWDC 24 上推出的，我们将在一些未来的工作中展示它。多函数支持基本上允许你将 LoRA 适配器打包到生成模型中，以便为不同的任务使用同一个模型（只需一小组额外的参数，称为适配器）。LoRA 是社区首选的大模型微调技术。例如，在扩散模型中，你可以使用 LoRA 生成不同风格的图像，如照片般逼真或卡通风格。我们相信 LoRA 是驱动苹果 Genmoji 实现的解决方案的一部分。对于语言模型，LoRA 适配器可用于将通用 LLM 调整到特定任务或领域。
要了解更多关于 LoRA 的信息，你可以查看这篇文章。
要了解更多关于多函数支持的信息，你可以查看苹果 coremltools 用户指南这里。
将 Mistral 7B 转换为 Core ML
高效运行大型语言模型最重要的单一组件是 kv-cache。如上所述，这是 WWDC' 24 发布的新有状态模型特性的绝佳候选者。transformers 库中的模型已经使用了高效的注意力实现，这些实现严重依赖于 kv-caching。然而，默认实现是为 Nvidia GPU 优化的，而这种硬件与 Apple Silicon 有着不同的约束。对于 Core ML，我们需要预先分配完整的缓存缓冲区，并确保每次调用模型时，我们都能就地更新缓冲区。这避免了低效的内存分配和张量拼接，同时也是 Core ML 有状态缓冲区的要求。
为了实现这个目标，我们必须使用一种考虑了这些因素的不同注意力实现。这需要修改 Mistral 架构的 transformers 建模代码，这是在这段代码片段中完成的。
注意：如果你想跟着做并复现转换过程（或转换另一个基于 Mistral 的模型，例如不同的微调版本），你可以使用这个脚本来运行所有的转换步骤。
追踪与转换
第一步是加载模型。我们将使用带有就地缓存 (in-place cache) 方法的补丁实现。
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 文件。Tokenization 是将文本数据转换为模型使用的数值表示的过程，每个模型都不同。
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 分支在本地运行它。苹果工程师向该项目贡献了代码，包括以下重要功能
完整的
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 分支。本文开头展示了运行转换后的 Mistral 模型的 swift-chat 示例。
使用 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 可以使用神经网络引擎，这在 iPhone 上尤其高效。要从神经网络引擎中获得最佳性能，需要一些额外的适配，我们计划在一些示例模型上进行这些工作。这项工作将基于这篇 2022 年（至今仍然非常相关）的苹果文章中讨论的经验。我们不会在 iPhone 上运行 Mistral 7B，但有几个较小的模型，如苹果的 OpenELM 或 DCLM，是探索的绝佳候选者！
这里展示的代码是高度实验性的。随着夏天的进行，我们计划采用这些方法并将其整合到
exporters中，这是一个旨在将 transformers 模型转换为 Core ML 的 Python 工具。希望你很快就能非常轻松地转换许多有趣的模型架构。我们将继续在
swift-transformers的preview分支上工作，以整合新功能或 API 变更。如果你感兴趣，请关注它！
你能如何帮助我们？
苹果在 WWDC 上发布的工具帮助我们实现了让 AI 对所有人来说都简单易用的长期目标，我们很乐意看到你将它们带向何方。我们展示的例子是实验性的，但你可以用它来将任何 Mistral 微调模型转换为 Core ML——如果你这么做了，请告诉我们！如果你想尝试其他模型架构，请随时向 swift-transformers 的 preview 分支提出问题或 PR——我们会尽力帮助你开始！
现在是运用你的创造力来解决你感兴趣的问题的最佳时机！去尝试，去享受，并告诉我们如何能帮助你。
