第8章 Autograd Engine：多线程后向调度器

“The engine drives backward like an event loop drives a server: tasks come in, get scheduled to the right thread, get executed, produce more tasks, and the cycle continues until the graph is done.”

—— torch/csrc/autograd/engine.cpp 顶部注释（节录改编）

8.1 问题：怎么并行跑反向 DAG

第 7 章我们看到 PyTorch 在前向时建好一张反向 DAG。loss.backward() 调用之后，要做的事：

1. 从 loss.grad_fn（DAG 的 root）出发，遍历所有可达的 Node
2. 对每个 Node 调 apply(inputs) 计算梯度
3. 把每个 Node 的输出按 next_edges 路由到上游 Node 的 InputBuffer
4. 当某个 Node 的所有输入都到齐时，把它放入 ReadyQueue 等待执行
5. 多线程并行执行，最终所有 leaf 张量的 grad_ 字段被填好

听起来简单，但要做到正确 + 快有几个挑战：

• 正确：DAG 上多条边汇入同一个 Node 时，必须等所有边的梯度都到齐才能执行（否则计算错误）
• 跨设备：CUDA Node 应该在 CUDA worker 上跑，CPU Node 在 CPU worker 上跑 —— 不能让 CPU 调 CUDA kernel 触发隐式同步
• stream 安全：同一 GPU 上多个 Node 在不同 stream 上可能并发，引擎必须正确同步
• 可重入：用户的反向规则里可能再调 torch.autograd.grad（二阶导）—— 引擎必须支持嵌套
• 错误处理：某个 Node 抛异常时，要让所有 worker 优雅停下、把异常传回用户线程

这些约束下设计出的引擎就是 PyTorch 的 autograd Engine。本章拆它的实现。

8.1.0.5 关键概念辨析

进入正文前先把几个易混淆术语理清：

• 反向图 (backward graph)：DAG，由 Node 和 Edge 组成。第 7 章建好的就是它
• GraphTask：Engine 为一次 backward 调用创建的”运行时上下文”对象，含 dependencies、not_ready、exec_info 等运行期状态
• NodeTask：派给 worker 的”工作单元”，含一个 Node 指针 + 它的 InputBuffer
• ReadyQueue：thread-safe 优先队列，存”已经可以执行”的 NodeTask
• worker thread：循环跑 thread_main 的常驻线程

这些名字容易混。一个粗略心智模型：反向图是静态结构（前向时建好），GraphTask 是动态实例（每次 backward 创建），NodeTask 是工作单位（图里每个 Node 一次跑就是一个 NodeTask）。

8.1.1 一个直觉：为什么不能”前向倒着跑”

新手有时候会问：反向不就是把前向反着跑一遍吗？为什么需要这么复杂的引擎？

答案是 梯度计算和前向是两套完全不同的逻辑：

• 前向是 y = f(x)，反向是 dx = g(dy, x) —— 反向函数 g 通常需要前向时的中间值（如 mm 反向需要 mat2.T）
• 前向是树状（一个输入生几个输出），反向是网状（多条边汇入同一节点要求和）
• 前向算子的实现已经在 ATen 里，反向需要专门的 XxxBackward 类
• 前向的 dependency 顺序由 Python 控制流自然确定，反向需要拓扑排序

所以反向不是”前向的 mirror”，而是基于”反向 DAG”的全新执行 —— 这就是为什么需要一个专门的 Engine。

更具体的对比：前向 c = a + b 是一行 Python，dispatcher 调度后命中 add kernel 就完事；而反向 loss.backward() 是一行 Python，背后是一整个图遍历 + 多线程调度 + 跨 device 路由。前向的复杂度藏在每个算子内部，反向的复杂度藏在图调度外部。这是同一份数学计算的两套不同工程实现。

8.2 全局架构

Engine 是单例，定义在 torch/csrc/autograd/engine.h:130：

struct TORCH_API Engine {
    static Engine& get_default_engine();    // 全局单例
    static Engine& get_base_engine();       // 基础版（不含 Python）

    virtual variable_list execute(
        const edge_list& roots,
        const variable_list& inputs,
        bool keep_graph,
        bool create_graph,
        bool accumulate_grad,
        const edge_list& outputs = {});

    void evaluate_function(
        std::shared_ptr<GraphTask>& graph_task,
        Node* func,
        InputBuffer& inputs,
        const std::shared_ptr<ReadyQueue>& cpu_ready_queue);

    auto thread_main(const std::shared_ptr<GraphTask>& graph_task) -> void;
    ...
};

它管理一组 device-bound worker 线程：

graph TB
    subgraph User["用户调用线程"]
        U[Python: loss.backward]
        U --> EX["Engine::execute(roots, inputs)"]
    end

    subgraph Engine["Engine 单例"]
        WCPU[CPU worker thread<br/>thread_main loop]
        WC0[CUDA:0 worker thread]
        WC1[CUDA:1 worker thread]
        WC2[CUDA:N worker thread]
    end

    subgraph Queues["每 device 一个 ReadyQueue"]
        RQ_CPU[CPU ReadyQueue]
        RQ_C0[CUDA:0 ReadyQueue]
        RQ_C1[CUDA:1 ReadyQueue]
        RQ_C2[CUDA:N ReadyQueue]
    end

    EX --> RQ_CPU
    WCPU <--> RQ_CPU
    WC0 <--> RQ_C0
    WC1 <--> RQ_C1
    WC2 <--> RQ_C2

    style Engine fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style Queues fill:#dbeafe,stroke:#3b82f6

每个 worker 线程有自己的 ReadyQueue（thread-local，藏在 tls_local_ready_queue）。当 Engine 要给某个 Node 派活儿时，它根据 Node 输出的 device 把 NodeTask push 到对应 device 的 ReadyQueue —— 那个 device 的 worker 线程被 condition_variable 唤醒，pop 出 task 执行。

