将机器学习和 AI 模型移植到 Apple 芯片

大家好 我是 Qiqi Xiao 我是 Core ML 团队的工程师 今天我将介绍 Core ML 工具迎来的 几项激动人心的更新 这些更新将帮助你 更好地将机器学习 和 AI 模型移植到 Apple 芯片 模型部署工作流程中 有三个重要阶段 我将重点关注准备环节 在这方面分享一些优化措施 确保你可以充分利用相关功能 更高效地在设备上 运行你的模型

在本次讲座中 我假设大家 已经有了一个机器学习模型 无论是预训练模型、经过微调的模型 还是从头开始训练的模型 都可以 如果你想进一步 了解模型训练 建议你观看今年的讲座视频： “基于 Apple GPU 训练 机器学习和 AI 模型”

今天的讲座不会介绍 如何编写代码 将你的模型整合到 App 中 建议大家观看下面的视频 进一步了解模型整合

现在 让我们开始吧

Core ML 工具是 一个开源 Python 软件包 包含一些实用程序 可优化和转换模型 以便与 Apple 框架配合使用 利用这个工具 可将 PyTorch 模型 转换为 Core ML 格式 这个格式已针对 Apple 芯片的执行能力进行优化

Apple 芯片的统一内存、CPU、 GPU 和神经网络引擎可为设备端 机器学习工作负载 提供低延迟、高效的计算 默认情况下 只需将模型 转换为 Core ML 格式 同时使用 Apple 的推理框架 你的 App 就能充分利用 Apple 芯片的澎湃动力

Apple 芯片为我们的 所有平台提供了强大支持 不过 每个平台都有自己 独有的特性和优势 因此 你需要综合考虑 目标平台和设备的 存储、内存和计算处理能力 这些属性需要与你的用例中 模型的大小、 准确度和延迟要求相匹配 在模型准备工作中 我们需要探索可以选择哪些方案 然后进行各种优化 来找到完美平衡点 让我们一起来探讨如何实现 首先 我会介绍模型压缩方面的内容 在本部分中 你将学习如何使用 不同的压缩技术 和工作流程 之后 我将有请 Junpei 为大家详细介绍 对有状态模型的支持 他还将介绍 转换器架构的优化

最后 他将介绍 对多函数模型的支持

现在 我们先从模型压缩开始

如今 高性能机器学习模型 往往非常大 这类模型需要更多的存储空间 会导致更高的延迟和更大的能耗 因此更难 在设备端高效部署 面对这些挑战 我们必须 通过模型压缩 减小模型大小

让我们先回顾一下 去年推出的 权重压缩技术

第一种技术 名为堆积 堆积会将具有相似值的权重 聚类到一组 并用聚类中心点的值来表示 聚类中心点 存储在一个查找表中 查找表的 N 位索引映射 就是压缩权重 在本例中 权重 被压缩为 2 位

第二种技术 名为量化 要进行量化 需要将浮点权重值 线性映射到整数范围 整数权重 与一对标度和偏置一起存储 称为量化参数 这些参数之后可用于 将整数映射回浮点数

第三种技术是剪枝 剪枝可以帮助你 以稀疏表示形式 高效地打包模型权重 首先对于权重矩阵 我将最小值设为 0 现在 我只需要存储位掩码 和非零值

现有技术 在很多模型上都很有效 现在让我们将其中 一些技术应用于某个大模型 看看我们如何扩展这些技术 实现进一步压缩

这里是一个 稳定扩散模型的例子 这个模型采用自然语言描述 也就是提示 可生成与这个描述 相匹配的图像

转换后 我发现最大模型的模型大小 在 Float16 精度下超过 5 GB 如果要在 iPhone 或 iPad 上运行 首先一点就是太大了 为了部署这个模型 我们需要对它进行压缩

让我们试试堆积 通过灵活地 选择位数 来实现不同的压缩比

例如 应用 8 位堆积 可将模型大小减小到 Float16 模型的一半左右 使用提示语 “穿燕尾服的猫 油画” 可在 Mac 上得到 很棒的图像 不过 我要先让这个模型 小于 2 GB 然后再考虑 与 iPhone 和 iOS 整合 6 位可满足这一要求 我终于可以在 iPad 上 运行这个模型了

