LinkedIn 严重依赖人工智能为其超过 5.75 亿会员提供内容和创造经济机会。随着深度学习技术的迅速发展,人工智能工程师已经开始在 LinkedIn 的关联驱动产品中采用深度神经网络,包括反馈和智能回复。这些用例中的许多都构建在由谷歌编写的深度学习框架 TensorFlow 上。
一开始,我们内部的 TensorFlow 用户在小型的、非托管的“裸机”集群上运行框架。但我们很快意识到需要将 TensorFlow 与基于 Hadoop 的大数据平台的庞大计算和存储能力结合起来。我们的 Hadoop 集群中存储了数百 pb 的数据,这些数据可以用于深度学习,因此我们需要一种可伸缩的方式来处理所有这些信息。幸运的是,TensorFlow 支持分布式训练,这是一种处理大型数据集的有效技术。然而,部署分布式 TensorFlow 并不是一个简单的工作,也不是所有数据科学家和相关工程师都具备的专业知识或能力——尤其是因为它必须手工完成。我们想要一种灵活和可持续的方式来弥补分布式 TensorFlow 分析能力和 Hadoop 伸缩能力之间的差距。
开源 TonY
为了满足我们的需求,并且我们知道还有许多其他对运行分布式机器学习感兴趣的人也在运行大型的 Hadoop 部署,我们已经在 YARN(TonY)上构建了 TensorFlow,我们今天将其开源。你可以查看 GitHub 上的 TonY 项目,了解如何使用它的详细信息。欢迎大家提意见并为其贡献代码!
开源项目传送门:https://github.com/linkedin/TonY
在本文的剩下部分中,我们将介绍 TonY 的内部细节、我们在 Hadoop 上实现并利用的扩展分布式 TensorFlow 的特性以及实验结果。
现有解决方案
在对 Hadoop 上运行分布式 TensorFlow 的初步调查中,我们发现了一些现有的解决方案。然而,我们最终确定没有一个满足我们的特定需求,导致我们决定构建 TonY。
基于 Spark 的 TensorFlow 是一个开源的解决方案,它允许你在 Apache Spark 计算引擎上运行 TensorFlow。我们能够在这个框架上装载两个内部的深度学习应用程序,但是遇到了一些问题,最明显的是 GPU 调度和异构容器调度都存在不足。此外,我们将来想要进行的任何调度和应用程序生命周期增强都必须在 Spark 中完成,这比在独立的 YARN 应用程序中进行更改要困难得多。
TensorFlowOnYARN 是另一个作为独立库运行的开源解决方案。不幸的是,这个项目中的容错支持和可用性并不符合我们的需求。此外,这个项目已经不再被维护。
基于这些原因,我们决定构建 TonY,以便完全控制 Hadoop 集群中的资源。另外,由于 TonY 是直接在 YARN 上运行,并且作为轻量级依赖程序运行,因此我们可以很容易地将它与 YARN 中的堆栈的低级部分和 TensorFlow 中的堆栈的高级部分结合起来。
TonY 的工作原理
类似于 MapReduce 提供的可在 Hadoop 上运行 Pig/Hive 脚本的引擎,Spark 提供运行调用 Spark APIs 的 scala 代码的引擎。TonY 旨在提供在 Hadoop 上运行 TensorFlow 任务的一流支持, 比如资源调度和容器环境设置。
TensorFlow 运行在 YARN 上的 TonY 上
TonY 有三个主要组件:Client、ApplicationMaster 和 TaskExecutor。下面是运行 TonY 作业的端到端过程:
用户向 Client 提交 TensorFlow 模型训练代码、提交参数及其 Python 虚拟环境(包含 TensorFlow 依赖项)。
Client 设置 ApplicationMaster(AM)并将其提交给 YARN 集群。
AM 根据用户的资源需求(参数服务器和任务量、内存和 GPU)与 YARN 的资源管理器进行资源协商调度。
一旦 AM 接收到分配,它就会在分配的节点上生成 TaskExecutors。
TaskExecutors 启动用户的训练代码并等待其完成。
用户的训练代码开始运行,TaskExecutors 和 AM 之间定期进行心跳检查。
Tony 架构
除了支持在 Hadoop 上运行分布式 TensorFlow 作业的基本功能外,TonY 还实现了各种功能来提高运行大型训练的体验:
GPU 调度:最近,Hadoop 增加了对 GPU 调度和隔离的原生支持。对于用户来说,这意味着他们可以确保一旦从 Hadoop 获得了他们的容器分配,他们就可以可靠地获得他们请求的 gpu 的数量。TonY 还可以获得 GPU 资源信息,因此它能够利用 Hadoop 的 API 从集群请求 GPU 资源。
细粒度的资源请求:由于 TonY 支持将不同的实体(例如参数服务器和任务量)请求为单独的组件,因此用户可以对每种类型发出不同的资源请求。例如,你的参数服务器和任务量可能有不同的内存需求。或者,你可能希望在 gpu 或其他特定硬件上运行训练,但是在参数服务器上使用 cpu 就足够了。对于用户来说,这意味着对应用程序的资源需求有更多的控制,对于集群管理员来说,这有助于避免昂贵的硬件资源浪费。
TensorBoard 支持:TensorBoard 是一个工具,它可以使你更容易地理解、调试和优化 TensorFlow 程序。由于 TensorBoard 进程是由任务启动时不确定的应用程序进程启动的,通常我们无法从 Hadoop UI 中看到 TensorBoard。最近,我们为 YARN 编写了一些代码,允许我们将 Hadoop 应用程序的跟踪 URL 重定向到 TensorBoard 上,这样只需点击一下就能看到 TensorBoard 了。
容错:使用大量的机器,TensorFlow 训练可能需要几个小时或几天的时间。因此,长期运行的 TensorFlow 作业比短期作业更容易受到瞬时错误或抢占的影响。TensorFlow 包含容错 API,可以将检查点保存到 HDFS,并从以前保存的检查点恢复训练状态。TonY 通过提供有弹性的分布式基础设施来从节点故障中恢复,从而简化了这个过程。如果一个任务没有心跳发送到 AM 或发送超时,TonY 将重新启动应用程序并从以前的检查点恢复训练。
实验结果
我们在 TonY 上运行了 Inception v3 模型,分配了 1 到 8 个任务,每个任务有一个 GPU(在 8 个任务中有 1 个同时也使用 CPU 训练),使用异步训练。该模型是一个著名的 ImageNet 深度神经网络,该数据集包含数百万张用于训练图像分类模型的图像。正如在 Inception v3 分布式训练示例中一样,我们测量 10 万次,批处理大小为 32。结果如下:
产生这些结果的测试环境为 40G RAM / 1 CPU,Tesla K80 GPU, RHEL 6.6 和 TensorFlow 1.9。通过 GPU 训练的 8 个任务进行 100,000 次后,最后的前 5 名错误率为 26.3%。
由于 TonY 位于协调分布式 TensorFlow 的层中,并且不干扰 TensorFlow 作业的实际执行,所以理论上它不产生开销。实际上,对于 GPU 训练,运行时是线性扩展的。我们还看到,在 CPU 训练时运行 GPU 训练,速度提高了大约 4 倍,考虑到模型的复杂性和深度,这一结果基本符合预期。