关于人工智能:大模型狂欢背后AI基础设施的老化与改造工程

作者|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 模型甚至更大。

图源 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 算子图（用于多种框架）、中级运行时原语和低级机器码生成。

多级两头示意 (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 解决。

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 新增适合的二进制格局（https://discourse.llvm.org/t/…）。

二进制格局须要思考很多因素，况且二进制格局决定了编译器的稳定性。MLIR 须要高度的灵活性能力高效应对各种各样的用例，须要实现高速度，而且 MLIR/LLVM 不能依赖第三方编码库。

不过，MLIR 的一大益处是编码难度极低。因为 MLIR 中所有操作的构造都雷同，所有操作都能够应用雷同的编码方式。上述的种种简单要求都是为了保障 MLIR 外围概念的紧凑和高效。思考到这些限度，咱们决定为 MLIR 定制编码方式（https://mlir.llvm.org/docs/By…）。

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 的稳定性（https://discourse.llvm.org/t/…）。

这些根底技术的改良最终都会融入产品中为用户服务。以上只是 Modular 致力晋升的有数核心技术之一。Modular 一方面极力适应大模型，一方面致力欠缺模型在设施上的部署，指标都是大幅晋升 AI 基础设施的性能和易用性。Modular 十分看好 AI 的将来以及 LLVM 和 MLIR 的倒退。

（本文由 OneFlow 社区翻译，译文转载请分割 OneFlow 取得受权。原文：1. https://www.modular.com/blog/…；2.https://www.modular.com/blog/…）

欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In…