libai icon indicating copy to clipboard operation
libai copied to clipboard

Dev add dalle2

Open ofhwei opened this issue 2 years ago • 12 comments

添加dalle2,后续持续完善中并记录遇到的一些问题。

python: 3.6.9 oneflow: 0.7.0+cu112

代码主要参考dalle2_pytorch, 里面还有一部分没太看明白:( ,基本上是 import torch -> import oneflow as flow, 然后把代码中的'torch'全部替换为'flow', 少部分接口改一改入参格式,基本可以跑通 :) 。

ofhwei avatar Aug 04 '22 02:08 ofhwei

1、在4卡机器上,这样启动的时间有点久,有没有什么快一点的方法呢 截屏2022-08-04 11 16 38

ofhwei avatar Aug 04 '22 03:08 ofhwei

projects/DALLE2/dalle2/data/bpe_simple_vocab_16e6.txt

这个文件应该不用提到 git 吧,超大

strint avatar Aug 04 '22 03:08 strint

projects/DALLE2/dalle2/data/bpe_simple_vocab_16e6.txt

这个文件应该不用提到 git 吧,超大

这个大约3M,感觉不是很大。。

ofhwei avatar Aug 04 '22 03:08 ofhwei

projects/DALLE2/dalle2/data/bpe_simple_vocab_16e6.txt 这个文件应该不用提到 git 吧,超大

这个大约3M,感觉不是很大。。

代码变更数量:+266,903

strint avatar Aug 04 '22 03:08 strint

在4卡机器上,这样启动的时间有点久,有没有什么快一点的方法呢

启动时间比较久是指通过指定环境变量的方式启动多个进程比较繁琐,还是指import时间比较长。 如果是觉得这种启动方式繁琐的话可以参考 https://docs.oneflow.org/master/parallelism/04_launch.html 使用launch模块。

shangguanshiyuan avatar Aug 04 '22 03:08 shangguanshiyuan

projects/DALLE2/dalle2/data/bpe_simple_vocab_16e6.txt 这个文件应该不用提到 git 吧,超大

这个大约3M,感觉不是很大。。

代码变更数量:+266,903

ok, 已经删除了

ofhwei avatar Aug 04 '22 03:08 ofhwei

在4卡机器上,这样启动的时间有点久,有没有什么快一点的方法呢

启动时间比较久是指通过指定环境变量的方式启动多个进程比较繁琐,还是指import时间比较长。 如果是觉得这种启动方式繁琐的话可以参考 https://docs.oneflow.org/master/parallelism/04_launch.html 使用launch模块。 好的,我试试。 前面这个是import太久了,现在还没完成。。

ofhwei avatar Aug 04 '22 03:08 ofhwei

好的,我试试。 前面这个是import太久了,现在还没完成。。

import 时间这么长很不合理,你看看在服务器裸机上能不能复现这个问题。 另外,我看你用的是python3.6,已经计划不支持3.6了。

shangguanshiyuan avatar Aug 04 '22 04:08 shangguanshiyuan

我看这个截图里WORLD_SIZE是4,是不是没有启动其他的3个进程,此时在等待。

shangguanshiyuan avatar Aug 04 '22 04:08 shangguanshiyuan

我看这个截图里WORLD_SIZE是4,是不是没有启动其他的3个进程,此时在等待。

刚试了一下,还真是这样(之前不太清楚要开4个shell窗口)。

ofhwei avatar Aug 04 '22 04:08 ofhwei

2、相关问题/进展记录

进展

  • 把模型从torch迁到oneflow还是比较方便的, 其中flow存在一些函数不支持或不一致的情况(如 F.pad 不支持bool类型,0.7版本缺unbind/full_like等 ),稍微修改一下就行,这一块工作量不是很大。然后是添加(copy)了GroupNorm, Conv2D, ConvTransposed2D等网络层进行适配。
  • 主要耗时是在寻找训练好的模型权重,主要来自这里,我试了若干个,根据文本生成的图片都噪声特别大,或者说感觉像是随机生成的🤔,所以现在dalle2推理部分可以跑通,但是效果远不及预期,还需要多试试。

其它问题 服务器上有4块卡,但是我目前都是使用单机单卡来跑,原因有两点

  1. 因为我在尽量利用libai.layers.xx等网络层,如Embedding, LayerNorm等,这些内置的参数plcament/sbp貌似都是一致(如layer_idx=0, sbp=broadcast); 另一方面,dalle2有太多参数了,目前不太想一个个去设置placement和sbp😂, 所以我基本都放在在第0块卡上。
  2. 参考

启动时间比较久是指通过指定环境变量的方式启动多个进程比较繁琐,还是指import时间比较长。 如果是觉得这种启动方式繁琐的话可以参考 https://docs.oneflow.org/master/parallelism/04_launch.html 使用launch模块。

要使用多卡的话每次都需要起四个进程,感觉还是稍微有点麻烦😂(可能是我的使用方式有问题),所以就直接使用单卡了。

后续计划

  1. 优化代码,排查问题,看看是不是有什么地方有遗漏,导致迁移的模型达不到预期效果。
  2. 现在单卡跑一次推理要40min左右,后续尝试多卡并行跑(如流水并行等-这块我暂时不是很熟悉),看看能不能提速。

ofhwei avatar Aug 05 '22 11:08 ofhwei

根据文本生成的图片都噪声特别大,或者说感觉像是随机生成的🤔

这里在对比输出的时候, 不要通过观察图片的结果, 而是在去除所有随机性的干扰以后, 用相同的输入 放到flow和torch的网络下 看看是否能得到相同的输出的tensor. 可以参考 https://github.com/Oneflow-Inc/OneTeam/issues/779#issue-1054706350 下的模型loss对齐经验 下的第一个板块.

