[Feature]: dw/dx 分离方案

### 🚀 背景描述 ### 需求来源 在流水想并行场景下，传统的1F1B并行策略虽然进一步优化了流水过程中的Bubble，但是依然有较多Bubble待优化，为了进一步的优化性能，提出了基于dw分离的Zero Bubble并行方案，将计算输入的梯度和计算权重的梯度作了分离， 调整调度方案，增加并行度，提升整体的吞吐量。如图所示，通过分离dx和dw计算，减少流水线并行的空闲时间。 ### 目标 给定一个计算函数，实现一种通用的，业务无关的dx/dw 图分离的方案，增强dw/dx计算的易用性，采用更灵活的调度策略，快速实现dw计算编排。  ### 设计思路 通过这张正向图为列：  ### 算法实现 核心流程 1. 反向计算图翻转 2. BFS获取dx计算子图闭包 3. 计算每一个权重与dx闭包的最邻近公共节点及节点的路径信息 4. 对权重作分组合并，同时合并相同节点及其节点的连边关系 5. 对中间节点注册prehook 6. 执行dx反向子图，获取dx值，同时获取中间激活值的梯度 7. 根据中间节点的连边使用信息，剪枝路径，分组执行dw计算子图 通过上述流程获取dx，dw的计算子图如下： dx计算图  dw计算图  #### dx路径闭包计算 1 通过翻转反向计算图，首先获取翻转后每个节点的出度信息，包括{'next_edge': [(grad_node, input_index)]}, 2. input_index即获取父节点的边的信息，用于路径选择。 3. 以输入节点为截止条件，根节点为首节点，BFS遍历反向计算图，得到dx的计算子图。 获取sub_graph的set集合，用于后续计算最邻近公共节点。 #### 计算最最邻近公共子节点 通过BFS算法，以每一个权重为起点，找到每一个权重的最邻近公共节点，以及连边信息，用于路径剪枝。 #### 相同公共节点合并 使用并查集思想，合并公共父节点路径，保证最小子树无交叉 #### 最近邻公共节点注册prehook，计算dx子图 1. 通过计算dx子图，获取dx梯度的同时，拿到中间节点的梯度值，用于dw计算 #### 路径剪枝，分组计算dw 1.获取到所有的 子图后，需要进行路径剪枝，避免子图多个根节点之间存在父子关系导致梯度重复累加，具体实现上，根据合并后的 param_group {"w":{w1}, "immediate": {a1, a2}, "edge_index": {a1:{0}, a2:{1}}} 以上述有重叠场景为例，根据param_group 的信息，我们需要把a2中间节点的第0条边置空，使用grad_node._set_next_edge(index, None)实现剪枝。分组计算和不分组差异 主要在显存可以及时释放，避免出现显存峰值上升。 ### 涉及到的对外API ### 接口设计 #### `forward_and_gradfn(fn, *inputs, weights=None, has_aux=False, grad_position=0, **kwargs)` - 返回：`forward_out, grad_fn` - grad_position=-1 表征对所有输入求导， grad_position=None表征不对输入求导， grad_position=[0, 1] 表征对输入的第0，1个位置求导。 - 返回值形式：按照inputs,kwargs的形状返回梯度 ### 6.2 `GradFunction` - `compute_input_grad(sens=None) -> dx` - `compute_weight_grad(keep_graph=False) -> dw` - `__call__(sens=None, keep_graph=False)`：按配置返回 dx/dw 或组合结果 ### 约束 - `compute…