作者｜River Riddle、Eric Johnson、Abdul Dakak
翻译｜胡燕君、杨婷

机器学习模型逐渐开展成人们口中的“庞然大物”。全球顶尖的科技公司纷繁踏上“军备竞赛”之路，立志训练出规划最大的模型（MUM、OPT、GPT-3、Megatron），而其他专注于出产体系的公司也相继扩展其原有模型，并获得杰出效果。

一切如火如荼，但是，鲜少有人提及，巨大的模型给现有的AI根底设施和开发流程带来了诸多实践性挑战。

大模型的权重可达100+GB，而目前的开发东西却还没跟上，运用起来十分费力，部署时往往要等上好几分钟乃至好几小时，这已经成为AI工程师的隐痛，不光糟蹋工程师的时刻，下降作业功率，还会拖慢迭代速度。

努力于AI根底设施东西研制的Modular团队以为，开发人员的作业功率是训练和部署模型的最大本钱之一。因而需求不断优化东西链，进步早期用户的体会，也方便开发人员。本文探讨编译过程中办理海量数据的技能难点，以及Modular为处理这些难点在根底设施（以及MLIR编译器框架）方面所做的改善。由OneFlow社区（ID：OneFlowTechnology）编译。

1、AI模型配套东西的易用性缺乏

机器学习中的图转化（Graph Transformations）、优化和编译器等技能的作用是进步AI模型的功能和便携性，让模型能够部署在某些目标硬件上。

编译器中，有TensorFlow Lite Converter这样的高层次“编译器”，它能够将TensorFlow SavedModel模型转化为高度优化的程序格局（如FlatBuffer格局），让模型能够在边际设备上履行；也有XLA和TorchScript JIT Compiler这样针对特定范畴的编译器，它们为AI模型创建中心表明（或许是一张“图”），然后将其编译成另一种格局——例如机器码或特定范畴的运转时表明（如CUDA图）。

AI图的编译与传统的编译很不一样。AI图包含两部分：图拓扑（各层之间怎么连接）和模型权重（特定层的参数）。从巨细来看，图拓扑以KB为单位，权重则以MB乃至GB为单位。举个比方，Meta公司发布的Open Pre-trained Transformers模型，其参数量从300亿、660亿到1750亿不等，适当于100+GB权重。Gopher和Megatron模型乃至更大。

图1:图源DeepMind论文

AI生态体系中现有的东西尚不能很好地处理大型模型。比方，Protobufs约束了传输数据巨细不能超过2GB，因而模型假如运用Protobufs序列化格局，就会备受掣肘。最新版TensorRT的文档中写道，“对于BERT和GPT等依据Transformer的神经网络模型，TensorRT在编译时能够耗费10倍于模型巨细的CPU内存”，可见TensorRT不适合大型模型。假如运用ONNX文件格局存储大型模型，就必须将模型权重分成多个文件分别存储。

以上种种不光给AI开发作业流添加了不必要的杂乱环节，也使模型丧失“单一事实来历”（SSOT），还导致模型分发更加困难。

为了应对模型权重太大的问题，大家或许会采纳变通办法，终究却或许导致整个AI开发作业流变得更杂乱。比方，由于某些编译器阶段耗时长达2分多钟，打断开发人员的作业节奏，所以Modular构建了一种缓存临时文件的机制。

尽管这种缓存机制和其他变通办法一样，仅仅治标不治本：它既非100%可靠，也不能处理Cache Miss（缓存缺失）的问题，不过由于Modular十分注重进步开发人员的作业功率，所以仍是决议选用这种机制。

2、Modular编译栈中的MLIR

Modular的技能栈中，MLIR编译器架构担任表明和转化AI模型，包含AI算子图（用于多种框架）、中级运转时原语和初级机器码生成。

图2:多级中心表明 (MLIR)

MLIR是LLVM编译器根底设施项目的子项目，LLVM旨在供给现代东西包，用以构建针对特定范畴的编译器。MLIR供给一套中心组件，用于硬件规划、量子核算、人工智能等多种核算范畴的建模、剖析和转化。

