初识 Oneflow

第一次听说 Oneflow 是在知乎的一个问题上《如何评价 7 月 31 日一流科技开源的深度学习框架 OneFlow?》,其的口号是主打性能(说实话,当时我是不信的)。当时并没有认真关注这个问题,因为国产的深度学习框架开源也不是一个两个了,比如清华的 Jittor,旷视的 MegEngine,华为的 Mindspore和百度的 PaddlePaddle。不关注的原因是因为这些框架都是为了公司内部的需求而首先设计的,之后再开源的,感觉可能不是很符合科研需求;另外已经熟悉了 PyTorch 的 API 和生态了,再转到另一个框架太麻烦也没有必要。最近,Oneflow 宣布其会兼容 PyTorch 的 API,并且放出了其性能测试报告(真香),一下子就让我有了兴趣,毕竟白嫖的性能谁不喜欢呢。本文主要是 Oneflow 官方发布的 Benchmark 和 介绍的搬运,后面等到放出与 PyTorch 兼容的新版本后会上手体验一下。

框架性能横评

Oneflow 既然主打性能,那肯定要凭数据说话的。为了证明 Oneflow 对比其他框架在性能上的优越性,其在官方仓库 Oneflow-Inc/DLPerf 放出了《基于ResNet50与BERT模型的深度学习框架性能评测报告》。下面我挑重点说一下,感兴趣可以去阅读原报告。

首先是在 ResNet50 上的性能结果对比,Oneflow 对比了几大主流的深度学习框架,包括 Tensorflow、PyTorch、MXNet和PaddlePaddle。其中 Tensorflow、PyTorch、MXNet 三者均有 NGC 版本,这个版本主要是配备了 DALI 这个数据加速读取的库,能够极大地加速数据读取速度,另外还有一些其他优化,可以在官网了解更多细节。下表中展示了不同框架使用 ResNet50 在图片分类上的吞吐率(batch size为256,并且使用了自动混合精度)。

#node #device OneFlow NGC MXNet NGC TensorFlow 1.x w/XLA NGC TensorFlow 1.x NGC PyTorch MXNet TensorFlow 2.x PaddlePaddle w/DALI
1 1 1472.72 1393.87 1198.55 945.18 802.9 833.65 1075.27 887.17
1 8 10629.32 9621.31 8588.45 6903.42 6154.66 2908.88 6898.00 6862.17
2 16 17920.40 16219.03 14931.03 12021.09 11991.94 5451.27 9870.15 6018.46
4 32 33141.02 30713.68 29171.69 24734.22 22551.16 10565.55 19314.31 11617.57

首先看一下单卡的结果,除了 NGC MXNet 能够看到 Oneflow 的车尾灯,其他框架都被甩了一条街。普通的 MXNet 没了 DALI 之后吞吐率也是暴跌,说明数据读取是个大的瓶颈。
然后看一下单机8卡的结果,从绝对吞吐率来说,Oneflow 依然是一枝独秀,吊打其他框架。从加速比上看,Oneflow 达到了 7.2(8卡),已经非常不错了。但是低于 PaddlePaddle 的 7.73(8卡),我怀疑 PaddlePaddle 就是捣乱的,单卡太差了,有着极大的上升空间;和 NGC Tensorflow 的 7.16(8卡)旗鼓相当。
最后直接看4机32卡,绝对吞吐率已经没有什么好看了,依旧是第一。主要关注加速比,Oneflow 是 22.51(32卡),低于 NGC PyTorch 的28.11(32卡),这个就有点低了,期待后续的优化吧。

除了 ResNet50 网络,Oneflow 在 BERT 网络上对比了其他框架(batch size开到最大,开启自动混合精度),结果见下表。对比下来,绝对吞吐率依然遥遥领先,加速比分别为 7.24(8卡),25.99(32卡),和 ResNet50 上的结论类似:单卡已经做的很好了,但是多卡加速比(尤其是大于8卡)仍有提高空间。另外一点,就是 Oneflow 能够支持的 batch size 比其他框架大很多,达到了 160,远超第二名 NGC Tensorflow 1.x 的 96。这一点对于小实验室很重要,因为 Oneflow 框架非常节省显存。

#node #device OneFlow bsz=160 NGC TensorFlow 1.x W/O XLA bsz=64 NGC TensorFlow 1.x with XLA bsz=96 NGC PyTorch bsz=96 PaddlePaddle bsz=160 OneFlow W/O clip bsz=160 MXNet W/O clip bsz=128
1 1 605.11 183.25 468.10 462.35 309.68 613.93 544.31
1 8 4381.66 1452.59 3559.8 3287.12 1666.54 4683.36 3825.21
2 16 8075.16 2653.74 5960.14 5426.07 2969.85 8777.57 7327.50
4 32 15724.70 5189.07 11650.0 10349.12 5452.35 17210.63 14822.31