8.2.1 Engine 启动时机

Engine 不是 PyTorch 启动时就立刻初始化所有 worker 线程。它是 lazy 启动的 —— 第一次有 backward 调用时才创建对应 device 的 worker 线程。这种 lazy 模式让”只跑 inference 的进程”完全不付 worker 线程的代价（每个线程 ~2MB 栈空间 + 线程切换开销）。

具体的触发点是 Engine::start_threads()（engine.cpp 中），它由 execute() 在第一次调用时通过 std::call_once 触发。Engine 析构时（通常是进程退出）给每个 worker 队列 push 一个 shutdown task，让所有线程 break 出 thread_main 优雅退出。

8.2.1.5 worker 数量为什么是 device 数量

很多人会问：为什么不开更多 worker（比如 16 个）？

答案是 GPU 任务本质是异步的。每个 CUDA worker 的工作主要是 launch kernel（几微秒）+ 调度内存（几微秒），真正的计算在 GPU 上由 stream 并发。如果开两个 CPU worker 服务同一个 GPU，它们会抢占同一个 stream 的提交权，反而引入同步开销。

所以 PyTorch 选择 “1 worker per device” 这个最简模型：每个 device 的 worker 串行 launch kernel，stream 之间在 GPU 上并发。这避免了多 CPU worker 协调的复杂度，把”并发”完全交给 GPU 自己处理。

CPU worker 也只有 1 个 —— 因为如果用户的算子是 CPU bound 的（如 dataloader 处理），他们应该用 num_workers 在 dataloader 层面并行，而不是依赖 autograd engine 多线程。

8.2.2 没有 GIL 的多线程：C++ Engine vs Python GIL

值得多说一句：Engine 的 worker 线程是纯 C++ 线程，运行 C++ 反向 kernel，完全不受 Python GIL 限制。所以反向能真正并行 —— CUDA:0 worker 在跑 MmBackward 的同时 CUDA:1 worker 可以跑另一个 MmBackward。

只有当反向调用进入 Python 时（比如用户写的 autograd.Function.backward），那个 worker 临时 acquire GIL 跑 Python 代码，跑完释放。这种”C++ 线程偶尔进 Python”的模式让 Engine 即便在 Python 主导的项目里也能榨出多核性能。

8.3 NodeTask：工作单元

engine.h:51：

struct NodeTask {
    std::weak_ptr<GraphTask> base_;     // 这个 task 属于哪次 backward
    std::shared_ptr<Node> fn_;           // 要执行的 Node
    InputBuffer inputs_;                 // 累积的梯度输入
    bool isShutdownTask_;                // 关闭信号 (Engine 析构时用)

    int getReentrantDepth() const;
};

InputBuffer 是关键 —— 它为这个 Node 的每个输入位维护一个累积器。从不同的 next_edge 流来的梯度被累加到同一个位上：

// 简化的 InputBuffer 用法
input_buffer.add(input_nr=0, grad_from_path_A);
input_buffer.add(input_nr=0, grad_from_path_B);   // 自动加到 path_A 上
// 当所有 input 都到齐，integrated_inputs = input_buffer.move()

这就是第 7 章 §7.3.1 提过的”DAG 多边汇入同一节点的梯度自动求和” —— 实现就在这里。

InputBuffer 的 add 方法不是简单的 +=，它会处理几种 corner case：

• stream-aware 累加：如果两条边的梯度来自不同 stream，要插 event 等待保证依赖
• device 不一致检查：两条边的 grad device 必须一致，不一致就报错
• dtype 提升：如果两条边的 dtype 不同（罕见但可能），按规则提升

这套累加在torch/csrc/autograd/input_buffer.cpp 里（实测整文件 326 行，核心 accumulate 逻辑约 100 行）—— 看似简单的”梯度求和”在多 stream / 多 device 场景下其实是个精巧的小系统。

8.4 ReadyQueue：thread-safe 优先队列

engine.h:86：

struct ReadyQueue {
 private:
    std::priority_queue<NodeTask, std::vector<NodeTask>, CompareNodeTaskTime> heap_;
    std::condition_variable not_empty_;
    std::mutex mutex_;

 public:
    void push(NodeTask item, bool incrementOutstandingTasks = true);
    void pushShutdownTask();
    NodeTask pop();
    ...
};

它是 mutex + condvar 保护的 std::priority_queue。优先级排序规则（CompareNodeTaskTime）：

// 简化的比较
bool operator()(const NodeTask& t1, const NodeTask& t2) {
    if (t2.isShutdownTask_) return true;        // shutdown 优先
    if (t1.getReentrantDepth() == t2.getReentrantDepth()) {
        return t1.fn_->sequence_nr() < t2.fn_->sequence_nr();   // 较旧的 Node 先
    }
    return t1.getReentrantDepth() < t2.getReentrantDepth();     // 较浅的 reentrant 先
}

这条规则的目的：

• Shutdown 任务最高优先级：Engine 析构时立即让 worker 退出
• 较浅 reentrant 的优先：嵌套 backward 的内层不要饿死外层
• 较旧的 Node 先：sequence_nr 小的 Node 通常靠近 leaf，先做能让更多其他 Node 准备好

这是经典的 拓扑序 + 优先级 调度。普通 BFS 反向遍历是 FIFO，PyTorch 在 FIFO 之上加了启发式排序，让 backward 进度更平稳。

注意 ReadyQueue 用的是 std::priority_queue 而不是更高级的 lockfree queue。这是因为 backward 路径的 push/pop 频次相对算子调用低（每个 Node 一次 push、一次 pop），mutex + condvar 的开销完全摊薄。如果改成 lockfree，代码复杂度会爆炸但收益微小。这是 PyTorch 在性能与简洁度之间的典型权衡。

8.5 GraphTask：单次 backward 的”状态机”

