为 Metal App 优化机器学习

♪ ♪

Denis：大家好 我是 Denis Vieriu 一名 Apple GPU Graphics 和 Display Software 团队的 软件工程师 今天 我将向你展示 Metal 今年 在机器学习领域引入的 全新功能和改进 首先来回顾一下 现有的机器学习后端 Metal 机器学习 API 已通过 Metal Performance Shaders 框架公开 MPS 是一组高性能 GPU 图元 可适用于多种领域 例如图像处理、 线性代数、和机器学习等 MPSGraph 是一个 基于 MPS 框架的 通用计算图形 并可扩展支持多维张量 Core ML 等机器学习推理框架 构建在 MPSGraph 后端之上 此外 MPSGraph 还支持训练框架 例如 TensorFlow 和 PyTorch 欲进一步了解有关 MPSGraph 和 ML 框架的更多信息 欢迎观看屏幕列出的 此前有关 Metal 的 WWDC 讲座

本次讲座将重点介绍 PyTorch 和 TensorFlow Metal 后端中增加的 更新和改进 针对 JAX 的新 GPU 加速 以及今年为 MPSGraph 新增的 ML 推理功能

PyTorch 和 TensorFlow 的 Metal 加速可以让你 利用 MPS 中的高效内核 在 Mac 上获得最佳性能 自 PyTorch 1.12 起 MPS 后端便开始 提供 Metal 加速 去年 PyTorch 生态系统 引入了该功能 并且自那时起 大量针对优化内存使用 和视图张量的更新开始涌现 今年 PyTorch 2.0 的 MPS 后端取得了出色进展 并已进入 Beta 阶段 但以上并不是全部的改进内容 最新版的 PyTorch 还包含大量其他更新 例如 MPS 操作性能分析、 自定义内核 以及自动混合精度支持等 在介绍每夜构建版本 所有的功能之前 我先来介绍一下 PyTorch 2.0 中的新功能

PyTorch 2.0 现在支持 最常用的 60 个 Torch 运算符 包括 grid sampler、 triangular solve 以及 topk 等

并且 测试覆盖范围也更加广泛 其中包括 对大部分 Torch 运算符的测试 梯度测试以及 基于 ModuleInfo 的测试

自 PyTorch 2.0 发布以来 由于多个流行模型在 macOS 上 采用 MPS 作为官方后端 因此其网络覆盖范围也得以扩展 这些模型包括 WhisperAI 等基础模型 YOLO 等物体检测模型 以及 Stable Diffusion 模型等 接下来 我们使用最新的 PyTorch 2.0 来看看其中一个模型的实际应用 在这个例子中 我正使用 YoloV5 一个在 M2 Max 上 运行的物体检测网络 左侧是我使用 PyTorch MPS 后端 运行网络并生成的实时图像 右侧是我在 CPU 上 运行同一模型生成的图像 可以看到 使用 MPS 后端运行的 左侧图像帧速率明显更高

此外 开发者不仅 在外部网络中采用 PyTorch MPS 后端 同时还为许多新的运算提供代码 包括直方图、 组归一化、以及 signbit 等

接下来 我将对最新 PyTorch 版本构建中的 新功能进行介绍 首先介绍支持 MPS 操作的 性能分析 PyTorch 每夜构建版本 提供性能分析支持 并使用操作系统 Signposts 为操作执行、CPU 和 GPU 之间的数据传输 以及由于不受支持的运算导致的 CPU Fallback 显示具体运行时间 你还可以使用熟悉的 Metal System Trace 工具 来可视化性能分析数据 该工具 属于 Instruments 欲进一步了解使用 Metal System Trace 对 ML App 进行性能分析的更多信息 欢迎观看去年的讲座 “使用 Metal 加速机器学习” 使用分析工具非常简单 调用 MPS 分析工具包中的 start 方法启用跟踪 然后在脚本结束时 使用 end 方法结束性能分析即可 现在 我用分析工具 来调试一个示例 本示例网络使用顺序模型 其中包含线性变换以及 Softshrink 激活函数 该模型共七层 该模型当前的性能并不理想 在这种情况下 我们可以使用 分析工具找出瓶颈

