PR Category

Execute Infrastructure

PR Types

Bug fixes

Description

问题描述

SOT/AST+PIR 模型暴露问题，伪码描述如下：

def fn(x):
   mean1 = builtin.parameter("parameter_1")
   y, mean2 = pd_op.batch_norm_(x, mean1, ...)
   return y

对于这个 Program 来说，动转静前向输入包含原输入 x 和 parameter mean1，前向输出包含原输出 y 和反向所需的中间变量 mean2，问题就出在这个 mean1 和 mean2 是 inplace 的（实际上是 view 的，但 PIR 下是一样的）

由于 mean1 和 mean2 是 inplace 的，所以他们对应同一个 Variable，mean2 作为输出我们会添加 shadow_output 来确保能够取出来，但由于 mean2 实际上就是 mean1，它已经有作为输入的 parameter_1 这个名字了，实际上是不需要通过 shadow_output 来重命名的。而且，如果进行重命名，在第一个 step 我们在 share 到 scope 时需要使用 parameter_1 这个名字来 share，而第二个 step 因为已经 rename 成另一个名字了，这就会导致 share 失败。

解决方案

归根结底，问题出在一个输入 Value 经过若干 inplace 操作得到一个输出 Value，此时这个输出 Value 不应该添加 shadow_output，而是应该直接替换为输入 Value。为了达成这一点，我们需要在前反向拆分时，对 Program 中所有 inplace 链进行分析，比如下图：

这里仅标识出了 inplace Value，忽略非 inplace Value，并且按照时序拉直，比如这里的链路 A 中的 A1 作为输入，A4 作为输出，那么 A4 不应该添加 shadow_output。

具体替换结果如下图所示：

但是这里为什么不直接将 A4 从中间变量中删掉呢？因为考虑到动转静前向的输出侧并不仅仅包含中间变量，还包含原前向的输出，针对中间变量尚且可以删掉，但如果针对输出的话就不太好删掉了，因此我们对中间变量和输出变量采用相同的操作，就是替换为 inplace 的输入侧源 Value。

如果考虑前反向全部情况，根据下图动转静前反向的输入输出关系，实际上我们需要处理如下几种情况：

前向：
x | param -> outputs
x | param -> middles
反向：
x | param | middles | outputs | out_grads -> x_grads
x | param | middles | outputs | out_grads -> param_grads

但是由于前向输入 x 是用 data op 创建的，data op 在经过 lower 到 kernel pass 后会跟随一个 shadow_feed op，而 shadow_feed op 不是 inplace 的，会创建一个 inner var，也就是说前向输入 x 永远不会出现上述的同时作为输入和输出的问题，因为在一开始就有一个非 inplace 的操作，因此我们不需要处理前向输入，只需要处理 parameter。

类似地，反向全部输入都是通过 block kwarg 创建的，而 kwarg 也会在 lower 后跟随一个 shadow_feed，因此整个反向都不会出现上述问题，也就是说反向不需要处理。

因此需要处理的只有：

前向：
param -> outputs
param -> middles

但由于最开始是考虑这种情况的，即便将来 data/kwarg 后不插 shadow_feed 也能很容易实现。（因为目前看反向会有过多 shadow_feed，非常影响性能，也许之后这个 shadow_feed 会被删掉也说不定）