stateDiagram-v2
    [*] --> Init: loss.backward创建
    Init --> Running: 把 root Node push 到 ReadyQueue
    Running --> Running: NodeTask 完成 / 减依赖 / push 后继
    Running --> Error: 任一 Node 抛异常
    Running --> Done: outstanding_tasks==0
    Error --> Cleanup
    Done --> Cleanup
    Cleanup --> [*]: future_result_ 通知用户线程

每次用户调 loss.backward() 创建一个新的 GraphTask（graph_task.h）：

struct GraphTask : std::enable_shared_from_this<GraphTask> {
    std::atomic<uint64_t> outstanding_tasks_{0};      // 未完成的 NodeTask 计数
    std::atomic_bool has_error_{false};                // 任一 Node 出错就停所有
    bool keep_graph_;                                  // retain_graph?

    std::mutex mutex_;
    std::unordered_map<Node*, InputBuffer> not_ready_; // 还没集齐输入的 Node
    std::unordered_map<Node*, int> dependencies_;      // 还要等多少条边

    std::unordered_set<Node*> nodes_in_graph_;
    c10::SmallVector<Node*, 4> graph_roots_;

    // exec_info 用来过滤"不需要计算"的子图 (见 .grad(inputs=...) 用法)
    std::unordered_map<Node*, ExecInfo> exec_info_;

    std::vector<at::Tensor> captured_vars_;            // .grad() 接口的输出
    std::shared_ptr<ReadyQueue> cpu_ready_queue_;       // 这次 backward 的 CPU 队列
    int owner_;                                         // 启动这次 backward 的线程 device
    std::shared_ptr<at::ivalue::Future> future_result_; // 完成信号
    ...
};

理解几个核心字段：

8.5.0 GraphTask 与用户线程的协作模型

GraphTask 设计上有一个有趣选择：它不是 Engine 单例的字段，而是 per-backward 的对象。每次用户调 loss.backward() 创建一个新 GraphTask，跑完即释放。这种”每次新建”的模式让多用户线程能并发跑各自的 backward，互不干扰。

具体协作：

1. 用户线程进 Engine::execute，创建 GraphTask
2. 用户线程自己也变成 worker：在 thread_main(graph_task) 里循环，参与执行
3. 与此同时，device 线程被唤醒，从对应 RQ 拿任务跑
4. 当所有 outstanding_tasks_ = 0，graph_task->future_result_ 标记完成
5. 用户线程的 thread_main 看到 future 完成，退出循环返回

注意第 2 点：用户线程自己是 worker。这就是为什么”backward 调用是同步的” —— 用户线程参与执行直到全部完成。这种设计避免了”用户线程提交任务后空转等待”的浪费，而是让它直接参与算 CPU 部分。

8.5.1 dependencies_：每个 Node 还要等多少条边

dependencies_[node] = N 表示 node 还要等 N 条边的梯度才能执行。每次有边的梯度被累加到 not_ready_[node]，dependencies 减 1。当 dependencies_[node] == 0 时，node 被 push 到 ReadyQueue 等待执行。

这套机制就是经典的 kahn 拓扑排序。Engine 执行前用一次 BFS 算出所有 dependencies，执行时用计数器递减驱动调度。

具体的 BFS 在 Engine::compute_dependencies：从 graph_root 出发，用 BFS 遍历所有可达 Node，每条边让目标 Node 的 dependencies 计数 +1。这次 BFS 也填好 nodes_in_graph_（用于过滤）。BFS 复杂度是 O(节点数 + 边数)，对一个 70 层 transformer 反向图大约几百个节点几千条边，BFS 耗时通常 < 1ms。

值得注意：BFS 没必要为每个 backward 都重做 —— 如果你跑同一个模型反复 backward，反向图结构是固定的，理论上 dependencies 也固定。但 PyTorch 仍每次重做，因为：(a) eager 模式下没法保证模型结构不变（用户可能动态改）；(b) BFS 本身很快，缓存的复杂度不值得。torch.compile 走的是另一条路 —— 编译后整个反向连 BFS 都跳过，直接执行编译好的 binary。

8.5.1.5 dependencies_ 与正确性

为什么这个计数器一定要存在？考虑一个反例：如果 Engine 看到任何 Node 准备好就立即执行，会发生什么？

Node X 有两条 in-edge：来自 Node A 和 Node B
A 先跑完，X 收到一份梯度 grad_A
此时 X 立即跑 → 用 grad_A 算输出
后来 B 跑完，又给 X 一份 grad_B
此时 X 已经跑过了 → grad_B 被忽略！梯度错误

dependencies_ 计数器就是防这种 race。只有所有 in-edge 的梯度都到齐，X 才被允许执行。这是反向 DAG 正确性的核心保障。

8.5.2 not_ready_：还没集齐输入的 Node 的 InputBuffer 缓存

当某 Node 的部分（不全）输入到齐时，它的 InputBuffer 暂存在 not_ready_[node]。其他边的梯度到来时累加进去。当 dependencies 减到 0，从 not_ready 取出 InputBuffer，构造 NodeTask 入队。

8.5.3 exec_info_：子图过滤

torch.autograd.grad(loss, [w1, w2]) 让用户只对部分 leaf 求导。这种”只跑某些路径”的需求由 exec_info_ 处理 —— 在初始化阶段从 root 反向 BFS 标记”哪些 Node 在通往目标的路径上”，执行时只跑标记过的 Node。这种 mask 让”对部分参数求导”不浪费时间在无关 Node 上。

8.5.3.5 exec_info_ 在 backward(inputs=…) 上的应用

torch.autograd.backward(loss, inputs=[w1, w2]) 是另一种触发 exec_info 的接口：你只想累积梯度到 w1 和 w2，不动其他参数。Engine 会预先做一次反向 BFS，标记”哪些 Node 在通往 w1/w2 的路径上”，执行时跳过非标记 Node。