在 Metal System Trace 中 首先确保已启用 os_signpost 这样你才可以捕捉 PyTorch 运算的信息 接着 确认设备 和恰当的可执行文件已就绪 在本示例中 可执行文件为 Python 二进制文件 然后点击记录按钮 Instruments 便会开始记录 PyTorch 的执行 我会让 Instruments 运行几秒 以捕捉足够的数据 接着我点击停止

在 os_signpost 标签页中 展开 PyTorch 的间隔时间线 该时间线显示了运算的执行时间 以及 PyTorch 的元数据 例如字符串识别符、 数据类型以及副本长度等

放大时间线 我们便可以看到 本示例所使用的 PyTorch 运算 根据跟踪模式 可以轻松识别 由七层组成的自定义顺序模型

根据跟踪 我们可以清晰地看到 瓶颈在于 对 CPU 的 Softshrink Fallback 这个过程效率十分低下 因为该模型需要承受 Softshrink 运算符的 CPU 执行以及额外副本带来的开销 而 GPU 却始终处于空载状态 GPU 时间线中大部分的间隙 来自 Softshrink 激活函数 回退到 CPU 为了解决该问题 我将编写 一个自定义内核来提高性能 编写自定义操作共分四步 首先 在 Objective-C 和 Metal 中 实现操作 接着 为 Objective-C 代码 创建 Python 绑定 并编译扩展 最后 扩展编译完成后 将操作导入训练脚本并使用该扩展 接下来 我从执行操作开始

首先 导入 Torch 扩展标头 该标头包含了编写 C++ 扩展 所需的所有 PyTorch 内容

接着 定义运算函数 并使用 get_command_buffer MPS 后端 API 获取一个 指向 MPSStream 命令缓冲区 的引用 同样地 使用 get_dispatch_queue API 获取一个串行调度队列的引用 然后 使用命令缓冲区创建一个编码器 并定义自定义 GPU 内核

你需要在调度队列中 对内核进行编码 以确保来自多线程的提交被序列化

在对所有工作编码后 使用 synchronize API 进行等待 直到当前命令缓冲区运行完成 这样你就可以看到序列化的提交 或者 若你无需进行序列化 使用 commit API 即可 下一步 绑定自定义函数 你可以使用 PYBIND11 轻松将 Objective-C 函数 绑定到 Python 对于本示例中的扩展 仅需两行绑定代码

完成绑定后 你需要编译扩展 首先 导入 torch.utils.cpp_extension 你可以使用其中提供的 load 函数 来编译扩展 接着 传递需构建的扩展名称 然后将相对或绝对路径传递给 源代码文件 并且 你还可以选择传递 额外的用于编译的编译器标志 接着 load 函数会将源文件 编译为共享库 并将其作为模块 加载到当前 Python 进程中

最后一步 在脚本中导入运算符 并进行使用

首先 导入经过编译的库 并更改此前的顺序模型 以使用自定义 Softshrink 内核

让我们再次运行该模型并查看结果 因为添加了新的自定义运算符 该模型运行效率更高了

CPU Fallback 创建的 副本和中间张量均已消失 而且顺序模型运行速度 也得以大幅提升 接下来 我们来探讨更多方法 以进一步改善你的网络

PyTorch MPS 后端 现在支持自动混合精度 让你可以 使用更少的内存提高训练速度 并且质量不会受到影响 为了理解混合精度 首先 我们来回顾一下 受支持的数据类型 混合精度训练是一种训练模型 允许你混合使用单精度浮点数 以及半精度浮点数 来训练深度学习模型 从 macOS Sonoma 开始 MPSGraph 添加了对新数据类型 bfloat16 的支持 bfloat16 是一种用于深度学习的 16 位浮点格式 该格式由 1 个符号位、 8 个指数位 以及 7 个尾数位构成 与标准 IEEE 16 位浮点格式不同 该格式并不是专门为 深度学习 App 而设计 自动混合精度可以同时启用 float16 以及 bfloat16 通过测试默认精度下的网络性能 自动混合精度 为每层选择恰当的精度 接着使用该混合精度设定 再次运行网络 便可以在不影响精度的情况下 优化网络性能 神经网络的某些层 可以以较低的精度执行 例如卷积层及线性层 而某些其他层则通常需要 更高的精度 例如规约层

