一些AI系统的技术随思

•

最近除了日常工作以外，大量时间花在技术细节的catch-up上。一个感受是，越多接触到底层技术，越会感觉到自己当前技术积累的薄弱。比较早以前，周围同事朋友给自己的建议是，应该花更少精力在技术上，花更多精力在业务梳理，外部协同和行业趋势的跟进上。也曾经在360度调研里，收到的评价是自己足够hands-on。现在到了新的环境，接触到了更多之前没有接触过的技术信息，认识了之前没有打过交道的不同类型的新同事，会觉得自己hands-on的程度是远远不够的。而在不同行业里，对技术细节理解掌握深度的区别，对于业务机会的判断，外部协同的把握，以及行业趋势跟进，也会带来不同的影响。

前不久和一位市场朋友聊天，提到公司的一位Senior VP，有一次和一位国内大厂类似CTO角色的同行交流，做了一些high-level的技术分享。会后收到的反馈是，对方觉得这位Senior VP的分享过于技术，导致那位CTO并没有完全消化。而那位市场朋友觉得他自己作为非技术出身的背景，都觉得那些分享并没有过于技术的味道。

我想，如果国内的技术行业想从粗放的业务扩张的风格，转变为精耕细作的技术创造突破性的风格，也许会需要更多人能够在职业生涯里的较长周期内（我曾经共事过的一位老大哥在这方面给我留下了深刻的映象，他当时已经承担了非常重要的项目推进以及资源决策的职能，但还保持着对技术如饥似渴的追求，记得有好几次我去找他讨论问题，都看到在他的办公桌上堆满了他正在跟进的一些领域新进展的paper，以及我们曾经多次在会议室里和团队同学讨论代码的优化细节，他也成为了技术追求领域我所认同的一个榜样），保持在技术细节层面的足够hands-on，才可能加速这种转变吧。当然不同行业，不同阶段，对技术细节的要求程度也存在差异，这跟行业分工定位有关。以及每一个技术代际，每一个具体的个体，因为时代和个人际遇的差异，所适合承担的那个“定位”也有所不同，结合主客观条件，尽量去最大化自己可能为所在行业带来的帮助就好了。

上面发散性的思考结束，接下来是一些最近会有思考，但并没有获得答案，也不确定什么时候能够获得答案的技术随思了，聊作整理，也欢迎同行讨论碰撞。

关于Triton的随思

因为一些原因，关注了一下OpenAI的Triton的项目。第一眼看到这个项目，会在想这会不会可能颠覆CUDA的生态，再看了看一些文档细节，结合之前的一些工作经历，对这个项目的演进趋势又有了一些不同的理解：

    • 所有的抽象本质上都是存在泄漏的。这种泄漏从最底层的ISA指令集就开始了。Triton作为架设在PTX之上的抽象层，也存在类似的问题。包括CUDA这一层抽象，也同样存在类似问题。
    • 决定一个新增抽象是否可能被广泛adopt，甚至取代已存在的一个抽象层次的关键点，在于这个新增抽象带来的新增价值要足够大。落实到朴素的用户产品设计的原则，那就是 用户价值=新体验 – 旧体验 – 切换成本。
    • Triton这一层抽象，和CUDA编程模型，以及TVM/XLA/MLIR这种深度学习编译器（目前我理解Triton的核心 定位还是解决深度学习场景的计算需求，而非其他）是存在直接的竞争的。所以决定Triton未来adoption ratio的关键是，相较于这些现存事物，其能带来的新增价值。
      • Triton最早的工作是在C语言层面提供了类似DSL的API。后来主要的推广以Python为主，我想是因为深度学习这个领域，Python的使用实在是太pervasive了。某种程度上，我觉得PyTorch团队后来选择在Python的bytecode level加入TorchDynamo的功能，也是出于类似的考虑。
      • Triton的抽象粒度是在Python层面提供了类似类似CUDA的并行编程模型，同时做了一定的简化，将部分需要程序员手工处理的优化在Triton内部优化pass里完成。不过我个人的认识，至少针对NV GPU，这种抽象粒度相较于CUDA本体或TVM之类工作提供的差异化价值不够显著。反而是类似CUTLASS这种提供building block的方式提供了更多的add-on value。
      • AI硬件架构还没有收敛，仍然在不断演化，特别是访存系统的演进还比较激进，这对于上层抽象带来的压力是巨大的，包括 Triton，TVM，甚至包括CUDA本身。
      • Triton对于隔离不同硬件的差异化，从而对于改善程序的可移植性是否会带来比较大帮助呢？我不确定。从实现手段来说，Triton是直接codegen到LLVM IR上，确实可能通过复用LLVM现有的backend对接的生态来简化同一份Triton program跑在不同硬件后端上的负担。关键还是在于Triton是想提供一套对多套硬件都能通用的抽象，还是想针对某款特定的架构相近的硬件做好这层抽象。我个人的认识，TVM或基于MLIR的一些自动codegen的工作，目前对于Tensor类计算，也只能对NV GPU架构相近的架构做到比较好的通用性。原因也比较朴素，抽象需要有相对稳定的被抽象基础以及大量的实践样本，NV GPU从16年推出Tensor Core到现在，无论是社区自身积累，还是NV推出的一些开源项目和软件库及资料文档，都为抽象提供了比较好的基础。这也是坊间传言一些GPU start-up通过兼容NV GPU的编程模型可以相对于从头build一套硬件架构和编程模型更快完成软件库搭建的原因之一了。