这种 mask 让”对部分参数求导”非常高效，特别是 LoRA 这种”冻结大部分参数、只训少量 adapter”的训练场景 —— 通常 99% 的反向 Node 不需要跑，只算 LoRA 那少量参数对应的子图。

代价是初始化阶段多一次 BFS 算 exec_info（和 dependencies 同时做）。对大型反向图这次 BFS 也只是几十毫秒。

8.5.4 captured_vars_：捕获指定 leaf 的梯度

当用户用 torch.autograd.grad(loss, [w1, w2]) 接口（非 loss.backward()），梯度不写到 w.grad_ 而是返回。captured_vars_ 就是这个返回值的存储位置。每条 ExecInfo 的 Capture 记录”哪个 input 索引要写到 captured_vars 的哪个位置”。

8.6 thread_main：worker 主循环

engine.cpp:518 是 worker 线程的灵魂。简化版：

auto Engine::thread_main(const std::shared_ptr<GraphTask>& graph_task) -> void {
    while (graph_task == nullptr || !graph_task->future_result_->completed()) {
        std::shared_ptr<GraphTask> local_graph_task;
        {
            NodeTask task = local_ready_queue->pop();   // 阻塞等待 task
            if (task.isShutdownTask_) break;            // 收到关闭信号

            local_graph_task = task.base_.lock();
            if (!local_graph_task) continue;

            set_device(worker_device);

            if (task.fn_ && !local_graph_task->has_error_.load()) {
                at::ThreadLocalStateGuard tls_guard(local_graph_task->thread_locals_);
                try {
                    GraphTaskGuard guard(local_graph_task);
                    NodeGuard ndguard(task.fn_);
                    evaluate_function(
                        local_graph_task,
                        task.fn_.get(),
                        task.inputs_,
                        local_graph_task->cpu_ready_queue_);
                } catch (std::exception& e) {
                    thread_on_exception(local_graph_task, task.fn_, e);
                }
            }
        }
        --local_graph_task->outstanding_tasks_;

        if (local_graph_task->completed()) {
            local_graph_task->mark_as_completed_and_run_post_processing();
            // 唤醒 owner 线程
            ...
        }
    }
}

——这是一个经典的 producer-consumer worker loop：

1. local_ready_queue->pop()：阻塞拿一个 task
2. 检查是否 shutdown 或者 graph_task 已经无效
3. 设置 thread-local state（grad_mode、autocast 等都从 graph_task 恢复）
4. 调用 evaluate_function 执行
5. 减少 outstanding_tasks 计数
6. 如果整个 GraphTask 完成，做后处理 + 唤醒 owner

graph_task == nullptr 的情况是 device worker 线程：它常驻不退出，给所有 backward 调用服务。graph_task 非空是用户调用线程或 reentrant：跑完这个 graph_task 就退出。

engine.cpp:498-518 注释里专门解释了这两种模式：

thread_main is used by:

1. Long-running device autograd threads (graph_task is nullptr)
2. Owning thread of the backward call drives thread_main (graph_task non-null)

理解这两种模式是看懂 Engine 多线程设计的钥匙。

8.6.1 thread_locals_ 在 GraphTask 上的保留

注意 thread_main 里的这一行：

at::ThreadLocalStateGuard tls_guard(local_graph_task->thread_locals_);

GraphTask 在创建时捕获用户线程的 ThreadLocalState（grad_mode、autocast 状态、record_function 等），device worker 在执行 task 时把这些状态恢复。这让用户的 with 块语义跨线程一致：

with torch.cuda.amp.autocast():
    out = model(x)
    loss = out.sum()
loss.backward()    # autocast 上下文已被捕获到 GraphTask, device worker 也走 amp 路径

如果不做这个 capture，device worker 不知道用户开了 autocast，反向算出来的梯度精度会错乱。这种”跨线程上下文捕获”是 Engine 与 PyTorch 上下文系统协作的关键细节。

8.7 evaluate_function：真正的工作

evaluate_function 是每个 Node 真正被执行的地方。它的工作流：

sequenceDiagram
    autonumber
    participant E as evaluate_function
    participant N as Node.apply
    participant GT as GraphTask
    participant RQ as ReadyQueue

    E->>N: outputs = node.apply(inputs)
    Note over N: 真正的反向数学计算<br/>(MmBackward, AddBackward, ...)
    N-->>E: outputs
    E->>E: 对每个 next_edge[i]:
    loop 每条 next_edge
        E->>GT: 拿到 next_node = next_edge[i].function
        E->>GT: 累加 outputs[i] 到 not_ready_[next_node]
        E->>GT: dependencies_[next_node] -= 1
        alt dependencies_[next_node] == 0
            E->>GT: 从 not_ready_ 取出 InputBuffer
            E->>RQ: push NodeTask(next_node, buffer) 到 next_node 所在 device 的 RQ
        end
    end

简化的代码（来自 engine.cpp 中的 evaluate_function）：

void Engine::evaluate_function(
    std::shared_ptr<GraphTask>& graph_task,
    Node* func,
    InputBuffer& inputs,
    const std::shared_ptr<ReadyQueue>& cpu_ready_queue) {

    auto outputs = func->apply(inputs.move());     // 真的反向计算

    auto& fn_info = graph_task->exec_info_[func];
    // ...处理 captured_vars (.grad 接口)...

    int num_outputs = outputs.size();
    for (int i = 0; i < num_outputs; ++i) {
        auto& output = outputs[i];
        const auto& next = fn->next_edge(i);
        if (!next.is_valid()) continue;

        Node* next_func = next.function.get();
        std::lock_guard<std::mutex> lock(graph_task->mutex_);

        auto& not_ready = graph_task->not_ready_[next_func];
        if (!not_ready.has_value()) {
            not_ready = InputBuffer(next_func->num_inputs());
        }
        not_ready->add(next.input_nr, std::move(output), ...);

        auto& dep = graph_task->dependencies_[next_func];
        if (--dep == 0) {
            // 全部输入到齐, 入队
            auto queue = ready_queue(cpu_ready_queue, output_device);
            queue->push(NodeTask(graph_task, shared_from(next_func), not_ready->move()));
        }
    }
}