将自动混合精度支持 添加到网络十分简单 首先 添加 autocast Autocast 同时支持 float16 和 bfloat16 Autocast 可以作为上下文管理器 允许脚本中的某段代码 以混合精度运行

在该代码中 MPS 操作会根据 autocast 选择的数据类型运行 从而在保持精度的同时 提高网络性能

MPS 后端也会得到显著优化 借助 PyTorch 2.0 和 macOS Sonoma 与此前版本相比 MPS 后端运行速度提高了 5 倍 以上就是 PyTorch 的全部内容 现在 我们开始讲解 TensorFlow TensorFlow Metal 后端已经 发展到稳定的 1.0 发行版本 在该版本中 插件中添加了一个 GrAppler 重映射优化过程 Metal 插件具有 混合精度支持 并且安装过程比以往更加简单

TensorFlow Metal 后端在 添加自动融合已识别计算模式的 功能之后 性能得到改善 这些计算包括融合卷积、 矩阵乘法、优化器操作以及 RNN 单元等 在计算图创建过程中 该优化会通过 grAppler pass 自动进行

在这里 我来举一个 二维卷积运算中的常见计算实例 在卷积神经网络中 卷积后跟着一个加法函数 的模式十分常见 GrAppler pass 在识别到该模式后 便会重新映射计算

从而 你便可以使用更优化的内核 实现相同的输出 并提高网络性能 与 PyTorch 一样 TensorFlow 也具备混合精度支持 并且 TensorFlow 允许对全局设置混合精度 这样 使用请求数据类型策略 所有的网络层便可以自动创建 只需要对当前代码进行少量修改 便可在标准工作流程中启用该更改

此外 全局策略可以设置为使用 float16 或 bfloat16

除了性能的提升 启用 Metal 加速获取的 用户体验也得到简化 自此 通过程序包管理器 按照通常的方法 安装 TensorFlow wheel 和 TensorFlow-Metal 插件 你便可以轻松启用 Metal 加速 如果你希望保持在 TensorFlow 开发的最前沿 当前 TensorFlow 每夜构建版本 便可为你提供 Metal 加速支持 接下来 我们讲解 针对 JAX 的新 GPU 加速 今年 JAX GPU 加速 同 PyTorch 和 TensorFlow 一样 也将得到 Metal 后端支持 JAX 是一个用于高性能数值计算 和机器学习研究的 Python 库 它基于流行的 NumPy 框架 用于处理大型数组 在机器学习研究中有三个重要扩展

首先 JAX 支持使用 grad 函数进行自动微分 它可以对 大部分 Python 要素子集进行微分 甚至可以计算高阶导数 并且 JAX 还支持 快速高效的矢量化 在给定函数 Apply_matrix 中 你可以循环遍历 Python 中的一批维度 但该函数运行的性能可能不太理想 在这种情况下 你可以使用 vmap 来自动添加批处理支持 此外 JAX 还可以借助 名为 jit 的 API 将函数即时编译为优化的内核 在相同的情况下 jit 可以 在 vmap 的基础上对函数进行转换 从而使其运行得更快

在配备 M2 Max 的 MacBook Pro 上 JAX Metal 加速 可以提供出色的加速效果 并且在相同的网络上 平均来说 速度是 CPU 的 10 倍 欲了解有关 JAX 环境设置和安装的更多信息 欢迎参考 Metal Developer 资源网页

接下来 我们再来讲讲 ML 推理 首先 我会介绍一种用于 MPSGraph 并可以优化加载时间的 新序列化格式 该新序列化格式可以从 其他框架现有的序列化网络中生成 最后 我会向你展示 如何利用 8 位整型量化 来优化网络的内存占用 我们开始吧 MPSGraph 可以使用 高度灵活的高级 API 逐层创建 欲了解更多细节内容 欢迎观看“使用 Metal Performance Shaders Graph 构建自定义 ML 模型” 在你定义并编译自定义图形后 该图形便会通过 MPSGraphExecutable 执行并获取结果 通常情况下 这个过程非常顺利 但是 若图形较为复杂并有很多层 那么在初始编译时 就需耗费很长时间启用 App