MLIR能够协助构建单个涵盖全栈的完好体系，比常规的技能栈功用更强大、模块化程度和可拓展性更高，也更易于维护。运用一致的根底设施让咱们得以便捷地将每一项改善迁移到自己的东西栈，使开发作业流完成更高的模块化和可拼装性。

除了Modular以外，TensorFlow、XLA、PyTorch和ONNX等也在运用MLIR进行模型表明和转化。跟着MLIR的用户生态不断扩展，在赞美MLIR长处的一起，也必须继续进行改善和完善。

3、MLIR办理权重的办法还有待进步

MLIR的基本组成部分之一是特色机制（Attribute），能够把它理解为被unique（或被memoize、intern）的常量数据。特色是用户可拓展的，也就是说，能够依据不同用例运用不同的特色类型。许多类型的值都能够被赋予特色，比方常量表达式值（如“5”、“10.0”等）、字符串字面量、枚举值（如“小于”、“大于”、“等于”等），数据组等等。大多数依据MLIR的AI东西都运用特色来保存AI模型的权重。

但是，问题呈现了：模型权重有或许极其巨大，但MLIR存储2 GB权重的办法和存储4 B权重的办法并没有差异——都运用同一特色，该特色包含一组被unique的元素。但对GB级的巨大数据运用unique办法显然不合理。

这个办法的难点在于：在MLIR中，当某个东西被unique，它就会被分配（allocated）、被hash 、然后被贮存到MLIRContext中。这些东西具有和MLIRContext相同的生命周期，只要当MLIRContext被毁掉，它们才会一起被毁掉。对于小的数值而言，这种机制带来许多好处，能够把数值传入传出，能够经过指针对unique后的值进行比较，还能够共享特色的内存分配（十分常见）等等。

但对数量巨大的权重而言，上述种种好处就变成了下风：咱们不期望对权重进行重新分配、仿制或运用unique办法，咱们只需求它们时刻短存在——当核算不再需求引证这些权重时，就要答应开释分配。例如，当运转模型量化东西时，需求对算子图进行转化，并生成新的权重，终究这些权重或许会被仿制多份，很多权重副本在编译结束前都将一向占用内存。

ML东西的另一个问题是MLIR怎么序列化至文件体系。一开始，MLIR没有二进制序列化格局，只要文本格局。对数量巨大的权重来说，这就造成问题，由于每个字节的二进制数据都会被转化为十六进制，后者占用的空间为前者的2倍。这样一来，咱们不光耗费了适当长的时刻进行进制转化（一个中等的GB级模型大约需求20秒），并且转化后的中心文件仍是本来的2倍大——2倍但是一个不小的数字！

4、内存占用：比拖慢开发功率更严重的影响

这一规划机制原意虽好，但它有或许下降编译功率，即运用最好的编译器也无济于事。最显着的问题是它会导致编译、监控和转化模型的时刻变长。但凡你曾用过“我的代码还在编译”作为日常摸鱼的借口，你就理解等待编译是多么痛苦的作业。选用这一机制，就意味着处理器不得不对GB级数据继续进行分配、仿制和hash处理。

图3:XKCD漫画 – 《还在编译》

比编译时长更严重的问题是内存占用，它会影响Modular技能栈中的其他架构功用的完成。例如，由于咱们的编译器和技能栈自身都高度并行，并且运用线上搜索等高档功用，内存占用会直接导致一些作业不能并行展开，也导致不能获得最高质量的成果。

Modular的价值中心是构建用户喜爱的东西。高档功用假如不好用，或许会影响功率，又或许附带一些注意事项（比方，“该功用对某些情况不适用”），那么用户就底子不会用。因而，Modular努力于处理巨大权重带来的根底性问题，简化用户的运用流程和开发人员的作业流程。

5、MLIR的中心改善

Modular团队是MLIR项目的重要贡献者，Modular企业文化的一大关键是“做对的产品”，Modular参加的一切项目都遵从这一要义。在推进MLIR开展的一起，Modular竭力确保MLIR项目的每一步路都正确，也加强与MLIR社区的协作，为所采纳的办法争夺认可。

Modular团队列出了大型模型东西应该具有的特色：

• 非必要不分配内存：对大型数据（比方权重）而言，从磁盘中实行内存映射比将数据仿制到已分配内存的block中更高效。

