机器学习部署落地:底层编译器那些事


我们为什么要了解编译器?

随着对机器学习模型投入生产的了解加深,编译器的话题不断出现。在许多应用情况下,特别是在边缘运行机器学习模型时,模型的成功仍然取决于其运行的硬件,这使得那些在生产中处理机器学习模型的人了解如何编译和优化模型以在不同的硬件加速器上运行变得重要。

理想情况下,编译器应该是隐形的,一切都能“自动运行”。然而,我们离实现这一目标还有很多年的路要走。随着越来越多的公司希自将机器学习应用于边缘计算,并且为机器学习模型开发了越来越多的硬件,为了弥合模型和硬件加速器之间的差距,正在开发越来越多的编译器,比如MLIR、TVM、XLA、PyTorch Glow、cuDNN等。

根据PyTorch的创始人Soumith Chintala的说法,随着机器学习的采用逐渐成熟,企业将会竞相展开竞争,看谁能够更好地编译和优化他们的模型。了解编译器的工作原理可以帮助选择合适的编译器,将模型部署到所选择的硬件上,最终帮助诊断性能问题。


云计算与边缘计算

想象一下,你已经训练出一个令人难以置信的机器学习模型,其准确性超出了你最疯狂的期望。你很兴奋地想要部署这个模型,以便用户可以访问它。

最简单的方法是将你的模型打包并通过托管的云服务(如AWS或GCP)部署,这是许多公司在开始进行机器学习时采用的方式。云服务在使公司将机器学习模型投入生产方面做得非常出色。

然而,云部署也有许多不足之处。首先是成本。机器学习模型可能需要大量的计算资源,而计算资源是昂贵的。即使在2018年,像Pinterest、Infor、Intuit等大公司已经每年在云服务上花费数亿美元的云费用。对于中小公司来说,这个数字可能在每年5万到200万美元之间,对云服务处理不当可能导致创业公司破产。

机器学习部署落地:底层编译器那些事

随着AWS的使用率飙升,公司对云费用感到惊讶。随着云费用不断攀升,越来越多的公司正在寻求将计算推向消费设备(边缘设备)。

在边缘上执行的计算量越多,云上所需的计算量就越少,他们就不必为服务器支付更多费用。其次,当模型已经部署在用户设备上时,网络延迟是一个比推理延迟更大的瓶颈。例如将ResNet50的推理延迟从30毫秒减少到20毫秒,但网络延迟可能会达到几秒。

将模型部署到边缘设备上还可以在处理敏感用户数据时具有吸引力。在云上进行机器学习意味着您的系统可能必须通过网络发送用户数据,从而使其容易被截取。云计算通常也意味着将许多用户的数据存储在同一个地方,边缘计算使得更容易遵守关于如何传输或存储用户数据的法规(如GDPR)。


编译兼容性

由于边缘计算在许多方面优于云计算,各个公司正竞相开发针对不同机器学习用例进行优化的边缘设备。包括谷歌、苹果、特斯拉在内的知名公司都宣布了推出自己芯片的计划。与此同时,机器学习硬件初创公司已筹集了数十亿美元来开发更好的人工智能芯片。

由于有许多新的硬件供应来运行机器学习模型,一个问题浮现出来:我们如何使使用任意框架构建的模型能够在任意硬件上运行?

要使框架在一种硬件上运行,它必须得到硬件供应商的支持。例如,即使TPU在2018年2月已公开发布,但直到2020年9月,PyTorch才支持TPU。在那之前,如果想使用TPU,就必须使用TensorFlow或JAX。

为一个框架在一种类型的硬件(平台)上提供支持是耗时且需要工程师大量投入的。将ML工作负载映射到硬件上需要理解并能够充分利用硬件的基础设施。然而,一个基本的挑战是不同类型的硬件有不同的内存布局和计算原语。

例如,CPU的计算原语以前是一个数字(标量),GPU的计算原语以前是一维向量,而TPU的计算原语是二维向量(张量)。然而,现在许多CPU都有向量指令,一些GPU也有张量核心,这是二维的。在一批256张图像x 3通道x 224宽x 224高上执行卷积运算符将会与1维向量相比,与2维向量相比,差异非常大。

