[Auto Parallel] Support gradient_merge pass for PIR #66641

AndSonder · 2024-07-26T15:58:56Z

PR Category

Auto Parallel

PR Types

Not User Facing

Description

Pcard-76459

为 PIR 适配 gradient_merge pass

依赖 PR：

【AutoParallel】Support to split program into sub-program in PIR #66130

before auto_parallel_gradient_merge_pass {
    (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"learning_rate_1",op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[]}},op_role:(Int32)0,persistable:[true],place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[],stop_gradient:[true]} : () -> pd_dist.tensor<f32, mesh_shape:[1],process_ids:[0],dims_mappings:[]>
    (%1) = "builtin.parameter" () {is_distributed:[false],is_parameter:[true],need_clip:[true],op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,parameter_name:"linear_0.w_0.dist",persistable:[true],stop_gradient:[false],trainable:[true]} : () -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%2) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"input0",op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,persistable:[false],place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,16],stop_gradient:[false]} : () -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%3) = "pd_op.relu" (%2) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%4) = "pd_op.matmul" (%3, %1) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false],transpose_x:false,transpose_y:false} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%5) = "pd_op.relu" (%4) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    () = "pd_op.send_v2" (%5) {dynamic_shape:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,peer:(Int32)1,ring_id:(Int32)22,stop_gradient:[],use_calc_stream:true} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> 
    (%6) = "pd_op.recv_v2" () {dtype:(pd_op.DataType)float32,dynamic_shape:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,out_shape:[(Int32)4,(Int32)16],peer:(Int32)1,ring_id:(Int32)22,stop_gradient:[true],use_calc_stream:true} : () -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%7) = "pd_op.relu_grad" (%5, %6) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%8, %9) = "pd_op.matmul_grad" (%3, %1, %7) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(2):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false,false],transpose_x:false,transpose_y:false} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%10) = "pd_op.relu_grad" (%3, %8) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%11, %12) = "pd_op.sgd_" (%1, %0, %9, <<NULL VALUE>>) {multi_precision:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[]},operand(2):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(3):{mesh_shape:[],process_ids:[],dims_mappings:[]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(1):{mesh_shape:[],process_ids:[],dims_mappings:[]}},op_role:(Int32)2,stop_gradient:[false,false]} : (pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<f32, mesh_shape:[1],process_ids:[0],dims_mappings:[]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, <<NULL TYPE>>) -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, <<NULL TYPE>>
}
after auto_parallel_gradient_merge_pass {
    ^kw:%kwarg_linear_0.w_0.dist@GRAD@MERGE
    (%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"learning_rate_1",op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[]}},op_role:(Int32)0,persistable:[true],place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[],stop_gradient:[true]} : () -> pd_dist.tensor<f32, mesh_shape:[1],process_ids:[0],dims_mappings:[]>
    (%1) = "builtin.parameter" () {is_distributed:[false],is_parameter:[true],need_clip:[true],op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,parameter_name:"linear_0.w_0.dist",persistable:[true],stop_gradient:[false],trainable:[true]} : () -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%2) = "pd_op.data" () {dtype:(pd_op.DataType)float32,name:"input0",op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,persistable:[false],place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,16],stop_gradient:[false]} : () -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%3) = "pd_op.relu" (%2) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%4) = "pd_op.matmul" (%3, %1) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false],transpose_x:false,transpose_y:false} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%5) = "pd_op.relu" (%4) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    () = "pd_op.send_v2" (%5) {dynamic_shape:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)0,peer:(Int32)1,ring_id:(Int32)22,stop_gradient:[],use_calc_stream:true} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> 
    (%6) = "pd_op.recv_v2" () {dtype:(pd_op.DataType)float32,dynamic_shape:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,out_shape:[(Int32)4,(Int32)16],peer:(Int32)1,ring_id:(Int32)22,stop_gradient:[true],use_calc_stream:true} : () -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%7) = "pd_op.relu_grad" (%5, %6) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%8, %9) = "pd_op.matmul_grad" (%3, %1, %7) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(2):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false,false],transpose_x:false,transpose_y:false} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%10) = "pd_op.add_" (%kwarg_linear_0.w_0.dist@GRAD@MERGE, %9) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%11) = "pd_op.relu_grad" (%3, %8) {op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>) -> pd_dist.tensor<4x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%12) = "pd_op.full" () {dtype:(pd_op.DataType)float32,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1]}},place:(pd_op.Place)Place(cpu),shape:(pd_op.IntArray)[1],stop_gradient:[true],value:(Double)1} : () -> pd_dist.tensor<1xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1]>
    (%13) = "pd_op.scale_" (%10, %12) {bias:(Float)0,bias_after_scale:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]}},op_role:(Int32)1,stop_gradient:[false]} : (pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<1xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1]>) -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>
    (%14, %15) = "pd_op.sgd_" (%1, %0, %kwarg_linear_0.w_0.dist@GRAD@MERGE, <<NULL VALUE>>) {multi_precision:false,op_dist_attr:{mesh:{shape:[1],process_ids:[0]},operand(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(1):{mesh_shape:[1],process_ids:[0],dims_mappings:[]},operand(2):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},operand(3):{mesh_shape:[],process_ids:[],dims_mappings:[]},result(0):{mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]},result(1):{mesh_shape:[],process_ids:[],dims_mappings:[]}},op_role:(Int32)2,stop_gradient:[false,false]} : (pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, pd_dist.tensor<f32, mesh_shape:[1],process_ids:[0],dims_mappings:[]>, pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, <<NULL TYPE>>) -> pd_dist.tensor<16x16xf32, mesh_shape:[1],process_ids:[0],dims_mappings:[-1,-1]>, <<NULL TYPE>>
}