那么 4 位呢？ 这时就无法生成符合描述的图像了 这意味着模型并不准确 这是怎么回事？

为方便理解 我们需要先重温一下 堆积表示法 iOS 17 只支持 按张量堆积 其中所有值都集中 在一个查找表中 使用 4 位堆积时 只有 16 个聚类中心点 可以映射回整个张量 由于粒度较低 因此会给大型矩阵带来 更多潜在误差 现在 我们来思考 如何增加粒度

我们先退一步来说 我想解释一下 机器学习领域 常用的几个术语 权重可以用矩阵表示 比如这里我用矩阵行 来表示输出通道 用矩阵列表示输入通道

我们可以 对整个张量进行压缩 也可对每组通道进行压缩 或者对每个通道单独压缩

我们甚至可以将每行 分割成更小的块 在子通道级别 应用压缩

图中从左到右 粒度逐渐变大

回到堆积话题 在 iOS 17 中 我们可以 将单个查找表分配给 整个权重张量 而在 iOS 18 中 我们现在可以存储 多个查找表 每组通道一个 这样 就能获得 更高的精度

对于线性量化 iOS 17 允许 按通道提供标度和偏置 而在 iOS 18 中 你可以 逐块提供这些量化参数

即使是剪枝技术 我们现在也可以进一步压缩模型

之前 密集权重经过稀疏化处理 变成了浮点精度的非零值 现在 浮点非零值 可使用堆积或量化 进一步压缩 这意味着模型可以更小 我们可以同时利用 两种不同压缩技术的优势 稀疏堆积 和稀疏量化 现在都可以实现

在前面讨论的 三种扩展中 让我们以堆积为例 看看如何在一些 Python 代码中 加以应用 首先 创建一个 OpPalettizerConfig 对象 描述你希望 如何压缩模型 在这里 我选择 4 位 使用“kmeans”模式进行聚类 我想为一组 16 个通道 分配 1 个查找表 因此 我将粒度设置为 “per_grouped_channel” 并将 group_size 设置为 16

定义好配置后 就可以将它放入 OptimizationConfig 然后使用 palettize_weights 来压缩模型

我们可以看到 这个过程非常简单

现在让我们看看 这个模型的表现如何 将 group_size 设为 16 时 可以 看到生成了另一只符合提示的猫 生成的图像 比以前好得多 这意味着我重新获得了 大部分的精确度 但大小只从 1.29 GB 稍微 增加到了 1.3 GB 增加的原因是 额外的查找表 占用了一点额外的空间

综上所述 在 iOS 18 中 按分组通道堆积 可通过多个查找表 来增加粒度 这种技术在具备神经网络引擎的 Apple 芯片上表现出色 除了逐块量化外 我们还将支持从 8 位 扩展到 4 位量化 这种模式特别针对 Mac 上的 GPU 进行了优化 最后 现在你可以将稀疏性 与其他压缩模式结合起来 这种方法在神经网络引擎上效果很好

在结束模型压缩内容之前 我想讨论一下 压缩工作流程 这是可以帮助你获得 更小模型的另一途径 到目前为止 我们一直使用的 压缩工作流程 是一种后训练方法 首先需要一个 预训练 Core ML 模型 在没有任何数据的情况下 非常便捷地进行压缩

然而 这种无数据 工作流程可能无法 在精确度和压缩量之间 实现最佳平衡 压缩率越高 精确度下降得越快

解决这个问题的 一种办法是使用训练时间压缩

你可以用训练数据 对 PyTorch 模型进行微调 同时压缩权重 然后 你可以将它 转换为 Core ML 格式

虽然这可以让你 获得更精确的模型 但这个过程可能会很耗时 而且 对大模型来说 需要的数据量也很大

现在 我们有一个新的工作流程 那就是使用校准数据 进行后训练压缩

这是一种介于无数据和 微调方法之间的折衷办法 在这种情况下 对模型压缩效果进行校准 所需的数据量不多

也完全不需要微调 因此设置更轻松 花费的时间也更少

今年 coremltools.optimize 推出了新的 API 执行这个工作流程会非常轻松

让我们举一个例子加以说明

