libai
libai copied to clipboard
Add support for dynamic control flow of auto placement.
修复之前auto placement不支持dynamic control flow的问题。
目前仍然有以下几个点需要注意或以后解决:
- transformers代码里有许多对于
ones,zeros的调用,但oneflow 0.9.0对它们的实现是调用了Ones().__call__(),这样会导致fx自动对__call__进行包装从而使其出错。目前oneflow master分支已经修改成了forward,所以onefx暂时没有做处理。目前要跑通需要用到oneflow的master分支。 - 目前
compile_auto_placement函数改成返回一个interpreter,在外面则是接收返回值之后调用run。之前直接返回GraphModule的方式不行,这可能是一个BUG,需要以后修复。 - 原本的测试代码中对auto placement之后的module调用了
generate函数,这样是不成立的,因为fx就是针对Module的forward进行trace的。这个和2无关,不论是否使用interpreter都是这样。这个可能可以实现,但是肯定会需要自己指定要trace的函数名,这部分后续需要再讨论一下。 - 在调用
compile_auto_placement的时候,不用指定concrere_args,但是有一个约定就是如果模型实际运行的时候会指定某个参数,那就需要对这个参数以keyword arg的形式传入compile_auto_placement。这个方式后续可以讨论修改一下。 - onefx也做出了少量必要的修改,这边的onefx版本应该升级到0.0.3。
原本的测试代码中对auto placement之后的module调用了generate函数,这样是不成立的,因为fx就是针对Module的forward进行trace的。这个和2无关,不论是否使用interpreter都是这样。这个可能可以实现,但是肯定会需要自己指定要trace的函数名,这部分后续需要再讨论一下。
@BBuf 这个我们可以一起讨论
看上去这个代码是个临时版本, 等ready了以后我再来review吧