[翻译] TensorFlow 分布式之论文篇 Large-Scale Machine Learning on Heterogeneous Distribute-一一网

本系列我们开始分析 TensorFlow 的分布式。之前在机器学习分布式这一系列分析之中，我们大多是以 PyTorch 为例，结合其他框架/库来穿插完成。但是缺少了 TensorFlow 就会觉得整个世界（系列）都是不完美的，不单单因为 TensorFlow 本身的影响力，更因为 TensorFlow 分布式有自己的鲜明特色，对于技术爱好者来说是一个巨大宝藏。

读论文有一种原则是：本领域最经典的论文，近5年最热的论文，近1年最新的论文。按照这个原则，本文主要介绍一篇 TensorFlow 经典论文 TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems。大家如果读了下面论文就会发现 TensorFlow分布式的博大精深。

本文图来自原始论文。

1. 原文摘要

基于我们在 DistBelief 方面的经验，以及对期望的系统特性和训练/使用神经网络需求的更全面的理解，我们构建了 TensorFlow ，这是我们的第二代系统，用于实现和部署大规模机器学习模型。 TensorFlow 采用类似数据流的模型来描述计算，并将其映射到各种不同的硬件平台上，从在 Android 和 iOS 等移动设备平台上运行推理，到使用包含一个或多个 GPU 卡的单机的中等规模训练和推理系统，再到在数百台具有数千个 GPU 的专用主机上运行的大规模训练系统。

拥有一个能够跨越如此广泛的平台的单一系统可以显著简化机器学习系统的实际使用，因为我们发现，如果大规模训练和小规模部署都分别有自己的独立系统，则会导致巨大的维护负担和漏洞。 TensorFlow 计算被表示为有状态数据流图，我们致力于使系统具有足够的灵活性，以便用户可以快速试验新模型，系统同时也具有足够高的性能和鲁棒性，可以被用于机器学习模型的训练和部署。

为了将神经网络训练扩展到更大规模的部署， TensorFlow 允许客户机通过复制和并行执行核心模型数据流图来轻松表达各种并行性，这样可以使用许多不同的计算设备来更新一组共享参数或其他共享状态。对计算描述的适度更改允许用户实现多种不同的并行方法。

TensorFlow 允许在参数更新的一致性方面具有一定的灵活性，这些宽松的同步要求允许我们可以在一些较大的部署中更加轻松。与 DistBelief 相比， TensorFlow 的编程模型更加灵活，性能显著提高，支持在更广泛的异构硬件平台上训练和使用更广泛的模型。

2. 编程模型和基本概念

TensorFlow 计算由一组节点组成的有向图来描述。该图表示一个数据流计算，也允许让某些类型的节点来维护和更新持久状态，并以类似于 Naiad 的方式在图中实现分支和循环控制结构。客户端通常使用一种前端语言（C++或Python）构建计算图。下图显示了使用 Python 前端构建并执行 TensorFlow 计算图的示例片段。

图 1. TensorFlow 计算图示例片段

图 2，计算图

在 TensorFlow 图中，每个节点表示操作的实例，其具有零个或多个输入和零个或多个输出。在计算图中沿普通边流动的值（从输出到输入）被称为张量。张量是任意维数组，其基本元素类型在计算图构造时被指定或推断出来。特殊的边（被称为控制依赖关系）也可以存在于图中：这些边上没有数据流，但它们表明控制依赖关系的源节点必须在控制依赖关系的目标节点开始执行之前完成执行。因为我们的模型包含可变状态，因此客户可以直接使用控制依赖关系来强制执行。我们的实现有时也会在内部插入控制依赖项，以强制某些独立操作之间的顺序，例如，可以控制峰值内存使用。

2.1 算子（Operations）与核（Kernels）

算子（Operation）表示一个抽象计算（例如，”矩阵乘法”或”加法”）。一个算子可以拥有属性，但是所有属性必须在计算图构造时被提供或推断出来，这样才能实例化一个执行该算子的节点。属性的一个常见用途是使算子在不同的张量元素类型上具有多态性（例如，加法算子即支持两个类型为 float 的 tensors 相加，也支持两个类型为 int32的张量相加）。