这就是反向 DAG 执行的全部 —— 每次执行一个 Node，把它的输出按边路由到上游 Node 的 InputBuffer，dependency 减到 0 时上游入队。整套调度就这一段。

8.7.1 evaluate_function 里的细节

实际的 evaluate_function 比简化版复杂得多（约 200 行），处理几个边界情况：

• anomaly mode：开启时每个 Node 执行前后都要打 stack 快照、检测 NaN/Inf
• stream sync：跨 stream 时插 event，跨 device 时插 H2D/D2H sync
• CheckpointValidGuard：activation_checkpoint 内部的反向不能再触发 checkpoint，用这个 guard 标记
• captured_vars 的 hooks：分布式相关的 deprecated hook
• leaf 检测：next_func 是 AccumulateGrad 时直接调用，不走 ReadyQueue（小优化避免 enqueue 开销）

这些细节平时用户感受不到，但每一条都对应过去几年里发现的某个 bug 或者性能问题。读这 200 行，能看到 Engine 在五六年里持续打磨的痕迹。

8.7.2 一个反直觉的优化：leaf Node 直接执行

PyTorch 在 evaluate_function 里有这么一段（简化版）：

if (next_func->is_accumulate_grad()) {
    next_func->apply(buffer.move());   // 直接执行, 不入队
} else {
    queue->push(NodeTask(...));
}

为什么 leaf 张量的 AccumulateGrad 直接同步执行？因为 AccumulateGrad 的工作是”把梯度加到张量 grad_ 字段”—— 几纳秒就完成，没必要 push 入队再 pop 出来执行（每次 push/pop 几百纳秒）。这种”特殊路径优化叶子节点”在百万次反向调用里能省可观时间。

这个优化也意味着：所有 AccumulateGrad 都在最后一个调用它的 worker 线程上执行，不在 leaf 张量原本所在的 device 线程上。但因为 AccumulateGrad 本身的 apply 实现就是”in-place 加到 grad_“，无所谓在哪个线程跑，正确性不受影响。

8.8 跨设备执行：CPU + CUDA worker 协作

考虑 MlpModule.cuda() 的反向：所有 Node 都在 CUDA:0。但 loss.cpu().backward() 时 SumBackward 在 CPU 上 —— 这是 PyTorch 必须处理的混合 device 场景。

Engine 的解法是 按 Node 输出 device 路由 task：

• 当 evaluate_function 决定上游 Node 入队时，根据上游 Node 的输入张量 device 选 ReadyQueue
• 反向链上 device 切换的地方（如 .cpu()），引擎会插入 stream sync 保证数据完整

// 简化路由逻辑
auto output_device = guess_device(next_func);
auto queue = ready_queue(cpu_ready_queue, output_device);
queue->push(NodeTask(graph_task, ...));

这种”task 跟着数据走”的设计避免了”CPU 线程提交 CUDA kernel”这种隐式同步陷阱。每个 device 的 worker 都在自己的 device 上跑 kernel，stream 内顺序天然正确，跨 stream / 跨 device 时 Engine 显式插 event 同步。

实践上，“task 跟着数据走”的对偶是”用户应该让数据集中在一个 device” —— 训练循环里反复 tensor.cpu() / .cuda() 不仅自身有同步开销，还让 Engine 在反向时被迫多次切 worker，整个 pipeline 的并发性下降。生产代码里这种”位置漂移”几乎都是 bug，不是 feature。

worker_device 是 thread-local 字段，告诉每个 worker 它服务的是哪个 device。set_device(worker_device) 在每个 task 执行前确保当前 CUDA 上下文是正确的 device —— 因为某些 CUDA API 隐式依赖”当前 device” 上下文（如 cudaMalloc），不显式 set_device 会导致 kernel 跑到错误的 GPU。这是 PyTorch 多 GPU 安全的基础。

第 17 章 DDP 章会展开多卡训练里 Engine 与 NCCL 的协作 —— DDP 通过 backward hook 在 grad 计算完成时立刻发起 AllReduce，这套机制完全依赖 Engine 的精准 device 路由。

8.8.1 跨设备的 stream 同步代价

跨 device 的反向有个隐藏代价：每次 D2H 或 H2D 拷贝都要插 CUDA event 等待对方 stream 完成。在反向 DAG 里如果 Node A (CUDA) → Node B (CPU) → Node C (CUDA)，会有两次跨 device 同步：

• A 输出在 CUDA stream 上 → 插 event 让 CPU 等 → CPU 拷贝 → CPU 跑 B
• B 输出在 CPU 上 → 拷贝回 CUDA → 插 event 让后续 stream 等 → C 跑

两次同步加起来可能几十微秒到几毫秒（取决于数据大小）。这是为什么生产代码里避免在反向路径上做 .cpu() 或者 .numpy() —— 这些操作会强制 GPU 等 CPU，让训练吞吐崩塌。

调试这种”跨 device 卡顿”的方法：用 torch.profiler.profile() 抓一段反向，看 chrome trace 里 CUDA stream 上的 idle gap —— 长 gap 通常对应跨 device 同步。

8.9 Reentrant Backward：嵌套反向

考虑这段代码：

y = f(x)
g = torch.autograd.grad(y.sum(), x, create_graph=True)[0]   # 一阶导
g2 = torch.autograd.grad(g.sum(), x)[0]                      # 二阶导（在 grad_fn 内部触发新 backward）

