(1) 一个人用AI如何写出比PyTorch更快的自研深度学习框架?

——系列文章之一

深度学习框架,是人工智能技术的基石。

小到入门者必学的 MLP,大到千亿参数的LLM,都是在某种深度学习框架上训练和推理的。深度学习框架是连接算法、开发者与底层硬件的关键基础设施:它负责高效实现各种基础算子,管理张量、显存、数据流、计算图和通信流程,同时向开发者提供相对友好的编程接口。

换句话说,深度学习框架是人工智能系统的“底座”。

它把运行不同模型所需的共同基础部件统一实现,免去了开发者重复造轮子的工作。你可以用较高层次的抽象来描述模型,比如”一个全连接层,输出神经元为512个”或”一个ReLU激活函数层”,而不需要自己去写矩阵乘法和max(0,x)。

很多可能不属于你研究兴趣范围的东西,比如数据集的预处理、GPU间的通信、算子的底层实现,深度学习框架会帮你包办,好让你把精力花在关键的方向上。

如果你是初学者,可以简单地把深度学习框架理解成积木。你选择你想要的部件,稍加调整,就能搭建出你要的模型,实现你要的功能。积木块之间的衔接,是深度学习框架为你做好的。

那么,今天的深度学习框架格局是怎样的?

如果把时间拨回十年前,那还是一个群雄逐鹿的年代:Theano、Caffe、Torch(Lua版)、MXNet、CNTK、TensorFlow……一大批框架同台竞技,各有各的拥趸,谁也说不准最终花落谁家。但十年过去,格局已经尘埃落定——如今真正占据主流的,基本只剩下两家:PyTorchTensorFlow,而PyTorch在学术界和LLM领域的统治力尤其明显。当年那些名噪一时的框架,大多已停止维护,或退居到极小众的角落。

这个高度收敛的格局,容易给人一种错觉:深度学习框架这件事,已经被巨头们做到头了,没有普通人插手的余地了。剩下的,无非是跟着主流框架调调API、训训模型而已。

真的是这样吗?

如果你刚要入门人工智能领域,想接触深度学习框架,周围的人一般会推荐你学PyTorch。因为PyTorch易于上手、功能强大、生态广泛,是当今最热门的主流深度学习框架。PyTorch 最初是由 Meta(当时的 Facebook)的 AI 研究团队开发的,其首个公开版本于 2017 年 1 月发布。当时TensorFlow和Theano是主流,但PyTorch凭借简明易懂的API、Pythonic的风格、易于调试的动态图特性,成功在学术界站稳脚跟,并迅速拥有广大的用户群体。我本人算是PyTorch的最早的一批用户。

从当年ILSVRC竞赛的经典模型,到如今各种神通广大的LLM,PyTorch都能提供良好的支持。PyTorch在LLM开发者当中备受青睐,是多数LLM的基础训练框架,是具有绝对统治力的主流,即使说PyTorch是当今人工智能世界的核心驱动也不为过。

然而本文并不是用来给PyTorch献上溢美之词的,也不是教你入门PyTorch的。

十年过去了,PyTorch进化成了一个足够完美的深度学习框架了吗?凭借先进的AI编程技术,PyTorch的开发团队把他们的框架优化得足够快、足够好、足够稳健了吗?答案是否定的。

截至2026年7月,你依然可以从最新版本的PyTorch看到各种显而易见却悬而未决的问题,包括但不限于:DDP训练中的权重同步一致性问题;模型反序列化的安全隐患;苹果MPS后端的”静默失效”等等。

但问题最集中的,还是torch.compile。首先,这个被寄予厚望的功能在 Python 3.14 上迟迟无法支持,你必须退回到 Python 3.13 或更旧的版本。其次,它的编译结果可能存在静默的正确性问题——代码不报错、程序照跑,但输出悄悄错了,不少人反映过梯度异常。而一旦把 torch.compile 与 CUDA Graph 搭配使用,还可能直接报错,甚至出现性能不升反降的情况。这并非某个最新版本的偶发问题,而是自 PyTorch 2.x 引入 torch.compile 以来就长期伴随的老毛病。事实上,即便在 torch.compile 看似正常工作的场景里,你也常常能看到它的性能反而输给 eager 模式——这本身就是不合理的。PyTorch 团队既没能保证 torch.compile 拥有与 eager 数学一致的行为,有时甚至连它的正常运行和应有性能都难以保障。

