使用 Core ML 运行 Mistral 7B

Hugging Face 详细介绍了一种在 Mac 硬件上运行 Mistral 7B 模型的工作流，利用了 WWDC 24 推出的新 Core ML 功能。通过结合状态缓冲区、新的 Swift Tensor 类型以及高级量化技术，7 千亿参数的模型可以在不到 4 GB 的内存下运行。

来自 WWDC 24 的新 Core ML 功能

Apple 最近对 Core ML 的更新为在设备上运行的大语言模型（LLM）提供了若干关键优化。

Swift Tensor（MLTensor）

全新的 MLTensor 类型为 Swift 中的多维数据结构提供了高级抽象，类似于 numpy 数组或 torch 张量的功能。它取代了对 MLMultiArray 或 MLShapedArray 的低层操作，后者以前需要将底层存储视为不透明指针进行访问。MLTensor 内置了诸如 softmax 的操作，简化了语言模型的前后处理流水线。

状态缓冲区

传统上，Core ML 模型以无状态函数的形式运行。状态缓冲区允许模型在 GPU 上预留一块内存，以在多次迭代中保持状态。对于 LLM，这用于实现键值（KV）缓存。将 KV 缓存保留在 GPU 上，可避免在每次生成 token 时在 CPU 与 GPU 之间传输大量状态数据的开销，从而通过降低内存带宽瓶颈显著提升性能。

块状量化

为在不显著降低质量的前提下降低模型体积，Core ML 现已支持块状量化。该技术为同一张量的不同区域创建多个查找表（LUT），而不是为整个张量使用单一表。这样可以将模型压缩至 4 位精度，与 float32 相比体积缩小 8 倍，与 float16 相比缩小 4 倍。

多功能支持

多功能支持使开发者能够将 LoRA（低秩适配）适配器打包进生成模型。这样，只需切换少量额外参数（适配器），即可让同一基础模型用于不同任务或风格，而无需重新加载整个模型。

将 Mistral 7B 转换为 Core ML

高效运行 Mistral 7B 需要自定义注意力实现，预先分配完整的 KV 缓存缓冲区并原位更新，以满足 Core ML 的状态缓冲区要求。

跟踪与转换流程

  1. Tracing: 模型使用 StatefulMistralForCausalLM 的补丁实现加载，并通过 torch.jit.trace 使用示例输入进行跟踪。
  2. Input Definition: 通过 coremltools 使用范围维度，使输入序列可以从单个 token 扩展到最大上下文长度 2048。
  3. State Preparation: 使用 ct.StateType 定义 keyCache 和 valueCache 缓冲区，以确保它们被转换为状态化的 Core ML 缓冲区。
  4. Conversion: 将模型转换为 Core ML，最低部署目标为 iOS 18 或 macOS 15，以使用新的状态化特性。

模型压缩

使用 OpLinearQuantizerConfig，模型通过 4 位线性对称量化（块大小为 32）进行压缩。这将最终的 mlpackage 大小从约 14 GB（float16）降低至约 3.8 GB。

执行与实现

用户可以使用 swift-transformers 仓库的 preview 分支运行已转换的 Mistral 7B 模型。

使用 Swift 运行

preview 分支已完整支持 MLTensor，并提供了 Swift 版的 Stateful API。可通过命令行执行推理：

git clone -b preview https://github.com/huggingface/swift-transformers
swift run transformers "Prompt text" --max-length 128 Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage

使用 Python 运行

也可以使用 Python 通过 coremltools 进行推理：

python3 generate.py Examples/Mistral7B/StatefulMistral7BInstructInt4.mlpackage --prompt "Prompt text"

未来路线图

Hugging Face 计划将这些实验性方法整合到 exporters 中，这是一款用于将 transformers 模型转换为 Core ML 的 Python 工具。此外，他们还希望通过探索 OpenELM 或 DCLM 等更小的模型，优化 iPhone 上 Apple Neural Engine（ANE）的性能，因为当前的 Mistral 7B 实现主要针对 Mac GPU 进行优化。

Sources