比 TensorFlow 云快 46 倍！IBM 用 Snap ML 和 Tesla V100 秀肌肉

2018 年 3 月 28 日 AI研习社 杨晓凡

AI 研习社按：TensorFlow 的机器学习库可以说大家都已经很熟悉了，但 IBM 的研究人员们表示这都是小意思。

今年 2 月的时候，谷歌的软件工程师 Andreas Sterbenz 曾在谷歌官方博客撰写文章，介绍如何使用谷歌云机器学习（Google Cloud Machine Learning）和 TensorFlow 对大规模广告投放和推荐系统做点击率预测。根据介绍，他用了大小为 1TB、包含了 42 亿条训练样本、一百万个特征的 Criteo Terabyte Click Logs 数据训练了机器学习模型，用来预测未来显示的广告可能被点击的概率如何。

在数据预处理之后，实际训练过程用到了 60 台计算服务器和 29 台参数服务器。训练模型花费了 70 分钟，训练损失最终降低到了 0.1293（训练损失可以看作最终预测准确率的粗略近似）。

Sterbenz 也尝试了不同的建模技术，看看能否继续降低训练损失。不过各种好方法基本都会带来更长的训练时间。最终他选用了深度神经网络，训练时有 3 个 epoch，一共花费了 78 小时训练完毕。

不过这种结果对 IBM 来说毫无吸引力，他们想要借机证明自己的训练框架有多么高效。他们认为，自己的带有 GPU 的 POWER9 服务器同样做前一项训练任务的话，要比谷歌云平台的这 89 台服务器快多了。

Power9 服务器架构特点

在苏黎世 IBM 研究院工作的 Thomas Parnell 和 Celestine Dünner 也找来了同样的 1TB 训练数据，训练的模型也是和前面 70 分钟完成训练的一样的逻辑回归模型。但不同的是，他们用的不是 TensorFlow 的机器学习库，而是 Snap Machine Learning（https://arxiv.org/pdf/1803.06333.pdf ）。

IBM 两人用的服务器是 Power System AC922，一共有八台 POWER9 服务器，每台服务器搭载两块 NVIDIA Tesla V100 GPU。训练只经过 91.5 秒就结束了，比 Sterbenz 在谷歌云机器学习平台上的 70 分钟快 46 倍。

他们也把自己的结果和其它各个系统的结果做成了图表进行对比：

可以看到，相比于运行在 TensorFlow 上，IBM 不仅通过 Snap ML 得到了很短的训练时间，训练损失也要稍低一些。

46 倍的速度提升毕竟不是一个小数目，那么其中的改进具体有哪些呢？

根据 Thomas Parnell 和 Celestine Dünner 两人介绍，Snap ML 中引入了许多不同层级的并行化计算设计，可以在同一个集群内的不同节点之间分配任务量、可以发挥出加速计算单元（比如 V100 GPU）的计算能力，而且在单个计算单元中也可以利用到多核心的并行计算能力。

其中的并行化设计可以概述如下：

首先把数据分配给集群内的各个计算节点
在单个节点中，数据一部分分配给 CPU、一部分分配给 GPU，CPU 和多张 GPU 可以同时进行计算
计算时，GPU 中的多个核心同时参与运算，CPU 的运算负载也是多线程的，可以更好利用多核心 CPU

Snap ML 中也有内置的层级化算法，可以让这各个级别的并行化手段高效协同运作。

IBM 的研究人员们并不是借此指责 TensorFlow 中没有好好利用并行化，但他们确实表示：「我们设计了专用的求解器，以便能够完全利用这些 GPU 的海量并行计算能力；同时我们还保证了数据在 GPU 内存中的局部性，避免让大量数据传输带来额外的开销。」

另外，AC922 服务器和 V100 GPU 之间的连接总线是 NVLink 2.0，而传统英特尔至强服务器（比如使用 Xeon Gold 6150 CPU @ 2.70GHz）只能使用 PCI-E 总线连接到 GPU。前者的有效数据传输带宽达到 68.1GB/s，后者仅有 11.8GB/s。看起来，PCI-E 总线速度可能也是系统性能的瓶颈之一，传输一个数据包需要 318ms，而 NVLink 2.0 只需要 55ms。

IBM 团队还表示：「我们还为系统中的算法设计了一些新的优化手段，可以更适合处理稀疏的数据结构。」

以上总总因素汇合起来，IBM 通过更好地利用 GPU 性能打败了谷歌的云服务器似乎还挺合理。不过据我们所知，IBM 目前并没有公开提供过任何 POWER9 与英特尔至强服务器之间的直接性能对比。另一方面，关于 Snap ML 到底有多好也只有在同样的硬件环境上运行 Snap ML 和 TensorFlow 才能知道。

（关于 Tesla V100 GPU，更多信息参见雷锋网(公众号：雷锋网) AI 科技评论此前报道一文详解英伟达刚发布的 Tesla V100 究竟牛在哪？）

via theRegister，雷锋网 AI 研习社编译