首页
学习
活动
专区
工具
TVP
发布

大规模分布式模型训练(02):TensorFlow分布式模型训练原理

TensorFlow分布式模型训练原理

distributed tensorflow可以利用分布式来更好的并行利用大量的计算资源来实现数据并行,提高数据吞吐率以及training的速度。如下图所示,用分布式tensorflow来并行计算和训练模型,从25个GPU扩展到200个GPU,数据的吞吐量也从500 Images/s显著提高到2000 Images/s以上。

TensorFlow的并行计算形式分为单台服务器内多块GPU(简称“单机多卡”)并行和若干台多卡服务器(简称“多机多卡”)并行。

分布式TensorFlow涉及以下几个方面的内容:

Replicating your model(复制你的模型)

Device placement for Variables(如何为参数放置设备资源)

定义分布式Tensorflow集群

Fault tolerance(容错设计)

模型复制 (replication)图内复制 (In-graph replication)

TensorFlow只使用1个CPU核心,即CPU0。绝大部分计算量均由GPU完成。

为了使所有GPU同时工作,需要将计算图划分成多个子图,每个子图对应一块GPU。CPU将各个子图上的计算任务指派给其对应的GPU,全程集中调度,调度过程大致如下:

计算图初始化:组装出计算图,给变量赋初值、生成随机数等。

划分子图:原则是子图之间的相互依赖尽量弱,各子图上包含的计算量尽量均衡。【实现原理是什么?】

作业指派:CPU将子图中的操作指派给对应GPU进行运算,少量不能在GPU上完成的图节点由CPU完成;指派步骤包括先将所需数据制到GPU上,再启动CUDA kernel来在这些数据上计算。

数据汇总交换:GPU使用自己显存上的数据更新图节点的值;CPU将所有子图的最新值从显存复制回CPU内存,汇总得到整个计算图的新值。然后将新数值分发给各个GPU。

各个GPU使用新值进下一轮计算。

图内复制适合于单机多卡。对于多机多卡情况,如果采用图内复制,分发数据带来的数据瓶颈就会比较大,数据传输会产生瓶颈。

图间复制 (Between-graph Replication)

图内复制只能使用一台服务器,GPU数量有限,导致大规模训任务练动辄消耗几天几周的时间,效率低下。

但是,如果将计算图剖分成多个子图,将一个子图指派给一台服务器,那么将面临几个难题:

计算图不够大,难以划分出足够多的子图。例如采用100台服务器,而深度神经网络只有1个输入层(784个神经元)、1个隐层(30个神经元)、1个输出层(10个神经元),这样的计算图几乎是不可剖分的。如果按神经元的层剖分,就是3级流水线,效率不高不说,关键是只能指派给3台服务器。

网络通信压力大:子图之间多多少少都有一些依赖关系,即使强行剖分成N个子图,子图之间的数据交换将带来巨大的网络压力,影响整体计算效率。

图内复制时不存在这两个难题,原因有两个:

需要子图数量少(通常2/4个,多的8/16个),从而剖分容易;

在服务器内存中汇总交换数据,PCIe的带宽和延时性能都比服务器之间的网络性能好很多。

当在运行多个client时,由中间的PS task来交互client之间的数据变化。实现原理是,当有两个worker task时,会创建两个同样名字的变量,然后放在PS中的内存中共享,当一个worker task更新了变量,那个对于另一个task也是可见的,所以这样可以训练得更快。

TensorFlow 封装了高级 API,会自动把参数部署到参数服务器上,把运算操作设定到 worker 上,从而实现一份代码可在多个 worker 上运行,简化客户端程序的编写。

多于1个的PS task时,如何进行参数共享、更新?

Device placement for Variables(如何为参数放置设备资源)

用Tensorflow中不同的设备函数来创建更多更加巧妙强大的参数放置策略。

Round-robin variables

最简单的一个方法叫方法,它会循环地分配变量。如下图所示,指定3个ps_tasks时,这个程序会把weights_1权重变量被放在到PS task0,biases_1偏置变量被放置到PS task1, weights_2权重变量被放置到PS task2,biases_2放到task0中。

优点:可以把整个模型建模代码用这个模块包起来。它只影响变量,把它们放在PS任务中,而其他的工作会自动到worker中执行。

缺点:并不是一个最优的平衡负载放置策略。

Load balancing variables

有一个简单的贪婪策略叫, 它可以根据参数的内存字节来完成类似在线垃圾收集的工作。根据weight和bias的字节数来放置到内存合适的task中,带来更好的负载平衡。

