https://zhuanlan.zhihu.com/p/321449610，介绍了autograd的各种能力

https://dev-discuss.pytorch.org/t/how-to-read-the-autograd-codebase/383
torch论坛的一个autograd的介绍

https://blog.ezyang.com/2019/05/pytorch-internals/:

中间这块确实太过复杂，有基本的概念就快速跳过，有需要再来单独读。
从python怎么写到kernel就行

这里有两个比较好的文章：
https://pytorch.org/blog/computational-graphs-constructed-in-pytorch/
https://pytorch.org/blog/how-computational-graphs-are-executed-in-pytorch/

Computational graph construction

这里以一个例子来串一下，看看在forward的时候，计算图是怎么构造的
以mul举例子，tensor乘tensor，对应反向的MulBackward0
代码来自26年比较新版本的pytorch重新编译的

• requires_grad=True
  • 设置后，对应调用TensorImpl::set_requires_grad

  • lazy分配AutogradMeta，包含三个核心字段：

    • Variable grad_，保存计算出来的梯度

    • shared_ptr grad_fn_，用来计算grad的函数。对应的就是tensor中的那个grad_fn

    • weak_ptr grad_accumulator，对应叶子节点，用来把多路的梯度做累计

• forward时，调用到mul_Tensor代码中

  • 会判断是否有输入的tensor设置了require grad。从而判断是否需要构造计算图

  • 如果需要计算图，则会构造一个新的node，为MulBackward0

  • grad_fn->set_next_edges()

    • 作用是把输入的tensor的AutogradMeta中的grad fn取出来。表示当前这个grad fn计算完之后，后续的计算节点

    • 会处理上面的grad_accumulator，如果是叶子节点，这里取出来的后续的节点就是grad accumulator。

  • 准备backward需要的变量。对应mul_tensor来说。两个参数self, other，如果self需要grad，则需要保存other来算。如果other需要grad，则需要保存self。

    • 保存到grad_fn中

    • 可以看MulBackward0的定义，成员变量就包含两个SavedVariable

  • 然后会做正式的前向操作：

    • at::redispatch::mul()，这里的dispatch key set做了 & after_autograd_keyset，应该是把autograd层的dispatch key去掉了，从而保证可以正常走到正式的forward
  • set_history
    • 作用就是把forward出来的tensor记录一下刚刚的grad_fn。表示输出的tensor接入到这个计算图中

然后再看一些细节的实现：

• 计算图通过Node/Edge来表示。
  • Node就是具体的grad_fn，有一个虚接口 apply()，表示计算梯度。以及edge_list，记录下游节点。（里面还有一堆hook）

  • Edge是单向边，指向forward时的上游，也就是grad时的下游

整体流程：

前向 x * y
  └─ dispatch → mul_Tensor (VariableType_0.cpp:7877)
       1. compute_requires_grad → 要建图
       2. new MulBackward0                         [Functions.h:4503 定义的节点]
       3. set_next_edges(collect_next_edges(x, y)) [连到 x/y 的上游或 AccumulateGrad]
       4. 存 self_=x, other_=y                      [SavedVariable，反向要用]
       5. redispatch::mul → 真正算出 result
       6. set_history(result, grad_fn)              [result.grad_fn = MulBackward0]

           这张图在前向结束时就已建好 ↓

反向 v.backward()
  └─ 引擎调 MulBackward0::apply(grads)              [Functions.cpp:19662]
       grad_self  = grad * other(=y)
       grad_other = grad * self(=x)
       └─ 沿 next_edges 把梯度送回 x、y 的节点/累加器

Function

上面这个是aten层生成的代码，然后来看一下自定义的autograd function，在autograd中是怎么生效的

class Exp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, i):
        result = i.exp()
        ctx.save_for_backward(result)
        return result
    @staticmethod
    def backward(ctx, grad_output):
        (result,) = ctx.saved_tensors
        return grad_output * result

• 注册阶段
  • 在定义class Exp的时候，在Function的init里面，会自动注册一个新的类，叫做ExpBackward，类型是BackwardCFunction。对应上面的MulBackward，是计算图中的节点，用来充当ctx。

  • 并且apply对应上面的backward调用，会调用用户写的backward函数。

  • 代码看BackwardCFunction即可

• 然后用户正常使用就会调用这个Exp.apply()，开始进行forward

  • 这个apply会进入到c层的THPFunction_apply，是组织这块的核心逻辑，会做：

  • unpack_input，解析输入的tensor，并且调用collect_next_edges，把输入tensor的grad fn取出来。

  • 创建ctx，对应代码中cls->_backward_cls，就是上面创建的ExpBackward

  • 创建一个PyNode，这里是对应计算图中的节点，并把ctx放进去。

  • 然后PyNode.set_next_edges()，连上上游的节点。此时autograd相关的meta就准备好了

  • 下一步是设置AutogradMode(false)，然后调用cls的forward函数。就是用户的forward

    • 没有细看实现，这里autograd mode false应该就是改一下tls中的dispatch key set什么的
  • 最后一步process_outputs
    • 调用set_history，把输出的tensor的grad fn记录上。

    • _get_tensors_to_save/_save_variables，把save_for_backward中的tensor取出来，转化成SavedVariable，保存到ExpBackward上，给反向使用

    • save_for_backward对应的是ExpBackward继承的FunctionCtx，会把save的tensor记录下来。放到to_save中，然后_get_tensors_to_save会把他取出来。