框架开发人员往往会专注于为一些服务器级别的硬件(如GPU)提供支持,而硬件供应商则倾向于为一些狭窄范围的框架提供自己的内核库(例如,英特尔有支持Caffe、TensorFlow、MXNet、Kaldi和ONNX的OpenVino。NVIDIA有CUDA和cuDNN)。

▎中间表示(IR)

与其为每一种新的硬件类型和设备开发新的编译器和库,为什么不创建一个中间层来连接框架和平台呢?框架开发人员将不再需要支持每种类型的硬件,只需要将他们的框架代码翻译成这个中间层。然后,硬件供应商可以支持一个中间框架,而不是支持许多个不同的框架。

这种“中间层”被称为中间表示(IR)。中间表示在编译器的工作中起着核心作用。从模型的原始代码开始,编译器会在生成适用于特定平台的硬件本地代码之前,生成一系列高级和低级中间表示。

要从IR生成机器本地代码,编译器通常会使用代码生成器,也称为codegen。ML编译器中最常用的代码生成器是LLVM,由Vikram Adve和Chris Lattner开发(他们通过创建LLVM改变了我们对系统工程的看法)。TensorFlow XLA、NVIDIA CUDA编译器(NVCC)、MLIR(用于构建其他编译器的元编译器)和TVM都使用LLVM。

这个过程也被称为“降级”,因为你将高级框架代码“降级”为低级的硬件本地代码。这并不是“翻译”,因为它们之间没有一对一的映射关系。

高级IR通常是ML模型的计算图。对于熟悉TensorFlow的人来说,这里的计算图类似于在TensorFlow 1.0中遇到的计算图,在TensorFlow切换到即时执行之前。在TensorFlow 1.0中,TensorFlow首先构建了模型的计算图,然后再运行它。这个计算图使TensorFlow能够理解模型并优化其运行时。

高级IR通常是与硬件无关的(不关心将在哪种硬件上运行),而低级IR通常是与框架无关的(不关心使用哪种框架构建的模型)。


性能优化很难

当代码“降级”以便在所选择的硬件上运行模型后,可能会遇到性能问题。代码生成器非常擅长将IR转换为机器代码,但是根据目标硬件后端,生成的机器代码可能性能不佳。生成的代码可能无法充分利用数据局部性和硬件缓存,也可能无法利用矢量或并行操作等高级功能,以加速代码执行。

典型的机器学习工作流包含许多框架和库。例如使用pandas/dask/ray从数据中提取特征、可能会使用NumPy进行矢量化、可能会使用像LightGBM这样的树模型来生成特征,然后使用使用各种框架(如sklearn、TensorFlow或transformers)构建的模型集合进行预测。

尽管这些框架中的各个函数可能已经进行了优化,但在这些框架之间几乎没有优化。 在计算中将数据从这些函数移动的天真方法可能会导致整个工作流程的几个数量级的减速。

通常在生产环境中发生的情况是,数据科学家/机器学习工程师会pip安装所需的软件包。在开发环境中,一切似乎都运行良好,所以他们将模型部署到生产环境。当他们在生产环境中遇到性能问题时,他们的公司通常会聘请优化工程师来为他们运行的硬件优化模型。

性能优化工程师很难找到,也很昂贵,因为他们需要在机器学习和硬件架构领域都有专业知识。优化编译器(同时优化代码的编译器)是一个替代方案,因为它们可以自动化模型的优化过程。在将机器学习模型代码降级为机器代码的过程中,编译器可以查看机器学习模型的计算图和其中的运算符(如卷积、循环、交叉熵)并找到加速它的方法。


优化机器学习模型

优化机器学习模型有两种方法:局部优化和全局优化。局部优化是指优化模型的一个操作或一组操作符。全局优化是指对整个计算图进行端到端的优化。

有一些标准的局部优化技术被认为可以加速模型,其中大多数技术都是通过并行运行或减少芯片上的内存访问来实现的。以下是常见的四种技术。

矢量化:对于一个循环或嵌套循环,不是逐一处理每个项目,而是使用硬件原语在内存中连续的多个元素上进行操作。