Partitioned variables

以上讨论的都是很小字节的参数,每个PS task都可以单独处理一个变量。但当遇到超大字节,比如可能是几万MB的数据该如何处理?为解决这个问题,提出了一个分割变量的方法。假设你用分隔符创建了一个变量,tensorflow会把这个变量分割成3个部分,分发到3个PS task中。

定义分布式Tensorflow集群Tensorflow集群基本概念

分布式TensorFlow Framework

gRPC

分布式Tensorflow底层的通信是gRPC 。

gRPC是一个RPC,即远程过程调用。通俗的解释是:假设在本机上执行一段代码num=add(a,b),它调用了一个过程 call,然后返回了一个值num,这段代码看上去只是在本机上执行的, 但实际情况是,本机上的add方法是将参数打包发送给服务器,然后服务器运行服务器端的add方法,返回的结果再将数据打包返回给客户端。

Cluster

一个TensorFlow Cluster有一个或多个jobs组成,每个job又由一个或多个tasks构成。Cluster的定义是通过tf.train.ClusterSpec来定义的。

TensorFlow一般将任务分为两类Job:一类叫参数服务器,parameter server,简称为ps,用于管理参数的存储和更新工作;一类就是普通任务,称为worker,用于执行具体的计算。如果参数的数量太大,一台机器处理不了,这就要需要多个Tasks,即是的集合。

Client

Client用来build一个TensorFlow Graph,并构建一个tensorflow::Session用来与集群通信。一个Client可以与多个TensorFlow Server交互,一个Server能服务多个Client。

Job

一个Job由tasks list组成,Job分ps和worker两种类型。ps即parameter server,用来存储和更新variables的,而worker可以认为是无状态的,用来作为计算任务的。workers中,一般都会选择一个chief worker(通常是worker0),用来做训练状态的checkpoint,如果有worker故障,那么可以从最新checkpoint中restore。

Task

每个Task对应一个TensorFlow Server,对应一个单独的进程。一个Task属于某个Job,通过一个index来标记它在对应Job的tasks中的位置。每个TensorFlow均实现了Master service和Worker service。Master service用来与集群内的worker services进行grpc交互。Worker service则是用local device来计算subgraph。

TensorFlow集群定义

上面的语句提供了一个TensorFlow集群定义,集群有两类Job,一类是ps,另一类是worker;ps由2个任务组成,worker由3个任务组成。

可以借助Kubernetes来进行TensorFlow集群定义。

定义完集群信息后,使用创建每个任务:

分布式任务启动,在ip:port对应的机器上运行每个任务(2个PS,3个Worker),共需执行五次代码,生成五个任务:

完整例子:

代码中,会根据job名,将内的 op放到ps tasks,将其他计算op放到worker tasks。默认分配策略是轮询。

容错性(Fault tolerance)定期保存checkpoint

使用来不断把参数检查点存储到磁盘上。

的参数设置为, 也就是设置分片数据。如上图所示,我们这个例子中有3个PS任务,所以分片为3份来并行写入。

如果你使用的是图间复制,你可以选择一个或多个worker任务来回应写入要求。一般来说我们用woker task0来执行一些额为的任务,因为worker是从0开始计数,所以总有worker0存在。这个任务我们称为, 主要负责一些重要的维护任务,比如写入检查点,初始化参数,以及在tensor board中记录一些总结。

现在支持分布式文件系统,比如Google ML, 或者如果你跑在一个Hadoop集群上,你可以写入一个HDFS。

容错恢复

的worker task失败

因为非 Chief worker task实际上是无状态的,恢复的时候,它会重新与它的PS task中建立连接,并且重新开始之前崩溃过的进程。

失败

因为PS task是有状态的,所有的worker task需要依赖PS task来发送它们的梯度并且取得新的参数值。负责监测这种错误,如果发生了这种错误,就打断整个训练,并从上一个检查点恢复所有的PS tasks。

失败

因为 chief worker task 负责其他所有的任务,需要确保它失败时能够让集群里所有的任务都回到一个很好的状态。如果发生了这种错误,打断训练,并在当它恢复了时候从上一个好的检查点恢复。

可以使用一个配置比如Hadoop ZooKeeper 或Etcd来选举而不是简单地定义为task0。

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20181127G0CMOM00?refer=cp_1026
  • 腾讯「腾讯云开发者社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 cloudcommunity@tencent.com 删除。

扫码

添加站长 进交流群

领取专属 10元无门槛券

私享最新 技术干货

扫码加入开发者社群
领券