search

Oneflow-Inc / swin-transformer

0 stars 0 forks source link

Swin-Transformer 显存问题跟踪 #6

Open Ldpe2G opened 2 years ago

Ldpe2G commented 2 years ago

Swin Transformer 8卡 DDP 显存占用对比

实验设置

实验机器 V100 16G Pytorch 版本： 1.10.1 Oneflow 版本： master nightly 0.6.0.dev20220102+cu112 Oneflow swin 测试分支： swin_clean_ldp Pytorch swin 测试基于官方 swin 仓库

测试结果

Framework Swin Model Batch Size per GPU GPU0 显存占用 (单位MiB)
Pytorch Tiny 28M 64 10264
Oneflow Tiny 28M 64 11418
Pytorch Tiny 28M 96 14054
Oneflow Tiny 28M 96 15698
Pytorch Tiny 28M 128 Out of Memory
Oneflow Tiny 28M 128 Out of Memory
Pytorch Small 50M 64 15054
Oneflow Small 50M 64 Out of Memory

测试结论

注

daquexian commented 2 years ago

Pytorch 在训 swin base 的时候，会用上 accumulation steps 和 gradient checkpointing 等技术

如果 pytorch 会用这些技术的话，就不是公平的比较了吧（因为我们没有用）

Ldpe2G commented 2 years ago

Pytorch 在训 swin base 的时候，会用上 accumulation steps 和 gradient checkpointing 等技术

如果 pytorch 会用这些技术的话，就不是公平的比较了吧（因为我们没有用）

这个测试没用到，我只是提一下，torch如果完整训练 base 会用上这些技术

Ldpe2G commented 2 years ago

实验设置

实验机器 V100 16G Pytorch 版本： 1.10.1 Oneflow 版本： dev_layer_norm_backward_param_grad Oneflow swin 测试分支： swin_clean_ldp Pytorch swin 测试分支： swin_clean_torch_ddp

都运行 train_with_real_data_ddp.sh

测试结果

Framework Swin Model Batch Size per GPU 记录最大显存占用 (单位MiB)
Pytorch Tiny 28M 64 9992
Oneflow Tiny 28M 64 11566
Pytorch Tiny 28M 96 13806
Oneflow Tiny 28M 96 Out of memory
hjchen2 commented 2 years ago

下面是我这边在分支dev_layer_norm_backward_param_grad_merge_refactor_release_tensor上跑的一些结果，看单卡显存已经和pytorch差不多了，但DDP多卡会比pytorch多1.5G，建浩后续可以帮忙一起看一下 。

注：分支dev_layer_norm_backward_param_grad_merge_refactor_release_tensor是dev_layer_norm_backward_param_grad的基础上合并了新奇重构release tensor的PR。

实验设置

实验机器 V100 16G Pytorch 版本： 1.10.1 Oneflow 版本： dev_layer_norm_backward_param_grad_merge_refactor_release_tensor

Framework Swin Model Batch Size per GPU DDP单卡显存占用 (单位MiB) DDP八卡显存占用 (单位MiB)｜
Pytorch Tiny 28M 64 9568MiB 9984MiB
Oneflow Tiny 28M 64 9578MiB 11686MiB
hjchen2 commented 2 years ago

下面是我这边在分支dev_layer_norm_backward_param_grad_merge_refactor_release_tensor上跑的一些结果，看单卡显存已经和pytorch差不多了，但DDP多卡会比pytorch多1.5G，建浩后续可以帮忙一起看一下 。

注：分支dev_layer_norm_backward_param_grad_merge_refactor_release_tensor是dev_layer_norm_backward_param_grad的基础上合并了新奇重构release tensor的PR。

实验设置

实验机器 V100 16G Pytorch 版本： 1.10.1 Oneflow 版本： dev_layer_norm_backward_param_grad_merge_refactor_release_tensor

Framework Swin Model Batch Size per GPU DDP单卡显存占用 (单位MiB) DDP八卡显存占用 (单位MiB)｜ Pytorch Tiny 28M 64 9568MiB 9984MiB Oneflow Tiny 28M 64 9578MiB 11686MiB

@daquexian @Ldpe2G depeng后续也可以把DDP单卡的结果跑一下看看

Ldpe2G commented 2 years ago

实验设置

实验机器 V100 16G Pytorch 版本： 1.10.1 Oneflow 版本： dev_layer_norm_backward_param_grad_merge_refactor_release_tensor

跑10次记录显存的波动范围

Framework Swin Model Batch Size per GPU DDP单卡显存占用 (单位MiB)
Pytorch Tiny 28M 64 9568 ~ 9568 MiB
Oneflow Tiny 28M 64 9516 ~ 10926 MiB
hjchen2 commented 2 years ago

继昨天发现DDP多卡时会出现显存异常的问题之后，今天修改了一下DDP的实现，让DDP直接返回原module，这样每个rank都独立执行，rank之间无通信，但这样改完之后跑DDP8卡显存仍然异常，所以就排除了DDP的嫌疑。

除去DDP引入的通信问题，那一个rank和起8个rank最大的区别就是CPU占用了，因为每个rank的data reader都开了8个子进程，那8个rank就至少有64个子进程。我尝试在8个rank时把data reader的num wokers设置为1，发现这样起8个rank的显存就和单卡是一致了。

Ldpe2G commented 2 years ago

datareader 里面的逻辑会用 gpu 吗，测试脚本的逻辑是从dataloader取出数据再 to 到 cuda 上的

Flowingsun007 commented 2 years ago

这个多出来的显存会是之前那个flow.Generator类似的产生的cuda context吗？

Flowingsun007 commented 2 years ago

datareader 里面的逻辑会用 gpu 吗，测试脚本的逻辑是从dataloader取出数据再 to 到 cuda 上的

datareader的逻辑都是cpu op没有用到gpu

hjchen2 commented 2 years ago

layer norm后向显存优化的PR（https://github.com/Oneflow-Inc/oneflow/pull/6996 ）已经合到master了，后续测显存可以更新一下代码再测。

hjchen2 commented 2 years ago

昨天新奇提出一个现象，release tensor指令执行时并没有真正触发tensor storage的释放，这个问题确实是view引起的。release tensor指令只处理了原tensor的eager blob object，把它自己的那份tensor storage给reset了，但其他view tensor的eager blob object还引用着tensor storage。view tensor的eager blob object不能通过release tensor指令主动reset tensor storage，只能等到触发虚构函数去reset。

有个问题是release tensor指令已经在执行了，证明所有与该tensor storage有关联的tensor都消亡了，与它们相关的指令都早已发送并执行了，那还有什么地方会引用到eager blob object呢？

我初步怀疑是指令本身引用了eager blob object，一条指令执行结束还需要等到这个指令被回收完成才会释放其引用的eager blob object。在测试过程中，由于某种原因导致vm回收指令变慢了，vm线程一直在忙于调度指令，导致显存不能及时释放，且一直又在申请新的显存，最终显存占用变大，甚至out of memory。

https://github.com/Oneflow-Inc/oneflow/pull/7287 这个PR里让tensor storage提供了一个直接释放的接口，eager blob object的tensor storage在reset的时候先执行一下这个接口，强制释放其管理的内存。