核（Kernel）是可以在特定类型的设备（例如CPU或GPU）上运行的算子的具体实现。 TensorFlow 通过注册机制定义了一系列算子和核，这样意味着用户可以通过链接其他算子和/或内核来进行扩展。下图显示了 TensorFlow 库中内置的一些算子。

表 1. 算子

2.2 会话（Sessions）

客户端程序通过创建会话与 TensorFlow 系统交互。会话先构建一个空计算图，为了创建计算图，会话接口支持Extend 方法，该方法可以让客户端附加节点和边来扩充这个空图，然后进行计算。

2.3 变量（Variables）

在大多数计算中，一个图会被执行多次，而大多数张量在图的单次执行之后都不会存在。然而，有一些张量是在计算图的执行过程之中始终存在的，位置也是固定的，其不能正常流动但是可以更新，比如模型的参数，这就引出了变量这个概念。

变量是一种特殊的操作，它返回持久可变张量的句柄，这些句柄可以被传递给少量特殊的操作，例如 Assign 和AssignAdd（相当于+=），通过这些操作就可以改变这些变量引用的张量。

3. 实现

TensorFlow 系统中的主要组件是客户端，它使用会话接口与主机以及一个或多个工作进程进行通信。每个工作进程负责协调对一个或多个计算设备（如 CPU 内核或 GPU 卡）的访问以及按照主设备的指示在这些设备上执行计算图节点。 TensorFlow 接口有本地和分布式实现两种。

当客户端、master 和 worker 都在单个机器上单个进程的上下文之中运行时（如果机器安装了多个 GPU 卡，则可能使用多个设备），将使用本地实现。
分布式实现与本地实现共享大部分代码，但支持通过一个环境对其进行扩展，在该环境中，客户端、master和 worker 都可以在不同机器上不同的进程中。在我们的分布式环境中，这些不同的任务是 job 之中的一些容器，这些 job 由集群调度系统来管理。这两种不同的模式如下图所示。

图 3. 调度模式

3.1 设备（Devices）

设备是 TensorFlow 的计算核心。每个工作者负责一个或多个设备，每个设备都有一个设备类型和名称。设备名称由以下几部分组成：

设备类型。
设备在工作者中的索引。
分布式设置中对于工作者所在作业和任务的标识（如果设备是进程本地的，则为 localhost）。

示例设备名称的样例如下：

“/job:localhost/device:cpu:0″。
“/job:worker/task:17/device:gpu:3″。

PyTorch 有针对 CPU 和 GPU 的设备接口的实现，其他设备类型可以通过注册机制提供新设备实现。每个设备对象负责管理设备内存的分配和释放，以及执行 TensorFlow 下发的核方法。

3.2 张量

在我们的实现中，张量是一个类型化的多维数组。我们支持多种张量元素类型，包括大小从 8 位到 64 位的有符号和无符号整数、IEEE 浮点和双精度类型、复数类型和字符串类型（arbitrary byte array）。张量所在设备的分配器负责管理张量的存储区，张量存储缓冲区是引用计数的，在没有引用保留时会进行释放。

3.3 单设备执行

让我们首先考虑最简单的执行场景：一个拥有单个设备的工作者进程。计算图中的节点按照节点之间依赖关系的顺序来执行。我们将跟踪每个节点尚未执行的依赖项数量的计数。一旦此计数降至零，该节点就有资格执行，并被添加到就绪队列中。就绪队列以某种未指定的顺序进行处理，其将节点的核方法执行委托给设备对象。当节点完成执行时，依赖于此已完成节点的所有节点的计数都将减少。

3.4 多设备执行

一旦一个系统有多个设备，就有两个主要的复杂问题：如何决定将每个节点的计算放在哪个设备上，如何管理这些放置（Placement ）所带来的跨设备数据通信。本小节讨论这两个问题。

3.4.1 决定设备（Node Placement）

给定计算图之后， TensorFlow 实现的主要职责之一是将计算映射到可用设备集。本文给出了该算法的一个简化版本。

布局（placement）算法的一个输入是成本模型（cost model），该模型包含每个计算图节点的输入和输出张量的大小（字节）估计，以及每个节点在获得输入张量之后所需的计算时间估计。该成本模型要么基于与不同操作类型相关的启发式静态估计，要么基于计算图早期执行的实际布局决策集进行测量/决定。

