利用 Metal 加快机器学习

大家好 我叫 Kamal Ramamoorthy 是 GPU, Graphics and Display 团队的软件工程师 在这个视频中 我将和我的同事 Sam 为你演示如何借助 Metal 为机器学习模型加速 训练是在 Apple 平台上 部署模型的第一步 第二步是在设备上为部署准备模型 最后 就可以将模型 整合到你的应用程序中 这就是我在本视频中 将重点介绍的内容 如果你使用 Core ML 来部署模型 那么 MPSGraph 会使用 Metal 为 GPU 加速 要了解 Core ML 的更多信息 请观看“在 Apple 平台上 探索机器学习”视频 你可能还需要观看 “利用 Core ML 在设备端 部署机器学习和 AI 模型”视频 你也可以利用 PyTorch、TensorFlow 和 JAX 等框架来训练模型 请观看“使用 Apple 芯片 GPU 训练 ML 模型”视频 了解更多内容 所有这些框架都基于 Metal Performance Shaders Graph 构建 这是一个使用 Metal 构建和运行通用计算图表的 框架 MPSGraph 提供了 对 GPU 同步和内存的低级别控制 所以在某些情况下 你可能需要直接使用 MPSGraph 例如 如果你的应用程序使用 Metal 那么你可以使用 MPSGraph 来安排 ML 任务和其他 GPU 工作的执行顺序 你也可以共享低级别 Metal 资源 如缓冲区 如果你刚刚开始使用 Metal 为机器学习加速 请观看往年 WWDC 上的视频 在本视频中 我和 Sam 将探讨三个主题 首先是 MPS 和 MPS Graph 的 改进功能 很多改进功能都旨在提升 转换器模型的效率 所以我会以这些模型为例 其次 可为基于 FFT 的 ML 模型加速的新功能 最后 Sam 将向大家介绍 MPSGraph Viewer 这个工具让你能够 直观查看 ML 模型 我们先来了解一下 专注于提高转换器模型性能的新功能 转换器是语言模型中常用的工具 用于翻译、预测和生成文本 输入是一系列标记 比如“The quick brown” 这个短句是包含 3 个标记的输入 语言模型通过预测 下一个标记“fox”来进行响应 这个新句子被不断地反馈回模型中 以生成新标记

MPS 和 MPSGraph 新增功能 可帮助你改进转换器模型 这些改进功能分为三类 首先是计算性能的提升 其次是内存带宽的节省 最后是转换器模型的质量改进 首先我将介绍计算性能的提升

基于转换器的模型 由多层转换器模块组成 我们来深入了解一下转换器模块的内部结构 组成部分包括多头注意力、归一化以及前馈模块 多头注意力模块 是计算最密集的模块 这个模块可计算大量多维矩阵乘法 属于计算密集型运算 输入矩阵通过矩阵乘法层进行投影 生成一个查询矩阵 Q 键矩阵 K 和值矩阵 V 这些矩阵随后被输入到 缩放点积注意力模块中 这就是转换器模型的核心所在

有两种方式可以优化 注意力模块的性能 如果你查看注意力模块的内部结构 会发现它是由几个运算组成的

MPSGraph 现在有一个运算 它将把这一系列运算 合并到一个内核中 执行起来更加高效

要使用这个运算 请对 MPSGraph 对象 调用 scaledDotProductAttention 方法 这个方法接受查询张量、 键张量和值张量作为参数

使用融合的 SDPA 运算应该能帮助你 提升转换器模型的性能 我们再来看多头注意力模块 这样我可以向你展示 利用这些张量来提高计算性能的其他机会 在这里你可以看到查询投影、 键投影和值投影 针对第一个标记的工作方式 矩阵乘法运算 将投影查询、键和值的嵌入矢量

要生成下一个输出标记 我们必须将之前生成的所有标记 输入到矩阵乘法中 这会导致重新计算之前迭代中 已经计算过的 键 (K) 和值 (V) 投影