• 无需进行hash或unique处理：咱们不期望费力气去查看2 GB Blob数据的持平性；要辨别权重，期望能够经过称号辨别，而不是看详细内容有没有被unique。

• 答应内联改变（Inline Mutation） ：假如数据只需求在一处运用，应当答应在原方位量化、转化和操作数据，而不是先仿制数据。

• 答应开释内存（deallocation） ：由于大模型的数据量十分巨大，因而当对某一数据的一切引证都不存在时，应当答应开释内存。

• 快速序列化：无论是即时编译，搜索优化参数，仍是本地迭代，都需求缓存IR，所以这一步必须快。

上述观点并不新颖，但传统编译器（比方适用于典型CPU编程言语的编译器）却还没有完成这些要求。

6、调整权重特色

上述前四点要求处理了咱们应该怎么运用MLIR这一基本问题：权重尽管是常量数据，但对它的办理应该差异于其他MLIR特色。一向以来，咱们的权重办理办法都很不适宜，这就好比试图将一块方钉挤进圆孔中，不仅糟蹋了空间，下降了咱们的开发速度，一起也添加了用户本钱。

所以Modular决议换一种办法来办理权重数据，这促成了MLIR的第一个基本扩展机制——“Resource机制”，在核算中将数据和对数据的引证区分开来。

在Resource机制中，序列化MLIR的每个Blob都或许包含额外的信息段，称为Resource。Resource要么是dialect（扩展MLIR时运用的相似namespace的抽象），要么是用于特定东西链数据的“外部（external）”资源。Resource中的数据用简略的键值对表明，创造出如下图所示的相似json的结构。

