先说一个小内存的部署 然后在说一下刚出的pytorch2.0 100%向后兼容,一行代码将训练提速76%
前段时间,PyTorch 团队在官方博客宣布 Pytorch 1.13 发布,包含 BetterTransformer 稳定版等多项更新。在体验新特性的同时,不少人也在期待下一个版本的推出。
出乎意料的是,这个新版本很快就来了,而且是跨越式的 2.0 版!
新版本的重要进步体现在速度和可用性,而且完全向后兼容。PyTorch 团队表示,PyTorch 2.0 是他们向 2.x 系列迈出的第一步,其稳定版预计在 2023 年 3 月初发布。
首先,PyTorch 2.0 引入了 torch.compile,这是一种编译模式,可以在不更改模型代码的情况下加速模型。在 163 个涵盖视觉、NLP 和其他领域的开源模型中,该团队发现使用 2.0 可以将训练速度提高 38-76%。
其次,PyTorch 2.0 是 100% 向后兼容的:代码库一样,API 一样,写模型的方式也一样。团队之所以称它为 2.0,是因为它有一些标志性的新特性,包括:
-
TorchDynamo 可以从字节码分析生成 FX 图;
-
AOTAutograd 可以以 ahead-of-time 的方式生成反向图;
-
PrimTorch 引入了一个小型算子集,使后端更容易;
-
TorchInductor:一个由 OpenAI Triton 支持的 DL 编译器。
PyTorch 2.0 将延续 PyTorch 一贯的优势,包括 Python 集成、命令式风格、API 简单等等。此外,PyTorch 2.0 提供了相同的 eager-mode 开发和用户体验,同时从根本上改变和增强了 PyTorch 在编译器级别的运行方式。该版本能够为「Dynamic Shapes」和分布式运行提供更快的性能和更好的支持。
在官方博客中,PyTorch团队还公布了他们对于整个2.0系列的展望:
以下是详细内容。
PyTorch 2.X:速度更快、更加地 Python 化、一如既往地 dynamic
PyTorch 2.0 官宣了一个重要特性——torch.compile,这一特性将 PyTorch 的性能推向了新的高度,并将 PyTorch 的部分内容从 C++ 移回 Python。torch.compile 是一个完全附加的(可选的)特性,因此 PyTorch 2.0 是 100% 向后兼容的。
支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。
-
TorchDynamo 使用 Python Frame Evaluation Hooks 安全地捕获 PyTorch 程序,这是一项重大创新,是研究团队对快速可靠地获取图进行 5 年研发的结果;
-
AOTAutograd 重载 PyTorch 的 autograd 引擎作为一个跟踪 autodiff,用于生成 ahead-of-time 向后跟踪;
-
PrimTorch 将约 2000 多个 PyTorch 算子规范化为一组约 250 个原始算子的闭集,开发人员可以将其作为构建完整 PyTorch 后端的目标。这大大降低了编写 PyTorch 特性或后端的障碍;
-
TorchInductor 是一种深度学习编译器,可为多个加速器和后端生成快速代码。对于 NVIDIA GPU,它使用 OpenAI Triton 作为关键构建块。
TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的,并支持 dynamic shapes(即能够发送不同大小的张量而无需重新编译),这使得它们具备灵活、易于破解的特性,降低了开发人员和供应商的使用门槛。
为了验证这些技术,研发团队在各种机器学习领域测试了 163 个开源模型。实验精心构建了测试基准,包括各种 CV 任务(图像分类、目标检测、图像生成等)、NLP 任务(语言建模、问答、序列分类、推荐系统等)和强化学习任务,测试模型主要有 3 个来源:
-
46 个来自 HuggingFace Transformers 的模型;
-
来自 TIMM 的 61 个模型:一系列 SOTA PyTorch 图像模型;
-
来自 TorchBench 的 56 个模型:包含来自 github 的精选流行代码库。
然后研究者测量加速性能并验证这些模型的准确性。加速可能取决于数据类型,研究团队选择测量 float32 和自动混合精度 (AMP) 的加速。
在 163 个开源模型中,torch.compile 在 93% 的情况下都有效,模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 float32 精度下,它的平均运行速度提高了 21%,而在 AMP 精度下,它的运行速度平均提高了 51%。
目前,torch.compile 还处于早期开发阶段,预计 2023 年 3 月上旬将发布第一个稳定的 2.0 版本。
TorchDynamo:快速可靠地获取图
TorchDynamo 是一种使用 Frame Evaluation API (PEP-0523 中引入的一种 CPython 特性)的新方法。研发团队采用数据驱动的方法来验证其在 Graph Capture 上的有效性,并使用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。TorchScript 等方法大约在 50% 的时间里都难以获取图,而且通常开销很大;而 TorchDynamo 在 99% 的时间里都能获取图,方法正确、安全且开销可忽略不计(无需对原始代码进行任何更改)。这说明 TorchDynamo 突破了多年来模型权衡灵活性和速度的瓶颈。
TorchInductor:使用 define-by-run IR 快速生成代码
对于 PyTorch 2.0 的新编译器后端,研发团队从用户编写高性能自定义内核的方式中汲取灵感:越来越多地使用 Triton 语言。此外,研究者还想要一个编译器后端——使用与 PyTorch eager 类似的抽象,并且具有足够的通用性以支持 PyTorch 中广泛的功能。
TorchInductor 使用 pythonic define-by-run loop level IR 自动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码和 CPU 上的 C++/OpenMP。TorchInductor 的 core loop level IR 仅包含约 50 个算子,并且是用 Python 实现的,易于破解和扩展。
AOTAutograd:将 Autograd 重用于 ahead-of-time 图
PyTorch 2.0 的主要特性之一是加速训练,因此 PyTorch 2.0 不仅要捕获用户级代码,还要捕获反向传播。此外,研发团队还想要复用现有的经过实践检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展机制来跟踪 Autograd 引擎,使其能够「ahead-of-time」捕获反向传递(backwards pass)。这使 TorchInductor 能够加速前向和反向传递。
PrimTorch:稳定的原始算子
为 PyTorch 编写后端具有挑战性。PyTorch 有 1200 多个算子,如果考虑每个算子的各种重载,则有 2000 多个。
在 PrimTorch 项目中,研发团队致力于定义更小且稳定的算子集,将 PyTorch 程序缩减到这样较小的算子集。目标是定义两个算子集:
-
Prim ops:约有 250 个相当低级的算子。这些算子适用于编译器,需要将它们重新融合在一起以获得良好的性能;
-
ATen ops:约有 750 个规范算子。这些算子适用于已经在 ATen 级别集成的后端或没有编译功能的后端(无法从较低级别的算子集(如 Prim ops)恢复性能)。
用户体验
PyTorch 2.0 引入了一个简单的函数 torch.compile,它会返回一个编译后的模型。
compiled_model *=* torch.compile(model)
compiled_model 保存对模型的引用,并将 forward 函数编译为一个更优化的版本。在编译模型时,PyTorch 2.0 给了几项设置来调整它:
def torch.compile(model: Callable,
*,
mode: Optional[str] = "default",
dynamic: bool = False,
fullgraph:bool = False,
backend: Union[str, Callable] = "inductor",
# advanced backend options go here as kwargs
**kwargs
) -> torch._dynamo.NNOptimizedModule
-
「mode」指定编译器在编译时应该优化的内容。
-
default 是一种预设模式,它试图在不花费太长时间或使用额外内存的情况下高效编译。
-
其他模式,如 reduce-overhead,可以大大降低框架开销,但要消耗少量额外内存。max-autotune 编译很长时间,试图为你提供它所能生成的最快的代码。
-
「dynamic」模式指定是否为 Dynamic Shapes 启用代码路径。某些编译器优化不能应用于动态形状的程序。明确你想要一个带有动态形状还是静态形状的编译程序,将有助于编译器提供更好的优化代码。
-
「fullgraph」类似于 Numba 的 nopython。它将整个程序编译成一个图,或者给出一个错误提示,解释为什么它不能这样做。大多数用户不需要使用这种模式。如果你非常注重性能,那么你可以尝试使用它。
-
「backend 」指定使用哪个编译器后端。默认情况下使用 TorchInductor,但还有其他一些可用的工具。
编译体验想要在默认模式中提供最大的好处和最大的灵活性。
常见问答
1、什么是 PT 2.0?
2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供相同的 eager 模式开发体验,同时通过 torch.compile 添加编译模式。这种编译模式有可能在训练和推理期间加速模型。
2、为什么是 2.0 而不是 1.14?
PyTorch 2.0 就是计划中的 1.14 。我们发布了大量新特性,我们相信这些新特性会改变用户使用 PyTorch 的方式,因此称其为 PyTorch 2.0。
3、如何安装 2.0?有什么额外要求吗?
安装最新的 nightlies:
CUDA 11.7
pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117
CUDA 11.6
pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116
CPU
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu
4、2.0 代码是否向后兼容 1.X?
是的,使用 2.0 不需要修改 PyTorch 工作流程。一行代码 model = torch.compile(model) 就可以优化模型,以使用 2.0 堆栈,并与其余 PyTorch 代码一起顺利运行。这是一个可选择项,不需要使用新的编译器。
5、默认启用 2.0 吗?
不,必须通过使用单个函数调用优化模型从而在 PyTorch 代码中显式启用 2.0。
6、如何将 PT1.X 代码迁移到 PT2.0?
代码应该按原样工作,无需任何迁移。如果想使用 2.0 中引入的新编译模式特性,那么可以从优化模型开始:
model = torch.compile(model)
虽然加速主要是在训练期间观察到的,但如果你的模型运行速度比 eager 模式快,也可以将它用于推理。
import torch
def train(model, dataloader):
model = torch.compile(model)
for batch in dataloader:
run_epoch(model, batch)
def infer(model, input):
model = torch.compile(model)
return model(\*\*input)
7、是否有任何不该使用 PT 2.0 的应用程序?
当前版本的 PT 2.0 仍处于实验阶段,并且处于 nightlies 版本。其中的动态形状支持还处于早期阶段,所以可以等到 2023 年 3 月稳定版发布后再使用该功能。
8、运行 PyTorch 2.0 时,代码有何不同?
开箱即用,PyTorch 2.0 与 PyTorch 1.x 相同,模型以 eager 模式运行,即 Python 的每一行都逐个执行。
在 2.0 中,如果用 model = torch.compile(model) 将模型打包,则模型在执行之前会经过 3 个步骤:
1)图获取:首先将模型重写为子图块。可由 TorchDynamo 编译的子图被「压平」,其他子图(可能包含控制流代码或其他不受支持的 Python 结构)将回退到 Eager 模式。
2)Graph lowering:所有 PyTorch 操作都被分解为特定于所选后端的组成内核。
3)图编译,内核调用其相应的低级设备专用操作。 whaosoft aiot http://143ai.com
9、2.0 目前支持哪些编译器后端?
默认和最完整的后端是 TorchInductor,但是 TorchDynamo 有一个不断增长的后端列表,可以通过调用 torchdynamo.list_backends(). 找到
10、2.0 版本的分布式训练能力如何?
Compiled 模式下的 DDP 和 FSDP ,比 FP32 中的 Eager 模式快 15%、AMP 精度快 80%。PT2.0 做了一些额外的优化,以确保 DDP 的通信 - 计算 overlap 与 Dynamo 的部分图创建良好协作。想要确保使用 static_graph=False 运行 DDP,更多细节参见:https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860
11、为什么我的代码用 2.0 的 Compiled Model 运行变慢?
性能下降最可能的原因是 graph break 太多。例如,类似模型前向 trigger 中的输出语句这样的东西会触发 graph break。详见:https://pytorch.org/docs/master/dynamo/faq.html#why-am-i-not-seeing-speedups
12、以前运行的代码在 2.0 中崩溃了,该如何调试?
参见:https ://pytorch.org/docs/master/dynamo/faq.html#why-is-my-code-crashing
参考链接:https://pytorch.org/get-started/pytorch-2.0/
## 然后在说说另一个话题哦
说到神经网络训练,大家的第一印象都是 GPU + 服务器 + 云平台。传统的训练由于其巨大的内存开销,往往是云端进行训练而边缘平台仅负责推理。然而,这样的设计使得 AI 模型很难适应新的数据:毕竟现实世界是一个动态的,变化的,发展的场景,一次训练怎么能覆盖所有场景呢?
为了使得模型能够不断的适应新数据,我们能否在边缘进行训练(on-device training),使设备不断的自我学习?在这项工作中,我们仅用了不到 256KB 内存就实现了设备上的训练,开销不到 PyTorch 的 1/1000,同时在视觉唤醒词任务上 (VWW) 达到了云端训练的准确率。该项技术使得模型能够适应新传感器数据。用户在享受定制的服务的同时而无需将数据上传到云端,从而保护隐私。
该研究提出了第一个在单片机上实现训练的解决方案,并且系统协同设计(System-Algorithm Co-design)大大减少了训练所需内存。
-
网站:https://tinytraining.mit.edu/
-
论文:https://arxiv.org/abs/2206.15472
-
Demo: https://www.bilibili.com/video/BV1qv4y1d7MV
-
代码: https://github.com/mit-han-lab/tiny-training
设备上的训练(On-device Training)允许预训练的模型在部署后适应新环境。通过在移动端进行本地训练和适应,模型可以不断改进其结果并为用户定制模型。例如,微调语言模型让其能从输入历史中学习;调整视觉模型使得智能相机能够不断识别新的物体。通过让训练更接近终端而不是云端,我们能有效在提升模型质量的同时保护用户隐私,尤其是在处理医疗数据、输入历史记录这类隐私信息时。
然而,在小型的 IoT 设备进行训练与云训练有着本质的区别,非常具有挑战性,首先, AIoT 设备(MCU)的 SRAM 大小通常有限(256KB)。这种级别的内存做推理都十分勉强,更不用说训练了。再者,现有的低成本高效转移学习算法,例如只训练最后一层分类器 (last FC),只进行学习 bias 项,往往准确率都不尽如人意,无法用于实践,更不用说现有的深度学习框架无法将这些算法的理论数字转化为实测的节省。最后,现代深度训练框架(PyTorch,TensorFlow)通常是为云服务器设计的,即便把 batch-size 设置为 1,训练小模型 (MobileNetV2-w0.35) 也需要大量的内存占用。因此,我们需要协同设计算法和系统,以实现智能终端设备上的训练。
方法与结果
我们发现设备上训练有两个独特的挑战:(1)模型在边缘设备上是量化的。一个真正的量化图(如下图所示)由于低精度的张量和缺乏批量归一化层而难以优化;(2)小型硬件的有限硬件资源(内存和计算)不允许完全反向传播,其内存用量很容易超过微控制器的 SRAM 的限制(一个数量级以上),但如果只更新最后一层,最后的精度又难免差强人意。
为了应对优化的困难,我们提出了 Quantization-Aware Scaling (QAS) 来自动缩放不同位精度的张量的梯度(如下左图所示)。QAS 在不需要额外超参数的同时,可以自动匹配梯度和参数 scale 并稳定训练。在 8 个数据集上,QAS 均可以达到与浮点训练一致的性能(如下右图)。
为了减少反向传播所需要的内存占用,我们提出了 Sparse Update,以跳过不太重要的层和子张的梯度计算。我们开发了一种基于贡献分析的自动方法来寻找最佳更新方案。对比以往的 bias-only, last-k layers update, 我们搜索到的 sparse update 方案拥有 4.5 倍到 7.5 倍的内存节省,在 8 个下游数据集上的平均精度甚至更高。
为了将算法中的理论减少转换为实际数值,我们设计了 Tiny Training Engine(TTE):它将自动微分的工作转到编译时,并使用 codegen 来减少运行时开销。它还支持 graph pruning 和 reordering,以实现真正的节省与加速。与 Full Update 相比,Sparse Update 有效地减少了 7-9 倍的峰值内存,并且可以通过 reorder 进一步提升至 20-21 倍的总内存节省。相比于 TF-Lite,TTE 里经过优化的内核和 sparse update 使整体训练速度提高了 23-25 倍。
结论
本文中,我们提出了第一个在单片机上实现训练的解决方案(仅用 256KB 内存和 1MB 闪存)。我们的算法系统协同设计(System-Algorithm Co-design)大大减少了训练所需内存(1000 倍 vs PyTorch)和训练耗时(20 倍 vs TF-Lite),并在下游任务上达到较高的准确率。Tiny Training 可以赋能许多有趣的应用,例如手机可以根据用户的邮件 / 输入历史来定制语言模型,智能相机可以不断地识别新的面孔 / 物体,一些无法联网的 AI 场景也能持续学习(例如农业,海洋,工业流水线)。通过我们的工作,小型终端设备不仅可以进行推理,还可以进行训练。在这过程中个人数据永远不会上传到云端,从而没有隐私风险,同时 AI 模型也可以不断自我学习,以适应一个动态变化的世界!