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

"The engine operates by having a single worker thread per work queue, and every work queue is pinned to a specific device where the operation is executed."

—— torch/csrc/autograd/engine.cpp:156-158，Note [Reentrant backwards]（该 Note 起于 torch/csrc/autograd/engine.cpp:147，v2.11.0）

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 由调用线程跑 —— 不同 device 的活儿各走各的队列
• stream 安全：同一 GPU 上多个 Node 在不同 stream 上可能并发，引擎必须正确同步
• 可重入：用户的反向规则里可能再调 torch.autograd.grad（二阶导）—— 引擎必须支持嵌套
• 错误处理：某个 Node 抛异常时，要让所有 worker 优雅停下、把异常传回用户线程

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

8.1.1 关键概念辨析

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

• 反向图 (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.2 一个直觉：为什么不能"前向倒着跑"

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

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

• 前向是 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);

    ...
  protected:
    virtual void thread_main(const std::shared_ptr<GraphTask>& task);
    ...
};

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

graph TB
    subgraph User["用户调用线程"]
        U[Python: loss.backward]
        U --> EX["Engine::execute(roots, inputs)"]
    end
    subgraph Engine["Engine 单例：device_ready_queues_ + 常驻线程"]
        WC0[CUDA:0 worker thread]
        WC1[CUDA:1 worker thread]
        WC2[CUDA:N worker thread]
    end
    subgraph Queues["每个非 CPU device 一个共享 ReadyQueue"]
        RQ_C0[CUDA:0 ReadyQueue]
        RQ_C1[CUDA:1 ReadyQueue]
        RQ_C2[CUDA:N ReadyQueue]
    end
    subgraph PerTask["每个 GraphTask 自带一个"]
        RQ_CPU["cpu_ready_queue_<br/>由调用线程自己消费"]
    end
    EX --> RQ_CPU
    U <--> 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 线程都有一个 thread-local 的"我服务哪个队列"指针（tls_local_ready_queue）。当 Engine 要给某个 Node 派活儿时，Engine::ready_queue() 按 Node 的 device() 选队列：非 CPU device 走全局共享的 device_ready_queues_[device.index()]，那个 device 的常驻线程被 condition_variable 唤醒来执行；CPU（以及所有 should_run_in_cpu_ready_queue() 为真的 device 类型）则走这次 GraphTask 自己的 cpu_ready_queue_。torch/csrc/autograd/engine.cpp:128-138 的注释说得很直白："The CUDA, XLA threads are shared among all invocations of backwards via device_ready_queues_, while the caller thread is dedicated to processing work for devices returning true in should_run_in_cpu_ready_queue"。

8.2.1 Engine 启动时机

Engine 不是 PyTorch 启动时就立刻初始化所有 worker 线程。它是 lazy 启动的 —— 第一次有 backward 调用时才创建 device worker 线程。这种 lazy 模式让"只跑 inference 的进程"完全不付这些线程的代价（线程栈 + 上下文切换）。

具体的触发点是 Engine::initialize_device_threads_pool()（torch/csrc/autograd/engine.cpp:1395），它内部用一个函数内静态变量的初始化保证只跑一次：

void Engine::initialize_device_threads_pool() {
  TORCH_CHECK(!in_bad_autograd_fork, "Unable to handle autograd's threading in combination with fork-based multiprocessing. ...");
  // Ensures device_ready_queues_ are initialized only once
  static bool start_device_threads_flag_ [[maybe_unused]] = [this]() {
    this->start_device_threads();
    return true;
  }();
}

注意开头那句 TORCH_CHECK(!in_bad_autograd_fork, ...) —— 这就是"父进程里 Engine 已经起过线程（device 线程或 reentrant 线程池线程）之后再 fork，子进程里 backward"会报错的地方：in_bad_autograd_fork 由 pthread_atfork 注册的回调在子进程里置位，而注册动作只发生在真要创建线程时（track_bad_autograd_forks()）。所以在纯 CPU 环境里，跑过普通 backward 再 fork 并不会触发它。Engine 析构（Engine::stop()）时，若各 device 队列都已空，就给每个队列 push 一个 shutdown task（ReadyQueue::pushShutdownTask()），让 device 线程 break 出 thread_main 退出，并最多等 TORCH_AUTOGRAD_SHUTDOWN_WAIT_LIMIT 秒（默认 10）；否则线程直接泄漏。

8.2.2 worker 数量为什么是 device 数量

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

一个角度是 GPU 任务本质是异步的：CUDA worker 线程的工作主要是 launch kernel，真正的计算在 GPU 上异步进行，提交线程本身并不是算力所在。

源码的做法就是 "1 worker per device"：Note [Allocating GPUs to autograd threads]（torch/csrc/autograd/engine.cpp:358）写的是 "a thread per CUDA device"，而且不同类型的设备同号共用一条线程（"colocate cuda device 0 with xla device 0"）—— start_device_threads() 取的是各类设备数的最大值，不是总和。注释自己也承认这个取舍"没有特别的理由"（"We don't have any good reason to prefer one or the other"）。每个 device 的线程串行 launch kernel，若前向用了多个 stream，这些 stream 之间在 GPU 上并发。