into develop

… fit_grad_merge_for_pir

paddle-bot · 2024-07-26T15:59:01Z

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

heavyrain-lzy · 2024-07-29T06:08:09Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+                break
+        for _, new_grad in new_params_to_grads:
+            new_grad = paddle._C_ops.scale_(new_grad, 1.0 / k_steps, 0.0, False)
+            new_grad.get_defining_op().op_role = int(OpRole.Backward)


这里的scale_算子应该归类到OpRole.Optimize中

heavyrain-lzy · 2024-07-29T06:13:54Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+def _pir_remove_cast_for_master_grad(main_program):
+    main_block = main_program.global_block()
+    for op in main_block.ops:
+        if _is_master_grad_cast_op(main_block, op):


_is_master_grad_cast_op函数不一定可以在pir下继续使用，需要结合PIR的amp功能进行调整。

heavyrain-lzy · 2024-07-29T06:14:43Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+    main_block = main_program.global_block()
+
+    for idx, op in list(enumerate(main_block.ops)):
+        if is_data_parallel_reduce_op(op):


is_data_parallel_reduce_op函数功能是否保留，需要进一步讨论。

heavyrain-lzy · 2024-07-29T06:18:58Z

test/auto_parallel/pir/mlp_demo.py

+
+        gradient_merge = strategy.gradient_merge
+        gradient_merge.enable = True
+        gradient_merge.k_steps = accumulate_steps


可以增加一个带有gradient_merge的单测，并对比loss结果

… fit_grad_merge_for_pir

JZ-LIANG · 2024-08-01T03:08:54Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+            "pd_op.c_allreduce_sum",
+            "pd_op.c_allreduce_avg",
+            "pd_op.c_reduce_sum",
+            "pd_op.c_reduce_avg",


有 op 遗漏，比如 sharding 的 reducescatter op。判断 dp 梯度同步 op 的逻辑会在很多 pass 中复用，建议把其写成一个公共的 utils，统一管理。当前 PR 逻辑影响不大，可以在后续 sharding pr 中升级修改。 @winter-wang

JZ-LIANG · 2024-08-01T03:11:30Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+            # NOTE(sonder): When "@RENAME@" is in the input name, it means that the op has been renamed.
+            # Such types input names are caused by shared parameter policy.
+            # Gradient merge should accumulate the gradient of ops without renaming.
+            if "@RENAME" in op_input_names[0]:


"@rename" 作为GM 逻辑中关键标志，最好不要裸写在代码的每一个地方。最好定义一个 GM_SUFFIX 全局变量统一管理。避免多处地方不同修改。

"@rename" 不是 grad_merge 引入的，是共享参数策略带来的

