模型并行是多机多卡分布式并行的一种典型技术。目前CANN软件栈模型并行模式支持将模型的部分层或者单层拆分到多个加速设备上并发计算,每个加速设备执行指定的一部分计算任务,在合适的位置通过通信操作完成模型计算结果的同步。
对于一些超大型规模参数的网络,单个昇腾AI处理器无法存储整个模型计算过程中的数据,模型并行将模型的参数分配到不同的节点上进行计算,可以降低计算对设备内存的需求。
需要注意的是,目前CANN软件栈不支持自动模型拆分,需要显式在计算图中完成层间或层内拆分,并通过集合通信算子完成数据传递。
模型并行方式要求算法人员根据模型特点和设备容量,将模型在合适的位置将需要切分的层分割开,由于每一层的计算依赖于前一层所有输出和当前层的所有参数,因此模型并行会在切分层后的每一步计算进行信息交换,通信负载较重。
每个模型的网络特点和计算行为的差异决定了每种网络适合的切分方式目前也需要单独设计。以ReID为例:fc的参数量非常大,原因是fc的output channel数非常大,如果有N块昇腾AI处理器供使用,将fc按照output channel的方向进行切分。正向计算至FC算子时,将由数据并行切换至模型并行,此时调用AllGather算子将每块卡上的数据并行结构汇聚至一块卡上,这样保证每块卡在进行模型并行的时候,能够获得所有的输入特征层。
反向计算至FC层后,调用ReduceScatter算子聚合数据,并按级别切分反向传播梯度至不同设备: