oneflow-inc / oneflow-xrt Goto Github PK

在测试 adam 的时候遇到1个问题：多个输入中，weight 和 weight_grad 是来自 non-cpu device，而 bias_correction_factor 来自 cpu device。导致在 compile 时，无法确定 input 到底是 device 上的还是 cpu 上的。

在oneflow的compile过程中，adam gen pass 中插入 bias_correction_factor op 时，一般插入的是 cpu 的 op。然后在之后的 build task graph 阶段，会再在 bias_correction_factor 和 adam optimizer 之间再插入一个 copy h2d 来完成转换。

但 xrt 的编译过程是发生在 build task graph 阶段之前，就无法检测到该 copy h2d 节点。其 adam xrt node 的输入就会来自不同的 device，编译成子图时，子图的 inputs 也就来自不同的 device 了。

oneflow 中 nn.Graph 的内部接口有一些 breaking update（譬如 finish_complie_and_init_runtime 被换成其他实现）。

XRTModule 需要随着一起更新。