这里我以剪枝为例 首先导入 必要的模块 现在可以指定 prune_config 这里我将 target_sparsity 设为 0.4 这意味着 40% 的权重值 将被剪掉 样本数设置为 128 然后 我用 LayerwiseCompressor 创建了剪枝器 使用 PyTorch 模型 和定义的配置 现在我开始处理 calibration_data_loader 你可以自行为你的 模型进行定义

然后 你可以调用 pruner.compress 使用 calibration_data_loader 来校准稀疏模型

现在可以将稀疏模型 转换为 Core ML 格式 或者像我们讨论过的那样 也可以进行稀疏堆积 你可以在剪枝后进行堆积 进一步压缩模型

为了传递稀疏元数据信息 我们继续使用 之前获得的稀疏模型 同样 你需要导入必要的模块 然后进行堆积配置 这里 我还将 4 位粒度 设置为“per_grouped_channel” group_size 设置为 16

你现在可以创建 PostTrainingPalettizer 对象 并使用 palettizer.compress 获取 sparse_palettized_model 这个 PyTorch 模型还可以无缝 转换为 Core ML 格式

在应用 40% 的稀疏度之后 稳定扩散模型的大小 进一步减小到 1.1 GB 在无数据后训练压缩流程下 模型只能产生噪声输出 而在数据校准工作流程下 模型还能生成 另一张猫图像 基于校准的工作流程效果很好 我喜欢这张新图像

总而言之 在 Core ML 工具 8 中 你可以尝试新的 压缩表示法

还可以选择 使用校准数据的 新压缩工作流程 今年还推出了许多其他新功能 由于时间有限 本视频无法逐一为大家介绍 请查看我们的文档 进一步了解 希望大家压缩成功

现在 我们有请 Junpei 来谈谈有状态模型 和其他优化

谢谢 Qiqi 大家好 我是 Junpei Zhou 是 Core ML 团队的工程师 我将谈谈有状态模型、 转换器模型的优化 以及对多函数模型的支持 帮助大家针对 Apple 芯片 准备机器学习模型

我们先来看看 对有状态模型的支持

对于任何机器学习模型 我们通常只有输入和输出 模型中的运算 处理输入后就会产生输出

通常 模型内部的数据流 只适用于特定的推理运行

然而 有时我们可能希望 数据有效期更长一些 这样就能在不同的 运行中存储信息 在这种情况下 通常会使用状态

模型可以从状态中读取数据 并将结果写回状态

状态中的信息是持久的 可以在不同的运行中使用

下面是一个累加器的例子 它会跟踪所有历史输入的总和

状态初始化为零

模型从状态和输入中读取数值 从而产生输出

同时 还会将结果写回状态

同样 出现另一个输入时 模型就会读取状态 产生输出 同时更新状态

在 Core ML 支持状态之前 在模型中处理状态的一种方法是 显式定义 额外的输入和输出

以累加器为例 我们需要累加器的 输入和输出字段

然后 需要将累加器输出 显式拷贝到 下一个累加器输入

这种方法的效率可能不高 尤其是在状态 较大的情况下

今年 Core ML 推出了 对有状态模型的支持

现在 模型会自动更新 状态张量 因此 你无需将状态张量 定义为输入或输出 相反 你可以就地更新状态 从而获得更好的性能

以累加器为例 在 PyTorch 中 我会这样定义模型

我会使用 torch 中的 register_buffer API 来表示模型的状态张量

完成这一步后 将有状态的 torch 模型 转换为有状态的 Core ML 模型 也很简单

只需使用新推出的 ct.StateType 来指定状态就可以 在这里提供的状态名称 应该与 register_buffer API 相同

然后 将状态传递给 ct.convert 用来转换有状态模型

现在 转换后的 Core ML 模型 将是有状态模型 在 torch 模型中 状态的读写 通常由就地运算 和切片更新运算完成

在转换过程中 Core ML 工具 会在转换后的 Core ML 模型中 生成相应的运算

如果你有兴趣进一步了解 如何在设备端部署 有状态 Core ML 模型 请观看这个视频

现在我们来谈谈 转换器优化

转换器是基础模型中 最常用的模型架构之一

它们由多个 多头注意力块构建而成 这些注意力块需耗费大量算力 通常是模型推理过程中的瓶颈

