【开源之夏】动转静支持子图高阶微分

[SigureMo]

背景

动转静可以跟踪动态图 API，自适应构建等价的静态子图，使用内核执行器，实现转静训练加速。目前动转静支持整图级别的高阶微分，即 paddle.grad 接口在@to_static 装饰函数内调用。我们期望动转静支持子图高阶微分，即 paddle.grad 接口在 @to_static 装饰函数外调用。

示例代码如下：

import paddle


def fn(x):
    return x**3


static_fn = paddle.jit.to_static(fn, full_graph=True)
x = paddle.to_tensor(1.0)
x.stop_gradient = False

# 使用动转静作为动态图中的子图
y = static_fn(x)
# 在子图外调用 paddle.grad
[y_grad] = paddle.grad(y, x, create_graph=True)
print(y_grad)

# 此时如果调用二阶导，就会有问题
[y_grad_grad] = paddle.grad(y_grad, x)
print(y_grad_grad)

# 或者不通过 paddle.grad 而是直接求反向，也是有问题的
# y_grad.backward()
# print(x.grad)

目标

1. 基于飞桨框架研发的动转静功能，升级 RunProgramOp（paddle/fluid/eager/to_static）模块，支持子图高阶微分功能，添加必要单测验证功能正确性，代码合入 PaddlePaddle 框架 repo
2. 在给定的一个科学计算高阶微分模型上，验证动转静子图高阶微分的训练精度对齐原生动态图

意义

• 支持子图高阶微分，可以实现科学计算模型的高阶微分训练加速
• 支持高阶微分模型动转静 SOT 自适应子图打断，可以在子图打断场景下实现任意阶高阶微分

参考

动转静是在动态图下通过 RunProgramOp 来运行静态图 Program，在 pir_run_program_ad_func 中，我们会构建反向 GradNode（PirGradNodeRunProgram）。目前我们 GradNode 的 operator() 会调用 PirRunProgramGradAPI 来执行反向计算。如果按照目前的思路继续向下开发的话，为了支持高阶微分，我们必然需要支持 PirRunProgramGradGradAPI、PirRunProgramGradGradGradAPI……

但如果我们可以抽离出来一个函数，这个函数的作用就是传入输入 x: list[Tensor]、以及所需 Program，并返回输出 out: list[Tensor]，对于这个函数来说，它只需要执行 Program 并取出输出即可，它不需要关心这个 Program 是前向还是反向。这样无论多少阶微分，我们都可以直接调用这个函数。

当然，我们还需要对当前流程进行若干重构，设想的新架构如下：

• run_program_ad_func(x: list[Tensor], program: Program, scope: Scope, x_values: list[Value], out_values: list[Value], **attrs) -> list[Tensor]（执行前向，构建反向）

  • （如果需要）调用 paddle.grad，添加反向，前反向拆分 + 前反向关联分析等
    • 如 skip_gc_vars 的计算和设置
    • 前反向中间变量获取
  • 执行 Program
    • 对 scope 产生影响，其中包含了前反向中间变量
    • 根据 out_values 获得当前 Program 的输出，构建输出 out: list[Tensor]
  • 构建反向 GradNode
    • 存储前向结果
      • forward_values = x_values + out_values
      • forward_out_grad_values = out_grad_values（out_grad_values 在前反向拆分时可以拿到）
      • forward_x_grad_values = x_grad_values（同上，前反向拆分时可以拿到）
      • forward_inputs = x + out
      • scope = scope
    • 设置反向 Program（前反向拆分后得到）

• GradNode::operator(out_grads: list[Tensor])（执行反向）

  • 调用 run_program_ad_func(forward_inputs + out_grads, backward_program, scope, forward_values + forward_out_grad_values, forward_x_grad_values, **new_attrs)，获取输出 x_grads: list[Tensor]，也即前向是 RunProgram OP 的反向是其自身，只不过输入不同，这便可实现任意阶微分

Note

• 该方案只包含大体流程，实际方案需要调研现有实现方式并细化
• 目前 Paddle 正处于老 IR 向 PIR 过渡的阶段，本任务所有开发工作仅针对 PIR 进行，无需考虑老 IR，可通过 FLAGS_enable_pir_api=True 开启 PIR 模式

其他资料

• 动转静原理文档

• 目前 PIR 下动转静前反向拆分关系可参考下图