g2 的计算在 g 的反向 Node 里又调用了 autograd.grad —— 这就是 reentrant backward。Engine 必须支持任意深度嵌套。

但 reentrant 有个坑：TSAN 死锁检测器最多允许一个线程持有 65 个 lock（参考 engine.h:31 注释）。每个 autograd Node 的执行会持有 graph_task->mutex_，深嵌套会很快越过这个限制。PyTorch 的解法是 MAX_DEPTH = 60：

static constexpr int MAX_DEPTH = 60;

当 reentrant depth > 60 时，Engine 把这次 reentrant backward 派给一个新的 worker 线程（reentrant_thread_init），把锁深度归零。这种”按深度切换线程”的策略让 PyTorch 支持几乎无限深的嵌套求导。

实战里多数代码不会触发 reentrant，但二阶优化器（如 Newton’s method）、meta-learning（MAML）、隐函数微分等场景会用到。理解这套机制能帮你诊断”backward 在二阶导数上抛 deadlock 错误”这类难找的 bug。

8.9.1 reentrant 的两种触发方式

reentrant backward 有两条产生路径：

1. 用户级显式 reentrant（最常见）：

y = f(x)
g = torch.autograd.grad(y, x, create_graph=True)[0]   # 一阶导, create_graph 让 grad 自身可微
g.sum().backward()                                     # 二阶反向

create_graph=True 告诉 Engine “保留 grad 的 grad_fn”，让 g 仍然是可微的。g.sum().backward() 触发对一阶 backward 节点的二阶反向 —— 这就是 reentrant。

2. C++ Node::apply 内部调用反向（罕见但 PyTorch 内部用）：

某些 Node 的 apply 实现里会显式调 at::sub::call(...) 这种带 autograd 的算子，触发新的反向 —— 同样产生 reentrant。

两种路径在 Engine 看来是统一的：每次进 execute 创建新 GraphTask，reentrant_depth 加 1。超过 MAX_DEPTH 就开新线程。

8.9.5 reentrant 与显存的微妙关系

reentrant backward 还有一个不常被讨论的副作用：深嵌套时显存翻倍。每一层 reentrant backward 都需要保留一张反向图（含 SavedVariable），如果嵌套 5 层，就有 5 张反向图同时活着。

在 MAML 这种 meta-learning 场景里，inner loop 跑 5 步、outer loop 求二阶导数，每步都保留 grad —— 总显存峰值可能是普通训练的 5-10 倍。这就是为什么 MAML 实现里大量用 activation_checkpoint + 梯度累积来扣显存。

理解这条副作用，你看到”二阶导数训练显存爆炸”就不会奇怪 —— 这是 reentrant + retain_graph 的固有代价。

8.10 异常传播

某个 Node 抛异常时，Engine 必须：

1. 停止整张 GraphTask 的所有进展（has_error_ = true）
2. 把异常传回用户线程（不是吞掉、不是 abort）
3. 让其他 worker 优雅退出循环

代码（engine.cpp:573 附近）：

try {
    evaluate_function(...);
} catch (std::exception& e) {
    thread_on_exception(local_graph_task, task.fn_, e);
}

thread_on_exception 把异常存到 graph_task->future_result_ 上。owner 线程在 wait() 时检查 future 状态，如果有异常就 rethrow。这种”future 跨线程传递异常”模式让 Engine 多线程异常处理对用户完全透明 —— 用户写的就是单线程感觉的 try/except，背后 Engine 替他做了所有跨线程协调。

8.10.1 Persisting PyErr 跨线程

C++ 异常跨线程传递相对简单，但Python 异常更麻烦：Python 用 thread-local 的 PyErr_Occurred 存储异常状态，跨线程时这个状态丢失。

PyTorch 在 engine.cpp 里有专门的 “Note [Persisting PyErr state across autograd engine threads]“（约 600 多行处）。它的解法：device worker 接到 Python 异常时，把 PyObject 异常对象存到 GraphTask，owner 线程在 wait 完成后从 GraphTask 取回这个对象、用 PyErr_Restore 设置回当前线程，让 Python 异常处理路径能正常工作。

这种”跨线程 Python 异常协调”是嵌入式 Python（任何 C++ 引擎对接 Python）必须解决的问题。PyTorch 的这套实现可以作为参考样本 —— 国内一些深度学习框架在自家 Python 绑定里就直接借鉴了它。

8.10.2 anomaly_mode：诊断”哪个 forward op 引发了反向 NaN”

很多用户遇到反向 NaN 时不知道是哪个 forward op 出的问题（因为反向时 stack trace 都在 Engine 内部）。torch.autograd.set_detect_anomaly(True) 的作用就是：

• 每个 Node 创建时记录”创建栈” (Python 端的)
• Node::apply 抛异常时把这个创建栈打印出来 —— 让用户看到”是哪个 forward op 触发的反向出错”

代价是 forward 时每个 Node 都要打 stack（慢一两倍）。所以默认关闭，仅在调试时开。

8.11 与 Tokio work-stealing 的对比

如果你读过本系列《Tokio 异步运行时》，会发现 Engine 的多线程模型与 Tokio 既相似又不同：

PyTorch 不严格 work-stealing 是因为 task 有强 device 亲和性：一个 CUDA Node 必须在 CUDA worker 上跑（不然 kernel launch 会失败或同步）。Tokio 的 task 是纯 CPU 的，可以自由迁移。

这条差异决定了两个 runtime 的核心设计：Engine 优化”路由正确性”，Tokio 优化”负载均衡”。两套思想没有优劣，是各自场景的最佳选择。

8.11.1 一个混合方案：CUDA Stream + 多线程

