大规模分布式训练--模型并行策略

应用场景

随着深度学习领域的发展，研究人员发现模型越大训练出来的效果越好，因此模型越来越大成为深度学习领域的一个显著特征。但是越大的模型对设备的要求越高，即需要单卡的算力更强，内存空间更大。当单卡运行不能满足模型的要求时，往往需要多卡甚至多台机器协调工作，共同完成训练工作。但如何协调多卡/多机来完成大模型的训练，是大规模分布式训练所需要解决的问题。

模型并行策略是大规模分布式训练很常见的策略之一。它通过将模型中特定子图中的权值均匀的分配到多张卡上，从而降低了模型对单卡的内存要求。帮助模型顺利运行起来。

什么是模型并行

模型并行是算子层面的并行，它利用某些算子的特性将算子拆分到好多个设备上进行计算。因此并不是网络中所有的算子都可以拆分计算，可以拆分的算子需满足如下特性：

• 可以并行计算的算子
• 算子其中一个输入来自于权值

综上， 模型并行中最主要应用的算子就是matmul算子。（conv算子也满足上述要求，不知道为啥没用，个人猜测可能是CNN网络规模不是很大，单卡可以cover, 没必要进行模型并行）。
理论上只要模型中有一个matmu算子就可以进行模型并行策略，算法示例图如下，左图为一个matmul算子的网络结构，右图为通过模型并行策略将一个matmul算子拆分到两张卡上进行计算，每张卡只需要保存原算子1/2的权值。最后通过通信算子allgather获取其他卡的计算结果，从而使每张卡都可以获取完整的计算结果。