这个开销会随着序列长度的 增长而增加 要缓解这个问题 可以在这些投影生成时进行缓存 以便在未来迭代中重复使用 你要做的第一件事是 创建 K 张量和 V 张量 来存储缓存的 K 值和 V 值 这些张量称为 KV 缓存 在第一次迭代中 计算第一个标记的 K 值和 V 值 并将它们插入到 KV 缓存中

现在你可以重复使用缓存的值 以便在第二次迭代中 只需计算第二个标记的 K 值和 V 值即可 这将矩阵-矩阵乘法 简化为矩阵-矢量乘法

通过为每次迭代创建新的张量 你可以将 K 投影和 V 投影 附加到 KV 缓存的末尾 但这会占用大量内存 相反 你也可以使用切片更新运算 就地更新现有的张量

然后使用切片运算 仅提取 KV 缓存中已计算过的部分

让我们来看看如何 用代码来完成这一过程

首先 创建一个表示缓存的占位符 这个张量的维度取决于 模型的具体细节 在这个示例中 我将只关注缓存的键部分 不过值部分的工作原理相同 为了能够就地更新 KV 缓存 你需要创建一个变量 这个变量代表缓存的当前状态 与常规图表运算的结果不同的是 你可以稍后更新这个变量 以引用另一个值

你需要将每个标记的键投影 插入到缓存中 你可以对 MPSGraph 对象 使用 sliceUpdateDataTensor 方法 来执行这一操作 开始数组和结束数组表示 要放置新值的位置 在本示例中 新值会被附加到 缓存有效部分的末尾 这个示例中的步长是一致的

现在 你可以将更新的缓存 赋值回原始变量 MPSGraph 会进行优化 以就地更新缓存分配 最后 你可以使用切片运算 仅提取 KV 缓存中已计算过的部分 这一部分是从缓存开头 到最新插入的键投影

然后 你可以将更新后的键缓存 传递给 SDPA 运算

当你完成这些计算改进后 内存带宽便成为新的瓶颈

存储大型语言模型权重所需的内存 可能达到数十千兆字节的量级 这些权重通常使用 16 位浮点数来表示 然而 MPS 支持将这些权重 量化为 8 位整数 从而将内存占用减半

作为今年的新增功能 MPS 还支持一种 4 位整数格式 这样你可以进一步减小权重的大小 MPS 支持用多种方法 将权重映射到这些量化格式上

这是一个示例张量 其中的元素分布在数轴上 对于 8 位量化 有 256 个可能的量化点 沿数轴线性分布 对于 4 位量化 存在 16 个点 在量化过程中 点将调整为最接近的量化值 量化缩放系数 可以使用右侧的公式确定 这会产生微小的误差 但最终我们能节省 2 倍或 4 倍的内存空间和带宽 对 MPSGraph 对象 使用 dequantize 方法 可反量化这些值

另一种量化技术使用查找表 当你的权重聚集在数轴的不同区域时 这种技术非常有用 利用仿射量化 量化值是均匀分布的 但输入值却不是 这会导致大多数量化数位闲置 因为大多数输入值都会聚集在 少数几个量化点周围 通过使用查找表 你可以更有效地利用量化数位 使用这项技术 你可以根据数据的分发设置 选择自己的量化点 你将这些量化值存储在查找表中 然后 将每个权重分配给 表中的 4 位或 8 位索引 这样可以大大提高灵活性 只是在查找表中查找值时 会牺牲一小部分性能

使用 dequantize 方法 将这些量化值 转换回 32 位浮点值 只需将量化权重 传递给 32 位查找表即可 然后 可以像往常一样 重复使用反量化张量 例如 将这些张量用作 矩阵乘法的输入 事实上 在这种情况下 MPSGraph 性能还会更进一步

如果你的图表包含 对权重的反量化运算 并且这一运算发生在矩阵乘法之前