实际上 PyTorch 在 CUDA 单 device 内部有”准 work-stealing”行为。如果一个 GraphTask 的 dependencies 里有多个独立 Node 都是 CUDA:0 的，它们都会被 push 到 CUDA:0 worker 的同一个 RQ，那个 worker 串行 pop 出来执行。但如果用户开了 multi-stream，多个 Node 的 launch 进不同 CUDA stream，CUDA driver 自己会让它们在 GPU 上并发。

所以 Engine 的”线程级并发”看起来只在多 device 场景才有意义，但其实单 device 场景的并发是通过 CUDA stream 实现的。Engine 提交 kernel launch 到 stream，stream 之间在 GPU 上并发。这是 PyTorch 多 stream 训练的工程基础。

8.11.2 与 jemalloc 对比的小启发

一个有趣的对比：jemalloc / tcmalloc 也是 per-thread cache + global pool 模型。Engine 的 per-device RQ + 全局 worker 思想异曲同工。这种”按局部性分片 + 全局协调”是高性能并发系统的通用模式 —— 任何”读多写少 + 强局部性”的工作负载都能从这套设计受益。

8.11.5 Engine 与 distributed engine 的关系

PyTorch 还有一个 torch.distributed.autograd 模块（torch/csrc/distributed/autograd/），用于 RPC 风格分布式训练 —— 跨进程的反向。它在普通 Engine 之上加了一层 DistEngine：

• 每个 RPC backward 创建一个 DistAutogradContext
• 跨进程的 Edge 通过 RPC 发回原节点
• 各进程的本地 Engine 各自跑本地子图，遇到跨进程 Edge 就 RPC 出去

DistEngine 是 PyTorch 早期分布式实验的产物（v1.4 引入），今天不是主流（DDP / FSDP 用的不是 DistEngine 而是 DDP hook + 普通 Engine）。但理解 DistEngine 的设计有助于看懂 PyTorch 怎么把”单进程引擎”扩展到”多进程协同”。第 17 章 DDP 章会展开 DDP 与普通 Engine 的协作关系，DistEngine 仅作背景知识简短提及。

8.12 一个具体场景：transformer 反向

把所有零件串起来。一个 70 层 transformer 的反向：

flowchart TB
    User["loss.backward 用户线程"]

    subgraph Setup["1. 初始化"]
        S1["GraphTask 创建"]
        S2["BFS 算 dependencies_"]
        S3["push graph_root NodeTask 到 cuda:0 RQ"]
    end

    subgraph Run["2. 多线程执行"]
        W0["CUDA:0 worker<br/>thread_main 循环<br/>处理大部分 Node"]
        WCPU["CPU worker<br/>处理 loss.cpu 之类的 Node"]
    end

    subgraph End["3. 完成"]
        E1["leaf 张量的 AccumulateGrad 跑完"]
        E2["所有 outstanding_tasks_ = 0"]
        E3["future_result_ 标记完成"]
        E4["唤醒 owner 线程"]
    end

    User --> Setup
    Setup --> Run
    Run --> End
    End --> User

    style W0 fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style WCPU fill:#dbeafe,stroke:#3b82f6

整个反向可能持续几百毫秒，期间 CUDA worker 处理 ~2000 次 NodeTask（每个 transformer 层约 30 次反向 Node 调用），每次执行一个 Node 的 apply 触发若干 ATen 算子的 dispatcher 调用 —— 整套协奏曲在用户看来就是一行 loss.backward()。

8.12.5 性能数字：Engine 自身开销

具体一些 Engine 自身的开销数字（实测 H100，PyTorch v2.11）：

对一个 70 层 transformer 反向（约 2000 次 NodeTask），Engine 自身调度开销约 2 ms。相对于反向数学计算（几十毫秒到几百毫秒）可以忽略。但对小模型反向（如 MNIST classifier），Engine 开销占比可能 30%+ —— 这是 torch.compile 在小模型上能拿到 50% 加速的部分原因（compiled autograd 把 Engine 调度开销也消掉）。

理解这些数字，你能精确估算”我的代码是 ATen kernel 慢还是 Engine 调度慢”，进而选择合适的优化方向。

8.12.6 一个总结性的”反向调用全旅程”

把整章串起来，看 loss.backward() 在 Engine 里走过的全旅程：

1. 用户线程调 loss.backward()，进 Python wrapper，下到 C++ Engine::execute(roots, inputs, ...)
2. Engine::execute 创建新 GraphTask，捕获用户线程的 ThreadLocalState
3. 触发 compute_dependencies BFS：从 graph_root 反向遍历 next_edges，填好 dependencies_ 和 nodes_in_graph_
4. 如果用户传了 inputs=...，再做一次 BFS 标记 exec_info_
5. 把 graph_root 对应的 NodeTask push 到 owner device 的 ReadyQueue
6. 用户线程调 thread_main(graph_task)，进入 worker 循环，自己也参与执行
7. 同时所有 device worker 被 condvar 唤醒，各自从对应 ReadyQueue 拿 task 执行
8. 每个 task 执行：set_device + 恢复 TLS → 调 Node::apply 算梯度 → 路由到上游 InputBuffer → dependencies— → 满足条件就入队
9. 直到所有 outstanding_tasks_ = 0，graph_task->future_result_ 标记完成
10. owner 线程被唤醒（如果在等），退出 thread_main 循环
11. 用户线程从 Engine::execute 返回，把 captured_vars（如有）作为返回值传给用户

整个过程对用户不可见，但每一步都精心设计。这就是”魔法”的真相 —— 一个精密协同的多线程系统，把数学链式法则在工程上实现到极致。

8.13 跨书关联