CPU 侧则更简单：Engine 根本不为 CPU 起线程。start_device_threads() 里那句 if (num_devices == 0) return; 意味着一台纯 CPU 机器上 Engine 一个 device 线程都不创建，反向全部在调用 backward() 的那个线程上跑完（唯一的例外是 §8.9 的 reentrant 嵌套超过 MAX_DEPTH，那时会从线程池拿或新开线程）。这也解释了为什么纯 CPU 训练时 loss.backward() 的耗时会完整地记在你的主线程上。

8.2.3 没有 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：工作单元

torch/csrc/autograd/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 方法不是简单的 +=，它会处理几种情况：

• 非加速器张量走捷径：CPU 梯度第一次到达直接放进槽位，之后的直接累加，不涉及 stream
• stream-aware 累加：加速器上的梯度要先确定"在哪条 stream 上累加"（消费者 stream、生产者 stream 或当前 stream 三选一），生产者与累加 stream 不同时插 event 等待，并记下 consumer 执行前要等的 ready_events
• 能就地累加就就地累加：accumulate() 在 grad mode 关闭、手里持有该张量的最后一个引用、且张量是稠密非重叠的普通张量时用 add_ 原地加，省一次分配；否则老老实实 old_var + var

这套累加在 torch/csrc/autograd/input_buffer.cpp 里（全文 326 行，其中 InputBuffer::add 一个函数就占了第 192–319 行）—— 看似简单的"梯度求和"在多 stream / 多 device 场景下其实是个精巧的小系统。

8.4 ReadyQueue：thread-safe 优先队列

torch/csrc/autograd/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）：

// torch/csrc/autograd/engine.h:88-105，注释与实现原文
// Returns true when t2 should be (weakly) BEFORE t1 in the queue.
// Shutdown tasks are first and then empty NodeTask are next.
struct CompareNodeTaskTime {
  bool operator()(NodeTask const& t1, NodeTask const& t2) {
    if (t2.isShutdownTask_) {
      return true;
    } else if (!t1.fn_ || t1.isShutdownTask_) {
      return false;
    } else if (!t2.fn_) {
      return true;
    } else if (t1.getReentrantDepth() == t2.getReentrantDepth()) {
      return t1.fn_->sequence_nr() < t2.fn_->sequence_nr();
    } else {
      return t1.getReentrantDepth() < t2.getReentrantDepth();
    }
  }
};

读这段要记住 std::priority_queue 的约定：比较器返回 true 表示 t1 排在 t2 后面（t2 先出队）。所以真实的优先级顺序是：

• Shutdown 任务最高优先级：Engine 析构时立即让 worker 退出；空 NodeTask 次之
• reentrant depth 更大的先跑：t1.depth < t2.depth 返回 true → depth 大的 t2 先出队。嵌套 backward 的内层先做完，外层才能继续往下走，否则外层会一直卡在等内层结果上
• sequence_nr 更大的先跑：t1.seq < t2.seq 返回 true → seq 大的先出队。sequence_nr 是前向创建顺序，越大表示越靠近 loss，启发式地先把靠近 loss 的反向做完（与第 7 章 §7.3.2 的说法一致）。AccumulateGrad 的 sequence_nr 被设成 UINT64_MAX（torch/csrc/autograd/functions/accumulate_grad.cpp:88-91），所以叶子一旦就绪总是最先执行

这是经典的 拓扑序 + 优先级 调度：普通 BFS 反向遍历用 FIFO 队列，PyTorch 换成了按上述规则排序的优先队列。在 CPU 上可以直接观察：b = a*2; c = a*3; d = (b+c).sum() 反向时，后创建的 c 上的 tensor hook 先于 b 上的触发（CPU torch 2.11.0 实跑）。

注意 ReadyQueue 用的是 mutex + condvar 保护的 std::priority_queue，不是 lockfree 队列。每个 Node 恰好一次 push、一次 pop，这把锁的开销是按 Node 计的固定成本 —— 它是否显著取决于单个 Node 的计算量，见 §8.13。

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（torch/csrc/autograd/graph_task.h:18，节选，字段顺序有调整）：

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<Variable> captured_vars_;              // .grad() 接口的输出
    std::shared_ptr<ReadyQueue> cpu_ready_queue_;       // 这次 backward 的 CPU 队列
    int owner_{NO_DEVICE};                              // 启动这次 backward 的线程 device
    const int reentrant_depth_;                         // 嵌套在几层父 GraphTask 之内
    c10::intrusive_ptr<at::ivalue::Future> future_result_; // 完成信号
    ...
};

理解几个核心字段：

8.5.1 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.2 dependencies_：每个 Node 还要等多少条边

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

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

