实现 TensorFlow 架构的规模性和灵活性

TensorFlow 是为了大规模分布式训练和推理而设计的，不过它在支持新模型和系统级优化的实验中的表现也足够灵活。

 

本文对能够同时兼具规模性和灵活性的系统架构进行了阐述。设定的人群是已基本熟悉 TensorFlow 编程概念，例如 computation graph, operations, and sessions。有关这些主题的介绍，请参阅 

https://tensorflow.google.cn/guide/low_level_intro?hl=zh-CN。如已熟悉 Distributed TensorFlow，本文对您也很有帮助。行至文尾，您应该能了解 TensorFlow 架构，足以阅读和修改核心 TensorFlow 代码了。

 

概览

TensorFlow runtime 是一个跨平台库。图 1 说明了它的一般架构。  C API layer 将不同语言的用户级代码与核心运行时分开。

 

 

图 1

 

本文重点介绍以下几个方面：

客户端：

将整个计算过程转义成一个数据流图

通过 session，启动图形执行

 

分布式主节点

基于用户传递给 Session.run() 中的参数对整个完整的图形进行修剪，提取其中特定子图

将上述子图划分成不同片段，并将其对应不同的进程和设备当中

将上述划分的片段分布到 worker services 工作节点服务上

每个 worker services 工作节点服务上执行其收到的图形片段

 

工作节点服务 Worker Services (每一任务一个)

使用内核实现来计划图形表示的计算部分分配给正确的可用硬件（如 cpu，gpu 等）

与其他工作节点服务 work services 相互发送和接收计算结果

 

内核实现

执行单个图形操作的计算部分

 

图 2 说明了这些组件的相互作用。“/ job：worker / task：0” 和 “/ job：ps / task：0” 都是工作节点服务 worker services 上执行的任务。“PS” 表示 “参数服务器”：负责存储和更新模型参数。其他任务在迭代优化参数时会对这些参数发送更新。如果在单机环境下，上述 PS 和 worker 不是必须的，不需要在任务之间进行这种特定的分工，但是对于分布式训练，这种模式就是很常见的。 

 

 

图 2

 

请注意，分布式主节点 Distributed Master 和工作节点服务 Worker Service 仅存在于分布式 TensorFlow 中。TensorFlow 的单进程版本包含一个特殊的 Session 实现，它可以执行分布式主服务器执行的所有操作，但只与本地进程中的设备进行通信。

 

下面，通过逐步处理示例图来详细介绍一下 TensorFlow 核心模块。

 

 

客户端

用户在客户端编写 TensorFlow 程序来构建计算图。该程序可以直接组成单独的操作，也可以使用 Estimators API 之类的便利库来组合层和其他更高级别的抽象概念。TensorFlow 支持多种客户端语言，我们优先考虑 Python 和 C ++，因为我们的内部用户最熟悉这些语言。随着功能的日趋完善，我们一般会将它们移植到 C ++，以便用户可以从所有客户端语言优化访问。尽管大多数训练库仍然只支持 Python，但 C ++ 确实能够支持有效的推理。

 

客户端创建会话，该会话将图形定义作为 tf.GraphDef 协议缓冲区发送到分布式主节点。当客户端评估图中的一个或多个节点时，评估会触发对分布式主节点的调用以启动计算。

 

在图 3 中，客户端构建了一个图表，将权重（w）应用于特征向量（x），添加偏差项（b）并将结果保存在变量中。

 

 

图 3

 

代码：

tf.Session

 

分布式主节点 Distributed master

分布式主节点：

基于客户端指定的节点，从完整的图形中截取所需的子图

对图表进一步进行划分，使其可以将每个图形片段映射到不同的执行设备上

以及缓存这些划分好的片段，以便在后续步骤中再次使用

 

由于主节点可以总揽步骤计算，因此它可以使用标准的优化方法去做优化，例如公共子表达式消除和常量的绑定。然后，对一组任务中优化后的子图或者片段执行协调。

 

 

图 4

 

图 5 显示了示例图的可能分区。分布式主节点已对模型参数进行分组，以便将它们放在参数服务器上。

 

 

图 5

 

在分区切割图形边缘的情况下，分布式主节点插入发送和接收节点以在分布式任务之间传递信息（图 6）。

 

 

图 6

 

然后，分布式主节点将图形片段传送到分布式任务。

 

 

图 7

 

代码：

MasterService API definition

Master interface

 

工作节点服务 Worker Service

在每个任务中，该部分负责：

处理来自 master 发来的请求

为包含本地子图规划所需要的内核执行

协调与其他任务之间的直接信息交换

 

我们优化了 Worker Service，以便其以较低的负载便能够运行大型图形。当前的版本可以执行每秒上万个子图，这使得大量副本可以进行快速的，细粒度的训练步骤。Worker service 将内核分派给本地设备并在可能的情况下并行执行内核，例如通过使用多个 CPU 内核或者 GPU 流。

 

我们还特别针对每对源设备和目标设备类型的 Send 和 Recv 操作进行了专攻：

使用 cudaMemcpyAsync() 来进行本地 CPU 和 GPU 设备之间的重叠计算和数据传输

两个本地 GPU 之间的传输使用对等 DMA，以避免通过主机 CPU 主内存进行高负载的复制

 

对于任务之间的传输，TensorFlow 使用多种协议，包括：

gRPC over TCP

融合以太网上的 RDMA

 

另外，我们还初步支持 NVIDIA 用于多 GPU 通信的 NCCL 库，请参阅：

tf.contrib.nccl

（https://tensorflow.google.cn/api_docs/python/tf/contrib/nccl?hl=zh-CN）。

 

 

图8

 

代码：

WorkerService API definition

Worker interface

Remote rendezvous (for Send and Recv implementations)

 

 

内核运行

该运行时包含了 200 多个标准操作，其中涉及了数学，数组，控制流和状态管理等操作。每个操作都有对应各种设备优化后的内核运行。其中许多操作内核都是通过使用 Eigen::Tensor 实现的，它使用 C ++ 模板为多核 CPU 和 GPU 生成高效的并行代码；但是，我们可以自由地使用像 cuDNN 这样的库，就可以实现更高效的内核运行。我们还实现了 quantization 量化，可以在移动设备和高吞吐量数据中心应用等环境中实现更快的推理，并使用 gemmlowp 低精度矩阵库来加速量化计算。

 

如果用户发现很难去将子计算组合，或者说组合后发现效率很低，则用户可以通过注册额外的 C++ 编写的内核来提供有效的运行。例如，我们建议为一些性能的关键操作注册自己的融合内核，例如 ReLU 和 Sigmoid 激活函数及其对应的梯度等。XLA Compiler 提供了一个实验性质的自动内核融合实现。

 

代码：

OpKernel interface

 

声明：文章收集于网络，如有侵权，请联系小编及时处理，谢谢！