反向的时候，通过PyNode，会调用到ExpBackward的apply，然后调用到用户写的backward函数。

```C++
定义 class Exp(Function)
└─ FunctionMeta.__init__ → 自动生成 Exp._backward_cls = ExpBackward [注册]

Exp.apply(x)
└─ Function.apply → _C._FunctionBase.apply → THPFunction_apply (C++)
1. unpack_input → collect_next_edges(x) [收边：x 是叶子 → AccumulateGrad]
2. ctx = ExpBackward() [ctx 兼作节点的 Python 身体]
3. cdata = PyNode(ctx); set_next_edges(...) [通用节点连进图]
4. AutoGradMode(false) 下调 Exp.forward [算 exp，save_for_backward]
5. _wrap_outputs → set_gradient_edge [y.grad_fn = cdata(PyNode)]
6. _save_variables [saved tensor 挂到节点]

y.backward()
└─ 引擎 → PyNode::apply
to_py_args(grads) → ctx.apply → Exp.backward(ctx, grad_output)
→ 返回梯度 → 沿 next_edges → x.grad

```

Computational graph execution

然后这里来看一下backward的执行流程，简略看一下有一个基本的概念
可以在pytorch仓库内让AI解释一下这篇博客https://discuss.pytorch.org/t/what-is-the-scenario-of-reentrant-backwards-in-pytorch-source-code/19330/2，个人认为讲的还比较清楚

• Python层
  • 有tensor.backward()，对应torch.autograd.backward()

  • 有torch.autograd.grad()，平常不知道哪里会用，我看到过的就是Megatron的PP中会用这个接口

  • 两者的区别就是backward是把梯度放到了叶子节点的grad属性中，而grad()则是直接返回叶子节点的grad，不会保存在属性中。

  • 最终都会调用到：Variable._execution_engine.run_backward

• THPEngine_run_backward

  • 核心的输入就是前向的输出tensor，以及对应的梯度

  • 这里是一个wrapper，作用是根据输出的tensor，把计算图中的节点取出来。作为本次backward的根节点

  • 调用到engine.execute，执行backward。

  • 如果是grad()模式的话，会返回对应的梯度，这里会做一下打包成python对象返回出去

• 执行backward用的数据结构：

  • GraphTask，一次engine.execute()会创建一个，对应一次backward操作
    • outstanding_tasks_：还剩多少任务，减到 0 表示这次 backward 完成。atomic int

    • dependencies_：每个 Node 还有几个前驱没执行完，减到 0 就 ready。

    • not_ready_：还没凑齐输入的 Node 及其正在累加的 InputBuffer。

      • std::unordered_map<Node*, InputBuffer> not_ready_;

      • Input buffer就是维护中间节点grad的buffer

    • exec_info_：只有指定了 inputs（或用 grad()）时才填，用来剪掉不需要的路径。

    • captured_vars_：grad() 模式下临时存放要返回的梯度。

    • cpu_ready_queue_：本 GraphTask 绑定的 CPU ready queue。

  • NodeTask

    • fn_，对应的是backward的node，用来执行一个节点的grad计算

    • input_，InputBuffer，积累这个buffer对应的grad

  • ReadyQueue

    • 实现上是一个priority queue，主要是为了pytorch插入一些特殊的task。正常认为就是一个普通的queue，做拓扑序的执行就行

    • 实现上应该是做了一些启发式的改进，根据Node的sequence nr来做的排序，forward过程中，越靠后构造的grad fn，他的sequence nr越大，backward执行的顺序就越靠前

• 然后到真正的实行engine.execute()

  • 构造GraphTask

  • compute_dependencies，记录每个节点的前驱数量。就是做一次bfs。然后这里会做一些剪枝，比如用户指定了只需要部分变量的grad，其他地方的grad就不需要流动了。

  • 引擎多线程，每个设备一个线程；调用方线程（跑 Python 解释器的那个）负责 CPU；每个线程用 TLS 记 worker_device 和 local_ready_queue

• thread_main，每个worker线程执行

  • 循环 pop 任务 → 检查是否 shutdown / GraphTask 是否还有效 → evaluate_function → 减 outstanding_tasks_ → 若 completed() 则做收尾并唤醒 owner 线程

  • 执行evaluate_function

    • 执行节点 → 把输出累加进后继的 InputBuffer → 减后继依赖计数、归零就入队

    • 这里还有一些更细节的处理，比如处理exec_info，调用hook，检查NaN等

    • 执行具体的grad_fn，得到output对应的grad，然后做累加。这里会判断下游节点的input buffer是否已经存在，然后把grad累加到下游节点的input buffer中

    • 前驱没了之后，对应的node为ready，则会push到队列中，把累加好的input buffer也带过去。下一次就会用这个input buffer作为输入了

  • 这里累加有两个，input buffer和grad accumulator。input buffer处理的是一次backward，tensor被多次使用时候的grad累加。另一个则是负责把梯度累加到grad这个属性上，是跨多次backward的

  • Node的meta中会记录输入tensor的device，用来在这里选择对应device的队列。每个device一个线程可能是希望多个device做并行。

• 博客中还有可重入backward的设计，感觉场景很少就先跳过了。这里有相关的讨论：https://discuss.pytorch.org/t/what-is-the-scenario-of-reentrant-backwards-in-pytorch-source-code/19330/2

autograd的逻辑简单梳理就这样，后面来看看常用的autograd相关的工具，比如各种的hook，retrain_graph等