布局算法首先运行计算图的模拟执行，然后使用贪婪启发式为图中的每个节点选择一个设备。此模拟生成的”节点到设备放置关系”最终也用作实际执行的放置。布局算法从计算图的源开始，并在前进过程中模拟系统中每个设备上的活动，在此遍历中：

当到达了一个节点，就考虑此节点的可使用设备集（如果设备不提供用户希望的实现特定操作的内核，则设备就不使用）。
对于具有多个可用设备的节点，布局算法使用贪婪启发式算法，看看将节点放置在每个可能设备上对节点完成时间会造成怎样的影响。该启发式方法不仅考虑了在成本模型中，该类设备上操作的估计或测量执行时间，还包括为了将输入从其他设备传输到该节点而引入的通信成本。然后选择最快完成节点操作的设备作为该操作的设备，
布局（placement）过程继续进行，以便为图中的其他节点（包括现在准备好模拟执行的下游节点）做出放置决策。

3.4.2 跨设备通信（Cross-Device Communication）

一旦决定了节点如何放置到设备之上（node placement），图就被划分成一组子图，每个设备一个子图。任何跨设备的从 x 到 y 的边将被删除，并替换为两条新边:

一条边是在 x 子图中，从 x 到新的 Send 节点。
一条边是在 y 的子图之中，从对应的 Receive 节点到 y。

请参见下图。

图 4 插入发送/接收节点之前和之后

在运行时，发送和接收节点将会彼此协调如何在设备之间传输数据。这使我们能够把发送和接收的所有通信隔离出来，从而简化运行时（runtime）的其余部分。

当我们插入发送和接收节点时，我们规范如下：特定设备上特定张量的所有用户都使用同一个接收节点，而不是特定设备上的每个下游用户都拥有一个自己的接收节点。这确保了所需张量的数据在”源设备→ 目标设备对”之间只传输一次，并且目标设备上张量的内存只分配一次，而不是多次（例如，参见上图中的节点 b 和 c）。

通过以这种方式处理通信，我们还允许将不同设备上的图的各个节点的调度分散到工作者之中：发送和接收节点在不同的工作者和设备之间传递必要的同步，这样就把主节点从调度任务之中解放出来。主节点只需要向每个具有计算图的任何节点的工作者发出单个 Run 请求（每次计算图执行），而不需要参与每个节点或每个跨设备通信的调度。这使得系统更具可伸缩性可扩展性，并且和主节点强制执行调度相比，可以执行更细粒度的节点执行策略。

3.5 分布式执行

计算图的分布式执行与多设备执行非常相似。在决定设备如何放置之后，将为每个设备创建一个子图。发送/接收节点对在跨工作进程通信时候使用远程通信机制（如 TCP 或 RDMA）来跨机器边界移动数据。

3.5.1 容错（Fault Tolerance）

我们可以在许多地方检测到分布式执行中的故障，目前主要依靠如下两种方式：

检测到发送和接收节点对之间的通信错误。
从主进程到每个工作进程的定期健康检查。

当检测到故障时，整个计算图执行将中止并从头开始。因为变量（Variable）节点指的是在图的执行过程中持续存在的张量，所以我们支持设置一致性检查点，以此来在重新启动时恢复这些状态。具体来说，每个变量节点都连接到一个 Save 节点。这些保存节点定期执行，例如每 N 次迭代执行一次，或每 N 秒执行一次。当它们执行时，变量的内容被写入持久存储，例如分布式文件系统。类似地，每个变量也都连接到一个恢复节点，该节点仅在重新启动后的第一次迭代中启用。

4. 扩展（Extensions）

在本节中，我们将介绍基本编程模型的几个更高级的特性。

4.1 计算梯度

许多优化算法，包括常见的机器学习训练算法（如随机梯度下降），会使用一组输入来计算一个成本函数（cost function）的梯度。因为这是一种常见的需求，所以 TensorFlow 内置了对自动梯度计算的支持。如果一个 TensorFlow 计算图中的张量 C 可能通过一个复杂的操作子图依赖于一组张量{ $X_k$ }，那么一个内置函数将返回张量集{ $dC/dX_k$ }。梯度张量和其他张量一样，通过使用以下步骤扩展 TensorFlow 图来计算。