并行化:对于一个输入数组(或n维数组),将其分成不同的、独立的工作块,并在每个块上进行操作。

循环切片:更改循环中的数据访问顺序,以利用硬件的内存布局和缓存。这种优化方式依赖于硬件。在CPU上的良好访问模式在GPU上可能不是一个良好的访问模式。以下是由Colfax Research制作的可视化效果。

操作符融合:将多个操作符融合成一个,以避免冗余的内存访问。例如,对同一数组进行的两个操作需要对该数组进行两次循环。在融合的情况下,只需一个循环。以下是由Matthias Boehm制作的示例。

机器学习部署落地:底层编译器那些事
循环切片


机器学习部署落地:底层编译器那些事
操作符融合

根据Weld的创建者Shoumik Palkar的说法,这些标准的局部优化技术可以预期给模型带来约3倍的速度提升。当然,这个估计是高度依赖上下文的。

要获得更大的速度提升,需要利用计算图的更高级结构。例如,对于一个卷积神经网络,计算图可以在垂直或水平方向上进行融合,以减少内存访问并加速模型。以下是由NVIDIA的TensorRT团队制作的可视化效果。

机器学习部署落地:底层编译器那些事
图优化


手工设计 VS 机器学习编译器

▎手工设计的规则

正如上一节中通过卷积神经网络的垂直和水平融合所示,执行给定计算图的许多可能方式。例如,给定3个操作符A、B和C,可以将A与B融合,将B与C融合,或者将A、B和C全部融合。

传统上,框架和硬件供应商会雇佣优化工程师,根据他们的经验,制定出如何最佳执行模型的计算图的启发式方法。例如,NVIDIA可能会有一个工程师或一个工程师团队,专门致力于如何使ResNet-50在他们的DGX A100服务器上运行得非常快。

在某种类型的硬件上运行得非常快的流行模型并不意味着任意模型在该硬件上都会运行得非常快。这可能只是因为该模型被过度优化了。

手工设计规则有一些缺点。首先,它们是非最优的。没有保证工程师提出的启发式方法是最佳解决方案。其次,它们是非自适应的。在新的框架或新的硬件架构上重复这个过程需要大量的努力。

这还受到模型优化依赖于组成其计算图的一组操作符的影响。优化卷积神经网络与优化循环神经网络是不同的。NVIDIA和Google专注于优化像ResNet和BERT这样的流行模型在他们的硬件上。

但是,作为机器学习研究人员,如果提出了一种新的模型架构,可能需要自己进行优化,以先证明它是快速的,然后再由硬件供应商进行优化。

▎使用机器学习加速机器学习模型

使用机器学习加速机器学习模型的目标是找到执行计算图的所有可能方式中最快的方式。如果尝试所有可能的方式,记录它们运行所需的时间,然后选择最佳方式,会怎么样?

问题是有太多可能的路径要探索,尝试所有这些路径将被证明是不可行的。如果我们使用机器学习:

  • 缩小搜索空间,以便不必尝试那么多路径。
  • 预测路径所需的时间,以便我们不必等待整个计算图完成执行。

估计整个计算图中通过一条路径所需的时间非常困难,因为这需要对该图做出许多假设。目前的技术可能仅限于关注计算图的一小部分。

如果在GPU上使用PyTorch,可能会看到torch.backends.cudnn.benchmark=True。当设置为True时,将启用cuDNN自动调优。cuDNN自动调优在预定的一组选项上执行卷积运算,并选择最快的方法。cuDNN自动调优在每次迭代中运行相同的卷积形状时非常有用。第一次运行卷积运算时速度会很慢,因为cuDNN自动调优需要时间来运行搜索。但在后续运行中,cuDNN将使用自动调优的缓存结果来选择最快的配置。

尽管cuDNN自动调优非常有效,但只适用于卷积运算符,据我所知,它只在PyTorch和MXNet中公开。更通用的解决方案是autoTVM,它是开源编译器TVM的一部分。autoTVM使用子图而不仅仅是一个运算符,因此它处理的搜索空间更加复杂。

autoTVM的工作方式非常复杂,但以下是要点:
1. 首先将计算图分解为子图。
2. 预测每个子图有多大。
3. 为搜索每个子图的最佳路径分配时间。
4. 将每个子图的最佳运行方式连接在一起以执行整个图。

