Oneflow-Inc / libai

LiBai(李白): A Toolbox for Large-Scale Distributed Parallel Training
https://libai.readthedocs.io
Apache License 2.0
391 stars 55 forks source link

Add support for dynamic control flow of auto placement. #476

Open AsakusaRinne opened 1 year ago

AsakusaRinne commented 1 year ago

修复之前auto placement不支持dynamic control flow的问题。

目前仍然有以下几个点需要注意或以后解决:

  1. transformers代码里有许多对于oneszeros的调用,但oneflow 0.9.0对它们的实现是调用了Ones().__call__(),这样会导致fx自动对__call__进行包装从而使其出错。目前oneflow master分支已经修改成了forward,所以onefx暂时没有做处理。目前要跑通需要用到oneflow的master分支。
  2. 目前compile_auto_placement函数改成返回一个interpreter,在外面则是接收返回值之后调用run。之前直接返回GraphModule的方式不行,这可能是一个BUG,需要以后修复。
  3. 原本的测试代码中对auto placement之后的module调用了generate函数,这样是不成立的,因为fx就是针对Module的forward进行trace的。这个和2无关,不论是否使用interpreter都是这样。这个可能可以实现,但是肯定会需要自己指定要trace的函数名,这部分后续需要再讨论一下。
  4. 在调用compile_auto_placement的时候,不用指定concrere_args,但是有一个约定就是如果模型实际运行的时候会指定某个参数,那就需要对这个参数以keyword arg的形式传入compile_auto_placement。这个方式后续可以讨论修改一下。
  5. onefx也做出了少量必要的修改,这边的onefx版本应该升级到0.0.3。
CPFLAME commented 1 year ago

原本的测试代码中对auto placement之后的module调用了generate函数,这样是不成立的,因为fx就是针对Module的forward进行trace的。这个和2无关,不论是否使用interpreter都是这样。这个可能可以实现,但是肯定会需要自己指定要trace的函数名,这部分后续需要再讨论一下。

@BBuf 这个我们可以一起讨论

CPFLAME commented 1 year ago

看上去这个代码是个临时版本, 等ready了以后我再来review吧