具体的遍历在 Engine::compute_dependencies(Node* root, GraphTask& task, uint64_t min_topo_nr)（torch/csrc/autograd/engine.cpp:1248）：用一个 std::vector 当栈（back() / pop_back()，所以严格说是深度优先的顺序，但顺序不影响计数结果），从 graph_root 出发遍历所有可达 Node，每条边让目标 Node 的 dependencies 计数 +1，同时填好 nodes_in_graph_。注意第三个参数 min_topo_nr —— 它由 compute_min_topological_nr(outputs) 算出（所有目标 output 边里最小的 topological_nr），遍历时 topological_nr 比它还小的节点不可能通向任何目标，直接剪掉。这就是第 7 章 §7.3.2 那个拓扑号在引擎侧的实际用途。复杂度是 O(节点数 + 边数)，相对反向本身的浮点计算是可忽略的一项。

值得注意：这趟遍历每次 backward 都重做。eager 模式下每次前向都会新建一张反向图（Node 都是新对象，上一轮的图默认在 backward 后已释放），dependencies_ 以 Node* 为键，本来就无从跨迭代复用。torch.compile 的做法是减少要遍历的节点：AOTAutograd 把整段编译区域的反向合成一个 CompiledFunctionBackward 节点（CPU torch 2.11.0 + backend="aot_eager" 实跑可见），Engine 看到的图随之小得多；即便开了 compiled autograd，Engine::execute 也是先跑完 compute_dependencies（torch/csrc/autograd/engine.cpp:1345）再把控制权交给编译后的反向（torch/csrc/autograd/engine.cpp:1352-1364）。

8.5.3 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.4 not_ready_：还没集齐输入的 Node 的 InputBuffer 缓存

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

8.5.5 exec_info_：子图过滤

torch.autograd.grad(loss, [w1, w2]) 让用户只对部分 leaf 求导。这种"只跑某些路径"的需求由 exec_info_ 处理 —— 在初始化阶段由 GraphTask::init_to_execute（torch/csrc/autograd/engine.cpp:1664）从 root 出发做一趟迭代式深度优先遍历，标记"哪些 Node 在通往目标的路径上"（needed_），执行时只跑标记过的 Node。这种 mask 让"对部分参数求导"不浪费时间在无关 Node 上。按 Note [Exec info]，exec_info_ 只有在".backward() 且不传 inputs"时为空；torch.autograd.grad 总会填它。

8.5.6 exec_info_ 在 backward(inputs=...) 上的应用

torch.autograd.backward(loss, inputs=[w1, w2]) 是另一种触发 exec_info 的接口：你只想累积梯度到 w1 和 w2，不动其他参数。Engine 会预先做同样的标记遍历，找出"哪些 Node 在通往 w1/w2 的路径上"，执行时跳过非标记 Node。与 .grad() 的区别在于：backward(inputs=...) 把目标节点的 needed_ 直接置真、梯度照常累加进 .grad；.grad() 则在目标节点上登记 captures_，梯度被截下来返回（GraphTask::init_to_execute 里的 accumulate_grad 分支）。

被跳过的只是不通往 w1/w2 的分支。要注意别高估省下的量：比如各层都挂着 adapter 的 LoRA，梯度仍要穿过每一层才能到达较浅层的 adapter，主干反向省不掉；实践中冻结参数通常直接设 requires_grad=False，让它们根本不进反向图。

代价是初始化阶段要多做一趟遍历：init_to_execute 是在 compute_dependencies 之后单独跑的（torch/csrc/autograd/engine.cpp:1345-1350），同样是 O(节点数 + 边数) 的纯指针遍历，相对反向本身的计算通常可以忽略。

8.5.7 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 主循环

torch/csrc/autograd/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（autocast、profiler 的 RecordFunction 状态、saved tensor hooks 等从 graph_task 恢复；grad_mode 是 GraphTask 构造时按 create_graph 设好的）
4. 调用 evaluate_function 执行
5. 减少 outstanding_tasks 计数
6. 如果整个 GraphTask 完成，做后处理 + 唤醒 owner

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

torch/csrc/autograd/engine.cpp:502-517 的注释原文列的其实是三种用法：

thread_main is used by: 1). autograd threads for devices (i.e. CUDA, XLA) 2). the caller/owning thread of the backward call on CPU (sync mode) 3). Renetrant backward that invoked by either 1) or 2)

——第 1 种在整个 Engine 生命周期里自旋，只有 Engine 析构 push shutdown task 时才终止；第 2 种同步地跑到自己那个 graph_task 完成为止；第 3 种（reentrant）同样是"完成即退出、并通知 owner 线程"，见 §8.9。（注释里的 "Renetrant" 是源码里的一处拼写笔误，照录。）

8.6.1 thread_locals_ 在 GraphTask 上的保留

注意 thread_main 里的这一行：

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

