张量模型并行详解 | 深度学习分布式训练专题

随着模型规模的扩大，单卡显存容量无法满足大规模模型训练的需求。张量模型并行是解决该问题的一种有效手段。本文以Transformer结构为例，介绍张量模型并行的基本原理。

---

模型并行的动机和现状

我们在上一篇《数据并行：提升训练吞吐的高效方法|深度学习分布式训练专题》详细介绍了利用数据并行进行大规模分布式训练。读者不禁要问：仅仅使用数据并行，是否可以完全满足大模型训练的要求？随着技术的发展，业界内训练的模型越来越大，模型朝着更深和更宽的方向发展。以自然语言处理（NLP）领域为例，模型从Bert发展到GPT，模型规模从数亿参数量增加到数百亿甚至是数千亿。当参数规模为千亿时，存储模型参数就需要数百GB的显存空间，超出单个GPU卡的显存容量。显然，仅靠数据并行无法满足超大规模模型训练对于显存的需求。为了解决这个问题，可以采用模型并行技术。与数据并行在不同设备都有完整的计算图不同，模型并行是不同设备负责单个计算图不同部分的计算。

模型并行从计算图的切分角度，可以分为以下几种：

1、按模型的layer层切分到不同设备，即层间并行，我们称之为流水线并行，如下左图[1]。

2、将计算图中的层内的参数切分到不同设备，即层内并行，我们称之为张量模型并行，如下右图[1]。本文主要讲述张量模型并行。

张量模型并行原理

张量模型并行需要解决两个问题：参数如何切分到不同设备（切分方式）；以及切分后，如何保证数学一致性（数学等价））。本文以NLP中的Transformer结构为例，介绍张量模型并行的切分方式和随机性控制。

切分方法

Transformer结构主要由嵌入式表示（Embedding）层、矩阵乘层（MatMul）和交叉熵loss计算层（CrossEntropy）构成。以上三种类型的组网层有较大的特性差异，需要设计对应的张量模型并行策略，但总体上看核心思想都是利用分块矩阵的计算原理，实现其参数切分到不同的设备[2]。下面详细介绍这三种层的切分方式。

嵌入式表示（Embedding）<