大模型Infra这些年，从黑铁时代到黄金时代再到白银时代
知乎原文 方佳瑞 腾讯 Llamafia 编辑推荐理由：带着历史视角看问题，会让理解更加深刻，更有沉淀。能够跨越五年尺度，解析一个领域为什么会发展到现在也体现出了作者深厚的功力。本文的从 MLSys 的视角出发，对于 LLM 研究有着很好的互补作用，也对一线的模型开发有很多启示。作者在文章末尾提到他的 research vision 是 Democratizing LLM，这一点也与 Llamafia 的使命相符合。因此我们取得了作者的授权，将文章从知乎 cross post 到 Llamafia Tech Review 越来越多人开始关注大模型，很多做工程开发的同学问我怎么入门大模型训练推理系统软件（俗称大模型Infra）。 作为一名曾经后台开发岗位同学，大家的心情我感同身受。近年来大模型崛起，每天新闻可谓繁花似锦，环球同此凉热。但需要提醒大家的是大模型Infra正在从从黄金时代进入白银时代。我21年开始写大模型训练系统，算是周期完整的亲历者，谈谈我这些年观察到大模型Infra经历的几个阶段： 2019-2021，黑铁时代 17年那篇著名的《Attention is All you Need》论文发表，以Transformers为基础零件的不同模型结构相继出现，Decoder-only的GPT-1（2018）、Encoder-only的Bert（2018），Encoder-Decoder的T5（2019）开始在NLP领域大杀四方。互联网企业中的翻译、对话、推荐等应用场景相继被Transformers占领。 顺应潮流，有人开始研究Scale Transformers模型。把模型变大在CNN时代是反直觉的，彼时大家正在通过NAS和AutoML等手段千方百计把模型变小，放到汽车、摄像头、手机里。Scaling Law的信徒主要是OpenAI和Google。在18年有用LSTM预测下一token的语言模型ELMo，OpenAI就用Transformers替换LSTM做了GPT-1，在Dota AI积累的哲学开始Scaling，随后就出了GPT-2和GPT-3。Google作为Transformers的策源地，17年就有提出了MoE架构的Transformers，20年已经能搞出用2K TPU训练的600B的大模型GShard，为此开发了基于TensorFlow的分布式训练框架MeshTensor。但和OpenAI不同，Google没有押宝Decoder-only的结构，而是大力发展Encoder-Decoder结构（也是Transformers最早被提出的结构）。DeepMind作为一个隶属Google但相对独立的英国研究机构，此时也有和GPT类似的Decoder结构模型Gopher和Chinchilla。 尽管，20年175B参数GPT-3的few-shot learining能力给业界带来了一些震撼，但国内对大模型技术路线持怀疑态度居多，拥趸寥寥，大部分算法同学对大模型并不感冒。我观察原因来源于两方面。 一方面，大多数人没有预训练的概念，当时NLP算法开发范式是搞私有数据+微调Bert，根据小模型时代的经验，应该专注在数据质量，而不是一个更大底座模型，哪怕先把模型变大两倍，也申请很多预算去训模型，ROI存疑。 另一方面，Infra没有准备好。训练一个大模型可不是一般人可以玩得起的，现在训模型是算法+工程同学相互配合，那时候没有这种兵种搭配概念，业务团队里只有算法工程师，搞GPU对他们是头疼事情，尤其是上线一个需要两张GPU才能运行的模型，简直是灾难。而隶属于中台部门的Infra团队，又不了解大模型的风向，跨部门的信息是有壁的。 国内最有动力去Scale模型是做算法刷榜的团队，但大多刷榜的模型大多无法逃脱PR+束之高阁的命运，不过这个过程锻炼了很多队伍，他们也成为后来ChatGPT爆火之后参与国内大模型训练的主力。 彼之砒霜，我之蜜糖。模型变大反而对AI Infra人是新机会。那个时代AI Infra的主旋律应用还是推荐系统，NVIDIA在押宝元宇宙作为新增长点，大模型对Infra同学也是新鲜事物。有OpenAI和Google前面开路，美国有些机构开始了将模型变大的探索性工作，主要是沿着把Encoder结构的Bert变大，不过对于训练来说Encoder和Decoder的差别不大。也正是借着这波机会，大模型训练框架Megatron-LM和DeepSpeed开始有了原型。20年微软搞了17B大Bert的Turing-NLG，其训练代码成为了DeepSpeed的原型。19年，NVIDIA搞了8.3B的Megatron-LM，没错Megatron-LM是一个大Bert的名字，Megatron-LM仓库里也是放训练模型的代码脚本，这些脚本实现了张量并行，后面逐渐发展成了最流行的训练框架。前面两者都是在PyTorch之上的搭建的，而且都是开源的。 这一时期，美国大厂训模型主要目的是试试水和秀肌肉，因而DeepSpeed和Megatron-LM最开始就是开源的，这也是一件好事。试想如果大模型出场就是巅峰，大家都像22年之后OpenAI那样藏着掖着，Infra技术扩散必然没这么快了。 用大规模GPU训模型，小模型时代的数据并行无法胜任了。一些训练优化的基本概念开始定型，比如ZeRO，张量并行，流水并行，Offloading，混合并行等。这些技术之前也都有前身，比如18年工作混合并行的FlexFlow，18年自给NAS出来的AmoebaNet训练的GPipe，ZeRO是一种Parameter Server特殊形式，等等。只不过这些技术在Transformers架构和高带宽互联网络中可以更加针对性地优化使用。 那个时候国内对大模型这种新鲜事物了解有限。袁老师的Oneflow是最早做大模型分布式训练的尝试的，SBP方式做自动化的模型并行理念还是比较超前的，不过可惜国内没有业务需求来支撑他们做PMF。华为的MindSpore也开始做自动并行的尝试。行动比较早团队的有阿里的M6，智源的GLM和华为的Pangu-alpha。 除了少数大模型Believer，国内大厂对大模型不感冒，各大厂的云/中台团队刚刚完成一轮X minutes训练ImageNet，还在思索用那么大规模的GPU去训练一个模型有什么商业价值。有一些NLP技术创业公司，技术路线还是是给不同的业务做定制的微调Bert。一些机构的投资人会打电话咨询技术人员，大模型是否是骗局。总体来说，22年之前，大模型对国内来说还是太超前，以至于共识非常薄弱，因此我称之为大模型Infra黑铁时代。 本人在黑铁时代的尾声21年初开始研究大模型训练系统的。在腾讯的WeChat AI，研究氛围比较宽松自由，而且很幸运我的+1和+2Leader都非常的技术远见，在没有落地前景的前提下支持我做这件事情。做大模型训练最大的障碍是计算资源，我当时能访问的GPU硬件主要是单机8卡V100，因此找的切入的方向是设计一个更好CPU-Offloading策略来让低端机器跑大模型，也就是PatrickStar。当时公司内有几个团队尝试联合起来申请几百张卡训模型，因为ROI不明确没有批下来。 2022-2023，黄金年代 经历了黑铁年代储备期之后，以Meta开源OPT-175B模型为标志，22年开始大模型Infra开始迎来黄金年代。 一方面，伴随NVIDIA芯片的计划迭代，一台DGX SuperPOD（2021）已经可以很轻松搞定百亿参数的大模型训练了，NVIDIA为了推这种超级计算机，也把大模型列为一种杀手应用来宣传。除了大模型还真不知道什么应用能把它填满，21年11月腾讯弄了一台SuperPod样机，拉全公司各种Infra团队一起测试两个月，本来列了一个密集测试排期表，最后发现只有我们调PatrickStar的在用机器，后面我们直接独占了全机。 另一方面，有人开始看到大模型的威力了，硅谷春江水暖鸭先知。22年5月份，Meta把OPT-175B的权重开源出来，它是为复现GPT3训练的，用的是PyTorch FairScale，虽然模型效果不敢恭维，但是真的是造福了广大大模型科研人员，做AI Infra的人也终于有一个真实模型可以做实验了。22年6月，HuggingFace也组织三十多个国家机构多国部队，开源出来Bloom-176B模型，用的DeepSpeed框架。 硅谷创业公司中，Character.ai和Anthropic.ai这时已经成立一段时间了，22年伊始，国外已经好几路人马在紧锣密鼓训练GPT-3级别的模型。印象比较深的是22年4月份，Transformers论文有两个作者离开Google成立了一个叫Adept.ai的大模型公司，用大模型帮人完成复杂任务，他们twitter的demo展示大模型能根据输入文字让Python画出一个柱状图表，我看到之后非常不可思议，那应该是我的ChatGPT时刻，现在再看类似东西已是见怪不怪了。 直到2022年11月，ChatGPT强势出圈，引爆了大模型。这一事件，显著加速了模型变大有价值共识的形成。大模型团队快速聚集了大量人力、物力、财力。大模型Infra也迎来了跨越式发展。 大模发展，训练先行。在训练系统领域，Megatron-LM和DeepSpeed在黄金年代快速迭代。不过说实话，前ChatGPT时代这两个软件质量堪忧，我感觉是处于无架构师状态，疯狂缝合各种research idea，系统bug很多，而且用户接口很不友好。但是，大模型训练框架，先发优势很重要。大家选型时不是考虑好不好用，而是看这个框架之前训出过什么模型。毕竟训练一次，小则几十万，大则几千万美金的投资，稳妥最重要。后ChatGPT时代，由NVIDIA品牌和技术背书的Megatron-LM的优势开始滚雪球般变大，到现在基本被其商业版本Nemo一统江湖。 在技术创新上，也是遍地是黄金。并行策略和算子优化等方面很多机会，比如Sequence Parallelism、Flash Attention之类simple but effective的作品都获得巨大影响力。预训练之后，还有 SFT，RLHF的需求，也训推混合调度和S-LoRA之类工作的诞生。 大模型Infra也深刻影响了上层算法的发展轨迹。比如，很长时间大家不敢增大context length，因为Attention计算中间QK^T矩阵的内存是序列的平方项开销，导致有一段时间Linear Attention，Approximate Attention用精度换长度的研究之风盛行。Memory Efficient Attention工程优化出现之后，最著名的是Flash Attention，直接把内存平方项干没了，大家又回到了老老实实用标准Attention的正规。...