支持数据并行（Allreduce）

简介

大规模AI训练集群中，通常采用数据并行的方式完成训练。数据并行即每个设备使用相同的模型、不同的训练样本，每个device计算得到的梯度数据需要聚合之后进行参数更新。

图1 数据并行方式训练的示意图

如果按照梯度聚合方式进行分类，数据并行的主流实现有PS-workers架构和AllReduce集合通信两种。在AllReduce架构中，每个参与训练的Device形成一个环，没有中心节点来聚合所有计算梯度。AllReduce算法将参与训练的Device放置在一个逻辑环路（logical ring）中。每个Device从上行的Device接收数据，并向下行的Device发送数据，可充分利用每个Device的上下行带宽。

AllReduce架构是为了解决了PS-workers架构无法线性扩展问题而提出的改良架构。各个节点按照算法协同工作，算法的目标是减少传输数据量，并充分利用硬件通信带宽。一般适合训练算力要求高、设备规模大的场景。

实现原理

图2 Allreduce模式
点击放大

以Ring算法为例介绍allreduce模式（称为Ring-allreduce），如图2所示，在Ring-allreduce架构下，每个设备都是worker，并且形成一个环，不需要中心节点来聚合所有worker计算的梯度。在一个迭代过程，每个worker完成一份mini-batch样本数据的前向计算、反向计算，得到梯度数据，然后使用Ring-allreduce算法完成梯度数据的同步。Ring-allreduce算法包括scatter-reduce和allgather两部分，梯度数据分多个步骤传递给环中的下一个worker，同时它也多次接收上一个worker的梯度数据。对于一个包含N个worker的环，每个worker需要从其它worker接收2*（N-1）次梯度数据（每次接收1/N的数据），并向其他节点发送2*（N-1）次梯度数据（每次发送1/N的数据）。

使用的接口

在TensorFlow中，一般使用tf.distribute.Strategy进行分布式训练，具体请参考https://www.tensorflow.org/guide/distributed_training。而昇腾AI处理器暂不支持上述分布式策略，TF Adapter提供了分布式接口npu_distributed_optimizer_wrapper，对传入的optimizer梯度函数添加npu的allreduce操作，最终返回输入的优化器，从而支持单机多卡、多机多卡等组网形式下，各个Device之间计算梯度后执行梯度聚合操作。用户调用该函数后，在生成的训练图中，梯度计算和更新算子之间插入了allreduce算子节点。

因此，对于原始TensorFlow训练脚本，需要经过修改后，才可在昇腾AI处理器上支持分布式训练。

数据集切分

分布式训练时，用户可以使用TensorFlow接口进行数据集切分。如果数据集切分时需要获取处理器资源信息，用户可以通过集合通信接口get_rank_size获取昇腾AI处理器数量，通过get_rank_id获取处理器id，例如：

  dataset = dataset.shard(get_rank_size(),get_rank_id())

Estimator模式下脚本迁移

TensorFlow会将策略对象传递到Estimator的Runconfig中，但是TF Adapter暂不支持这种方式，用户需要将相关代码删除。例如：

迁移前：

mirrored_strategy = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(
    train_distribute=mirrored_strategy, 
    eval_distribute=mirrored_strategy,
    session_config=session_config,
    save_checkpoints_secs=60*60*24)

迁移后：

config = tf.estimator.NPURunConfig(
    session_config=session_config,
    save_checkpoints_secs=60*60*24)

然后调用npu_distributed_optimizer_wrapper，对传入的optimizer梯度函数添加npu的allreduce操作，最终返回输入的优化器，从而在昇腾AI处理器上实现分布式计算。具体方法为：

def cnn_model_fn(features,labels,mode):    
  #搭建网络   
  xxx    
  #计算loss
  xxx    

  #Configure the TrainingOp(for TRAIN mode)    
  if mode == tf.estimator.ModeKeys.TRAIN:      
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) # 使用SGD优化器
    optimizer = npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度
    train_op=optimizer.minimize(loss=loss,global_step=tf.train.get_global_step()) # 最小化loss
    return tf.estimator.EstimatorSpec(mode=mode,loss=loss,train_op=train_op)

