模型训练

如何提高大型 AI 模型的开发速度

2024 年 11 月 1 日

机器学习模型越来越大——有些人甚至说，非常 Suuuper 庞大。世界上最先进的技术公司一直在进行军备竞赛，看谁能训练最大的模型（几个巨头有目共睹的在推陈出新），而其他专注于生产系统的公司则将现有模型进行了大规模扩展，取得了巨大的效果。在所有因为大模型产生的兴奋和不安中，更大的模型给现有的 AI 基础设施和开发人员工作流程带来的无数实际挑战，这一点倒是没什么人讨论。

开发大型模型的众多挑战之一，是使用无法处理具有巨大权重的模型的工具。这是一种痛苦经历，毕竟这些模型可能高达 100+ GB。一个不言而喻的事实是，AI 部署工程师通常需要等待几分钟甚至几小时才能让工具处理大型模型。这对生产力来说并不好，也不是我们希望 AI 专家花时间的方式，而且会降低快速迭代的能力。

我们认识到开发人员的生产力是训练和部署模型成本的重要组成部分。我们不断优化我们的工具链，以改善我们早期客户和内部开发人员的生活。这篇文章讨论了在编译过程中管理数千兆字节数据的技术挑战，以及我们为解决这些问题而对基础设施（和 MLIR 编译器框架）所做的更改。

使用 AI 模型具有独特的挑战性

如果您不熟悉机器学习中的图形转换、优化和编译器，则无需害怕。这些技术用于提高 AI 模型的性能和可移植性，或使其能够部署到某些目标硬件。有高级“编译器”，如 TensorFlow Lite Converter，它将 TensorFlow SavedModel 转换和优化为高度优化的程序格式（即 FlatBuffer），以便在边缘设备上执行。还有领域特定的编译器，如 XLA 和 TorchScript JIT 编译器。它们创建或使用 AI 模型的中间表示（例如“图”）并将其编译为另一种格式 - 例如机器代码或领域特定的运行时表示（例如 CUDA 图）。

编译 AI 图实际上与传统的编译问题完全不同。AI 图包含两件事：1）图拓扑（各层如何互连）和 2）模型权重（与特定层相关的参数）。就大小而言，图拓扑为千字节数量级，而权重为兆字节和千兆字节数量级。例如，看看 Meta 发布的一些更大的模型。Open Pre-trained Transformers 有 30B、66B 甚至 175B+ 个参数，相当于 100+ 千兆字节的权重。还有更大的模型，如 Gopher 或 Megatron。

训练计算优化大型语言模型 (Deepmind)

AI 生态系统中的现有工具无法很好地处理大型模型。例如，protobuf 具有固有的 2 GB 限制，这给以此序列化格式支持的模型格式带来了问题。在最新版本的 TensorRT 中，“基于 Transformer 的网络（例如 BERT 和 GPT）在编译期间可能会消耗高达模型大小 10 倍的 CPU 内存”，这使其难以用于大型模型。在 ONNX 中，用户必须将模型权重拆分到多个文件中才能支持大型模型。所有这些都给 AI 开发工作流程带来了不必要的复杂性，可能会失去模型的“单一事实来源”，并且通常使分发模型变得更加困难。

更糟糕的是，这些权重的重量可能会导致您添加变通方法，从而使您的整体 AI 开发工作流程变得复杂。例如，我们构建了一种缓存临时文件的机制，因为某些编译器阶段需要 2 分钟以上的时间，这使我们的开发人员无法进行交互式工作。

与其他解决方法一样，我们意识到这种缓存是一种“胶带解决方案”：它不是 100% 可靠的，并且在缓存丢失时没有帮助。由于我们非常关心开发人员的工作效率，因此我们决定解决问题的核心。

Modular 编译堆栈中的 MLIR

Modular 堆栈利用 MLIR 编译器基础架构来表示和转换 AI 模型，包括 AI 运算符图（用于多个框架）、中级运行时原语和低级机器代码生成。我们的团队拥有许多 MLIR 的基础架构师，他们深入参与了向世界发布 MLIR，并且我们今天继续积极维护核心 MLIR 的大部分内容。

多级中间表示 (MLIR)

MLIR 是 LLVM 编译器基础架构项目的一个子项目，该项目提供了一个用于构建特定领域编译器的现代工具包。它提供了一组核心构建块，这些构建块对于建模、分析和转换各种计算领域（包括硬件设计、量子、人工智能等）必不可少。

MLIR 使我们能够构建一个涵盖整个堆栈的单一内聚系统，该系统比传统堆栈更强大、分层、可扩展且更易于维护。使用统一的基础架构使我们的改进能够轻松地在我们的工具堆栈之间传输，并提高整个开发工作流程的可组合性和模块化程度。

Modular 并不是唯一一个利用 MLIR 的系统——许多其他系统都使用 MLIR 进行表示和转换，包括 TensorFlow、XLA、PyTorch、ONNX 等。随着这个生态系统的不断发展，我们都可以庆祝 MLIR 的好处，但我们也必须继续投资于它的发展。

MLIR 是个好东西，但其管理权重的方法却不尽如人意！