GraphTask 在创建时（也就是调用 backward() 的那一刻，在 Engine::execute 里）捕获用户线程的 ThreadLocalState（dispatch key set、autocast dtype、RecordFunction、saved tensor hooks 等，见 aten/src/ATen/ThreadLocalState.h），grad_mode 则被强制设成 create_graph 的值（GraphTask 构造函数里的 thread_locals_.set_grad_mode(grad_mode)）。device worker 在执行 task 时把这些状态恢复，让包在 backward() 外面的 with 块语义跨线程一致：

with torch.profiler.profile() as prof:
    loss.backward()    # RecordFunction 状态被捕获到 GraphTask，device worker 上跑的 Node 也会被记录

注意捕获的是 backward 调用时的状态，不是前向时的。常见的 AMP 写法把 loss.backward() 放在 autocast 块外面，这时反向里 autocast 是关着的；在 CPU torch 2.11.0 上用一个打印 torch.is_autocast_enabled('cpu') 的自定义 Function.backward 实测：backward 在 autocast 块外调用时打印 False，块内调用时打印 True；grad_enabled 在普通 backward 里为 False、create_graph=True 时为 True。这种"跨线程上下文捕获"是 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

简化的代码（来自 torch/csrc/autograd/engine.cpp:1063-1233 的 evaluate_function；省略了开头的 stream 等待、exec_info / captures 处理、anomaly 检查和 exec_info 过滤）：

void Engine::evaluate_function(
    std::shared_ptr<GraphTask>& graph_task,
    Node* func,
    InputBuffer& inputs,
    const std::shared_ptr<ReadyQueue>& cpu_ready_queue) {
    // ...stream 等待、exec_info 过滤与 captures（.grad 接口）...

    auto outputs = call_function(graph_task, func, inputs);  // hooks + 真正的反向计算 + validate_outputs

    auto& fn = *func;
    if (!graph_task->keep_graph_) {
        fn.release_variables();                  // retain_graph=False：释放 SavedVariable
    }

    auto num_outputs = outputs.size();
    // ...anomaly mode 下的 NaN 检查...

    std::lock_guard<std::mutex> lock(graph_task->mutex_);   // 整个循环在这把锁里
    for (const auto i : c10::irange(num_outputs)) {
        auto& output = outputs[i];
        const auto& next = fn.next_edge(i);
        if (!next.is_valid()) continue;

        bool is_ready = false;
        auto it = graph_task->dependencies_.find(next.function.get());
        if (--it->second == 0) {                 // 最后一条边到了
            graph_task->dependencies_.erase(it);
            is_ready = true;
        }

        auto not_ready_it = graph_task->not_ready_.find(next.function.get());
        if (not_ready_it == graph_task->not_ready_.end()) {
            InputBuffer input_buffer(next.function->num_inputs());
            input_buffer.add(next.input_nr, std::move(output), ...);
            if (is_ready) {                      // 只有一条入边：直接入队
                auto queue = ready_queue(cpu_ready_queue, next.function->device());
                queue->push(NodeTask(graph_task, next.function, std::move(input_buffer)));
            } else {                             // 先存进 not_ready_
                graph_task->not_ready_.emplace(next.function.get(), std::move(input_buffer));
            }
        } else {
            auto& input_buffer = not_ready_it->second;
            input_buffer.add(next.input_nr, std::move(output), ...);
            if (is_ready) {                      // 全部输入到齐，入队
                auto queue = ready_queue(cpu_ready_queue, next.function->device());
                queue->push(NodeTask(graph_task, next.function, std::move(input_buffer)));
                graph_task->not_ready_.erase(not_ready_it);
            }
        }
    }
}

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

8.7.1 evaluate_function 里的细节

实际的 evaluate_function（约 170 行，另有它调用的 call_function）比简化版复杂，处理几个边界情况：

• anomaly mode：开启且 check_nan 为真时，Node 跑完后对每个输出做一次 NaN 检查（只查 NaN，不查 Inf）；前向的"创建栈"是建 Node 时就记下的，出错时由 GraphTask::set_exception_without_signal 打印，见 §8.10.2
• stream sync：执行前，对每个来自加速器的输入，若它 ready 的 stream 与本 Node 的 stream 不同，就让本 Node 的 stream 等对应 event。跨 device 的拷贝本身不在这里做，而是由 .to() / .cpu() 对应的反向节点（如 ToCopyBackward0）完成
• CheckpointValidGuard（在 call_function 里）：记录"当前嵌套的每一层 backward 是不是都是不带 inputs 的 .backward()"（can_checkpoint() 即 exec_info_.empty()）。use_reentrant=True 的 checkpoint 在反向里据此判断，若为假就报错 "When use_reentrant=True, torch.utils.checkpoint is incompatible with .grad() or passing an `inputs` parameter to .backward()"（torch/utils/checkpoint.py:273-279）
• retain_graph：keep_graph_ 为假时，call_function 在调用 Node 前先 will_release_variables()，evaluate_function 在 Node 跑完后立刻 release_variables() 释放它的 SavedVariable。第二次 backward 解包已释放的 SavedVariable 就会报 §8.16 第 8 条那个错。retain_graph 只影响释放，不影响调度
• captured_vars 的 hooks：分布式相关的 deprecated hook

