(5) 静态图 Vs 动态图:深度学习框架的两种世界观

——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之五

上一篇我们把深度学习框架拆解成了张量、算子、自动微分、计算图、数据管线、分布式训练这几个核心部件。其中有一个概念被我当作”背景”一笔带过,但它实际上决定了一个框架的基本气质——那就是计算图到底是静态的,还是动态的。

你需要明白的是,这是两种不同的深度学习框架世界观,各有长处,但依然有办法兼具两者的优点,以至于在当下的深度学习框架领域里,两者的差距正在缩小、边界正在变得模糊。

一、计算图:框架里的”施工图纸”

在讨论静态图和动态图之前,先统一一下什么叫计算图(Computational Graph)

一个深度学习模型,无论多复杂,本质上都可以拆成一系列基本运算:矩阵乘法、卷积、加法、激活函数、损失函数……如果把这些运算画成节点,把张量数据流动画成有向边,就得到一张有向无环图——这就是计算图。节点是算子,边是数据。

计算图的价值在于它把”模型长什么样”和”怎么执行”分离开了。有了这张图,框架就可以做很多事:

  • 自动微分:从损失节点反向遍历图,沿链式法则传播梯度;
  • 图优化:常量折叠、死代码消除、公共子表达式消除;
  • 算子融合:把相邻的 conv、bias、batchnorm、relu 合并成一个 kernel;
  • 内存规划:分析张量生命周期,复用显存;
  • 并行调度:把不相关的子图扔到不同流上并发执行。

换句话说,计算图是框架做全局优化的抓手。没有图,框架就只能在局部”见一个算子优化一个算子”;有了图,框架才能”放眼全局”。

二、静态图与动态图:两种世界观

历史上,框架围绕”这张图什么时候建、怎么用”分成了两派。

2.1 静态图:先画图,再施工

静态图(Static Graph,define-and-run)的思路是:用户先用某种方式把整张计算图描述出来,框架对这张图做编译、优化、内存规划,然后再喂数据执行。

TensorFlow 1.x 的 tf.Graph + tf.Session、Theano 的符号图,都是这个路线的代表。用户写的代码更像是在”声明”一张图,而不是在”执行”指令。优点和缺点都很鲜明:

优点:

  • 全局优化空间大。因为图在运行前已经完整可见,编译器可以做跨算子的常量折叠、算子融合、布局转换、内存复用规划。
  • 部署友好。图可以导出为 SavedModel、ONNX、TensorRT Plan 等格式,脱离 Python 运行。
  • 更容易做确定性训练。因为图结构和 kernel 序列在运行前已经确定,如果再配合静态内存规划、固定调度和确定性算子选择,就更容易让同一程序多次运行得到一致结果。

缺点:

  • 对初学者不友好。静态图的思路对于尚不理解计算图为何物的初学者而言可谓相当诡异,很多人无法理解“为什么我要先定义session再run”、“为什么我写了一个运算却没有执行”。
  • 调试体验相对较差。你写的不是”普通 Python 程序”,而是”描述图的程序”。出错时很难直接 print 中间张量,也很难用 pdb 单步跟进。
  • 动态逻辑表达困难。Python 里的 iffor、递归、数据依赖控制流,在静态图里需要 special 处理(比如 tf.condtf.while_loop),学习曲线陡峭。
  • 图构建有开销。复杂模型构图和编译可能要几十秒甚至更久。

2.2 动态图:边施工,边画图

动态图(Dynamic Graph,define-by-run)则反其道而行之:你写一行代码,框架就执行一行、建一个节点;中间张量立刻可见,调试点哪里打哪里。

PyTorch Eager、TensorFlow Eager、Chainer、DyNet 都属于这一派。PyTorch 的成功很大程度上归功于动态图带来的”所写即所得”体验。研究者可以像写 NumPy 一样写模型,可以用原生 Python 控制流,可以随时随地 print(x.shape) 或者 import pdb; pdb.set_trace()

优点:

  • 符合直觉。和写普通 Python 程序几乎没有区别。
  • 调试方便。任何中间结果都能立即查看。
  • 动态模型友好。变长序列、动态结构、条件分支、稀疏计算都可以自然表达。

缺点:

  • 缺乏全局视野。框架看到的是一次一次独立的算子调用,很难做跨算子的全局优化。
  • 调度开销大。每个算子都要走一遍 Python 派发、C++ dispatcher、CUDA kernel launch,小算子密集时 CPU 调度会成为瓶颈。
  • 中间结果常驻显存。动态图为了反向传播,通常要把前向中间结果都保留下来,内存占用更大。