JZ-LIANG · 2024-08-01T03:12:42Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+            op.op_role = OpRole.Optimize
+            main_block.move_op_to_block_end(op)
+
+            if op.name() in ["pd_op.c_allreduce_sum", "pd_op.c_reduce_sum"]:


如何区分 TP 中在反向的 allreduce？

后续通过在 dp 中添加 allreduce 算子的 attr 解决这个问题

JZ-LIANG · 2024-08-01T03:18:15Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

@@ -526,12 +618,48 @@ def parse_program(
    return grad_to_gradient_merge


+def _pir_parse_program(


感觉缺少一个关键逻辑：《GM 的子图调度编排》
PIR 版本的 GM 实现需要相对于原有 programIR 的 GM 升级最关键的点：
programIR 的 GM：控制流实现 GM 调度
PIR 的 GM：通过类似 PP 子图调度编排实现 GM 的调度，最好 PP 和 GM 复用同一套框架逻辑。

将计算图切分成 FW-BW-OPT 三个子图，通过编排不同子图的调度顺序实现 GM 逻辑(类似流水线的调度编排)
如：FW-BW-FW-BW-FW-BW-FW-BW-OPT

这个这一版的 pr 先不考虑，后续方案确定后再实现

JZ-LIANG · 2024-08-01T03:19:58Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+            gard_defining_op = grad.get_defining_op()
+            paddle.pir.set_insertion_point_after(gard_defining_op)
+
+            new_gradient_merge_var = main_block.add_kwarg(


需要考虑和 AMP 中的 master grad 逻辑兼容，如果AMP 中开启了 master grad，累加的 persistable 的 grad 应该和 amp 复用同一个 master grad

JZ-LIANG

LGTM

AndSonder · 2024-08-01T07:43:54Z

python/paddle/distributed/passes/auto_parallel_gradient_merge.py

+def _pir_remove_cast_for_master_grad(main_program):
+    main_block = main_program.global_block()
+    for op in main_block.ops:
+        if _is_master_grad_cast_op(main_block, op):
+            main_program.remove_op(op)
+
+


这里的适配逻辑还有些问题，下一个 pr 修复并添加 amp 情况下的单测

heavyrain-lzy and others added 15 commits July 17, 2024 17:12

split program and polish the executor

da84e28

merge

e6fee47

change unit test

403db76

fix 1f1b test

23eed1a

polish

9870779

polish

f2a63b4

polish

0ea4320

merge develop

03c5d20

fix conflict

2125aca

Merge branch 'support_pp_pir' of https://github.com/heavyrain-lzy/Paddle

0f0cc15

into develop

add pp in pir

e3a68f1

fix ut

391af96

fit grad_merge for pir

68557a3

Merge from heavyrain-lzy/support_pp_pir

c7b228f

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

16ae000

… fit_grad_merge_for_pir

paddle-bot bot added the contributor External developers label Jul 26, 2024

AndSonder and others added 2 commits July 29, 2024 12:52

Merge branch 'develop' into fit_grad_merge_for_pir

b58dab2

update test

3c42563

heavyrain-lzy reviewed Jul 29, 2024

View reviewed changes

AndSonder added 6 commits July 29, 2024 14:11

apply suggestions from code review

24644fa

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

5b8b788

… fit_grad_merge_for_pir

fit master grad

2a59204

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

c4871d3

… fit_grad_merge_for_pir

fix unset op_role

0983f08

fix unset op_role

d3020d8

JZ-LIANG reviewed Aug 1, 2024

View reviewed changes

AndSonder requested a review from JZ-LIANG August 1, 2024 07:15

JZ-LIANG approved these changes Aug 1, 2024

View reviewed changes

AndSonder commented Aug 1, 2024

View reviewed changes

JZ-LIANG merged commit 8caa713 into PaddlePaddle:develop Aug 2, 2024
31 checks passed

		@@ -526,12 +618,48 @@ def parse_program(
		return grad_to_gradient_merge


		def _pir_parse_program(

[Auto Parallel] Support gradient_merge pass for PIR #66641

[Auto Parallel] Support gradient_merge pass for PIR #66641

Uh oh!

Conversation

AndSonder commented Jul 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Category

PR Types

Description

Uh oh!

paddle-bot bot commented Jul 26, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JZ-LIANG left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

AndSonder commented Jul 26, 2024 •

edited

Loading