对于PyTorch这样一个极度热门的框架、对于torch.compile这种非常重要、使用频率很高的基础功能,问题长期悬而未决,几乎是不可思议的。

更重要的是,我觉得PyTorch很慢

LLM的庞大容量和数以月计的训练时间会让人觉得”慢”是理所当然的;背靠Meta的强大技术团队、十年的技术积淀、庞大的用户群体和生态,也会让人对PyTorch的性能深信不疑。但这就如同”名牌的衣服质量一定好”一样,是一种未经检验的迷信。

支撑了无数LLM训练的PyTorch,真的已经把性能压榨到极限了吗?

未必。

你可知道,一个人在业余时间用AI写出的自研深度学习框架,都有可能在训练吞吐量上胜过PyTorch?

你可知道,你很多年前就用PyTorch跑过的VGG16,其实完全可以更快——不是快一点点,而是在同等硬件、同等超参数、公平对比的前提下,每个训练epoch的耗时被显著压缩?

今天已不再是十年前那个深度学习框架风起云涌的时代。但今天是AI编程技术、AI Agent快速普及、深刻改变人们生产生活方式的时代。

你手中握着的AI工具,其实能让你创造不可思议的奇迹,只是你可能想都没想过。如果你只把AI当成一个新型的搜索引擎,或者一个长期在线的聊天伙伴,那可真是暴殄天物。

本系列文章,就介绍如何用AI写出比PyTorch更快的自研深度学习框架。

你将会看到,做一个深度学习框架需要哪些基本部件,一个模型完整训练的流程是怎样跑通的,一个自研深度学习框架从架构设计、技术选型,到AI开发,再到性能测试和开源发布的完整项目生命周期。

你将可以亲手运行这个完全开源的自研深度学习框架。至于它是不是真的更快、到底能快多少,你完全可以自己去见证——所有数据都经过7个不同GPU平台的交叉验证,代码全部开源,任何人都可以复现。

你将会看到设计中的各种思考和权衡,看到如何打通一个庞大的技术栈,看到如何借助AI去攻克一个本来看似不可能的技术工程。

如果你是一个想要研究AI加速硬件、学习AI算法的初学者,这个系列应该能让你了解AI编程、深度学习框架、深度学习算法,以及GPU这类AI加速硬件;如果你已经是AI领域的老手,这个系列或许也能给你提供一个新的视角,助你打开一些新的思路,让你看到一些新的可能。

系列文章持续更新中,欢迎关注。

系列文章列表:

标题
(1) 一个人用AI如何写出比PyTorch更快的自研深度学习框架?
(2) 深度学习框架是什么?
(3) AI编程技术及原则
(4) AI做深度学习框架有哪些难点?
(5) 静态图 Vs 动态图:深度学习框架的两种世界观
(6) Tech-Renaissance全景图:七大模块与一条样本的完整旅程
(7) 张量、NHWC布局与256字节对齐:框架里的”砖块”
(8) 数据加载管线:一套抽象,兼容 MNIST、CIFAR 与 ImageNet
(9) DTS数据格式:为高速训练定制的存储方案
(10) 多线程预处理、NUMA 感知与异步双缓冲:让 GPU 不再等待 CPU
(11) FusedNormalization:把数据增强的最后一步融进一次内存遍历
(12) BluePrint DSL:用声明式语法定义神经网络
(13) ArchPlan与编译管线:从模型定义到可执行图
(14) ComputationGraph与GraphAtlas:一份图纸,多处复用
(15) MemoryPlan与显存分区:框架的灵魂设计
(16) DTensor:一张图纸,多卡共享的分布式张量抽象
(17) 多流并发架构:计算流、传输流、更新流的协同
(18) CUDA Graph全捕获:把训练循环变成一次GPU提交
(19) AMP自动混合精度训练:FP16的速度与精度平衡
(20) CBR融合算子:Conv + BatchNorm + ReLU的三位一体
(21) 融合优化器:SGD、AdamW与LARS的整区批量更新
(22) 损失函数、学习率调度与参数初始化:训练算法配置层
(23) Philox可复现随机数与确定性训练
(24) NCCL通信与分布式数据并行
(25) 集成测试与性能密码:一个人做的框架,凭什么能比PyTorch快?
(26) 附:Tech-Renaissance V4.20 的安装运行指南
(27) 系列后记:”技术”如何”觉醒”

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

ICP备案号:京ICP备2025133467号-1