分布式的随思

最近go through了Meta的OPT的工作，扫了一下里面引用的老文章，顺便又去翻了翻引用这篇老文章的Megatron-LM 2021年的文章 ，再结合之前一些对这个领域的关注和思考，有了一些新的思考

    • 现在能看到的关于分布式的工作，我个人的感觉还是集中在解决一个个E2E的问题，不管是DeepSpeed里的工作，还是Megatron-LM的工作，还是PyTorch Lightning的工作，也包括Google在XLA TPU上的一系列工作。其中Google XLA TPU上的工作，因为有了一个比较干净的表示层，我觉得具备了一定的生态属性。
    • 这些工作在具体某个模型或场景上，都还是蛮有意义的，相当于通过一个一个具体的问题把领域的认识推得更靠前了一些。
    • 但这些工作给我的一种感觉是“碎片化”，就是彼此之间的复用并不容易，工作A如果想应用在工作B里，往往需要reimplement一遍，而分布式策略的空间，随着模型（Dense or Sparse, CV or NLP, …），硬件(NVLink, NVSwitch, Multi-node NVLink, IB, Grace-Hopper…)、训练方法的变化(Vanilla optimizer, Second-order optimizer, Neighbor-aware loss optimizer)，其实呈现一个比较明显的NP-hard的性质，想在一个工作里把这些问题全解干净，看上去很吸引人，但可能会总存在一定的gap。
    • 所以另外一种解决这个问题的思路是构建一套可能复用，share的基础设施，让这个领域的从业人员都可以在上面贡献自己的想法，便于复现复用，而这套基础设施的关键模块可能包括 ：
      • 前端定义。比如如何将PyTorch的模型描述转换成一套适合分布式策略处理以及交给runtime执行的前端IR，这方面的工作，我理解不管是TorchDynamo，还是LazyTensor都是类似的考虑
      • 一些重要的可复用的基础脚手架的定义实现。比如常见的通信原语，allreduce/allgather/sendrecv/all2all/reduce/gather等等。这方面的工作，其实包括 NCCL以及PyTorch distributed里的primitive都是在提供类似的能力。
      • 一套或多套可复用的框架。比如cost function的抽象（具体 的cost function由分布式探索的人来填充)，策略的抽象（具体 的策略细节由具体 实施人员决定），runtime的抽象（同上)。将整个分布式执行的过程抽象成一套类“八股文”的形式，大家都按这套八股文来填充。因为有了相对统一的模板，复用性就变强了。
    • 如果结合技术发展历史来看，其实LLVM之于单机CPU编译器，就起到了这个作用，可能LLVM就是因为顺行当时的技术时代需要，提供了这种领域复用积累的基础能力，才会获得那么快的发展。MLIR则是Chris的另一个类似的野望了，但还需要时间检验。其实之前和TVM的同学讨论，我个人的理解，TVM的核心想法也是期望打造一套生态基础，吸引更多人加入进来，而不是自己做完所有事情，这也是当时TVM的核心 团队花了一年多时间重构底层TIR的原因（当时我其实有点不完全理解TVM团队做这个选择的考虑，现在反而更加理解其背后的考虑，也会很欣赏TVM的核心团队在当时的那个时间点能够做出这个判断，因为这种判断做好了能打开更大空间，做不好会因为失去了一些更具体 的能解决实际问题的feature的迭代机会，导致错失一些时间窗口，是一个有概率性的决策了，而不是一个确定性的工程决策）。