因为我在尽量利用libai.layers.xx等网络层,如Embedding, LayerNorm等,这些内置的参数plcament/sbp貌似都是一致(如layer_idx=0, sbp=broadcast); 另一方面,dalle2有太多参数了,目前不太想一个个去设置placement和sbp😂, 所以我基本都放在在第0块卡上。

这个我近期会提一个pr, 在libai下面不需要全部替换model中的tensor, 在model中只替换自己需要的layer即可, 其他的都默认全部转换为broadcast.

多卡的启动

在libai下其实提供了多卡启动的脚本. https://github.com/Oneflow-Inc/libai/blob/5d5acf9aa69ab5b5da8ae9d992dce4afe0d1964c/tools/train.sh#L13-L14.

但我看到这个pr下面目前还没有config.py之类的. 所以如果你想方便启动的话. 可以参考这个 重新写一个.sh的shell脚本

CPFLAME avatar Aug 09 '22 02:08 CPFLAME

nsys profile --stats=true -o dalle2 python3 test.py 利用nsight system工具打开上述结果,发现ComputeVarUsingWelfordWrapper耗时占比非常大,这个ComputeVarUsingWelfordWrapper具体是指什么操作呢,另外如何避免/减少它的开销? py: 3.8 oneflow: 0.8.0 image

ofhwei avatar Aug 16 '22 09:08 ofhwei

@Flowingsun007 露阳在disco diffusion里有完全一样的问题,要增加group norm,改变计算方式,是不是已经搞定了

yuanms2 avatar Aug 16 '22 09:08 yuanms2

nsys profile --stats=true -o dalle2 python3 test.py 利用nsight system工具打开上述结果,发现ComputeVarUsingWelfordWrapper耗时占比非常大,这个ComputeVarUsingWelfordWrapper具体是指什么操作呢,另外如何避免/减少它的开销? py: 3.8 oneflow: 0.8.0 image

这个ComputeVarUsingWelfordWrapper就是var op的kernel实现,这个kernel实现的效率有问题,暂时还没有修复。如果是网络中用到了nn.GroupNorm的话,可以先按照这个pr:https://github.com/Oneflow-Inc/oneflow/pull/8905 中,修改一下GroupNorm中var的实现,可以绕过这个var kernel的问题,大大加速

Flowingsun007 avatar Aug 16 '22 09:08 Flowingsun007

这个ComputeVarUsingWelfordWrapper就是var op的kernel实现,这个kernel实现的效率有问题,暂时还没有修复。如果是网络中用到了nn.GroupNorm的话,可以先按照这个pr:https://github.com/Oneflow-Inc/oneflow/pull/8905 中,修改一下GroupNorm中var的实现,可以绕过这个var kernel的问题,大大加速

修改了GroupNorm,速度提升了很多👍

ofhwei avatar Aug 16 '22 10:08 ofhwei

oneflow版本从0.7->0.8时,对global tensor操作,Device2Device Memcpy耗时会大量增加 如下部分代码:

import oneflow as flow
from oneflow import nn

print(flow.__version__)

def rotate_half(x):
    d = x.shape[-1]
    x1, x2 = x[..., :d // 2], x[..., d // 2:]
    return flow.cat((-x2, x1), dim=-1)

class RotaryEmbedding(nn.Module):
    def __init__(self):
        super().__init__()
        freqs = flow.randn(81, 32, placement=flow.placement(type='cuda', ranks=[0]), sbp=flow.sbp.broadcast)
        self.register_buffer("freqs", freqs)

    def forward(self, t, start = 0, end = 32):
        t_left, t, t_right = t[..., :start], t[..., start:end], t[..., end:]
        t = t * self.freqs.cos() + rotate_half(t) * self.freqs.sin()
        return flow.cat((t_left, t, t_right), dim=-1)

rotaryEmbedding = RotaryEmbedding()
t = flow.randn(6, 32, 81, 64, placement=flow.placement(type='cuda', ranks=[0]), sbp=flow.sbp.broadcast)
for i in range(2000):
    print(i)
    t = rotaryEmbedding(t)

在python3.8+oneflow0.7时是OK的,但是把oneflow升级为0.8.0时,会产生大量的Device2Device Memcpy耗时操作, 如下图: 截屏2022-08-18 15 07 38

ofhwei avatar Aug 18 '22 07:08 ofhwei

感谢反馈问题。我们调查一下。

@clackhan

lixinqi avatar Aug 18 '22 08:08 lixinqi

oneflow版本从0.7->0.8时,对global tensor操作,Device2Device Memcpy耗时会大量增加

原因是支持所有 sbp 输入的 slice kernel 实现后,在 kernel 部分对最后一维切分的处理不够好, 后面在 https://github.com/Oneflow-Inc/oneflow/pull/8589 中已经优化,可以使用 nightly 版本再测一下。

wyg1997 avatar Aug 19 '22 02:08 wyg1997

可以加一下 requirements.txt, 在libai的环境下 还需要额外手动安装一下包 .

我这边尝试的时候发现需要手动安装的就有如下的包:

resize_right
ftfy
kornia
einops

CPFLAME avatar Sep 09 '22 03:09 CPFLAME

dalle2的可以加一个requirements.txt,

resize_right
ftfy
kornia
einops

再加一个下载权重的脚本, 或者有一个下载权重的readme说明也可以.

CPFLAME avatar Sep 16 '22 03:09 CPFLAME

可以在reviewer里面加一下ci

CPFLAME avatar Sep 23 '22 11:09 CPFLAME