则 MPSGraph 会将这两个运算 替换为单个量化矩阵乘法运算 这种运算会在需要时 即时反量化权重 而不是存储 反量化权重的临时副本

量化可以节省内存和内存带宽 但同时也可能引入数值误差 现在 我来演示 提升转换器模型质量的两种方法

当你量化权重时 每个权重都会映射到较低精度的值 你还可以选择一个缩放系数 以及一个偏移值 以便在反量化时应用于量化值 但是 对所有权重 应用单个缩放系数和偏移值 将会限制重构值的准确度 相反 你可以单独量化元素块 每个块都有自己的缩放系数和偏移值 这样你就可以为每个块更精准地 匹配缩放系数和偏移值

执行这个操作的代码与前面的示例类似 只有一点不同 你要为每个块传入 包含缩放系数和零点值的张量 而不是传入单个缩放系数和零点值

以上就是量化内容 下面我将为大家介绍另一种方式 可以使用适配器 来提高转换器模型的质量

适配器是小型层 你可以将它们插入到 仅包含少数运算和参数的模型中 在微调模型时 仅会更新适配器内部的参数 这可用于让预先训练的 基础模型适应新任务 也可用于补偿由量化引入的误差 你可以使用 MPSGraph 可调用对象 将适配器添加到你的模型中

它的工作原理是 每个适配器都是单独的 MPSGraph 可从主图中进行调用 首先 通过为每个适配器 指定唯一的名称 从基础图表中插入对适配器的调用 为了在代码中实现这一点 需要对输出的形状和类型进行定义 这些输出由对适配器的调用生成 然后 对主 MPSGraph 对象 使用 call 方法 将调用添加到适配器 这时你可以提供可调用对象的名称、 输入和输出类型

接下来 为每个适配器创建 MPSGraph 在本例中 我将创建一个占位符 将“输入”表示为未排序的张量 然后 将“输入”张量乘以 2 来创建“输出”张量

最后 将每个适配器的图表 编译为图表可执行文件 这些图表的编译方式 与其他 MPS 图表一样 首先 通过提供确切的形状 来定义图表的输入类型 然后 调用图表对象的 compile 方法 提供 Metal 设备、 输入类型和 outputTensor

现在 你已从主图表中 添加了对适配器的调用 并为每个适配器 编译了图表可执行文件 最后需要做的是 将适配器名称映射到 主图表中的实际图表可执行文件 使用 GraphCompilationDescriptor 为你的网络 编译主 MPSGraph 时 操作就完成了 首先 创建一个字典 将每个适配器名称 映射到图表可执行文件 并在描述符上设置字典 然后 只需在编译主图表时 提供编译描述符 以上就是设置适配器的全部步骤 总结一下 使用适配器和可调用对象可定制模型 以执行新任务并提升质量 因为你可以使用它们 来微调你的模型输出 下面 我来为大家介绍一下 今年 MPS 和 MPSGraph 中的 “傅里叶转换”的新功能 快速傅里叶转换简称 FFT 可将信号或图像等数据 从时间域或空间域转换到频率域 在处理音频的机器学习模型中 比如语音转文本模型 以及从单个音轨中 分离不同音频源的模型 FFT 是常见的预处理步骤 FFT 也可用于加速某些卷积层 并且在很多图像处理和 科学计算应用程序中也可加以利用 举例来说 要从音频信号中提取文本 首先需要对输入波形 进行短时傅里叶转换 (STFT) 然后 通过转换器模型分析频谱 以提取文本 我已经介绍了如何使用 MPSGraph 在 GPU 上高效执行 ML 模型 但你也可以利用 MPSGraph 对快速傅里叶转换的支持 将整个管道转移到 GPU 上进行处理 第一步是实现短时傅里叶转换 这可通过将输入波形 分成多个较短的视图或窗口来实现 这些视图或窗口可能会彼此重叠 每个窗口实际上是一个独立的信号 为了减少频谱泄漏 每个窗口都需要乘以一个窗口函数 最后 你可以使用常规的批量一维 FFT 运算来计算每个窗口的 STFT 为了将波形分成多个较短的窗口 你可以创建步长视图 首先 设置窗口视图的形状 在这个示例中 窗口的宽度将是 512 个元素 接下来 为每个维度设置步长 这个示例使用的值为 256 这意味着在第二维度中每跨一步 就表示在底层一维数组中 跳过 256 个元素 最后的批次维度设置为 1 但你可以使用更大的批次大小

