OneFlow源码解析：自动微分机制

撰文 | 郑建华
更新｜赵露阳、王迎港

深度学习框架一般通过自动微分（autograd）机制计算梯度并反向传播。本文尝试通过一个简单的例子，粗浅地观察一下OneFlow的autograd的实现机制。

1 自动微分基础

自动微分相关的资料比较多，个人感觉自动微分的原理介绍(https://mp.weixin.qq.com/s/Bw...)这个系列及其引用的资料对相关背景知识的介绍比较完整清晰。

下面分几种情况对梯度传播的原理做一些直观解释。

1.1 stack网络的梯度传播

以x -> f -> g -> z这个stack网络为例，根据链式法则：

∂z/∂x = ∂z/∂g * ∂g/∂f * ∂f/∂x

实际运行时，在梯度反向传播过程中：

z将∂z/∂g传给g。
如果节点g有权重w需要计算梯度，就计算∂z/∂w = ∂z/∂g * ∂g/∂w。
g需要计算∂g/∂f，再乘以z传过来的梯度，将结果传给f。g只需要给f传递链式乘积的结果，不需要传递各项明细。
在训练阶段的前向计算时，g需要保存∂g/∂f计算依赖的中间结果、以供反向计算时使用。
其它节点的传播情况依次类推。

1.2 简单graph的梯度传播

以下面这个简单的graph拓扑为例。

在继续之前，需要了解一下多元复合函数微分的基本公式。

下图中，u和v都是关于x和y的函数，z是关于u和v的函数。

根据这个公式可以知道，z对x的梯度分别沿两条链路传播，z -> u -> x和z -> v -> x，节点x将两个梯度之和作为z对x的梯度。

1.3 复杂graph的梯度传播

再看一个拓扑稍微复杂点的例子：

上图可以视为x -> U -> L，其中U是e -> ... -> h的子图。f -> g的子图可以视为V。

对于节点h来说，它需要把梯度传给g和k。对节点e来说，它需要对f和k传来的梯度求和，才是∂L/∂e。这样，L对x的梯度，仍可以按链路拆解，一条链路前后节点间的梯度是乘积关系，传入的多条链路梯度是加和关系。
这篇博客(https://blog.paperspace.com/p...)中有一个几乎一样的拓扑图，给出了部分权重参数的梯度公式。

2 autograd中tensor相关的一些基本概念

2.1 叶子节点

OneFlow的autograd文档(https://docs.oneflow.org/en/m...)中介绍了leaf node和root node的概念。只有输出、没有输入的是leaf node，只有输入、没有输出的是root node。

个人理解，如果把weight、bias、data视为计算图的一部分，这些节点就是叶子节点（op不是叶子节点）。尤其是从反向计算图的视角(https://discuss.pytorch.org/t...)看，这些节点的grad_fn是空，反向传播到这些节点就会停止。

is_leaf和requires_grad有比较密切的关系，但二者又是独立的。PyTorch是这样解释的：(https://pytorch.org/docs/stab...)

requires_grad=false的节点都是叶子节点。比如data。
requires_grad=true的节点如果是用户创建的，也是叶子节点。比如weight和bias。
在梯度的反向计算过程中，只有叶子节点的梯度才会被填充。对于非叶子节点，如果要填充梯度信息，需要显式设置retain_grad=true。
requires_grad=true才会计算、填充梯度。比如y = relu(x)，y是op创建的、不是叶子节点。但如果x需要计算梯度，则y.requires_grad==true。但不需要为y填充梯度。

关于叶子节点这个概念，目前找到的主要是直观描述，还没看到严格、清晰的定义。也可能是因为用户一般不会直接使用is_leaf(https://discuss.pytorch.org/t...)，这个概念只是在阅读代码的时候才会涉及到。

下面的资料可以供进一步参考：

What is the purpose of is_leaf? (https://discuss.pytorch.org/t...)
叶子节点和tensor的requires_grad参数（https://zhuanlan.zhihu.com/p/...）

2.2 tensor detach

Tensor的detach方法（https://github.com/Oneflow-In...）会创建一个新的tensor，新tensor的属性中

requires_grad = false
is_leaf = true

detach的意思是从grad的反向计算图中把tensor分离出来。新的tensor与原来的对象共享存储，但不参与反向图的拓扑构造。原有对象的requires_grad属性不变。

比如下面的代码，修改一个对象的数据，另一个对象的数据也会改变。

import oneflow as flow
y = flow.Tensor([1, 2, 3])
x = y.detach()
x[0] = 4
assert(y[0] == 4)

3 示例代码

本文通过如下代码来观察OneFlow的autograd机制。

import oneflow as flow
# y is scalar
x = flow.tensor([-1.0, 2.0], requires_grad=True)
y = flow.relu(x).sum()
y.backward()
print(x.grad)

# y is not scalar
x = flow.tensor([-1.0, 2.0], requires_grad=True)
y = flow.relu(x)
y.backward(flow.Tensor([1, 1]))
print(x.grad)

y.backward方法有两种接口：

如果y是一个标量（比如loss），不需要传递任何参数。
如果y是一个向量，需要传入一个与y的shape一致的向量作为参数。

为什么会有这种区别呢？下面几篇参考资料中对这个问题做了比较详细的解释。简单的说：

如果函数的输出是向量，在反向传播的过程中会造成梯度tensor shape的维度膨胀，实现复杂、性能差。
如果函数的输出是标量，反向传播梯度tensor的shape与参数变量的shape一致，不会出现维度膨胀，更容易实现。
对于向量版本的backward，可以假想存在某个loss函数，backward的参数是loss传播到y这里的梯度。因为前后节点间的梯度是乘积关系，所以用ones替代这个假想的梯度，这样计算结果x.grad就是y对x的梯度。

后续将以y.backward(flow.Tensor([1, 1]))为例观察一下autograd的机制。其反向图只有x <- y这一步。

参考资料

自动求梯度（https://tangshusen.me/Dive-in...梯度）
PyTorch 的 backward 为什么有一个 grad_variables 参数？（https://zhuanlan.zhihu.com/p/...）

3.1 梯度结果的存储

Tensor的grad属性（https://github.com/Oneflow-In...），在读取值时调用的是acc_grad()方法（acc应该是accumulate的缩写）。这样就知道梯度实际存储在哪里，读代码时可以重点关注相关部分。

调用流程如下：

注：图片中的MirroredTensor在最新源码中，已经更名为LocalTensor，其实是一样的。

4 autograd相关的类图关系

下图展示了autograd相关类的关系

在看autograd代码之前，可以参照这个类图，了解其中的结构和关系，有助于理解代码中各个部分的作用。

在eager模式下，用户通过op的组合逐步构建出前向计算图。在执行前向计算的过程中，引擎会为autograd需要的反向计算图记录必要的信息，在调用backward方法时执行这个反向计算图。

对照上面的类图

站在tensor的视角

前向op输出一个tensor y，即TensorIf <- ReluFunctor这部分。
从y可以找到反向计算图实际执行梯度计算的类，即TensorIf -> FunctionNode -> ReLU这个链路。

FunctionNode的backward_fn_包含了OpExprGradClosure。它只负责计算当前节点的梯度。
ReLU是执行梯度计算的类，它会调用ReluGradFunctor这个op来执行梯度计算。

站在反向图存储的视角

反向图相关的信息在FunctionNode中保存。
反向计算图的root是tensor（比如y或loss）的grad_fn_node_变量。
FunctionNode的next_functions_表示反向图的下游节点，当前节点把梯度结果传给这些下游节点。这些FunctionNode的连接就构成了反向图的拓扑结构。
tensor的梯度存储路径是TensorImpl.AutogradMeta.acc_grad_
AutogradMeta.current_grad_是反向图上游传递到当前节点的梯度合计。如果tensor t输入给op u和v，那么u和v反传的梯度会累加到current_grad_。current应该表示截至当前正在计算时的累加和。
FunctionNode虽然并不持有tensor实例，但它持有tensor的AutogradMeta成员变量指针。
基于上述relu的例子中的节点y
output_meta_data_即y.autograd_meta_
input_meta_data_即x.autograd_meta_
所以FunctionNode能获取到上下游的梯度数据并进行读写

AutoGradCaptureState可以存储一些梯度计算需要的状态信息，比如计算relu的梯度时需要用到它的前向输出结果y。

站在反向图执行的视角

GraphTask负责反向图的执行。
FunctionNode只保存必要的数据。
GraphTask基于这些数据，自己构造遍历需要的数据结构，遍历所有节点、执行梯度计算。

5 前向计算过程中为autograd所做的准备

反向图的执行过程是数据驱动的，数据的存储结构和内容决定了执行的具体动作。

以下讨论只针对eager模式。lazy模式下，反向图的构建是多轮优化passes的一部分（https://github.com/Oneflow-In...）。
之前在讨论Op、Kernel与解释器时已经了解Interpreter的作用。只是当时重点关注op的执行，忽略了grad相关的内容。
GetInterpreter（https://github.com/Oneflow-In...）返回的其实是一个AutogradInterpreter对象（https://github.com/Oneflow-In...），在它的Apply方法中（https://github.com/Oneflow-In...），调用内嵌Interpreter的同时，也会记录grad计算需要的信息。

AutogradInterpreter::Apply的主要流程如下：

Apply的第一步会先计算requires_grad。只要op的任一输入的requires_grad为true，op的输出的requires_grad也为true（https://github.com/Oneflow-In...）（前提是输出的数据类型支持梯度）。y的requires_grad就是在这里决定的。
比如y = relu(x)，如果数据类型支持梯度，y.requires_grad就等于x.requires_grad。

然后会调用内嵌的解释器internal_执行相关计算。在调用内嵌解释器期间，会临时禁止梯度模式，比如有些op可能会嵌套、多次调用解释器（ReluGradFunctor也会通过解释器执行），这些都不需要梯度逻辑。

需要说明的是，构造x时不会执行grad相关的逻辑，因为inputs的requires_grad都是false，x的requires_grad是在构造的最后才设置的（https://github.com/Oneflow-In...）。

下面重点看一下几个核心函数的逻辑细节。

5.1 梯度闭包的构建

前面对类图的说明中已经提到，OpExprGradClosure只负责当前节点的梯度计算。

GetOrCreateOpGradClosure函数（https://github.com/Oneflow-In...）的核心代码如下：

template<>
Maybe<OpExprGradClosure> BuiltinOpExprImpl<UserOpConf>::GetOrCreateOpGradClosure() const {
if (!op_grad_func_.get()) {
    ...
    op_grad_func_.reset(NewObj<std::string, OpExprGradFunctionIf>(proto().op_type_name()));
    JUST(op_grad_func_->Init(*this));
  }
return std::make_shared<OpExprGradClosure>(op_grad_func_);
}

NewObj会调用AutoRegistrationFactory（https://github.com/Oneflow-In...）获取预先注册的工厂、创建对象。之前在讨论Op指令在虚拟机中的执行时也看到过类似的注册机制。

这里op_type_name的值是relu，在代码中搜索"relu"，可以找到注册ReLU的宏（https://github.com/Oneflow-In...）。宏展开后的代码如下：

static AutoRegistrationFactory<std::string, OpExprGradFunctionIf>::CreatorRegisterTypeg_registry_var4("relu", ([]() { return new ReLU; }));

所以实际返回的对象是ReLU（https://github.com/Oneflow-In...）。其Init函数是个空操作。

OpExprGradClosure只是简单的把ReLU存下来供backward执行时调用。整个调用流程如下：

5.2 捕获梯度计算需要的数据

调用流程如下：

Capture函数（https://github.com/Oneflow-In...）的作用就是为后续的梯度计算保存必要的数据。
需要注意的是，OpExprGradFunction::CaptureIf（https://github.com/Oneflow-In...）中保存的是detach的tensor。这些tensor与原来的tensor共享数据；可以读写梯度数据，但不会参与反向图的拓扑构造。这个函数把Interpreter传过来的op的detached outputs传给ReLU::Capture（https://github.com/Oneflow-In...）（就是relu的前向输出y），ReLU::Capture就把output[0]存到ReLUCaptureState的saved_tensors_中（https://github.com/Oneflow-In...）。因为对于relu来说，根据y就可以计算梯度。

5.3 保存反向图结构信息

AutogradInterpreter::Apply中会构造一个lambada表达式backward_fn（https://github.com/Oneflow-In...），其核心逻辑只有一行grad_closure->Apply。

这个lambda的主要作用就是捕获grad_closure这个智能指针。lambda表达式最终会作为FunctionNode的backward_fn_变量。这样才有类图中FunctionNode到OpExprGradClosure这条线，才能从FunctionNode找到closue、执行节点的梯度计算。

GetThreadLocalAutogradEngine()->AddNode这个函数（https://github.com/Oneflow-In...）很关键，AddNode的主要任务（https://github.com/Oneflow-In...）是为inputs和outputs创建FunctionNode、并保存反向图遍历需要的数据。其输入参数中的inputs/outputs，是前向计算的op的inputs/outputs。对于relu来说，inputs就是x，outputs就是y。

在上述示例代码中，对于x，因为它是叶子节点、也需要梯度，在AddAccumulateFunctionNode会将grad_fn_node设置为一个空操作的函数（https://github.com/Oneflow-In...）。之所以是空操作，是因为叶子节点只需要存储梯度、不需要自己计算梯度；它所需要的梯度计算结果会由反向图的上游节点保存到x.autograd_meta_中。

之后会为y构造GraphFunctionNode并形成节点连接（https://github.com/Oneflow-In...）、并保存到grad_fn_node（https://github.com/Oneflow-In...）。需要注意的是，这里的backward_fn就是AutogradInterpreter::Apply中的lambda表达式（https://github.com/Oneflow-In...）。
需要注意的是，AddBackwardFuncPtr中的inputs/outputs是针对op而言，GraphFunctionNode构造函数中同名变量的是针对FunctionNode而言，二者的含义和指向的对象是不一样的。

构造完成后，x和y的grad_fn_node_字段数据内容如下：
x.grad_fn_node_

name_: accumulate_grad
next_functions_: 空
input_meta_data_: 空
output_meta_data_: size=1，x.autograd_meta_，requires_grad=true，is_leaf=true
output_tensor_infos_: 对应x, relu前向op的input
backward_fn_: 空函数，AddAccumulateFunctionNode中定义的

y.grad_fn_node_

name_: relu_backward
next_functions_: size=1, x.grad_fn_node, 空操作, AddAccumulateFunctionNode中构造的GraphFunctionNode
input_meta_data_: x.autograd_meta_, requires_grad=true, is_leaf=true
output_meta_data_: size=1, y.autograd_meta_, requires_grad=false, is_leaf=false
output_tensor_infos_: 对应y, relu前向op的output
backward_fn_: AutogradInterpreter::Apply中定义的lambda函数

backward就是根据这些数据，从roots出发，完成反向图的遍历。

6 backward的入口

在《OneFlow源码阅读4：tensor类型体系与local tensor》（https://segmentfault.com/a/11...）中提到过，Tensor类在Python端经过一层包装，通过Python机制为Tensor类注册一些方法，backward就是包装的方法之一。

7 autograd.grad

通常，我们都会通过tensor.backward或autograd.backward触发梯度计算和反向传播，但偶尔也会用到autograd.grad这个接口。autograd.grad和autograd.backward很相似，不同之处主要在于：

autograd.backward以outputs(Tensor)作为起点，计算每一个叶子节点的梯度，并且梯度可累积，且保存于对应inputs(Tensor)的tensor.grad上。
而autograd.grad 接口则是从指定的 outputs为起点，以指定的 inputs为终点计算梯度，并按 inputs 参数的顺序返回一个由inputs相对应的grads构成的TensorTuple。且梯度是直接获得的，不在inputs的tensor.grad中累积。

由于autograd.grad就只执行后向计算图中的一部分，在OneFlow 静态图模式下(lazy mode)TaskGraph 统计入度时就需要做一次剪枝，把不需要计算的结点去掉（参考 TaskGraph::ComputeDependenciesAndPruneNode(https://github.com/Oneflow-In...) 接口），同时记录每个 inputs 序号，在 FunctionNode::Apply(https://github.com/Oneflow-In...) 执行后，把需要保存的 grad 及时捕获，最后返回给用户。

8 反向计算中GraphAutogradEngine的调用流程

反向图计算的流程分析可以结合3类信息

流程代码
上述x和y的grad_fn_node_的值
类图以及类之间的关系

RunBackwardAndSaveGrads4LeafTensor函数的几个参数是：

outputs: relu的输出y
out_grads: 用户自己构造的ones [1, 1]

8.1 反向传递过来的梯度的累加

RunBackwardAndSaveGrads4LeafTensor函数中，PushPartialTensor的作用就是将loss传过来的梯度累加到 autograd_meta_.current_grad_.acc_tensor_。第4节中提到，TensorArg.acc_tensor_存储的就是loss传过来的梯度的合计。这就是roots（即y）接收到的梯度，要么是框架自动创建的ones，要么是用户提供的梯度（通常也是ones）。

这行代码的逻辑可以用如下伪码表示

outputs[i].impl_.autograd_meta_.current_grad_.acc_tensor_ += out_grads[i]

8.2 反向图计算任务的构造与执行

FunctionNode只是记录了反向图的基础信息。RunBackwardAndSaveGrads4LeafTensor中会再构造一个GraphTask对象来表示一次反向计算任务。

GraphTask的构造函数主要是初始化反向图的roots_节点，并将图中各个节点的依赖计数dependencies_置为0。根据示例代码，roots_就是y（通常是loss）。
ComputeDependencies会对反向图进行深度优先遍历、统计图中各个节点的依赖计数。
GraphTask::Apply中实现了反向图的遍历逻辑（传入的save_grad_for_leaf参数是true）。当FunctionNode的依赖为0时，节点才会被放入执行队列，后续会对反向图执行按拓扑序遍历。FunctionNode::Apply执行时，它的依赖都执行完毕了。GraphTack::Apply这个函数中，涉及梯度计算逻辑主要包括两部分：

8.3 节点的梯度计算

FunctionNode::Apply中，处理output_meta_data_的for循环的核心逻辑可以用如下伪码表示：

acc_tensor = output_meta_data_[i].current_grad_.acc_tensor_
if (acc_tensor != nullptr) {
  output_grads[i] = acc_tensor_
} else {
  output_grads[i] = zeros()
}

从中可以看出来，output_grads的作用就是拷贝上游传过来的梯度数据（指针），作为backward_fn_的参数。
后面可以看到，backward_fn的核心逻辑是：

// d(y)表示当前节点对y的梯度，比如relu对其输出y的梯度。
input_grads = d(y) * output_grads

input_grads就是当前节点传给下游节点的梯度，调用backward_fn时会对它进行赋值。

处理input_meta_data的for循环的核心逻辑可以用如下伪码表示。实质就是将当前节点传给下游节点的梯度，累加到下游节点的current_grad上，从而实现梯度的传播。如果tensor输入给多个op，每个op的梯度会加起来。

input_meta_data_[i].current_grad_.acc_tensor_ += input_grads[i]

8.3.1 梯度计算的执行：backward_fn

以下只考虑前述示例的root节点的执行。也就是y对应的FunctionNode。对于y来说，backward_fn就是AutogradInterpreter::Apply中定义的lambda表达式。对于relu来说，执行过程如下：

之前在5.1节已经确认，OpExprGradClosure::impl_就是ReLU。
如前所述，backward_fn的参数中，output_grads是上游传过来的梯度数据，backward_fn需要计算relu的梯度，二者的乘积赋值给in_grads。这些参数会一直传递到ReLU::Apply。
functional::ReluGrad的Functor名字是ReluGrad。对应的Functor是ReluGradFunctor（命名空间是oneflow::one::functional::impl）。

ReluGradFunctor之后，是基于Primitive kernel实现的计算逻辑。ReluGradFunctor中对应op名字是"relu_grad"，这个relu_grad的注册被包在一个宏定义中，实际上会返回一个BinaryPrimitiveKernel，这是一种稍显特殊的基于Primitive的kernel，其具体为ep::primitive下的一种BroadcastElementwiseBinary工厂，其对应的cpu和cuda注册分别位于：

oneflow/core/ep/cpu/primitive/broadcast_elementwise_binary.cpp
oneflow/core/ep/cuda/primitive/broadcast_elementwise_binary.cu

最终实现位于binary_functor.h：

template<DeviceType device, typename Src, typename Dst>
struct BinaryFunctor<device, BinaryOp::kReluBackwardWithDyY, Src, Dst> {
  OF_DEVICE_FUNC BinaryFunctor(Scalar attr0, Scalar attr1) {}

  OF_DEVICE_FUNC Dst operator()(Src dy, Src y) const {
    return static_cast<Dst>((y <= static_cast<Src>(0.0)) ? static_cast<Src>(0.0) : dy);
  }
};

至此，完成了梯度计算的逻辑。

8.4 梯度的存储

FunctionNode::Apply执行完毕后，GraphTask::Apply调用FunctionNode::AccGrad4LeafTensor为叶子节点拷贝梯度数据。
在上述例子中，因为y不是叶子节点，处理到y.grad_fn_node_时不会进行实质处理。对于x，会调用CopyOrAccGrad，这个函数逻辑的伪码形式如下

autograd_meta.acc_grad_ += autograd_meta.current_grad_

autograd_meta.acc_grad_就是Python端读到的x的梯度。

8.5 临时梯度的释放机制
上述第5.点中，描述了前向图构建过程中已经存放了对应的FunctionNode以及前向op所对应的反向backward_fn，实际求梯度、反向传播时，这一个个 backward_fn串联起来构成了反向计算图拓扑，对于其中的每个节点，backward_fn中都可以表示为output_grads、inputs/outputs(可选) -> inputs_grads的一个函数。
其中output_grads 就是链式法则中上游计算的累计梯度，当前节点backward_fn计算完成后，该节点的output_grads就不会再被使用到，从而变成了临时梯度。之后会调用 FunctionNode->ReleaseOutTensorArgs() 来及时释放该临时梯度。

参考资料

欢迎下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In...

OneFlow源码解析：自动微分机制

1

自动微分基础

2

autograd中tensor相关的一些基本概念

3

示例代码

4

autograd相关的类图关系

5

前向计算过程中为autograd所做的准备

6

backward的入口

7

autograd.grad

8

反向计算中GraphAutogradEngine的调用流程

OneFlow

引用和评论

SiliconCloud上线Reflection-Llama-3.1-70B

Light·技术公益创造营，开营了！

数据集汇总｜18个电影/音乐数据集汇总，覆盖影片/歌曲推荐、电影评价、歌词识别、音乐流派······

详解Diffusion扩散模型：理论、架构与实现

Looktech 发布 AI 智能眼镜：旋钮交互引领创新，智能体小程序开启无限可能

RAG+Agent人工智能平台：RAGflow实现GraphRA知识库问答，打造极致多模态问答与AI编排流体验

腾讯AI开放平台的接口调用指南