大模型狂欢背后：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 模型甚至更大。
图 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 的结构。
上面例子展示了如何调整 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/rfc-a-binary-serialization-format-for-mlir/63518）。
二进制格式需要考虑很多因素，况且二进制格式决定了编译器的稳定性。MLIR 需要高度的灵活性才能高效应对各种各样的用例，需要实现高速度，而且 MLIR/LLVM 不能依赖第三方编码库。
不过，MLIR 的一大好处是编码难度极低。因为 MLIR 中所有操作的结构都相同，所有操作都可以使用相同的编码方式。上述的种种复杂要求都是为了保证 MLIR 核心概念的紧凑和高效。考虑到这些限制，我们决定为 MLIR 定制编码方式（https://mlir.llvm.org/docs/BytecodeFormat/）。
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/mlir-generic-ir-stability-and-upgradability/65371）。
这些基础技术的改进最终都会融入产品中为用户服务。以上只是 Modular 致力提升的无数核心技术之一。Modular 一方面竭力适应大模型，一方面努力完善模型在设备上的部署，目标都是大幅提升 AI 基础设施的性能和易用性。Modular 非常看好 AI 的未来以及 LLVM 和 MLIR 的发展。
（本文由 OneFlow 社区翻译，译文转载请联系 OneFlow 获得授权。原文：1. https://www.modular.com/blog/increasing-development-velocity-of-giant-ai-models；2.https://www.modular.com/blog/increasing-development-velocity-of-giant-ai-models-part-2）
其他人都在看
欢迎下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-Inc/oneflow/
不至于成为世界上最快的深度学习框架。 2022-03-23 加入
★ OneFlow深度学习框架：github.com/Oneflow-Inc/oneflow ★ OF云平台：oneflow.cloud
评论