『TensorFlow』分布式训练_其一_逻辑梳理

1,PS-worker架构

将模型维护和训练计算解耦合,将模型训练分为两个作业(job):

  • 模型相关作业,模型参数存储、分发、汇总、更新,有由PS执行
  • 训练相关作业,包含推理计算、梯度计算(正向/反向传播),由worker执行

该架构下,所有的woker共享PS上的参数,并按照相同的数据流图传播不同batch的数据,计算出不同的梯度,交由PS汇总、更新新的模型参数,大体逻辑如下:

  1. pull:各个woker根据数据流图拓扑结构从PS获取最新的模型参数
  2. feed:各个worker根据定义的规则填充各自batch的数据
  3. compute:各个worker使用第一步的模型参数计算各自的batch数据,求出各自batch的梯度
  4. push:各个worker将各自的梯度推送到PS
  5. update:PS汇总来自n个worker的n份梯度,求出平均值后更新模型参数

分布式经典架构PS-worker会重复上面步骤,直到损失到达阈值或者轮数到达阈值。

2,数据并行模式分类

根据数据流图构建模式分类:

根据参数更新机制分类:

  • 异步训练:各个worker独立训练,计算出梯度后即刻更新参数,不需要等待其他worker完成计算
  • 同步训练:所有worker完成本轮计算后,汇总梯度,更新模型,计算能力强的worker需要阻塞等待其他worker

两种训练机制同时支持上面两周数据流图构建模式。一般来说同步机制收敛快,异步单步计算快,但易受单批数据影响,不稳定。

3,同步优化器

tensorflow进行同步(同步训练模式专用)各个worker梯度并进行优化时,会使用特殊的优化器即同步优化器,tf.train.SyncReplicasOptimizer,其第一个参数为普通优化器,我们可以定义一个普通的优化器传入,后续参数如下:

参数名称功能说明默认值
replicas_to_aggragate并行副本数num_workers
total_num_replicas实际副本数(worker数目)num_workers

并行副本数指期望的每一步中并行的batch数据数目,实际副本数指参与的workers数目,

  • 并行=实际:全民参与,一个worker领取一个batch数据
  • 并行>实际:能者多劳,先完成自己batch的worker会继续领取未训练数据,PS会等到梯度份数到达并行数后进行模型参数计算
  • 并行<实际:替补等位,存在空闲的worker,取代可能出现的异常worker,确保训练过程高可用

运算过程

  • 计算梯度过程同普通优化器,调用基类的Optimizer的compute_gradients成员方法
  • 更新参数时重写了Optimizer的apply_gradients方法,见tensorflow/python/training/sync_replicas_optimizer.py

讲解同步优化器工作逻辑之前,介绍两个概念,

梯度聚合器

每一个模型参数有一个自己队列,收集来自不同worker的梯度值,梯度聚合器包含M个队列对应M个模型参数,每个队列收集来自N个worker计算出来的N个梯度值。

同步标记队列

存储同步标记,实际上就是N个global_step值,每个worker领取一个,用于控制同步

以全民参与模式为例

worker工作模式如下:

  1. 从同步标记队列领取一个global_step,表示全局训练步数的同步标记
  2. 将同步标记值赋予worker的本地训练步数local_step
  3. 从PS获取最新模型参数
  4. 计算出M个梯度值
  5. 将M个梯度值推送到PS上的M个梯度队列中

PS工作模式如下:

  1. 从梯度聚合器上收集worker推送过来的梯度值,每个队列收集N份(对应N个global_step下训练值)后,计算均值,收集齐M个均值后,得到M对{模型参数,梯度值}的聚合元组
  2. 更新模型参数
  3. 向同步标记队列推送N个global_step+1标记

聚合器收集梯度值并校验local_step是否符合global_step,是则接收梯度值,计算能力强的worker提交梯度后由于没有同步标记可以领取所以被阻塞,PS集齐N份后更新参数,发布下次N个同步标记,开始下一步训练。

由于初始PS不会更新参数发布同步标记,所以需要初始化同步标记队列——sync_init_op,直接向队列注入N个0标记。

分布式模型训练需要的主要初始化操作如下(opt指tf.train.SyncReplicasOptimizer):

操作名称常用变量名功能说明
opt.local_step_init_oplocal_init_oploacl_step初始值
pot.chief_init_oplocal_init_opgobal_step初始值
opt.ready_for_local_init_opready_for_local_init_op为未初始化的Variable设置初始值
opt.get_chief_queue_runnerchief_queue_runner同步标记队列启动QueueRunner实例
opt.get_init_tockens_opsync_init_op同步标记队列初始化
tf.global_variables_initializerinit_op全局Variable设置初始值

如果使用模型管理类Supervsor,可以将大部分工作交由其代劳。

以能者多劳模式对比

模型参数个数M,worker个数N,并行副本数R(R>N),此时

梯度聚合器仍然有M个参数收集队列,每一个队列要收集R份才进行汇总,R>N所以会存在某个worker领取多份数据的情况。

同步标记队列存储R个同步标记,以确保每一步中梯度聚合器可以收集到R份数据。

4,异步优化器

异步优化器没有很多附加参量,和单机训练几乎一致,只是每个worker获取参数需要从另一个进程PS中得到而已。

5,模型管理类Supervsor

本质上是对Saver(模型参数存储恢复)、Coordinator(多线程服务生命周期管理)、SessionManager(单机以及分布式会话管理)三个类的封装

Coordinator会监测程序的线程是否运行正常,任何异常的出现都会向Supervisor报告,此时Coordinator讲程序的停止条件设置为True,Supervisor停止训练并清理工作(关闭会话、回收内存等),其他服务检测到True后会各自关闭服务,终止线程。

SessionManager帮助用户创建管理单机或是分布式会话,以便简化数据流图的生命周期和维护逻辑,同事负责将checkpoint文件中加载出的数据恢复到数据流图中。

流程逻辑如下:

  1. 创建Supervisor实例,构造方法需要传入checkpoint文件和summary文件存储目录(Supervisor的logdir参数)
  2. 调用tf.train.Supervisor.managed_session,从Supervisor实例获取会话实例
  3. 使用该会话执行训练,训练中需要检查停止条件,保证训练正确性

获取managed_session时,Supervisor会通过QueueRunner同时启动一下三个服务:

  • 检查点服务:将数据流图中的参数定期保存,默认10min保存一次,且会识别global_step(Supervisor的global_step参数)
  • 汇总服务:默认2min一次
  • 步数计数器服务:向汇总添加global_step/sec,2min一次

使用managed_session创建会话时,会自动恢复上一次的结果并继续训练。