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

1

自动微分基础

1.1 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的梯度传播

1.3 复杂graph的梯度传播

2 

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

2.1 叶子节点

• 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填充梯度。
  
  

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

2.2 tensor detach

• requires_grad = false
• is_leaf = true
  
  

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

3 

示例代码

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

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

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

• 自动求梯度 （https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter02_prerequisite/2.3_autograd?id=_233-梯度）
  
  
• PyTorch 的 backward 为什么有一个 grad_variables 参数？（https://zhuanlan.zhihu.com/p/29923090）

3.1 梯度结果的存储

4 

autograd相关的类图关系

• 站在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所做的准备

5.1 梯度闭包的构建

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_);}

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

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

5.3 保存反向图结构信息

name_: accumulate_gradnext_functions_: 空input_meta_data_: 空output_meta_data_: size=1，x.autograd_meta_，requires_grad=true，is_leaf=trueoutput_tensor_infos_: 对应x, relu前向op的inputbackward_fn_: 空函数，AddAccumulateFunctionNode中定义的

name_: relu_backwardnext_functions_: size=1, x.grad_fn_node, 空操作, AddAccumulateFunctionNode中构造的GraphFunctionNodeinput_meta_data_: x.autograd_meta_, requires_grad=true, is_leaf=trueoutput_meta_data_: size=1, y.autograd_meta_, requires_grad=false, is_leaf=falseoutput_tensor_infos_: 对应y, relu前向op的outputbackward_fn_: AutogradInterpreter::Apply中定义的lambda函数

6 

backward的入口

• python/oneflow/framework/tensor.py
• python/oneflow/autograd/__init__.py
• oneflow/python/oneflow/autograd/autograd.py
• oneflow/api/python/autograd/autograd.cpp
  
  

import oneflow as flowx = flow.tensor([-1.0, 2.0], requires_grad=True)y = flow.relu(x)y.backward(flow.Tensor([1, 1]))print(x.grad)

• outputs: y, relu输出的tensor
• out_grads: [1, 1]
  
  

• 如果y是一个向量，backward必须传入一个与y的shape一致的向量（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L72-L81）。
• 如果y是一个标量，backward不要参数，框架会自动构造一个全1的tensor（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L70）。

7 

反向计算中GraphAutogradEngine的调用流程

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

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

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

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

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

• GraphTask的构造函数（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L452）主要是初始化反向图的roots_节点，并将图中各个节点的依赖计数dependencies_置为0。根据示例代码，roots_就是y（通常是loss）。

• ComputeDependencies（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L321）会对反向图进行深度优先遍历、统计图中各个节点的依赖计数。

• GraphTask::Apply（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L405）中实现了反向图的遍历逻辑（传入的save_grad_for_leaf参数是true）。当FunctionNode的依赖为0时，节点才会被放入执行队列（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L439），后续会对反向图执行按拓扑序遍历。FunctionNode::Apply执行时，它的依赖都执行完毕了。GraphTack::Apply这个函数中，涉及梯度计算逻辑主要包括两部分：

• 调用node->Apply执行单个节点的梯度计算（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L421）
• 调用node->AccGrad4LeafTensor存储算好的梯度（https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L430）

7.3 节点的梯度计算

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

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

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

7.3.1 梯度计算的执行：backward_fn

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

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);  }};

7.4 梯度的存储

autograd_meta.acc_grad_ += autograd_meta.current_grad_

参考资料

• https://github.com/Oneflow-Inc/oneflow/tree/48e511e40e09551408c96722c09bd061ce320687
• OneFlow学习笔记：Autograd解析
• 自动微分的原理介绍
• 自动求梯度 (https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter02_prerequisite/2.3_autograd?id=_233-梯度)
• PyTorch 的 backward 为什么有一个 grad_variables 参数？(https://zhuanlan.zhihu.com/p/29923090)
• PyTorch 101, Part 1: Understanding Graphs, Automatic Differentiation and Autograd （https://blog.paperspace.com/pytorch-101-understanding-graphs-and-automatic-differentiation/）
  
  

• Vision Transformer这两年
  

• OneFlow源码解析：Global Tensor
  

• 更快的YOLOv5问世，附送全面中文解析教程
• 李白：你的模型权重很不错，可惜被我没收了
• Stable Diffusion半秒出图；VLIW的前世今生
• 大模型狂潮背后：AI基础设施“老化”与改造工程
• OneEmbedding:单卡训练TB级推荐模型不是梦

---