——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之四
现在AI编程很强了。
你可以看到Claude Code、Cursor、Copilot、Windsurf、Trae等琳琅满目的AI编程工具,看到网上的达人用它们实现五花八门的功能。你可以看到AI在各种人机编程竞赛中完爆人类,从速度到质量都无可挑剔。你可以看到各大厂商隔三差五地更新编程LLM,拿出令人心服口服的数据,以及令人望洋兴叹的价格。
现在AI编程很强了。所以有人会觉得用AI实现什么都不是问题,只要Token足够。想必让AI做一个深度学习框架也不难吧?
但是,事实并非如此。
Tech-Renaissance就是综合利用AI编程技术开发出来的自研深度学习框架。但是这整个过程,可能远比你想象的要困难,也远比我最初想象的要困难。下面就从我的切身体会来讲讲用AI做深度学习框架到底有哪些难点。
首先最基本的一点,你需要知道上下文长度(Context)的概念。LLM既有输入上下文长度的限制,也有输出上下文长度的限制。输出上下文长度通常远小于输入上下文长度。截至写这篇文章的时候,多数厂商提供的最大的LLM也只支持1M的上下文长度,而实际编程时可用的窗口往往还要更小。深度学习框架属于中大型项目,光是源码的文本量就极其庞大——PyTorch的项目总大小约200MB,其中源码大约占150MB;TensorFlow的项目总大小大约400MB,其中源码大约占220MB;即使是超轻量级的Tech-Renaissance,它的V4.20.692版的源码也达到了5MB。尽管LLM通常会只选取部分相关的片段来读取,但它无法保证阅读完这些片段就能解决问题。debug的时候,bug可能存在于代码的任何一个角落;优化的时候,它需要分析到底哪里有优化的空间;更新模块的时候,可能牵一发而动全身——深度学习框架的模块之间耦合极强,改动一个算子,可能同时牵涉API签名、算子注册、调度分发逻辑和反向传播规则;框架内部还存在大量”不成文的约定”,比如”临时缓冲区必须由这个模块来统一分配”、”这个张量必须放在这个分组”、“末尾必须预留16个字节”,这些隐式契约只读局部代码是看不见的,AI一旦违反,错误往往要到运行时才暴露。莱曼定律指出”系统越大,复杂度就越高,维护它的成本就越高”,这对于AI也是适用的。LLM可以分段读取信息,但快要超过Context上限的时候它会自动进行压缩(Compact)。压缩必定会丢失信息。对于代码这种非常需要精确的文本,一旦压缩,LLM极有可能就忘记它原本的样子。举例来说,我会告诉LLM:”本项目编译前必须执行XXX命令来配置环境”——这句话对绝大多数LLM来说一旦压缩就会遗忘。上下文长度对于人类团队来说,反而是个相对比较好解决的问题——人类记忆容量远超目前的LLM上下文长度,而且人类的记忆与协作机制、项目管理经验通常可以高效地锁定需要修改的代码。这对LLM来说,当然也有办法解决,比如RAG、代码索引、工具调用、自动测试、分层提示词和规范文档等,只不过当前解决得还不够好。上下文长度问题在深度学习框架开发的场景下会被放大——因为AI更擅长局部的补全,而框架要求的是全局的协调与一致。当你的并行化、流水线、双缓冲不能很好地被AI理解,它最终给你的实现就不会是你希望的样子。
第二点,LLM的知识是有限的。一来LLM本身的参数量有限,二来LLM的知识存在一个时效性问题。LLM的知识截止到它完成训练的那一刻,准确地说是截至它的训练语料被确定的那一刻。第三方库如果更新了最新的API,它是不知道的;最新的算法、架构、工具、设计理念,它是不知道的。虽然可以通过RAG技术来弥补,但它不一定能找到、也不一定能正确地理解和应用这些信息。有人统计过人类程序员对LLM说的话当中频率最高的内容,其中一个就是——”要使用已有的API啊!”更麻烦的是,深度学习框架大量依赖底层库和硬件生态:CUDA、cuDNN、cuBLAS、NCCL、Triton、oneDNN、Eigen、XNNPACK、LLVM、pybind11、构建工具链、不同平台的驱动版本……这些东西不仅更新快,而且很多关键知识并不集中出现在一本教科书里。我在开发的时候碰到的一个最基本的问题,就是我发现我的编程AI只会用cuDNN Legacy API,而不会用cuDNN Frontend API。这样它每实现一个功能就必须去查第三方库的示例代码或者官方网站、社区资料,查完之后再试试这样写对不对、那样写对不对,反复的查找、编程、调试就花费巨量的时间。还有一类知识的缺口更隐蔽:对于GPU极致优化所依赖的那些微架构细节——比如访存合并、通信分桶、引擎筛选等等——高质量的公开资料本来就稀缺,很多经验散落在底层库的源码和工程师的口口相传里,训练语料里严重不足。所以AI通常能生成”可以运行”的模块,但很难给你逼近硬件极限的性能。
第三点,LLM不一定能得到最优解。咱们具体讲debug和优化这两方面。理想情况下你可能认为,把代码扔给LLM,然后让它”给我改好”,它就会把它改好了,修复所有bug,并且优化到极致,因为你听说过AI的编程能力超过人类了。但事实上,LLM经常调试一个bug要一整天,因为它很可能无法理解bug的产生机制。这里有个深度学习框架特有的麻烦:它的很多bug不是”崩溃”这种明显的错误,而是编译通过、程序能跑,但数值悄悄错了。LLM会进行耗时极长的调试,然后告诉你,它也没办法,建议你自己查查什么什么地方,或换个方法实现。在框架开发早期,开发数据加载和预处理模块的时候,就在Linux服务器上碰到死锁。这个bug是一个架构级的问题,并不是某一两行代码引入的,当时调用了8个不同的LLM,花了三天三夜才把问题解决掉,因为绝大多数LLM在检查时都陷入了盲区,始终找不准问题、更给不出正确的方案。至于优化,那是一个更大、更没有标准解法的问题。LLM一来未必能找到有优化空间的点,二来未必能想出优化的方法,三来它并不知道优化到什么程度才算足够好,四来它无法确保优化不引入bug。性能优化本质上是一个”假设—测量—验证”的闭环,如今的AI Agent虽然已经可以跑profiler、看数据,但要从跨层级的线索中把因果链理清楚——性能瓶颈到底在kernel launch的开销、访存带宽,还是内存分配器的碎片化——这种系统性的性能归因,它目前还明显弱于资深的系统工程师。更重要的是,对于大型系统,优化时的搜索空间规模会随代码量呈几何级增长,尝试的次数再乘以每次编译运行和benchmark所需的时间,那将是天文数字。除了前面所说到的一切,还有一点应该很多人能想到,那就是LLM有幻觉(Hallucination)。LLM完全有可能在开发过程中未经核实或询问就自信地生成大量看似合理、却与事实不符的代码,也就是“想当然”。这一点非常致命,尤其是在开发初期,很多相关部件尚未实现、甚至有关规则尚未完善的时候,LLM完全有可能自作主张地替你实现,而且它不一定会在最终汇报中提及。这样以来别说是最优解,你得到的是stub、简化实现、或根本牛头不对马嘴的东西,更要命的是,它能运行。
第四点,关键决策方面,LLM通常倾向于常识而非创新。LLM的训练本质上讲可以理解为从大量的素材中发现规律。事实上人工智能领域举足轻重的贝叶斯流派的观点就是”把一切知识都表示为概率”。一个信息在训练中出现的频率越高,那么推理时这个信息激活它相应神经元的概率通常也就越高。而且,对于“编程实现一个功能”这样的任务来说,能编译运行成功才是最重要的,LLM会采用更常见、更稳健的方法来实现——结果就是,它能实现出来,但那是次优解。我举个简单的例子:算子里如果需要一个缓冲区,LLM通常会怎么做?它会在算子里malloc。算子是热路径,不停地malloc和free会造成巨大的延时。如果你不制定足够的规则、给它足够的提示,LLM可能会无视这个问题,或者即使发现了问题也认为”这是迫不得已的妥协”。如果没有足够的人工干预,你对它说”帮我写一个深度学习框架”,只要有足够的Token和足够的编程、调试时间,多半应该也是能做出来的,只不过性能会惨不忍睹,而且LLM并不知道到底是做错了什么才会让性能这么差。即使是最最理想的情况,LLM掌握了大量的方法、细节,对深度学习框架的设计了如指掌,它可能顶多设计出第二个PyTorch或第二个TensorFlow,因为据它所知这就是世界上最好的框架,它会觉得做到跟它们一样就是最好的。我并没有否认AI的创新能力和科研能力,我的意思是,如果你不在提示词上多下功夫,它通常就是给你平庸的答案。
最后一点,可能跟LLM的本质无关,而是跟这项任务有关——那就是,”做一个深度学习框架”跟”做一个软件”不是同一层次的概念,而“做一个拥有顶级性能的深度学习框架”跟”做一个深度学习框架”也不是同一层次的概念。为什么这么说呢?当你做一个普通软件的时候——比如一个网页、一个小游戏、一个小型办公软件,那都是有大量模板的。倒不如说,当下很多编程LLM的核心演示就是这些接地气的简单功能。但深度学习框架本质上是”运行时 + 编译器 + 数值系统”的混合体,它涉及的技术栈远比一般的软件要庞大:它本身要支持的算子就复杂多样,需要支持不同的数据集、不同的预处理、不同的训练算法、不同的数据类型、不同的操作系统、不同的硬件平台;再往深了走,还有自动微分、计算图优化、算子融合这些接近编译器研究级别的问题。这些内容对于基本运行来说可能缺一不可,但深度学习框架并没有一个简单的通用模板。再然后,深度学习框架如果光是要能正常运作并得到数学正确的结果,那倒还好说,因为公式大家都知道、操作系统会帮你管理内存、第三方库会帮你实现计算,网上又还有各路高手的经验之谈帮你避坑,只要按部就班地实现,结果与主流深度学习框架进行对齐即可;但你如果要性能很好、甚至超越主流深度学习框架,那就非常困难了——因为你在试图解决一个世界级的难题,而LLM连框架的百分之一的代码都读不完,更别提要做到突破极限的优化。现有的主流框架哪里还能优化?更快的框架是什么架构?算子实现上还能怎么创新?你指望它能给你答案,它也在指望你能给它答案。
说了那么多,我想表达的一个观点就是:AI很强,但目前还不够好。你要用它实现普通的功能,那很容易;但你要解决复杂系统的复杂问题,并且要把性能做到极致,那就有很大的挑战。你如果指望对AI吼一句”给我做出一个最好的深度学习框架来”它就老老实实地给你做出来,那基本是痴人说梦。不但你做不出,PyTorch的开发团队也做不出,因为这些年也不见得他们用AI把框架优化得有多完美。
但是,除了对当前AI的局限性的批评和抱怨之外,我也得出了两个建设性的经验。金无足赤,人无完人,AI无完AI,但你可以与AI共同进化。我的两个经验是:
第一,人不能比AI更糊涂。你可以不知道API,可以不知道语法,可以不知道算法,但你必须清楚地知道你在干什么、你想要得到什么。你必须知道你在做的这个东西在整个系统中发挥什么样的作用,它怎么跟其他部件交互,它要优化到什么程度才不成为瓶颈。简单来说,就是要有架构思维。AI可以只对眼前的模块负责,但那个维护全局一致性、把控核心权衡的人,必须是你。如果你比AI糊涂,那就会被AI糊弄,或者出了bug也不知道怎么回事;如果你比AI糊涂,那你可能无法掌握开发的真实情况,很多地方可能只是stub,但你以为已经实现了;如果你比AI糊涂,性能就会丢失在你意想不到的地方,到最后你只能接受”我没法做得更好”的事实;如果你比AI糊涂,项目做大以后就会变成一团乱麻,你东修一块西修一块,始终无法让它完美运转。在我的实践里,最有效的方式不是把AI当成“全自动工程师”,而是把它当成一个非常勤奋、知识面很广、动手能力很强、但需要明确任务边界和验收标准的协作者。你要给它架构文档,给它模块约束,给它代码规范,给它测试目标,给它性能指标。你不能只说“优化一下”,你要说清楚优化哪个路径、指标是什么、不能破坏哪些语义、用什么benchmark验证。
第二,人必须比AI更顽强。AI是会放弃的,AI是会妥协的,AI是会投降的。它把常用的方法都试过一遍之后就会不停地劝你”这样就好”、”接受现实”。但你要达到你的目标,就必须不停地push它,让它做得更好。我在优化数据加载器的时候和优化融合算子的时候,很快就实现了基本功能,但性能的优化却花了很长时间。直观地来说,大概就是实现只需一天,但优化却要一个月。AI换了一种又一种方法,每次尝试失败后都给我两个选项:(1)接受现实。(2)再试试别的方法。——我永远选第二个。