最后 你可以通过对输入张量 调用 arrayView 方法 来创建步长视图 最棒的是 通过为输入数组的内存 添加别名 视图运算即可执行 无需进行任何拷贝 从而节省内存和 GPU 时间

现在 你可以计算所有窗口的 FFT 首先 为步长视图数据创建占位符 你需要将数据 从步长视图 NDArray 中加载出来 然后在运行图表时再提供这一数据 接下来 乘以窗口函数 这通常是 Hann 窗或 Gaussian 窗 例如 你可以使用 MPSGraphConstant 张量 最后 你可以创建 FFTTensor 运算 以上就是“快速傅里叶转换” 的相关内容 接下来 有请 Sam 来为我们演示 如果你希望深入了解 MPSGraph 结构体 Sam 将为你带来一些好消息 感谢 Kamal 大家好 我叫 Sam Colbran 也是来自 GPU, Graphics and Display 团队 的软件工程师 你或许还不知道 Metal 包含了 Xcode 和 Instruments 中的高级 工具 可帮助你充分利用 Apple GPU 强大的功能触手可及 Metal 管道 以及在设备端运行的 AI 模型 会变得更大、更复杂 不过 虽然你可以借助 Xcode 中的依赖项查看器 直观查看 Metal 管道 但以前无法直观查看 MPSGraph 今天 我很高兴为大家介绍 Xcode 16 中 Metal 工具的新成员 MPSGraph Viewer 这是一款全新的工具 专为机器学习和 AI 而设计 现在 你可以直接在 Xcode 中 打开 MPSGraph 软件包 并直观查看你的运算是如何关联的 在开始演示之前 我们先回顾一下 如何实际创建 MPSGraph 软件包 无论你是直接使用 MPSGraph 还是一直在其他框架中开发 ML 模型 如果你已直接使用 MPSGraph 创建了模型 首先 将你的图表编译到 MPSGraph 可执行文件中 然后 使用可执行文件上的 serialize API 创建软件包

作为今年的新增功能 你也可以直接从 CoreML 软件包 创建 MPSGraph 可执行文件 和以前一样 接下来你可以将可执行 文件序列化为 MPSGraph 软件包

或者 如果你以前使用另一个框架 比如可以导出为 ONNX 格式的框架 你可以使用 mpsgraphtool 来转换模型 让我们一起来看一个示例

我正在使用 Mistral 的模型 其中有 70 亿个参数 这个模型已转换为 CoreML 格式 具体内容可观看今年的 “将机器学习和 AI 模型 移植到 Apple 芯片”视频 现在 可以使用命令行访问 mpsgraphtool 打开常用的“终端” 然后使用 convert 参数运行 mpsgraphtool

就是这样 新创建的 mpsgraph 软件包 即可开始使用 现在 使用新的 MPSGraph Viewer 可以轻松查看 我已经在 Xcode 16 中打开了 转换后的 Mistral 软件包 让我来介绍一下屏幕上的内容 从左上角开始 这些是编译选项 默认情况下 查看器按原样显示图表 也就是说 它没有针对任何特定设备进行优化 因此 无论你使用什么设备 运算显示应该是一样的

下方是运算导航器 这里会显示图表中使用的 所有运算的列表 中间就是图表 最后 屏幕的右边是运算检查器 我们稍后还会再谈到这一项

在当前级别 没办法看清图表 所以我要放大一些

