大模型狂欢背后：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/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）

其他人都在看

点击“阅读原文”，欢迎Star、试用OneFlow最新版本。