MLIR 的基本构建块之一是属性，您可以将其视为“唯一”（又称记忆化或实习化）的形式常量数据。属性是用户可扩展的，这意味着它们可能根据用例采用各种形式。属性用于诸如常量表达式值（例如“5”、“10.0”等）、字符串文字、枚举器（例如“小于”、“大于”、“等于”等）、数据数组……等等。大多数基于 MLIR 的 AI 工具使用属性来保存 AI 模型的权重。

然而，这是一个问题：模型权重可能非常大，MLIR 以与四字节张量相同的方式存储两千兆字节的权重张量——在包含唯一元素数组的属性中。鉴于我们刚刚如此紧密地使用了“unique’d”和“千兆字节”这两个词，这会产生一个明显的问题！

挑战就在这里：当 MLIR 中某些东西是 unique’d 时，它会被分配、散列并存储在“MLIRContext”中。这些对象的生命周期与 MLIRContext 相关联，并且它们在上下文被销毁之前不会被销毁。这对于较小的值非常有用，因为我们可以传递它们并通过指针比较 unique’d 对象、共享属性分配（非常常见）等等。

对于巨大的权重张量，这些好处变成了负担：我们不想重新分配、复制或唯一化它们。我们也不希望它们永远存在：当计算不再引用它们时，释放大权重很重要。例如，当我们运行量化模型的工具时，它需要转换运算符图并生成新的权重 — 并且最终可能会得到该数据的多个副本，这些副本在编译过程期间都有效。

ML 工具的另一个问题是 MLIR 如何序列化到文件系统。当我们开始时，MLIR 没有二进制序列化格式 - 只有文本格式。这对于大权重来说是一个问题，因为每个字节的二进制数据最终都以十六进制形式发出 - 占用的空间是其编码数据的 2 倍。这意味着我们不仅要花很长时间来创建十六进制（对于一个相当大的多 GB 的模型，大约需要 20 秒），而且我们的中间文件是它们应该大小的两倍 - 已经很大的数字的 2 倍！

不仅仅是开发人员生产力的影响

这种善意的设计机制甚至可能削弱最好的编译器。最明显的挑战是它增加了编译、检查和转换模型所需的时间。如果您曾经使用过“我的代码正在编译”这个借口，您就会意识到这会带来多大的麻烦。在这里，我们迫使处理器不断分配、复制和散列多 GB 的数据。

XKCD - 编译

比编译时更大的问题是内存使用会影响 Modular 堆栈中更大规模的架构功能。例如，由于我们的编译器和技术堆栈本身是高度并行的，并且利用了在线搜索等高级功能，因此内存使用直接影响我们可以并行完成的工作量。这对于获得最高质量的结果非常重要。

我们的核心理念是构建用户会爱上的工具。我们意识到，如果高级功能难以使用、影响生产力或有重大警告（例如，它们并非在所有情况下都有效），它们就不会被使用。我们很高兴能够用大权重修复这些基础问题，这使我们能够从用户的生活和工作流程中减少复杂性。

以正确的方式修复此问题：MLIR 的核心补充

Modular 团队是 MLIR 的杰出利益相关者和维护者，我们的文化很大一部分是“正确构建” - 这适用于我们参与的每个项目。在我们为 MLIR 的发展做出贡献并推动其发展的同时，我们有既得利益来确保这些步骤适合整个项目，并且我们与整个 MLIR 社区合作以就我们的方法达成共识。

我们退后一步来了解使用大型模型工具解决此问题所需的条件，并列出以下内容：

仅在必要时分配内存：我们知道从磁盘内存映射大数据（如权重）比将数据复制到 malloc 的块中更有效率。‍
没有哈希或唯一性：我们不要检查 2 GB 数据块的相等性；权重应该通过名称来标识，而不是通过内容来隐式地唯一化。‍
启用内联变异：如果数据只有一个用户，我们应该能够量化、转换和操作数据，而不是先复制它。‍
启用释放：我们正在处理的数据非常庞大，当对数据的最后一个引用被销毁时，我们需要释放它。‍
快速序列化：无论是 JITing、搜索优化参数，还是只是在本地迭代，我们都会出于多种原因缓存 IR，而且它应该很快。

这些见解都不是特别新颖的，但这些是传统编译器（例如典型的 CPU 编程语言）不会遇到的要求。

作为我们设计工作的成果，我们为 MLIR 添加了两个扩展，它们更好地支持我们的用例，并且适用于整个社区的许多其他用例。我们正在积极地将这些贡献回上游 MLIR 社区。

立即咨询 JOTO

JOTO 提供覆盖企业智能体规划与搭建、AI 平台私有化部署、RAG 知识工程、AI 安全治理、FDE 驻场共创及持续运营优化的全周期 AI 落地服务，帮助企业把验证中的 AI 能力转化为安全、可控、可持续迭代的生产力。联系 JOTO 获取 AI 落地咨询

想把这些做法用到你的业务里?

留下你的场景和痛点,我们帮你判断从哪一步开始。

联系我们

如何提高大型 AI 模型的开发速度

立即咨询 JOTO

想把这些做法用到你的业务里?

开启企业级 AI 落地

填写需求单