现在 我可以看到高层级结构 继续放大 细节会更多 现在 我可以滚动查看 每个运算的所有输入和输出 以及它们之间是如何关联的 这样能够直观查看并理解 图表的结构

Mistral 现在是一个转换器模型 正如 Kamal 之前解释的那样 它们由转换器模块层叠而成 让我们试着找出它们 首先 我会查找 新的缩放点积注意力运算符 这个运算符 在每个转换器模块中应该属于 多头注意力的一部分 我可以执行搜索 但我已经在左侧的 运算导航器中看到了 有 32 个运算符 我将展开这个群组并点按第一个群组 这样可跳转到图表中的相应部分

这个运算看起来有 5 个输入 你应该已经认识查询、键和值 Kamal 之前介绍过

我将画面缩小一些 可以更好地查看整个转换器模块

跟随这些关联 我能看到组成查询、键和值的 模块 即便是在这个级别 我也可以看到键和值中的变量 我将放大以显示键中的变量 由于这个模型是从 带有状态的 CoreML 中导出的 它使用 KV 缓存并利用 MPS 中的 新 assign-to-variable 和 read-from-variable 运算 正如 Kamal 展示的那样 计算性能会得到提升 现在 为了简化图表 查看器可能会在多个位置 显示一些运算和变量 就像这个例子一样 我在右侧的检查器中选择了变量之后 可以看到变量首次创建的位置 以及所有使用变量的位置

好了 刚才介绍的是一个转换器模块 其余的呢 正如检查器一样 实际上 我可以在运算导航器中 同时选择多个运算 就像这样 高层级结构就很清晰了 可以看到各个层 全都一遍又一遍地重复出现 非常美观 现在 我们来谈谈常量 你可能已经注意到 直接显示在图表内部的 绿色预览部分 但我还能在左侧的 常量导航器标签页中找到它们 并按大小排序 我会选择最大的常量 然后 连按两下打开常量查看器 你可以在这里检查经过训练的权重 并深入了解 模型学到的内容 这样有助于你发现机会优化模型 以实现更出色的设备端整合 但请记住 查看器按原样显示图表 它没有针对任何特定设备进行优化

实际上 执行的图表可能会有所不同 例如 通过将运算整合到 单个 Metal 内核中 MPSGraph 可以自动优化运算 可以使用查看器直观查看运算 我来演示一下具体方法 我打开了一个 MPSGraph 软件包 其中包含 ResNet50 和之前一样 我可以放大屏幕 以查看所有的运算和常量 但现在 让我们来看看 我的设备图表是什么样的

在左上方的编译选项中 我将选择我的设备

现在 放大来看 我可以看到运算已被分组到这些 Metal 拼接着色器区域 展开可以看到内部 由于这些运算融合到了 一个经过优化的 Metal 着色器中 从内部来看 没有内存开销 因而可显著提高性能 总的来说 了解图表最终如何在硬件上执行 对于真正理解图表运行时的 性能特性非常有用 以上就是 MPSGraph Viewer 新增功能的全部内容 现在 我们来总结一下今天的内容 正如 Kamal 先前所说 借助 Metal 可以为机器学习加速 你可以使用 Metal Performance Shaders Graph 这个工具已经集成到热门框架的内部 比如 CoreML 为你带来 Apple 芯片的出色性能 今年 转换器的新功能 可以帮助你改进计算 包括结合 KV 缓存的 新的缩放点积注意力运算 具有量化功能的内存带宽 还能通过适配器和可调用对象 提升质量 借助新的步长 NDArray API 在 MPSGraph 中可以更快地 进行傅里叶转换计算 最后 新的 MPSGraph Viewer 有助于理解和深入洞察 模型在 Apple 芯片上是如何执行的 请务必查看 MPSGraph 文档 和示例代码 当然 模型集成是最后一个板块 如果你还没有学习过 请务必观看相关精彩视频 了解更多关于训练和部署的知识 感谢观看 祝你在 WWDC 度过美妙的时光