2.3 边界正在模糊

有意思的是,这场”静态 vs 动态”之争并没有以一方的彻底胜利告终,而是走向了融合

TensorFlow 2.x 默认使用 Eager Execution,但提供了 @tf.function,在保留动态编程体验的同时把函数编译成静态图。PyTorch 2.0 则在 2023 年推出了 torch.compile:前端 TorchDynamo 捕获 FX 图,AOTAutograd 处理前向与反向,后端 Inductor 做算子融合、kernel 生成,甚至尝试 CUDA Graph 捕获。JAX 从设计之初就鼓励用户用 jax.jit 把动态 Python 函数编译成 XLA 静态图。

所以现在的局面不是”静态图框架”和”动态图框架”泾渭分明,而是动态框架纷纷长出编译模式,静态框架也纷纷补 eager 体验。这是一个大家都承认的事实:动态图的灵活性是以牺牲全局优化空间为代价的;而静态图的性能优势,则建立在”编译期确定一切”的前提之上。

三、为什么 Tech-Renaissance 选择静态图编译

讲清楚了通用背景,现在回到我们的框架。

Tech-Renaissance 的定位非常明确:它不是给研究者提供灵活调试玩具的框架,而是一个追求极致训练吞吐量的生产级训练框架。这个定位决定了我们必须从架构上优先保留全局优化空间。静态图不是”为了不同而不同”,而是我们后续一系列核心能力的必要前提

具体说,以下三件事,离开了静态图几乎不可能做好。

3.1 CUDA Graph 全捕获

CUDA Graph 是 NVIDIA GPU 提供的一种机制:把一段 GPU 工作(多个 kernel、memcpy、事件同步)一次性捕获下来,之后通过 cudaGraphLaunch 反复重放,从而消除每次 kernel 的 CPU 派发开销。

但 CUDA Graph 有几个硬性约束:

  • 输入输出张量的虚拟地址必须固定
  • 捕获期间不能动态分配显存
  • 控制流必须是静态的,不能有条件分支或动态循环;
  • 不能出现 host-device 同步(比如 .item() 取标量到 CPU)。

这些约束天然地指向静态图:只有在编译期就确定所有张量的形状、内存偏移、kernel 序列,才能在运行期放心地把整个训练循环捕获成 CUDA Graph。动态图框架如果想做这件事,就必须先把动态执行痕迹冻结成静态子图——这正是 torch.compile 在做的事,但它经常因为 graph break、动态形状、地址不稳定而被迫回退到 eager。

在 Tech-Renaissance 里,我们不是”事后捕获”,而是”从设计之初就为捕获而建”。

3.2 MemoryPlan 静态显存规划

Tech-Renaissance 有一个非常核心的设计:运行期不做动态显存分配。所有张量的偏移、大小、生命周期都在编译期由 MemoryPlan 确定。

MemoryPlan 把显存按语义划分为 68 个命名 Region(枚举槽位共 69 个,最后一个为数组边界哨兵):权重区、梯度区、一阶/二阶动量区、EMA 权重区、AMP FP16 权重区、输入缓冲区、特征图区、临时张量区、结果区等等。每个 Region 的偏移在 finalize() 之后固定下来,所有 DTensor 只是这个布局里的一个 (offset, stride, shape) 描述符,不持有实际内存。

这种设计带来了几个直接好处:

  • CUDA Graph 所需的地址稳定性天然满足
  • 消除动态分配延迟和内存碎片
  • 可以按 Region 做批量操作,比如一次性清零全部梯度、一次性切换 EMA 权重、一次性完成整区更新;
  • 多卡分布式布局一致,同一张 MemoryPlan 图纸可以在多张 GPU 上共享执行。

如果框架是动态图,显存通常由运行时分配器按需申请,地址每次可能不同,就很难做到这种静态分区。

3.3 确定性训练

Tech-Renaissance 强调可复现性:在相同硬件、相同种子、不调用非确定性算子的情况下,同一程序多次运行可以得到字节级一致的结果。

这件事也依赖于静态图。因为只有图结构、内存布局、kernel 序列、调度顺序在编译期就固定下来,我们才能:

  • 用 Philox 计数器随机数生成器精确控制每个随机点的 seed 和 offset;
  • 在数据增强、参数初始化、Dropout 等所有随机点保持一致;
  • 避免运行时动态调度带来的不确定性。

动态图框架理论上也可以做确定性训练,但由于 Python 执行路径、内存分配时机、异步调度顺序更容易波动,工程上要保证端到端一致需要付出更多代价。

四、代码层面看 Tech-Renaissance 的静态图实现