这些细节平时用户感受不到，但正是它们让这段调度在多 stream、hook、checkpoint 等场景下依然正确。

8.7.2 一个容易记错的点：AccumulateGrad 没有"直接执行"的快路径

有一种常见的误解是：既然 AccumulateGrad 的工作只是"把梯度加到张量的 grad_ 字段"，Engine 应该给它开一条不入队的快路径。v2.11 的 evaluate_function 里没有这样的分支。把上面那段代码对着 torch/csrc/autograd/engine.cpp:1161-1233 读一遍就能看到：无论上游是 AccumulateGrad 还是普通的 XxxBackward，只要 --it->second == 0，走的都是同一句

auto queue = ready_queue(cpu_ready_queue, next.function->device());
queue->push(NodeTask(graph_task, next.function, std::move(input_buffer)));

Engine 唯一"绕过 device 队列"的场合在 Engine::ready_queue() 里，条件也和叶子无关 —— 是 multithreading_enabled 被关掉、或者 should_run_in_cpu_ready_queue(device.type()) 成立时，直接返回调用方传进来的 cpu_ready_queue。换句话说：决定去哪个队列的是 device，不是节点类型。

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

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

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

• 当 evaluate_function 决定上游 Node 入队时，根据上游 Node 的 device() 选 ReadyQueue：Node::device()（torch/csrc/autograd/function.h:261）取它输入元数据里第一个非 CPU 的 device，全是 CPU 才算 CPU
• 反向链上 device 切换的地方（如 .cpu()），引擎会插入 stream sync 保证数据完整

// engine.cpp，evaluate_function 里真实的路由两行
auto queue = ready_queue(cpu_ready_queue, next.function->device());
queue->push(NodeTask(graph_task, next.function, std::move(input_buffer)));

路由依据是 上游 Node 自己的 device()（Node 在构造时从输入元数据里记下来的），而不是当前这条梯度张量的 device。

这种"task 跟着数据走"的设计让每个 device 的 worker 都在自己的 device 上提交 kernel，stream 内顺序天然正确，跨 stream 时 Engine 显式插 event 同步（Note [Streaming backwards]，torch/csrc/autograd/engine.cpp:177）。它是一种调度约定，不是硬性前提：torch.autograd.set_multithreading_enabled(False) 时，Engine::ready_queue() 对任何 device 都返回调用线程的 cpu_ready_queue，CUDA Node 也照样在调用线程上跑。

实践上，"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 在每个参数的 AccumulateGrad 上挂 post hook（torch/csrc/distributed/c10d/reducer.cpp:182-199），一个 bucket 里的梯度全部就绪就发起这个 bucket 的 AllReduce，从而让通信与剩余的反向计算重叠。

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 跑

每一次这样的跨越都是一次真实的 stream 等待加一次 PCIe 拷贝，而且它卡在反向的关键路径上 —— GPU 那边的后续 kernel 只能干等。这是为什么生产代码里避免在反向路径上做 .cpu() 或者 .numpy()。

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

8.9 Reentrant Backward：嵌套反向

考虑这段代码：

import torch