NPUDistributedOptimizer分布式优化器在当前版本依然兼容。
Estimator模式下，使用npu_distributed_optimizer_wrapper实现allreduce功能时，由于NPUEstimator中自动添加了NPUBroadcastGlobalVariablesHook，因此无需手写实现broadcast功能。

如果原始脚本使用TensorFlow接口计算梯度，例如grads = tf.gradients(loss, tvars)，需要在计算完梯度之后，调用npu_allreduce接口对梯度进行allreduce。

迁移前：

grads = tf.gradients(a + b, [a, b], stop_gradients=[a, b])

迁移后：

grads = npu_allreduce(tf.gradients(a + b, [a, b], stop_gradients=[a, b]))

sess.run模式下脚本迁移

Estimator模式下，使用npu_distributed_optimizer_wrapper实现allreduce功能时，由于NPUEstimator中自动添加了NPUBroadcastGlobalVariablesHook，因此无需手写实现broadcast功能。但sess.run模式的训练脚本还需要用户手写实现broadcast功能。具体方法为：

在变量初始化之后，训练之前，通过集合通信接口broadcast进行变量广播：

from npu_bridge.npu_init import *

def broadcast_global_variables(root_rank, index):
    """Broadcasts all global variables from root rank to all other processes.
    Arguments:
    root_rank: rank of the process from which global variables will be broadcasted
    to all other processes. 
    index: rank_id
    """
    op_list = []
    for var in tf.trainable_variables():
        # the input and out tensor of HCOMBroadcast interface are list
        if "float" in var.dtype.name:
            inputs = [var]
            outputs=hccl_ops.broadcast(tensor=inputs,root_rank=root_rank)
        if outputs is not None:
            op_list.append(outputs[0].op)
            op_list.append(tf.assign(var, outputs[0]))

    return tf.group(op_list)

...
bcast_op = broadcast_global_variables(root_rank, index)
sess = tf.Session()
...
sess.run(bcast_op)

此外，broadcast接口中有改图的操作，如果图无法修改（例如冻结了图或者使用tf.train.Supervisor创建session等），则在需要先取消图冻结：

with sv.managed_session() as sess:
  sess.graph._unsafe_unfinalize() # 取消冻结的Graph
  sess.run(bcast_op)

执行训练时，在使用梯度优化器计算完各Device数据后，直接调用npu_distributed_optimizer_wrapper进行梯度数据聚合：
```
from npu_bridge.npu_init import *
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001) # 使用SGD优化器
distributedOptimizer=npu_distributed_optimizer_wrapper(optimizer) # 使用NPU分布式计算，更新梯度
```
NPUDistributedOptimizer分布式优化器在当前版本依然兼容。
如果原始脚本使用TensorFlow接口计算梯度，例如grads = tf.gradients(loss, tvars)，需要在计算完梯度之后，调用npu_allreduce接口对梯度进行allreduce：
```
grads = npu_allreduce(tf.gradients(a + b, [a, b], stop_gradients=[a, b]))
```

Keras模式下脚本迁移

如需在Keras模式下进行分布式训练，需要在Keras模型编译时修改optimizer，调用npu_distributed_optimizer_wrapper，对传入的optimizer梯度函数添加npu的allreduce操作；并且在keras_model.fit的callbacks参数中增加NPUBroadcastGlobalVariablesCallback。

迁移前：

from npu_bridge.npu_init import *

data = xxx
labels = xxx

opt = tf.keras.optimizers.Adam(learning_rate=0.001)
keras_model.compile(optimizer=opt,loss='sparse_categorical_crossentropy')
keras_model.fit(data, labels, epochs=10, batch_size=32)

迁移后：

from npu_bridge.npu_init import *

data = xxx
labels = xxx

opt = tf.keras.optimizers.Adam(learning_rate=0.001)
opt = npu_distributed_optimizer_wrapper(opt)           # allreduce
keras_model.compile(optimizer=opt,loss='sparse_categorical_crossentropy')
callbacks = [NPUBroadcastGlobalVariablesCallback(0)]  # 变量进行广播
keras_model.fit(data, labels, epochs=10, batch_size=32, callbacks=callbacks)

KerasDistributeOptimizer分布式优化器在当前版本依然兼容。

父主题： 分布式训练脚本迁移