autoTVM测量沿着每条路径运行所需的实际时间,从而为其训练成本模型提供了基准数据,以预测未来路径所需的时间。这种方法的优点是,因为模型是使用运行时生成的数据进行训练的,所以它可以适应它运行的任何类型的硬件。缺点是,花费更多的时间来改进成本模型的性能。

机器学习部署落地:底层编译器那些事
autoTVM加速


机器学习部署落地:底层编译器那些事
autoTVM成本模型

像TVM这样的编译器是自适应的、灵活的,并且在想尝试新硬件时特别有用。一个例子是在2020年11月苹果发布M1芯片时。M1是基于ARM的片上系统,ARM架构多少是已知的。然而,M1仍然具有其ARM实现的许多新颖组件,并且需要进行大量的优化,以使各种机器学习模型在其上运行快速。在发布一个月后,OctoML的人员展示了autoTVM的优化几乎比Apple的Core ML团队手动设计的优化快30%。当然,随着M1的发展和手动设计优化的深入,自动优化很难击败手动设计优化。但是系统工程师可以利用像autoTVM这样的工具来加速优化过程。

虽然自动调优结果令人印象深刻,但它们也有一个限制:TVM可能会很慢。您需要遍历所有可能的路径并找到最优的路径。这个过程可能需要几小时,甚至对于复杂的机器学习模型可能需要几天。然而这是一次性的操作,优化搜索的结果可以被缓存并用于优化现有模型,同时也可以作为未来调整会话的起点。可以针对一个硬件后端对模型进行一次优化,然后在该后端的多个设备上运行它。


不同类型的编译器

最广泛使用的编译器类型是由主要框架和硬件供应商开发的领域特定编译器,针对特定的框架和硬件组合。毫不奇怪,最流行的编译器通常由最大的供应商开发。

NVCC(NVIDIA CUDA编译器):仅适用于CUDA。闭源。

XLA(加速线性代数,Google):最初旨在加速TensorFlow模型,但已被JAX采用。作为TensorFlow存储库的一部分开源。

PyTorch Glow(Facebook):PyTorch已经采用了XLA以在TPU上运行PyTorch,但对于其他硬件,它依赖于PyTorch Glow。作为PyTorch存储库的一部分开源。

第三方编译器通常非常雄心勃勃(例如,认为您可以比NVIDIA更好地优化GPU的信心必须非常足够)。但是第三方编译器很重要,因为它们有助于降低开发新框架、新硬件代、新模型性能的开销,使小型参与者有机会与已经为其现有产品大量调优的大型参与者竞争。

我看到的最好的第三方编译器是Apache TVM,它可以与各种框架(包括TensorFlow、MXNet、PyTorch、Keras、CNTK)和各种硬件后端(包括CPU、服务器GPU、ARM、x86、移动GPU和基于FPGA的加速器)一起使用。

我认为令人兴奋的另一个项目是MLIR,最初由Google的Chris Lattner(LLVM的创始人)发起。然而,现在它在LLVM组织下。MLIR实际上不是一个编译器,而是一个元编译器,允许您构建自己的编译器。MLIR可以运行多个IR,包括TVM的IR、LLVM IR和TensorFlow图。


编译器的未来展望

当模型准备部署时,尝试不同的编译器以查看哪一个能够为您提供最佳的性能提升是有意义的。可以并行运行实验。一个推断请求的小提速可以在数百万或数十亿个推断请求中积累成大的回报。

尽管在机器学习的编译器方面已经取得了巨大的进展,但在完全从一般的ML从业者中抽象出编译器之前,仍然有很多工作要做。想象一下像GCC这样的传统编译器。在C或C++中编写代码,GCC会自动将代码转换为机器代码。大多数C程序员甚至不关心GCC生成了什么中间表示。

在将来ML编译器可以是同样的方式。使用框架以计算图的形式创建ML模型,ML编译器可以为运行的任何硬件生成机器本地代码。


本文转自:Coggle数据科学,转载此文目的在于传递更多信息,版权归原作者所有。如不支持转载,请联系小编demi@eetrend.com删除。

最新文章