理论讲完,看一下代码是怎么落地的。

4.1 用户侧:先定义,后编译,再运行

tests/example/mlp_mnist.cpp 为例:

BluePrint mlp = seq(
    fc(1024, true), relu(),
    fc(512, true), relu(),
    fc(256, true), relu(),
    fc(10, true)
);

DeepLearningTask task;
task.model(mlp)
    .loss(CrossEntropyLoss().label_smoothing(0.1f))
    .total_epochs(kTotalEpochs)
    .optimizer(AdamW().weight_decay(1e-4f))
    .scheduler(CosineAnnealingLR().base_lr(0.001f).warmup(5));

task.compile();
auto result = task.run();

注意这里的 task.compile():它不是一个简单的”打印配置”,而是真正把 BluePrint 描述的模型编译成内部 IR、MemoryPlan、ComputationGraph,并捕获成可执行图。调用 task.run() 时,框架只是在重复执行已经编译好的图。

4.2 三阶段状态机

TaskBase 强制了一个三阶段状态机,定义在 include/renaissance/core/types.h

enum class Phase : uint8_t { PLANNING, MEMORY_LOCKED, COMPILED };
  • PLANNING 阶段:用户或编译器分配 DTensor、配置初始化策略;
  • MEMORY_LOCKED 阶段:内存布局锁定,注册计算图;
  • COMPILED 阶段:硬件初始化、CUDA Graph 捕获完成,进入运行期。

这个状态机在 include/renaissance/task/task_base.h 中有明确约束:

// PLANNING 阶段
[[nodiscard]] DTensor alloc(const Shape& shape, DType dtype, Region region);
void finalize_memory();

// MEMORY_LOCKED 阶段
void add_graph(const std::string& name, ComputationGraph graph,
               StreamKind stream = StreamKind::COMP_1);

// COMPILED 阶段图执行
void run(const std::string& name);

每个阶段能调用的 API 是受限的。这种强约束保证了我们不会在运行期突然”插入”一张新图或”动态”申请一块新内存——那不是这个框架的设计目标。

4.3 五阶段编译器

Compiler::compile() 是静态图编译的核心入口,定义在 include/renaissance/graph/compiler.h

static Result compile(const ArchPlan& arch,
                      const CompileSpec& base_spec,
                      const PlanConfig& plan_config,
                      Initializer& initializer,
                      const std::vector<CompileSpec>& variant_specs = {});

它的内部管线分为五个阶段:

  1. derive_all_shapes:对 6 个变体(base、train_last、train_lowres、train_lowres_last、val_base、val_last)分别做形状推导;
  2. compute_max_slot_bytes:对每个 (layer, tensor) 位置跨变体取最大 slot 字节数;
  3. create_memory_plans:用最大 slot 构造 6 个独立的 MemoryPlan,保证同一 DTensor 在不同变体下 offset 一致;
  4. build_computation_graph:遍历 LayerDescriptor 序列,构建训练图和推理图的拓扑;
  5. share_or_clone:shape-only 变体共享同一份 ComputationGraph,val 变体不填训练图。

这里的核心思想是:把”图拓扑”和”形状/内存布局”解耦ComputationGraph 是零形状信息的纯拓扑容器,一份图可以服务多个 shape-only 变体。

4.4 零形状信息的计算图与图集

include/renaissance/graph/computation_graph.h 中的注释说得很清楚:

核心特征:零形状信息。节点中只存 (GraphNode, OpParams, tensor_id);Shape / DType / Region 全部从 MemoryPlan 的 DTensor 获取;一份图供多个 shape-only 变体共享。

class ComputationGraph {
    // 每个 GraphId 桶各存一个节点序列
    std::array<std::vector<GraphNode>, static_cast<size_t>(GraphId::COUNT)> graphs_;
    std::vector<GraphNode> linear_nodes_;
};

运行期,哪个变体、哪个子图该执行,由 GraphAtlas 决定。GraphAtlas 是一张 6 变体 × 多子图的映射表,在 include/renaissance/graph/graph_atlas.h 中实现:

class GraphAtlas {
    static constexpr size_t kMaxVariants = 6;
    static constexpr size_t kMaxGraphIds = static_cast<size_t>(GraphId::COUNT);
    // Phase A 填入 cg / mp / shape_id
    // Phase B 预捕获后填入 captured_idx
    // Phase C 运行时 O(1) 数组访问
};

shape 无关的子图(比如 TRANSFER、COMM、OPTIMIZER、EMA_UPDATE)在所有变体间共享;shape 相关的训练/推理图按 ShapeId 去重。运行期只需要一次数组索引,零 hash、零分支、零动态查找。

