tensorflow nccl_op && 单机多卡 allreduce training

Tensorflow1.13.1中关于nccl的源码位于:tensorflow\core\kernels\nccl_ops.cc;tensorflow\core\nccl\nccl_manager.cc其中核心代码位于nccl_manager.cc文件，入口API为：AddToAllReduce

想要看懂这部分的代码要首先理解nccl的原理和基本用法，建议参考NCCL官方文档：https://docs.nvidia.com/deeplearning/sdk/nccl-developer-guide/docs/，文档中有一些基本的使用examples，安装nccl后可以写简单的例子进行实践（安装教程：https://github.com/NVIDIA/nccl），其中主要有两种用法（除了结合MPI的例子）：单线程多设备和多线程多设备.

其中多线程多设备的用法必须由应用程序控制好同步：

具体可能导致的问题参见：https://docs.nvidia.com/deeplearning/sdk/nccl-developer-guide/docs/troubleshooting.html#concurrency-between-nccl-and-cuda-calls-nccl-up-to-2-0-5-or-cuda-8及https://github.com/NVIDIA/nccl/issues/217.

 

1.Tensorflow单机多卡AllReduce训练：https://github.com/tensorflow/tensorflow/issues/22692

首先需要了解这个op的特点：

      tensorflow提供了调用nccl库的python API，位于python\ops\nccl_ops.py文件下，函数名为:all_sum(allreduce操作中的一种)，函数返回一个list tensor，与输入list tensor分布在相同设备.比如输入list=[gpu0_var0_grad, gpu1_var0_grad]，那么输出返回两个tensor，一个位于gpu0，一个位于gpu1，二者值相同，为输入所有tensor的和.所以利用这个op单机多卡训练的时候，要将不同设备对应的相同var的grad作为这个op的输入.这样就能求得梯度和，再除以设备个数就得到了均值，然后就可以更新参数了.

与PS架构不同的是，每个设备(节点)都持有一份参数，自己更新自己的参数，因为算出来的梯度相同，所以参数能保持一致，导出模型的时候，只导出单张卡的参数和前向计算图就OK了.

所以与PS架构的单机多卡训练相比，要改动的地方如下：1.每个设备有自己的参数.2.计算梯度均值的时候改用nccl的all_sum算子.3.导出模型的时候，只导出单张卡的参数.