class Sin(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.sin()

    @staticmethod
    def backward(ctx, grad_out):
        x, = ctx.saved_tensors
        with torch.enable_grad():
            x_ = x.detach().requires_grad_()
            y = x_.sin()                                   # 在反向里重建一小段前向
            gx, = torch.autograd.grad(y, x_, grad_out)     # 反向函数内部再启动一次 backward
        return gx

x = torch.randn(3, requires_grad=True)
Sin.apply(x).sum().backward()
print(torch.allclose(x.grad, x.cos()))   # True（CPU torch 2.11.0 实跑）

Sin.backward 是 Engine 正在执行的一个 Node，它在里面又调用了 torch.autograd.grad —— 这就是 reentrant backward：外层 GraphTask 还没跑完，同一个线程又进了一次 Engine::execute。use_reentrant=True 的 torch.utils.checkpoint 就是这种结构（CheckpointFunction.backward 里重算前向后调 torch.autograd.backward，torch/utils/checkpoint.py:325）。Engine 必须支持任意深度嵌套。

但 reentrant 有个坑，源码注释（torch/csrc/autograd/engine.h:31-36）写得很清楚：TSAN 的死锁检测器在一个线程同时持有超过 65 个 lock 时会报错（注释里给的出处是 google/sanitizers issue #950）。而 PyTorch 在每个自定义 C++ autograd Node 上都持一把 mutex，深嵌套很快就会越过这个限制。于是有了 MAX_DEPTH = 60（torch/csrc/autograd/engine.h:37）：

static constexpr int MAX_DEPTH = 60;

当本线程上的嵌套层数 current_depth 达到这个上限（current_depth >= max_recursion_depth_，torch/csrc/autograd/engine.cpp:1453）时，Engine 把这次 reentrant backward 交给线程池里的另一个线程跑（没有空闲线程时 std::thread t(&Engine::reentrant_thread_init, this) 新开一个，见 torch/csrc/autograd/engine.cpp:1636），当前线程在 Engine::execute 末尾的 fut->wait() 上阻塞等它；新线程的 current_depth 从 0 开始，锁深度就此归零。这种"按深度切换线程"的策略让 PyTorch 支持几乎无限深的嵌套。在 CPU torch 2.11.0 上把上面这种"反向里再 grad"嵌套 131 层实测：前 61 层（最外层的普通 backward 加 60 层 reentrant）都在主线程上执行，之后每 61 层换一条新线程 —— 纯 CPU 进程也会因此起线程。

实战里多数代码不会触发 reentrant。最常见的来源是 use_reentrant=True 的 activation checkpoint，以及自己在 autograd.Function.backward 里调 torch.autograd.grad 的代码。理解这套机制能帮你读懂"反向里突然多出几条线程"这类现象。

8.9.1 什么算 reentrant，什么不算

判断标准在 Engine::execute 里只有一行：bool not_reentrant_backward_call = worker_device == NO_DEVICE;（torch/csrc/autograd/engine.cpp:1321）。worker_device 是线程局部变量，只有在线程正在驱动某个 GraphTask 时才不是 NO_DEVICE（调用线程在 execute_with_graph_task 里把它设成 CPU_DEVICE，跑完再复位；device 线程在创建时就设成自己的设备号）。所以：

1. 在某个 Node 的反向函数里调 backward / grad：此时线程正在跑外层 GraphTask，worker_device != NO_DEVICE，新 GraphTask 的 reentrant_depth_ = 外层深度 + 1 —— 这是 reentrant。

2. 顶层的高阶导数不是 reentrant：

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

第一次 grad 返回时，外层 GraphTask 已经结束、worker_device 已复位为 NO_DEVICE；第二次 grad 是一次全新的顶层 execute，reentrant_depth_ 为 0。create_graph=True 的作用只是让第一次反向在 grad mode 下执行、给 g 建出自己的反向图，与 reentrant 无关。

8.9.2 别把高阶求导的显存账算到 reentrant 头上

"二阶导数训练显存爆炸"是真实现象，但根源是 create_graph=True，不是 reentrant。create_graph=True 的一阶反向会给梯度本身建反向图，而且 retain_graph 默认跟随 create_graph（torch/autograd/__init__.py:375-376），原图的 SavedVariable 也不会被释放。

在 MAML 这种 meta-learning 场景里，inner loop 每跑一步都要保留一张带 create_graph=True 的图，outer loop 再对它们整体求二阶导 —— 显存开销随 inner loop 步数线性叠加。

8.10 异常传播

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

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

代码（torch/csrc/autograd/engine.cpp:573 附近）：

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

thread_on_exception 经 GraphTask::set_exception 置 has_error_，并把异常设进 graph_task->future_result_。owner 线程在 Engine::execute 末尾 fut->wait() 等到完成，接着取 fut->value() 时，Future::value() 发现有异常就 rethrow（aten/src/ATen/core/ivalue_inl.h:1014-1019）。这种"future 跨线程传递异常"模式让 Engine 多线程异常处理对用户完全透明 —— 用户写的就是单线程感觉的 try/except，背后 Engine 替他做了所有跨线程协调。

8.10.1 Persisting PyErr 跨线程

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

PyTorch 为此专门写了 Note [ Persisting PyErr state across autograd engine threads ]，正文在 torch/csrc/autograd/python_engine.cpp:399，引用点分布在 torch/csrc/autograd/engine.cpp:581、torch/csrc/autograd/python_engine.cpp:96 与 torch/csrc/autograd/python_function.cpp:84。它的解法：worker 线程接到 Python 异常时先 persist() 把异常对象保存下来（PythonEngine::thread_on_exception），异常随 future 带回 owner 线程后，PythonEngine::execute 的 catch (python_error& e) 再 e.restore() 把它恢复到当前线程，让 Python 侧的异常处理路径能正常工作。

这种"跨线程 Python 异常协调"是嵌入式 Python（任何 C++ 引擎对接 Python）必须解决的问题，PyTorch 的这套实现可以作为参考样本。

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

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

• 每个 Node 创建时记录"创建栈" (Python 端的)
• 反向中某个 Node 出错时，GraphTask::set_exception_without_signal 把这个创建栈打印出来 —— 让用户看到"是哪个 forward op 触发的反向出错"。CPU torch 2.11.0 实测，一个故意返回 NaN 的自定义 Function 会先得到一条 UserWarning: Error detected in BadBackward. Traceback of forward call that caused the error: 加前向调用栈，随后抛出 RuntimeError: Function 'BadBackward' returned nan values in its 0th output.；返回 Inf 则不会被拦下

代价是 forward 时每个 Node 都要抓一次 Python 调用栈，反向每步还要做 NaN 检查（engine.cpp 里那句 isnan(output)._is_any_true().item<bool>() 是个会同步的 device→host 读取）—— 前向变慢、反向多一轮同步。所以默认关闭，仅在调试时开。

8.11 与 Tokio work-stealing 的对比

如果你读过《Tokio 源码深度解析》专栏（尤其是第 5 章多线程 Scheduler），会发现 Engine 的多线程模型与 Tokio 既相似又不同：

维度	PyTorch Engine	Tokio Runtime
工作单元	NodeTask	Task (Future)
队列	per-device ReadyQueue (priority)	per-worker 本地环形队列 + LIFO slot + 全局 inject 队列
调度	task 跟随 device 路由	worker 偷其他 worker 的队列
阻塞点	mutex + condvar	atomic + park
线程数	每个非 CPU device 一个常驻线程；CPU 部分由调用线程兼任	等于 CPU 核心数（可配）
异常传播	future_result_ + has_error_	JoinHandle 返回的 Result

PyTorch 不做 work-stealing，是因为它按 task 的 device 亲和性把 task 路由到固定队列：CUDA Node 进对应 device 的队列，由那个 device 的线程消费，别的线程不会来偷。这是调度约定而不是硬性限制 —— 如 §8.8 所说，关掉 multithreading 后 CUDA Node 也能在调用线程上跑。Tokio 的 task 没有这种亲和性，可以自由迁移。

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

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

严格说 PyTorch 在单 device 内部没有任何 work-stealing。如果一个 GraphTask 的 dependencies 里有多个互相独立的 Node 都在 CUDA:0，它们会被 push 到同一个 device_ready_queues_[0]，由那个 device 的常驻线程串行 pop 出来执行（reentrant 超深时接手的线程池线程也消费同一个队列，但此时原线程正阻塞等待）—— 线程这一层是纯串行的。真正的并发发生在下一层：这条线程只是把 kernel launch 提交出去，多个 Node 若落在不同 CUDA stream 上，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 按 device 分队列是为了把 task 路由到对应设备的线程。

8.11.3 Engine 与 distributed engine 的关系

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

• 分布式的前向和反向都包在同一个 DistAutogradContext 里（用户用 with dist_autograd.context() as context_id: 创建）
• 跨进程的 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["遍历图算 dependencies_"]
        S3["push graph_root NodeTask 到 cuda:0 RQ"]
    end
    subgraph Run["2. 多线程执行"]
        W0["CUDA:0 worker<br/>thread_main 循环<br/>处理大部分 Node"]
        WCPU["调用线程本身<br/>消费 cpu_ready_queue_<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

层数乘上每层的反向 Node 数，一次反向要跑的 NodeTask 是上千量级；每个 NodeTask 的 apply 里又是若干次 ATen 算子的 dispatcher 调用 —— 整套协奏曲在用户看来就是一行 loss.backward()。

8.13 Engine 自身开销由什么构成

与其背几个数字，不如把 Engine 每跑一个 Node 要付的固定成本列清楚 —— 这几项都能在前面的源码里逐条对上：

成本项	来源
一次 ReadyQueue::push / pop	std::mutex 加锁 + std::priority_queue 的 O(log n) 堆调整 + condition_variable 唤醒
dependencies_ / not_ready_ 的哈希表读写	都在 graph_task->mutex_ 这把 GraphTask 级、所有 worker 共用的锁内（evaluate_function 里那句 std::lock_guard<std::mutex> lock(graph_task->mutex_)）
GraphTask 创建（每次 backward 一次，不按 Node 计）	一次 compute_dependencies 全图遍历 + 捕获 ThreadLocalState
跨 stream 同步	evaluate_function 开头那段 opt_parent_stream->wait(event)
每个 Node 的 TLS 恢复	at::ThreadLocalStateGuard

关键的定性判断有两条。第一，这些成本是按 Node 计的固定开销，与张量多大无关；所以算子越小、Node 越多，Engine 调度在总时间里占比越高 —— 大模型的大 GEMM 反向里它可以忽略，小模型、小 batch、碎算子多的反向里它就会浮上来。第二，graph_task->mutex_ 是一把跨 worker 共享的锁，Node 数量越多、worker 越多，锁竞争越明显。compiled autograd（§7.14）之所以对小算子密集的反向收益大，正是因为它把这两条成本一起编译掉了。

想知道自己的模型落在哪一侧，唯一可靠的办法是 profile：抓一段反向的 chrome trace，看 worker 线程上"在 kernel 里"和"在队列/锁上"的时间比例。

8.14 一个总结性的"反向调用全旅程"

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

1. 用户线程调 loss.backward()，进 Python wrapper，下到 C++ Engine::execute(roots, inputs, ...)
2. Engine::execute 创建新 GraphTask，捕获用户线程的 ThreadLocalState
3. 触发 compute_dependencies 遍历：从 graph_root 沿 next_edges 遍历，填好 dependencies_ 和 nodes_in_graph_
4. 如果用户传了 inputs=...（torch.autograd.grad 总是如此），再单独遍历一次标记 exec_info_
5. 把 graph_root 对应的 NodeTask push 到 graph_root 自己 device 对应的 ReadyQueue（CPU 就是本次的 cpu_ready_queue_）
6. 用户线程调 thread_main(graph_task)，进入 worker 循环，自己也参与执行
7. 有 task 入队的 device 队列上，对应的 device worker 被 condvar 唤醒，从队列拿 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.15 跨专栏关联

• 《Tokio 源码深度解析》专栏第 5 章（多线程 Scheduler 与工作窃取）：Tokio 的全局 inject 队列 + 每 worker 本地环形队列（外加 LIFO slot）模型与 Engine 的 device-bound RQ 形成对照。值得看的细节是两者怎么处理"work-conservation"
• 《vLLM 推理内核深度解析》专栏第 6 章（Worker / Executor）：vLLM 的多 Worker 同样是每个进程绑一张卡，但它是"一个 writer 广播、所有 Worker 执行同一份 SchedulerOutput"的模式，没有 Engine 这种按 device 路由单个 task 的调度，模型简单得多
• 《Rust 编译器与运行时揭秘》专栏第 17 章（增量编译）：rustc 也维护一张依赖图，但用法正好相反 —— 它按需递归求值、用红绿算法判断缓存能否复用，而 Engine 用 dependencies_ 计数从 root 正向推动执行，两相对照很有意思
• 《MCP 协议设计与实现》专栏第 3 章（JSON-RPC 与消息格式）：MCP SDK 用 _responseHandlers 按消息 ID 记下每个待响应请求、收到响应再回调；Engine 的 outstanding_tasks_ 只记未完成任务的数量、不记身份 —— 都是"追踪未完成的工作"，粒度不同

8.16 几条工程经验

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

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

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

3. backward 报 "One of the differentiated Tensors does not require grad"：这条来自 torch/csrc/autograd/python_engine.cpp:335，触发点是你传给 torch.autograd.grad(outputs, inputs) 的某个 inputs 张量本身 requires_grad=False。检查那批输入张量的 requires_grad。同样的错误出在 backward(inputs=[...]) 上时，文案是另一条：can't retain_grad on Tensor that has requires_grad=False —— 因为 backward 路径会先对每个 input 调 retain_grad()（torch/csrc/autograd/python_engine.cpp:330-332，CPU torch 2.11.0 实跑两种文案均已确认）

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

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

5.1 zero_grad(set_to_none=True) 之后，反向结束时有的参数 .grad 还是 None：说明这次反向根本没跑到它的 AccumulateGrad —— 它不在这次的反向图里，或被 inputs= 排除了。若 .grad 是全 0 张量而不是 None，则说明它在图里、只是收到的梯度恰好为零。需要确认所有 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：retain_graph 默认跟随 create_graph（即通常为 False，传进 Engine 就是 keep_graph），反向跑一次后各 Node 的 SavedVariable 被 release_variables() 释放（§8.7.1），再 backward 解包它们就报这个（文案常量是 ERR_BACKWARD_TWICE，torch/csrc/autograd/saved_variable.h:18）。修法：第一次 backward 加 retain_graph=True，或重新 forward 一次

9. RuntimeError: Function XxxBackward returned an invalid gradient at index N - expected dtype/layout/device ...：通常是用户写的 autograd.Function.backward 返回了与对应输入对不上的梯度。校验在 engine.cpp 的 validate_outputs_impl 里，逐条比对梯度个数（文案是 "invalid number of gradients - expected " 后面拼上期望个数与实得个数）、dtype、layout、device，形状则先交给 metadata.maybe_reduce()（torch/csrc/autograd/engine.cpp:941）尝试 sum_to 归约；报错文案由 torch/csrc/autograd/engine.cpp:1055 那个 ostringstream 拼出："Function " << fn.name() << " returned an " << msg

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

import torch

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()

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

8.18 几条 Engine 设计的"通用启示"

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

第一：任务粒度要适配调度开销 —— 每次入队都要付一把锁、一次堆调整、一次条件变量唤醒。当单个工作单元本身的成本已经和这套调度开销同量级时，队列就从"并发手段"变成"净损耗"。PyTorch 自己的答案不是给小节点开后门（§8.7.2 说明了它并没有这么做），而是往上一层解决：用 torch.compile / compiled autograd 把一串小 Node 编成一段代码，从根上减少任务数

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

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

第四：用户线程参与执行避免空转 —— 让调用 backward 的线程也变成 worker（负责 CPU 部分），而不是单纯在 future 上等待

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

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

第七：嵌套深度上限来自 TSAN 等死锁检测器的现实约束 —— 60 / 65 这种数字不是任意选的，是工具链对"单线程同时持锁数"的限制。设计支持递归 / 嵌套的系统时要查清楚这类约束

把这七条记下来，写自己的并发执行引擎时能少踩不少坑。

下一章拆 nn.Module —— PyTorch 用户每天都用的最基础类，背后是 __setattr__ / __getattr__ 拦截这套 Python 元编程（nn.Module 本身是个普通类，没有元类）。