注意力块由 多个 MatMul 和 Softmax 运算 以及 3 种大型张量组成 分别称为 查询、键和值 为了生成标记 当前词的查询向量 会关注序列中先前 所有词的键向量 然后根据注意力得分 融合值向量

也就是说 在每个时间戳 先前标记的 键和值向量 已经在之前的步骤中计算过了

为避免这种重复计算 我们可以使用键值缓存 简称 KV 缓存 来存储每一步计算出的 键和值向量

以便在之后的步骤中使用

在大语言模型中 KV 缓存是一种非常有效的常用技术 可加快解码速度 这种情况非常适合使用 我们之前介绍过的 有状态模型

如果没有状态 KV 缓存只能用于 I/O 考虑到 KV 缓存的张量通常很大 因此这样做的效率很低

有了状态 就可以 更高效地进行就地更新

这种有状态的 KV 缓存非常有用 尤其是针对 Apple 芯片 准备转换器模型时

让我们回到 注意力块结构 前面介绍了 有状态 KV 缓存 如何提高注意力的 键和值准备效率 这让我们自然会想到另一个问题：

能否提高注意力计算的效率？

这就引出了另一种 优化转换器模型的方法

那就是缩放点积注意力计算的 融合表示法 缩放点积注意力也称 SDPA

在这之前 使用 SDPA 运算 转换 torch 模型时 Core ML 工具会将 SDPA 运算 分解成几个较小的运算 包括 matmul、softmax 和其他运算

今年 将最低部署目标 设为 iOS 18 后 Core ML 工具会在转换后的 Core ML 模型中使用 SDPA 运算

这个 SDPA 运算 会一次性接收所有输入 更高效地计算注意力 这种方法在 Apple 芯片 GPU 上 效果非常出色

如果你想进一步了解如何 使用 Metal 优化 SDPA 运算 强烈建议观看这个视频

介绍了这些精彩的功能之后 我来演示一下如何 使用 Core ML 工具 针对 Apple 芯片 准备 Mistral-7B 模型 Mistral-7B 是一个强大的 大语言模型 拥有 73 亿个参数 利用今年推出的 有状态 KV 缓存和 SDPA 运算 我们可以将它 转换为 Core ML 模型 以便在 Apple 芯片上 流畅运行

让我们使用 “Jupyter Notebook”这款工具 来展示如何将 Mistral-7B 模型 转换为有状态 Core ML 模型 我们采用的方法是 后训练逐块量化

我来导入一些 utils 让代码更简洁、更容易理解

第一步是 让原始 Mistral 7B 模型变为有状态

从原始 Mistral-7B 模型入手 我们需要修改模型 使它变为有状态模型 以便高效处理 KV 缓存

我们可以为 KV 缓存注册缓冲区 将缓冲区命名为 “keyCache”和“valueCache”

模型中还有其他一些修改 比如在前向传递过程中 对 KV 缓存进行就地更新

你可以在后面的 utils 中 找到详细信息

在前向函数中 由于我们已经 为 KV 缓存注册了缓冲区 这里我们不再需要 将它作为输入

完成所有这些改动后 就得到了有状态 Mistral-7B 模型 它将 KV 缓存作为持久状态

现在 我们可以使用 从 HuggingFace 自动下载的权重 来初始化这个模型 并将模型设置为 eval 模式

指定示例输入后 对于这里的 Mistral-7B 模型 我们 需要 input_ids 和 causal_mask 然后我们就可以跟踪模型了

要转换这个模型 我们同样要指定输入和输出 这里我们使用 ct.RangeDim 来指定输入形状 因为不同运行中的 提示长度可能不同

为确保转换后的 Core ML 模型 也具有 KV 缓存的状态 我们需要使用 今年新推出的 ct.StateType 来指定状态

这里 我们使用“keyCache” 和“valueCache”这两个名称 与我们在 torch 模型中 注册缓冲区时使用的名称保持一致

最后 我们使用输入、 输出和状态来调用 ct.convert 并将 minimum_deployment_target 设置为 iOS 18

将这些状态传递给 ct.convert 后 转换后的 Core ML 模型 就是一个有状态模型 具有 k 缓存和 v 缓存 可以就地更新

转换完成后 让我们把模型 保存到磁盘 检查模型大小