4.5 CUDA Graph 捕获

src/graph/capture_cuda.cpp 负责把 ComputationGraph 中的节点序列捕获为 CUDA Graph:

void CapturedGraph::capture_cuda(const ComputationGraph& cg,
                                  const MemoryPlan& mp,
                                  GraphId gid,
                                  StreamKind stream_kind,
                                  const DeviceContext& ctx,
                                  CapturedGraph& result) {
    cudaStream_t primary_stream = select_primary_capture_stream(stream_kind, ctx);
    const auto& nodes = cg.linear_nodes().empty() ? cg.nodes(gid) : cg.linear_nodes();

    cudaStreamBeginCapture(primary_stream, cudaStreamCaptureModeThreadLocal);
    // ... 遍历节点,处理多流同步,cudaStreamEndCapture, cudaGraphInstantiate
}

捕获前,我们会做 warmup;捕获时,MultiStreamCaptureState 会管理三条计算流、一条传输流、一条更新流之间的依赖和事件同步。这也是为什么 Tech-Renaissance 能做到训练循环全阶段 CUDA Graph 捕获——从 H2D 传输、首层前向、深层前向+反向、梯度通信、优化器更新、BN 统计量同步,全部可以在捕获范围内。

五、与 PyTorch Eager 和 torch.compile 的对比

说到这里,可能有人会问:PyTorch 不是也有 torch.compile 吗?它不也是在试图做静态优化和 CUDA Graph 吗?

是的,但路径完全不同。

PyTorch Eager 是动态图的原教旨主义者:每个算子都经过 Python dispatcher、Autograd、C++ kernel launch,逐个提交给 GPU。它的优势是灵活和易用,代价是调度开销和全局优化空间受限。

torch.compile 是在动态图之上叠加的一层JIT 编译器。它用 TorchDynamo 捕获 FX 图,用 Inductor 做融合和 kernel 生成,并尝试 CUDA Graph 捕获。这是一个非常有工程价值的方向,但它也面临几个结构性挑战:

  • Graph break。遇到 torch.save.item()、数据依赖控制流等不支持的操作时,Dynamo 会中断当前图,中间片段回退到 eager 执行。每个 graph break 都意味着一次 host-device 同步和优化机会的损失。
  • 动态形状。输入形状变化会触发重新编译。虽然 PyTorch 有符号化形状的支持,但复杂场景下仍会导致反复编译或 shape specialization。
  • 地址稳定性。CUDA Graph 要求捕获时的地址在重放时不变。PyTorch 的缓存分配器通常是”伪确定性”的,需要额外机制(如 private memory pool、CUDA Graph Trees)来保证,理论上并非 100% 稳定。
  • 正确性风险torch.compile 的编译结果与 eager 在数学上应当一致,但社区长期报告过静默错误、梯度异常、与 CUDA Graph 搭配报错等问题。这不是贬低 PyTorch,而是说明在动态图基础上事后补静态优化,天然比从静态图出发更困难

Tech-Renaissance 的选择是:不要事后捕获,而是从头静态。我们把所有形状、内存、图拓扑、调度关系都在编译期根据一定的规则确定下来,因此 CUDA Graph 全捕获、MemoryPlan 静态分区、确定性训练这些能力都能自然落地,不需要在灵活性和性能之间反复拉扯。

显然这存在代价:Tech-Renaissance 不支持运行期动态修改模型结构,不支持数据依赖控制流,调试体验更接近”先编译、再运行”的 C++ 程序,而不是 PyTorch Eager 的”所见即所得”。这是我们有意识地接受的 trade-off。但是这带来的性能上的收益,很快大家就可以见识到

六、小结

静态图和动态图没有绝对的优劣,只有场景取舍

  • 如果你的首要诉求是研究快速迭代、调试方便、动态模型,动态图是更自然的选择;
  • 如果你的首要诉求是训练吞吐、部署效率、确定性、端到端复现,静态图能给你更大的优化空间。

Tech-Renaissance 属于后者。它不是”静态图框架”四个字可以概括的,而是围绕静态图编译这一选择,把 Compiler 五阶段编译、68-Region MemoryPlan(含哨兵共 69 个槽位)、零形状 ComputationGraph、GraphAtlas 变体复用、多流 CUDA Graph 全捕获 串成了一条完整链路。后续几篇文章会分别展开这些部件,但它们的共同底座,就是本篇讲的静态图编译

下一篇,我们把视角拉高,用一条样本从磁盘到指标输出的完整旅程,串起 Tech-Renaissance 的七大模块,给你一张全局地图。

发表回复

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

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