/// Here we have some MLIR operations.
module {
  func.func @foo() {
    // Cool stuff here ...
  }
}
/// Here we have an `external_resources` section. The resource section's syntax is designed to be unique as to not conflict with other MLIR syntax (which is user extensible!).
{-#
  external_resources: {
    mlir_reproducer: {
      pipeline: "func.func(cse,canonicalize),inline",
      disable_threading: true
    }
  }
#-}

上面比方展现了怎么调整MLIR来用Resource进行复现。MLIR再生器（Reproducer）实践上是一种配置，它包含转化管道（Transformation Pipeline）等履行信息，用于复现某种毛病或失利。在运用Resource之前，咱们经过在MLIR文件顶部添加注释来表明这些履行信息。现在能够运用Resource将这些履行信息合并为第一类信息。

从前需求进行unique处理导致长时间占用内存的大型权重数据，现在能够运用Resource机制进行贮存。在IR中，咱们对特色选用轻量级引证而不再选用底层数据：

其他优势：

• 运用IR进行调试时更不简单出错，然后带来更好的开发体会：Resource是专门的信息段；咱们不必忧虑在调试时会不小心转储整整4GB的数据。

• 咱们能够在无需数据的情况下合理地处理IR：由于IR只保存对数据的引证，不保存数据自身，假如需求，咱们能够省略底层Resource数据。这样做的好处包含能够极大地简化再生器生成流程，再生器本来就不需求用到大型权重数据（想象一下，你曾经需求向同事发送1.2GB的再现器文件，现在的再生器文件只要20MB大）。

经过引进Resource这个新概念，咱们在程序和数据之间树立明晰的别离机制。现在，咱们不再将权重数据直接传递给某一特色。相反，咱们向特色传入一个弱引证，并将权重数据传给一个专门的办理器。这样，咱们就能更好地控制权重分配、改变和毁掉的时刻和办法。

7、新增MLIR二进制编码办法

有了更好的权重表明办法之后，下一步咱们只需找到更高效的权重贮存办法来完成MLIR表明的序列化。

到此为止，MLIR只要文本序列化格局，这种格局运用ASCII十六进制字符串来表明权重。但是，Modular的终极目标是尽或许加速本地开发流程，因而需求摒弃文本序列化格局，为MLIR新增合适的二进制格局（discourse.llvm.org/t/rfc-a-bin… ）。

二进制格局需求考虑许多要素，何况二进制格局决议了编译器的稳定性。MLIR需求高度的灵活性才能高效应对各式各样的用例，需求完成高速度，并且MLIR/LLVM不能依赖第三方编码库。

不过，MLIR的一大好处是编码难度极低。由于MLIR中一切操作的结构都相同，一切操作都能够运用相同的编码办法。上述的种种杂乱要求都是为了确保MLIR中心概念的紧凑和高效。考虑到这些约束，咱们决议为MLIR定制编码办法（mlir.llvm.org/docs/Byteco… ）。

8、用户收益

为MLIR添加Resource机制和二进制编码办法大大加速了东西链和开发流程，并大幅下降内存占用，进步了功能和速度体现，也全体改善了MLIR。

为了验证上述改善带来的功能变化，能够测验不同规划的模型上依据MLIR的图编译器中“降级”和“优化”过程的实践速度（将TensorFlow序列化模型转化为符合MLIR运转时输入格局的模型），以及该过程中的实践内存占用。

速度进步：编译作业流

测验成果发现，MLIR的速度大幅进步。从TensorFlow序列化模型（TensorFlow 2.10模型）转化为MLIR运转时输入格局的模型，这一过程触及很多底层表明转化，经过改善后，实践履行时刻缩短了1.8~2倍，履行速度随模型巨细按份额缩放。

详细而言，处理TensorFlow序列化模型耗时极短——生成MLIR时将很多权重数据写入磁盘这一过程是首要的耗时来历。经改善后，代码处理时刻比本来快10倍，全体履行时刻的快慢首要取决于固态硬盘（SSD）将 >1 GB数据写入磁盘的速度。

ML开发人员运用咱们的东西，能够加速模型编译速度，然后进步出产功率，削减迭代时刻。咱们的东西能够优化出产环境以及模型的加载和编译，包含依据流入数据的动态模型加载和卸载，以及各种个性化或经过精细化调整的用户模型。

速度进步：序列化

引进二进制编码办法不光能够加速编译作业流，还能加速序列化速度。经过外部东西与MLIR进行交互，包含运转时类型查看（Introspection）、缓存和再生器生成等，都需求对序列化MLIR进行读写。

经过对不同规划的模型进行了序列化测验，成果同样发现峰值功能大幅提速，且SSD写入过程仍然为首要耗时来历。详细而言，大型模型文本数据的读取耗时约5秒，而二进制数据的读取耗时仅不到10毫秒；二进制格局的写入速度则约是文本格局数据的5倍。

对Modular而言，引进二进制编码办法能够加速以MLIR为中心的根底设施和东西的开发速度，改善本来高本钱、低速度的开发情况。比方，调试器（Debugger）的功率很大程度取决于编译作业流中缓存模型表明的功率，而引进二进制编码办法能够进步调试器的功率，然后进步底层编译器的功能。

内存占用

二进制序列化格局的mmap（一种内存映射办法）功能以及经过Resource机制完成的IR和数据的相互独立功能够大幅削减内存占用。测验发现，各种规划的模型编译流程中的内存占用都大大下降——由于曾经需求为模型权重分配内存，现在不需求了。

9、晋级AI生态

Modular的愿景不仅仅为了方便咱们自己，而是晋级整个AI行业的生态。前文提及的新式Resource表明和二进制编码办法都已提交至上游的LLVM/MLIR仓库中。

Modular起初的研制动机是为了处理Modular的客户遇到的问题并进步自身内部根底设施，但这些改善产生的积极影响并不限于自己的用例，还能改善其他以MLIR为根底技能的产品。例如，由于二进制编码办法的引进，MLIR社区如今正在评论怎么确保MLIR的稳定性（discourse.llvm.org/t/mlir-gene… ）。

这些根底技能的改善终究都会融入产品中为用户服务。以上仅仅Modular努力进步的无数中心技能之一。Modular一方面竭力习惯大模型，一方面努力完善模型在设备上的部署，目标都是大幅进步AI根底设施的功能和易用性。Modular十分看好AI的未来以及LLVM和MLIR的开展。
（本文由OneFlow社区翻译，译文转载请联络OneFlow获得授权。原文：1.
www.modular.com/blog/increa… ；2.www.modular.com/blog/increa… ）

欢迎下载体会 OneFlow v0.8.0 最新版本：
github.com/Oneflow-Inc…