关于LLVM的随思

  • 最近花了点时间研究了一下modular AI的vision，顺便返回来又去读了一些相关的背景资料，包括这篇经典的老文章，很有共鸣的感觉。我对LLVM的关注比较早，大约在09年LLVM还不算非常有名（当时主要是Apple和一家叫AutoESL的EDA公司在主要支持它)就关注过，真正有实操经验是17年左右因为开始关注XLA的工作，后来因为MLIR也相应地有连续性的关注，也会经常琢磨LLVM带来的一些启发。每次琢磨，都会有一些不同的感觉。从之前把LLVM当作一个实现XLA的技术底座，到看到MLIR以后，有些朦胧地感觉LLVM模块化的设计思想可能不像看起来那么容易，有些“大巧不工”的味道，后来加入硬件公司以后接触了解到更多新硬件上基于LLVM开发device compiler(不是AI compiler)的便利性（配一些.td 文件，加一些pass，就能相对快的完成一个后端codegen模块的开发)，再到在NV之后，从不同视角接触不同类型的fusion codegen相关工作。对于这篇文章里提到的LLVM成功的复盘总结的认同感就变得更加强烈了。我现在的认识 里，LLVM 的成功是典型的工程技术和架构技术的成功（并没有非常原始创新的research idea，但是把一些从第一性原理上已经被验证的idea在工程架构上落实地非常到位），核心 在于其 “模块化”，“可复用”，“IR表示的完备性”，“支撑工具的完备”使得其具备非常强烈的自演化的“技术生态系统”的特质，这种技术特质甚至对于技术组织协同的效率都会带来很大的影响。这个特质一但完成了早期的momentum积累 之后，就会产生巨大的雪球效应。这些技术philosophy的东西，说起来似乎很快就能有认同感，但真正实际landing到生产系统里，结合一个具体 的方向场景落实到位，恐怕远没有那么简单了。稍微具体 一些，这里从实施角度来说比较难的问题有几个：
    • 针对AI系统这个领域，怎样"模块化"能够尽可能避免leaky abstraction以及减少模块化带来的系统性能开销和调用交互的开销，这需要对目标domain非常深刻的理解。更具体一些，需要对AI软硬件技术的全栈几乎都需要有深刻的理解。
    • 怎样才能做到真正的“可复用”？在我看来，“可复用”的核心点是要对一个功能模块可能使用的不同contexts有着清晰的边界isolation的定义。有些问题相对容易做到可复用，比如一个cse的pass或是dead node elimination的pass，但即便是这种可复用，要真正做到位，也需要有一个能够贯穿系统全栈的稳定的representation，得到这个representation，并不那么容易(想想倒出来碎了一地的ONNX吧)。至于更复杂的，怎样对Tensor计算的描述空间进行抽象复用，适用不同硬件，就更难了。
    • “IR表示的完备性”。比较早以前，和一个朋友讨论问题，他提到说编译器最核心的就是IR设计。在另外的一个场合，有人提到编译器设计最核心的就是IR设计和在这个IR设计之上的一系列变换。这几种说法都稍微有些过于简化的嫌疑，但从描述问题核心的角度，我觉得是合理的。也反映出来设计一个完备的IR表示远没有那么容易。因为IR设计反映的是对目标问题的抽象，抽象的价值在于它的普遍义，而复杂性也在于得到这个普遍义的获得过程。
    • "支撑工具的完备"。这个是我觉得有可能被忽略，但在生产系统里很重要，甚至可能是花费时间更多的部分。包括debug工具，测试工具(比如filecheck)，DSL工具（比如tablegen)，统一的序列化反序列化工具的支持等。因为生产系统会有大量research工作不会touch的dirty的细节，而缺失了对这些细节的有效的把控，生产系统的演化就可能看起来差了毫厘，实际上已经完全不是那么回事了（比如，一套codegen系统在真正布署到生产环境以后，JIT可能会面临在线编译的耗时过长，以及编译结果本身的执行性能抖动，AOT可能面临的binary size过大导致不得不引入一系列技术手段进行mitigate）。

