——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之二
本文先讲一些基础的东西。
但凡一个人要进入人工智能领域，任何课程都会在交代简单的数学基础之后，就让他去用深度学习框架运行简单的神经网络。很可能是PyTorch，或者TensorFlow。在关于深度学习的论文或博客文章里，这样的句子也经常出现：”我们的模型基于 PyTorch/TensorFlow 实现”。这两个名字已经变得像空气一样理所当然，以至于很多人在写下第一行 import torch 或import tensorflow as tf之前，从未真正问过一个问题：深度学习框架到底是什么？它解决了什么问题？为什么别的领域不搞框架，唯独深度学习领域要框架？
本文试图捋一捋深度学习框架的这些常识。不打算介绍得太系统、太全面，但至少对引出后面介绍本框架的博客文章有所帮助、也方便初学者朋友对我们这个系列文章所要讨论的主题有一个基本的认识。
一、没有框架的世界，是什么样子
想象一下，如果没有深度学习框架，你要如何训练一个神经网络。
你得自己写 CUDA 代码实现矩阵乘法，自己管理 GPU 显存的分配和释放，自己推导每一层的反向传播公式并手写梯度计算代码，自己处理数据从硬盘到内存再到显存的搬运，自己实现多 GPU 之间的通信同步。这还只是最基本的要求。
其中最麻烦的部分，往往不是前向计算，而是反向传播。对于一个包含几十层甚至上百层的现代网络——里面混杂着卷积、注意力机制、残差连接、批归一化——人工推导每一层的梯度公式几乎是不可能完成的任务。更糟的是，梯度计算一旦写错，程序通常不会报错，只会让模型”训练不动”，或者收敛到一个诡异的结果，排查成本极高。而且，你换一种网络结构，之前手推的梯度公式全部作废，得重新来一遍。研究者的大部分精力会被”如何求导”这件事本身吃掉，而不是”这个模型该长什么样”。
深度学习框架的出现，就是为了把这些”通用但复杂”的底层工作一次性做好，让所有开发者共享，从而把研究者的注意力从”怎么一步步算”解放出来，聚焦到”我要什么”上。打个比方：盖房子的人不会自己烧砖、自己炼钢；他们用现成的建材，按照设计图纸组装。深度学习框架就是这套”建材体系”——卷积、归一化、激活函数是基础构件，梯度的自动计算、数据的加载搬运、多卡的通信同步则是隐藏在幕后的施工队。
二、框架的心脏：计算图与自动微分
几乎所有主流框架的骨架，都建立在同一个技术之上：自动微分（Automatic Differentiation，简称 AD）。这也是深度学习框架区别于 NumPy 这类普通数值计算库的核心标志。
理解自动微分，先要区分几种不同的求导方式。手动微分是人工推导公式再写代码，精确但极其繁琐，只适用于简单模型。数值微分利用导数定义 f'(x) ≈ (f(x+h)-f(x))/h 做差分近似，实现简单，但计算量大、精度差，还有截断误差。符号微分（如 Mathematica 所做的）对表达式做符号求导，能得到精确的导数表达式，但对复杂函数会出现”表达式膨胀”问题，效率低下。
自动微分走的是另一条路：它把一次前向计算拆解成一系列基本运算（加减乘除、矩阵乘法、卷积等）构成的计算图（Computational Graph）——一个有向无环图，节点是算子，边是张量数据的流动。每个算子的求导规则（比如乘法的导数、ReLU 的导数）只需要写一次，之后无论你怎么组合搭建网络，框架都能沿着计算图，利用链式法则自动、精确地算出梯度。这个过程既没有数值微分的截断误差，又不会像符号微分那样表达式膨胀，是精度与效率的完美平衡。
自动微分有两种基本模式：前向模式在计算输出的同时同步传播导数，适合”输入少、输出多”的场景；反向模式先完整做一遍前向计算并记录计算图，再从输出端反向逐层传播梯度，适合”输入多、输出少”的场景。神经网络的训练恰恰是数千万甚至数千亿个参数（输入）对应一个标量损失值（输出），所以反向模式——也就是我们常说的反向传播（Backpropagation）——是深度学习框架的标准选择。当你调用 .backward() 时，框架就是从损失节点出发，沿着计算图的边反向遍历，逐节点应用链式法则，把梯度传回每一个参数。
计算图的构建方式，直接决定了一个框架的核心范式，历史上也由此分出两条路线。静态图（define-and-run）要求先用专门 API 定义好完整的计算图，框架对整张图做编译优化（常量折叠、算子融合、内存复用规划等），然后再喂数据执行——TensorFlow 1.x、Theano 是典型代表。这种方式给了框架充分的全局优化空间，也便于部署，但代价是调试体验差：你写的是”描述图的代码”，出错时很难用普通的 print 去定位。动态图（define-by-run）则是随着代码的实际执行边跑边建图，你写一行代码，它就立刻执行、立刻拿到结果——PyTorch 从设计之初就选择了这条路。这意味着你可以用 Python 原生的 if、for、递归去构建模型结构，可以随时打断点检查中间结果，调试体验和写普通 Python 程序几乎一致。这正是 PyTorch 能在学术界迅速反超 TensorFlow 的关键原因——研究工作需要频繁试错，动态图的即时反馈价值极高。
有意思的是，这场”静态 vs 动态”之争最终没有分出绝对赢家，而是走向融合。TensorFlow 2.x 默认启用了 Eager Execution（本质就是动态图）；PyTorch 2.0 引入的 torch.compile，则是反过来在动态图之上加一层”即时编译”，在运行时把实际执行过的图捕获下来做静态优化。两边都在往中间靠：研究阶段要动态图的灵活与可调试性，部署阶段要静态图的性能与可移植性。这背后有一个共同的技术思路，叫”图捕获（Graph Capture）”——框架追踪你的动态操作，然后把它冻结成一张可优化的静态图。这也说明动态图的灵活性终究是以牺牲部分编译优化空间为代价的，如何兼得两者，至今仍是框架设计中最大的未解难题之一。
三、框架”框”住的其他东西：张量与算子
如果自动微分是框架的心脏，那么张量（Tensor）就是流淌在整个系统里的血液。标量是零阶张量，向量是一阶张量，矩阵是二阶张量，而深度学习里常见的批量图像数据（batch × 通道 × 高 × 宽）则是四阶张量。一个框架里的张量远比 NumPy 的多维数组复杂——它至少要携带这些信息：形状（shape）、数据类型（dtype，如 FP32、FP16、INT8）、内存布局（NCHW 还是 NHWC）、所在设备（CPU 还是 GPU）、以及是否需要梯度、指向哪个计算图节点等。这些信息看似琐碎，实际上深刻影响着计算效率——比如 FP16 张量比 FP32 少占一半显存，在支持 Tensor Core 的 GPU 上算得更快；PyTorch 历史上默认使用 NCHW，早期 cuDNN 对 NCHW 路径支持最成熟；而近年来 cuDNN 对 NHWC / channels-last 的优化越来越充分，尤其在 FP16 Tensor Core 场景下，NHWC 往往能取得更好的卷积吞吐。Tech-Renaissance 选择统一使用 NHWC，正是为了顺应这一趋势。两种布局可以互相转换，但转换本身有开销，所以框架通常在整条训练管线里统一使用一种布局。
有了张量这个数据载体，算子（Operator）就是操作数据的基本计算单元——如果把神经网络比作一座建筑，算子就是砖块。卷积、全连接、批归一化、ReLU/GELU 等激活函数、池化、交叉熵损失，这些耳熟能详的名字，本质上都是一个个算子。框架需要为每种算子提供正确且高效的实现，并保证同一个算子在不同硬件（比如 A100 和 RTX 5090）上给出数学等价的结果。这里的工程量常常被低估：同一个卷积操作，在 CPU 上要用 SIMD 指令优化，在 GPU 上要写 CUDA Kernel 并利用 Tensor Core，输入尺寸、通道数、卷积核大小、数据精度的任何变化，都可能导致最优实现方案完全不同。因此现代框架通常会调用 cuDNN、cuBLAS 这类硬件厂商提供的高度优化库，而不是重新造轮子。
四、别让 GPU”挨饿”：数据管线
很多人第一次做性能优化时会惊讶地发现：训练慢，未必是模型算子慢，而是 GPU 在等数据。
一次完整的训练迭代，数据要经历这样的旅程：从硬盘读取原始文件 → 解码（如 JPEG 解压）→ 预处理（缩放、裁剪、归一化、数据增强）→ 打包成 batch → 从 CPU 内存拷贝到 GPU 显存 → GPU 计算。如果前面几步跟不上 GPU 的计算速度，昂贵的计算资源就会空转等待，这被称为”数据饥饿”。对于大规模数据集，这个问题尤为突出：一块高端 GPU 算完一个 batch 可能只需要几十毫秒，但单线程从硬盘加载并预处理一个 batch 可能需要几百毫秒。
为了解决这个瓶颈，现代框架普遍采用几种机制：用多个 worker 子进程并行做数据预处理（比如 PyTorch DataLoader 的 num_workers 参数，在 Linux 下通常是多进程）；利用 CUDA 的异步拷贝，在 GPU 计算当前 batch 的同时把下一个 batch 从 CPU 搬到 GPU，这需要用到传输带宽更高的锁页内存（Pinned Memory）；维护双缓冲区，让 CPU 填充与 GPU 处理两个过程交替重叠；以及设计更高效的数据存储格式（如 TensorFlow 的 TFRecord），把图像和标签预先打包为二进制记录以加速读取。这一整套机制的目标只有一个：让数据管线的速度追上 GPU 的计算速度，实现计算与传输的完全重叠。
五、从单卡到集群：分布式训练
单张 GPU 的显存和算力终归有限。当模型越来越大（比如千亿参数级别的语言模型，仅存储 FP32 权重就需要数百 GB 显存）、数据越来越多时，就必须让多张 GPU 协同工作。分布式训练主要有三种策略。
数据并行是最常用、最容易实现的策略：每个 GPU 持有完整的模型副本，处理不同的数据子集，各自独立完成前向和反向计算后，通过 AllReduce 通信把所有 GPU 上的梯度加和平均，保证参数始终一致。PyTorch 的 DistributedDataParallel 就是典型实现。它的优点是实现简单、扩展性好，缺点是要求每张卡都能装下完整模型。
模型并行则是把模型的不同部分放到不同 GPU 上，数据依次流过各卡——适合单卡装不下整个模型的场景，但由于卡与卡之间存在前后依赖，同一时刻只有部分 GPU 在工作，硬件利用率较低。
流水线并行是模型并行的改进版：把模型切成多个阶段分布到不同 GPU，但不再等一个 batch 完全处理完才开始下一个，而是像工厂流水线一样让多个 batch 同时注入管线，不同 GPU 同时处理不同 batch 的不同阶段，从而提高整体利用率。
在实际的大模型训练中，这几种策略往往被混合使用——例如在 Transformer 内部使用张量并行（模型并行的一种变体），在层之间使用流水线并行，在不同数据副本之间使用数据并行，这就是所谓的”3D 并行”。这些复杂的通信机制，最终都要依赖框架对底层通信库（如 NCCL）的封装，一个优秀的分布式实现还需要尽可能把通信和计算重叠起来，以隐藏网络延迟。
六、把这一切串起来：训练循环与优化器
有了张量、算子、自动微分、数据管线和分布式支持，框架还需要把它们组织成一个完整的训练循环：取一个 batch 的数据 → 前向传播得到预测 → 计算损失 → 反向传播算出梯度 → 优化器根据梯度更新参数。这个循环重复数万到数百万次，直到模型收敛。
优化器控制着参数更新的具体策略。最基础的随机梯度下降（SGD）沿梯度反方向更新参数；引入”动量”的 SGD 会累积历史梯度方向，让更新更平滑，有助于跳出局部最优；Adam 结合了动量与自适应学习率，对每个参数独立调整学习率，让训练对超参数不那么敏感，是目前最流行的优化器之一；AdamW 进一步修正了权重衰减的实现方式,将其与自适应学习率解耦,已成为当前大模型训练的标准选择。配合优化器的还有学习率调度策略——比如先用较大学习率”热身”（Warmup）再逐渐衰减，或让学习率按余弦曲线平滑下降（余弦退火），这些细节共同决定了模型能否又快又稳地收敛。
七、简史：从百家争鸣到双寡头
深度学习框架的发展只有十余年，但格局变化很大。Theano（2008 年前后，蒙特利尔大学 Bengio 团队）最早把符号计算图、自动微分和 GPU 加速系统化地结合在一起，是许多后来者的思想源头，但其 API 偏底层、编译速度慢，2017 年停止了主要开发。Caffe（2013 年前后，加州大学伯克利贾扬清主导开发）用配置文件描述网络结构，在计算机视觉领域一度极为流行，但难以表达循环、条件分支等动态逻辑，遇到 ResNet、Inception 这类复杂结构时局限性凸显。基于 Lua 的 Torch，以命令式、动态构图的编程体验著称，其设计理念后来被 PyTorch 直接继承，但 Lua 语言的小众性限制了它的传播。MXNet 曾同时支持命令式与符号式编程、并被亚马逊采用，但始终未能积累起足够大的社区，最终于 2023 年 9 月退役至 Apache Attic（项目归档）。
2015 年，Google 开源 TensorFlow，凭借强大的分布式能力和静态图带来的编译优化，迅速成为工业界的标杆；随后在 2019 年发布的 TensorFlow 2.0 中默认启用了动态执行，向易用性妥协，却也因此付出性能代价，迫使官方随后全力推广 @tf.function 图编译来救场。2017 年 1 月，Facebook（现 Meta）发布 PyTorch 首个公开版本，其动态计算图带来的”所写即所得”体验极大降低了调试门槛，到 2019 年前后已在学术论文中的使用率反超 TensorFlow，并在后续的大模型浪潮中巩固了统治地位。此外，Google 推出的 JAX 以函数式编程和 XLA 编译见长，凡是需要复杂梯度操作（如高阶导数、逐样本梯度）的前沿研究，都能看到它的身影；DeepSpeed、Megatron-LM 这类专用训练框架则不替代 PyTorch，而是作为插件提供更极致的分布式训练能力。如今，研究和生产中占据主流的基本只剩 PyTorch 与 TensorFlow 两家，PyTorch 的优势在学术界和大模型领域尤为突出，JAX 作为新锐力量份额也在稳步上升；而 Theano、Caffe、MXNet 等曾经的名字，大多已淡出主流视野。
八、框架的边界与未来
深度学习框架和 NumPy 这类科学计算库的边界在哪里？NumPy 提供的是无状态、即时求值的数值运算，没有自动微分、没有计算图、也没有设备管理——可以说它只是框架张量系统的一个”CPU 回退子集”。框架和 MLOps 平台（如 Kubeflow）的边界又在哪里？框架解决的是”如何高效正确地训练一个模型”，而 MLOps 解决的是”如何规模化、可复现地管理成百上千个模型的训练与上线”，两者互补而不重叠。
值得注意的是，Python 本身并不是框架真正的执行主体——当你调用一个卷积层时，Python 只是”发令枪”，真正干活的是框架的 C++ 后端和硬件厂商提供的底层库。Python 的全局解释器锁决定了它难以高效利用多核并行，因此计算密集任务最终都要下沉到编译后的底层代码去执行。这也是为什么”编译器化”正在成为各大框架的共同方向：PyTorch 的 torch.compile、TensorFlow/JAX 的 XLA，都试图在保留动态图开发体验的同时，通过即时编译捕获执行图并做静态优化，把性能损失降到最低。未来的框架很可能进一步走向自动化——自动做算子融合、自动生成硬件相关的计算核心、甚至自动权衡并行策略与内存布局，同时需要适配日益多样化的 AI 芯片。
结语
深度学习框架不是一个简单的”调库工具”，而是一整套连接数学抽象与物理硬件的复杂系统：它是自动微分引擎，是计算图的构建者与优化者，是硬件抽象层，也是分布式通信的调度者。正因为这套系统的存在，研究者才能用几十行 Python 代码描述一个数十亿参数的模型，并让它在数千张 GPU 上高效、正确地运行——而不必操心矩阵乘法该怎么在 GPU 上并行，梯度该怎么一层层手推。
但理解框架的价值，不只在于用好它，更在于知道它为什么这样设计。当你遇到”为什么这个操作不支持梯度””为什么分布式训练卡住了””为什么 torch.compile 能加速”这类问题时，如果你了解张量、算子、自动微分、计算图这几层骨架各自在做什么，你就有了一条清晰的排查思路，而不是把框架当成一个纯粹的魔法黑盒。这大概就是”深度学习框架到底是什么”这个看似基础的问题，值得被认真回答一次的原因。
下一篇文章，我们先聊聊 AI 编程技术本身——毕竟接下来要讲的框架设计，全靠 AI 编程完成。