MPSGraph 现推出了新的 序列化格式 MPSGraphPackage 用于解决这个问题 这种新的序列化格式可以让你 提前创建 MPSGraphExecutable 在创建完成之后 优化后的 MPSGraphExecutable 便可以直接 从 MPSGraphPackage 文件中进行加载 创建 MPSGraphPackage 非常简单

你只需创建一个序列化描述符 并将其传递给需要序列化的 MPSGraphExecutable 的序列化函数 同时 你还需要 传递一个路径以供存储 创建完程序包后 以下便是你将图形 加载到 App 中的方法 你需要一个编译描述符 和存储程序包的路径 接着 使用这些值来初始化 MPSGraphExecutable 如果你已经用过 MPSGraph 那么你便可以轻松使用 我们提供的 API 来使用这个新的序列化格式 但如果你此前使用的是其他框架 使用全新的 MPSGraphTool 你便可以轻松迁移至 MPSGraphPackage 对于 Core ML 的用户而言 你可以将 ML 程序传递给 MPSGraphTool 它会为你创建 MPSGraphPackage 使用 ONNX 也是如此 你可以将 ONNX 文件 作为输入 这个新工具可以让你 将现有模型快速包含到 MPSGraph App 中 而无需手动编写推理模型 以下是使用命令行工具的方法 为 MPSGraphTool 提供一个标志来声明输入模型的类型 在本例中为 Core ML Package 接着为 MPSGraphTool 提供输出目标的路径 以及输出模型的名称 最后 定义目标平台 和最低操作系统版本 转换完成后 生成的 MPSGraphPackages 便可以加载到 App 中 直接执行 接下来 我们来探讨如何使用 8 位整型量化来提高计算效率 通常我们在训练和推理中 会使用浮点格式 例如 16 位浮点格式 但是在推理中 这些模型 便需要更长的时间来预测结果 相反 在许多情况下使用降低精度 或 8 位整数 反而可以达到不错的效果 这样做不仅可以帮你节省内存带宽 还可以减少模型的内存占用

8 位整型格式具有两种量化类型： 对称量化和非对称量化 MPSGraph 现已支持 这两种类型的 API 相较于对称量化 非对称量化可让你 指定量化偏差 在这里我们用 zeroPoint 表示

我们通过一个示例来深入了解一下 量化计算的使用方法 首先 将激活函数和权重 以 Int8 格式作为输入 使用 MPSGraph 中的 dequantizeTensor 操作 将该输入反量化为浮点格式 现在 将浮点输入送入 卷积运算中 这样 得到的浮点张量 借助 quantizeTensor 操作 又可以量化为 Int8 MPSGraph 可以自动将这些内核 融入到单一操作中 从而节省内存带宽 甚至还可以提高网络性能 以上就是如何在 MPSGraph 中 使用量化支持的全部内容

除了此前介绍的新功能 MPSGraph 还可以支持 其他机器学习的运算符 从今年开始 大部分图形操作 都可以支持复数类型 你可以使用单精度复数 也可以使用半精度浮点格式复数 基于复数数据类型 MPSGraph 添加了用于计算 快速傅里叶变换的运算符 这样 你便可以应用 复数到复数、复数到实数 以及实数到复数等 多达 4 个维度间的转换 这些转换在音频、视频 及图像处理 App 中都非常常见 此外 你现在可以使用 MPSGraph 来执行三维卷积、 网格采样、Sort 和 ArgSort 以及累积操作 包括求和、乘积、最小值和最大值 以上便是有关 MPSGraph 新功能的全部内容 让我们回顾一下 本次讲座介绍的内容 我介绍了如何借助 Metal 加速 PyTorch 和 TensorFlow 等流行 ML 框架 现在你还可以使用新的 Metal 加速 JAX 框架 此外 我们还谈论了如何使用 新的序列化工具 将来自其他框架的现有模型 与 MPSGraph 实现无缝集成 以上便是本次讲座的全部内容 我非常期待看到你使用这些功能 创造出精彩的内容 感谢你的观看 ♪ ♪