转换完成 我们来检查一下模型的磁盘大小

模型相当大 有 13 GB

让我们运行模型预测 来测试模型 模型使用的是 Core ML 框架 这个提示是让模型 推荐六月份 去西雅图旅行的好去处

我们可以看到 尽管模型相当大 转换后的有状态 Core ML 模型 仍然可在 Mac 上流畅运行

这个 Float16 模型 有点太大了 让我们用一下 Qiqi 刚才提到的 一些模型压缩技术

我们可以使用 ct.optimize 模块 来指定压缩 config

在 OpLinearQuantizerConfig 中 我们将 dtype 指定为 int4 粒度为“per_block” block_size 为 32

我们将它设置为 global_config 因为我们要压缩 模型中的所有权重

如果需要 你可以根据运算类型 或运算名称指定 config

然后 我们可以使用 ct.optimize 模块中的 linear_quantize_weights 工具 根据 config 来量化模型

根据我们指定的 config 会进行“int4”“per_block”量化 其中每个块有 32 个元素 在这个工具中 Core ML 工具 会检查模型中的所有权重 并通过线性量化算法 对这些权重进行压缩

量化完成

我们来看看量化后的模型大小

不错！ 量化模型的磁盘大小 比最初的 Float16 模型 小得多 从 13 GB 减少到 4 GB 以下

让我们使用同一提示 运行这个量化模型

真棒！ 这个小得多的模型 运行速度甚至 比原来的 Float16 模型更快 而且生成的输出质量也不相上下

我们在本演示中可以看到 使用有状态 KV 缓存和量化 来准备 Mistral-7B 模型 非常简单 你可以按照类似的工作流程 将你喜欢的语言模型 移植到 Apple 芯片上

最后 还想谈谈 另外一项功能 那就是对多函数 Core ML 模型的支持

让我们考虑一个由 多个函数组成的模型 这些函数组合需要 一起运行来进行预测 例如 模型可能有一个特征提取器 然后是分类器层 用来产生分类输出

此外 还有一个回归层 用来生成回归输出

这两者可组合成一个模型 用于处理这两项任务

这是在 PyTorch 中的样子

分类器和回归器 共享同一个特征提取器 且两者都已转换为 Core ML 模型

今年 我们将推出 对多函数模型的支持 你可以使用 Core ML 工具 合并这两个模型

在合并过程中 Core ML 工具会计算 权重的哈希值 尽可能删除重复的共享权重

合并后的多函数模型 可通过共享特征提取器 完成这两项任务

这是使用 Core ML 工具 合并 Core ML 模型的示例代码片段

使用 ct.convert 转换两个模型 并将它们存储为单独的 mlpackages 后

可以创建 MultiFunctionDescriptor 来指定要合并哪些模型 以及合并模型中新的 function_name

然后 可以使用 save_multifunction util 生成合并后的 多函数 Core ML 模型

通过 Core ML 工具的 Python API 载入多函数模型时 你可以指定 function_name 来载入特定函数 然后像往常一样进行预测

多函数模型 在大模型中非常有用

在微调大模型时 对整个模型进行微调的 成本通常太高

相反 常见的做法 是在微调过程中使用适配器

适配器通常较小 参数更少 但表现出的性能 与全面微调后的模型相当

对于同一个基础模型 不同字段和下游任务 可能有不同的适配器

适配器块通常不只是 位于模型的开头或结尾 它们也可以与中间层交互

你可以使用 Core ML 新的 多函数模型 在单个模型中 捕获多个适配器 每个适配器都会 定义一个单独的函数

基础模型的权重 会尽可能共享

欢迎观看“利用 Core ML 在设备端 部署机器学习和 AI 模型”视频 进一步了解如何 将多个函数整合到 App 中

在那个视频中 你会看到 一个文本转图像的例子 其中一个模型 与多个适配器一起用于 生成不同风格的图像

总结一下今天的内容 可以使用 新的模型压缩技术和工作流程 大大减小模型大小 同时仍能实现良好的精度 还可以轻而易举地 让模型使用状态 并提供多个函数

将大型转换器模型 移植到 Apple 芯片从未如此简单

更多详细信息 请查阅 Core ML 工具文档 谢谢大家