大模型狂潮背后：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算子图（用于多种框架）、中级运行时原语和低级机器码生成。</