Infra: Tensor Parallel

对于 Column-wise TP：

X [\begin{matrix} W_{1} & W_{2} & \dots & W_{n} \end{matrix}] = [\begin{matrix} X W_{1} & X W_{2} & \dots & X W_{n} \end{matrix}]

反向传播的时候，令 $Y_{i} = X W_{i}$ ：

\frac{\partial L}{\partial W_{i}} = X^{⊤} \frac{\partial L}{\partial Y_{i}}, \frac{\partial L}{\partial X} = \sum_{i = 1}^{n} \frac{\partial L}{\partial Y_{i}} W_{i}^{⊤}

对于 Row-wise TP：

[\begin{matrix} X_{1} & X_{2} & \dots & X_{n} \end{matrix}] [\begin{matrix} W_{1} \\ W_{2} \\ ⋮ \\ W_{n} \end{matrix}] = \sum_{i = 1}^{n} X_{i} W_{i}

反向传播的时候，令求和结果为 $Y$ ：

\frac{\partial L}{\partial W_{i}} = X_{i}^{⊤} \frac{\partial L}{\partial Y}, \frac{\partial L}{\partial X_{i}} = \frac{\partial L}{\partial Y} W_{i}^{⊤}

对于一个 MLP

MLP (x) = σ (X W_{1}) W_{2}

我们可以将 $W_{1}$ 做 Column-wise 拆分， $W_{2}$ 做 Row-wise 拆分，这样子做完 $X W_{1}$ 之后不需要做 All-Gather，反向传播过完 $W_{2}$ 之后也不需要 All-Gather 完整梯度。

对于 attention，考虑到 Megatron (Shoeybi et al., 2020) 里没有考虑 number of attention heads 大于 TP 的情况（代码位置），我们也只讨论这一点

python

if self.num_attention_heads % self.tensor_model_parallel_size != 0:
    raise ValueError(
        f"num_attention_heads ({self.num_attention_heads}) must be a multiple of "
        f"tensor_model_parallel_size ({self.tensor_model_parallel_size})."
    )

对于 MHA 就非常好做了，对于

O = [\begin{matrix} O_{1} & O_{2} & \dots & O_{h} \end{matrix}] W_{o} O_{t} = Attention (X W_{q}^{(t)}, X W_{k}^{(t)}, X W_{v}^{(t)})

我们将 $W_{o}$ 进行 Row-wise 切分。然后每张卡单独计算一定量的 heads，这样子天然是一个 Column-wise 切分的状态。

References

Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., & Catanzaro, B. (2020). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arxiv.org

References

Cite this post