• 《Tokio 异步运行时》第 X 章 调度器内核：Tokio 的 inject_queue + LIFO 本地队列模型与 Engine 的 device-bound RQ 形成对照。值得看的细节是两者怎么处理”work-conservation”
• 《vLLM 内核探秘》第 6 章 Worker / Executor：vLLM 的多 Worker 执行也有类似的”按 device 路由 task”思想，但 vLLM 处理推理（无反向图）所以模型简单得多
• 《Rust 编译器之路》增量编译：编译器的依赖图遍历也用类似的 dependencies_ 计数 —— 节点等待依赖完成才能跑
• 《Serde 元编程》derive 宏的派生顺序：Serde 的 derive 也是依赖图驱动 —— 类型 A 的 Serialize 可能依赖类型 B 的 Serialize，编译器按依赖顺序处理。与 autograd Engine 的 dependencies_ 是同一思想的不同应用
• 《MCP 协议剖析》分布式状态：MCP 在多 server 场景做”待响应消息追踪”也用计数器 + 完成回调，与 outstanding_tasks_ 完全同构

8.14 几条工程经验

实战里 Engine 相关 issue 的诊断思路：

1. backward 慢，但单算子快：可能是 dependencies_ 不平衡，多线程退化成单线程。用 torch.autograd.profiler.profile 抓反向阶段每个 Node 的耗时分布

2. backward hang 不返回：检查是否 reentrant 死循环（A 反向调 B 反向、B 反向又调 A 反向）。用 py-spy dump 看用户线程 stack

3. backward 报 “RuntimeError: one of the differentiated Tensors does not require grad”：是 GraphTask 在 BFS 时发现 root 不连通到任何可求导的 leaf。检查输入张量的 requires_grad

4. 多 GPU 训练某个 backward 步骤特别慢：可能是隐式 H2D / D2H 拷贝。用 chrome trace 看每个 worker 线程的活跃时间，差异大的就是嫌疑

5. backward 显存峰值高：参考第 7 章 §7.5.4，主要是 SavedVariable 占用。Engine 自身的内存开销很小

5.1 加 set_to_none=True 但有的参数仍是 0 不是 None：因为这些参数在当前反向图里完全没收到梯度（exec_info 把它们排除了）。zero_grad 没发现它们，grad_ 保持上次的值。需要确认所有 forward 路径都覆盖了你期望的参数

6. backward 输出 grad 是 None 但前向没问题：通常是某条路径上的 Node 没有 grad_fn（如某个张量从 tensor.detach() 来）。检查反向链上每一步的 grad_fn 是否非空

7. distributed training 反向卡在某个 rank：某个 rank 的 grad 算完后立刻发起 NCCL AllReduce，如果其他 rank 的反向还没到那一步就会等。诊断方法：用 NCCL_DEBUG=INFO 看哪个 collective 在等

8. RuntimeError: Trying to backward through the graph a second time：默认 keep_graph=False，反向跑一次后图被释放，再 backward 会报这个。修法：第一次 backward 加 retain_graph=True，或重新 forward 一次

9. RuntimeError: Function XxxBackward returned an invalid gradient：通常是用户写的 autograd.Function.backward 返回了错误形状的梯度。Engine 在每个 Node 输出后会校验 grad 形状与对应输入是否匹配，不匹配就报错

8.15 一个练习：观察 backward 的多线程行为

import torch, threading, time

class TrackingMode:
    """简陋的线程跟踪 — 在 evaluate_function 调用时记录 thread id"""
    pass   # 真实实现要 hook engine, 此处仅演示思路

a = torch.randn(1000, 1000, device='cuda', requires_grad=True)
b = torch.randn(1000, 1000, device='cuda', requires_grad=True)
c = torch.randn(1000, 1000, device='cuda', requires_grad=True)

# 一个会触发多 worker 协作的反向
y = a @ b + a.cpu().sum() + c.exp().sum()
loss = y.sum() if y.dim() > 0 else y
loss.backward()

更好的实战：用 torch.autograd.profiler.profile() 跑这段代码，导出 chrome trace，能直接看到 CPU worker 与 CUDA worker 的活跃时间线。这是理解 Engine 多线程的最直观方式。

8.16 几条 Engine 设计的”通用启示”

把 Engine 思想抽象到任何”调度 + 工作”系统：

第一：任务粒度要适配调度开销 —— ReadyQueue 的 push/pop 几百纳秒。如果工作单元只值几十纳秒（如 AccumulateGrad），就要绕过队列直接执行（§8.7.2 的优化）。一般原则：工作 / 调度开销 ≥ 10x 才值得入队

第二：计数器驱动调度比扫描驱动更高效 —— dependencies_ 是 O(1) 决定”该谁跑了”，比每次扫描整张图（O(n)）找”谁的输入到齐了”快几个数量级。这种 reference counting 调度思想在编译器、垃圾回收、操作系统调度里随处可见

第三：线程亲和性优于工作迁移 —— 当工作单元有强 device 亲和性时（如 CUDA Node），按设备分区队列比让任意 worker 抢任意 task 更优。Tokio 这种 “work-stealing” 是为没有亲和性的纯 CPU task 设计的

第四：用户线程参与执行避免空转 —— 让调用 backward 的线程也变成 worker，避免它在 future 上 park 浪费一个 CPU。这种”caller is also worker”模式在 Rust async-task crate、Node.js libuv 都有体现

第五：捕获上下文跨线程恢复 —— GraphTask 把用户线程的 ThreadLocalState 捕获，让 device worker 能正确恢复 grad_mode、autocast 等状态。这是任何”工作派发到其他线程”系统的必修课

第六：异常通过 future 跨线程传递 —— 不要用 abort 或者吞掉，把异常存到 future / promise，让 owner 线程在 wait 时拿到 + rethrow

第七：栈深限制是 TSAN 等死锁检测器的现实约束 —— 60 / 65 这种数字不是任意选的，是工具链限制。设计支持递归 / 嵌套的系统时要查清楚目标平台的栈深约束

把这五条记下来，写自己的并发执行引擎能避开几乎所有大坑。

下一章拆 nn.Module —— PyTorch 用户每天都用的最基础类，背后是 Python metaclass 的精彩玩法。