AI编译技术的随思

    • 针对现行的NV GPU架构（以及未来一两代架构中）我的mental model里会把fusion分为两类，一类是计算密集算子为核心的fusion(GEMM/Conv-centric)，另一类是访存密集算子为核心 的fusion(elementwise + elementwise, elementwise + reduce, reduce + elementwise, reduce + reduce, …) 。在一些新的AI硬件公司里，也在尝试打破计算密集算子和访存密集算子的fusion边界，以期通过更大尺度的fusion获得更好的性能，这涉及到硬件架构上的更多工作。
    • Fusion不等于常规意义的AI编译。因为：
      • Fusion可能有不同的支持方法。手写，经典的AI编译(XLA/TVM的codegen)手段，AI编译codegen+手工提供building block的hybrid作法，还有比较tricky的对一个pre-built好的binary通过类似代码注入的方式来进行简易的fusion(一种泛化的LTO)。这几种手段各有其适用场景和优缺点。
      • 经典AI编译手段，自动化更强。我个人的观点，这种手段对于访存密集算子fusion已经没有本质上的技术可行性问题，主要是工程量。但对于计算密集算子codegen，目前还没有已知足够好的策略能够做到完全auto codegen。虽然有一些来自于社区的工作(MLIR和TVM社区都有 )，在Ampere架构上可能通过自动化的方式生成和SOTA library性能相近的kernel，但我觉得这是因为Ampere的编程模型已经稳定了比较久，留出了这种空间。随着Hopper,的推出，我的判断会比较快地继续拉开这种auto codegen作法和手写，以及codegen+手工building block作法的性能差距。完全自动化是一个迷人的技术方向，但是对于计算密集算子，因为调度空间过于庞大，又需要非常精细的处理，在硬件基础足够稳定之前，自动化codegen从生产落地来讲，我觉得不确定性还非常强。
      • codegen+手工提供building block的作法，我知道有若干AI硬件公司是这样来做的。相当于专家经验和自动化技术的结合。这样可以更快的迭代起来。这里要解决的一个核心问题，除了技术以外，还涉及到手工优化专家和自动化技术的配合协调。我听说过有些公司里，手工优化的同学和AI编译同学的关系理顺不是那么容易的事情。因为整个解空间那么大，哪些放在手工，哪些放在自动化，是一个此消彼长的过程，只有建立一个能够就事论事，从技术第一性原理迭代演进的文化，才可能以比较小内耗的方式推进这种迭代，否则就会有大量的内耗冲突。
      • LTO的作法，一个很大的优势是，能够确保对于GEMM/Conv这种基础kernel的性能不会引入太多variance。如果是常规的fusion&codegen作法，很可能会因为引入了fusion时新生成的逻辑，在和GEMM/Conv的代码作完融合以后，改变了shared mem/register等片上资源的消耗，于是再做一次编译会影响到register spill或occupancy，带来性能抖动。当然，LTO的作法，依赖于对device compiler的定制能力，另外对过于灵活的pattern支持起来会比较吃力。
      • 手写。我们有时候说fusion会潜意识里跟codegen划等号。其实fusion是目的，codegen是手段，也可能用手写作为手段来解决fusion的问题。对于足够高频稳定的pattern，手写是更高效的方式。
    • AI编译要解决的问题，实际上是GPU上kernel生成的问题。针对具体模型场景，这个问题其实是可能拆分成不同的子问题，每个子问题有其更适合的手段，有些适合手写，有些适合LTO，有些适合auto codegen，有些适合codegen+手写building block。No silver bullet。这种情况下，有一套能够把这些不同手段有效组织在一起的uniform的框架，就会非常有助于技术的打通，交互和共享流动了。因为此时适合用手段A解决的问题，也许随着软件，硬件技术演化，假以时日，更适合用手段B来解决了，如果底层的技术框架鼓励允许这种流动，就可以让整个系统向着更逼近系统设计理论上限的方向演化，否则就可能形成越来越多的系统legacy，导致出现软件系统的“创新者的窘境"。

最后，打个小广告，上面这些思考，都是我围绕日常工作中产生的，如果对上面这些问题感兴趣的朋友，可以关注一下这里的信息，感兴趣的朋友欢迎邮件联系 [email protected]，一起探索AI软硬全栈系统的未来技术演进路径。

来源：知乎 www.zhihu.com

作者：杨军

【知乎日报】千万用户的选择，做朋友圈里的新鲜事分享大牛。
点击下载