对于 Column-wise TP：
\[
X \begin{bmatrix}
W_1&W_2&\cdots&W_n
\end{bmatrix} = \begin{bmatrix}
XW_1 & XW_2 & \cdots & XW_n
\end{bmatrix}
\]
反向传播的时候，令 \(Y_i=XW_i\)：
\[
\frac{\partial\mathcal{L}}{\partial W_i} = X^\top\frac{\partial\mathcal{L}}{\partial Y_i}, \quad \frac{\partial\mathcal{L}}{\partial X} = \sum_{i=1}^n\frac{\partial L}{\partial Y_i}W_i^\top
\]
对于 Row-wise TP：
\[
\begin{bmatrix}
X_1&X_2&\cdots&X_n
\end{bmatrix}
\begin{bmatrix}
W_1\\W_2\\ \vdots\\W_n
\end{bmatrix} = \sum_{i=1}^nX_iW_i
\]
反向传播的时候，令求和结果为 \(Y\)：
\[
\frac{\partial\mathcal{L}}{\partial W_i} = X_i^{\top}\frac{\partial\mathcal{L}}{\partial Y}, \quad \frac{\partial\mathcal{L}}{\partial X_i}=\frac{\partial\mathcal{L}}{\partial Y}W_i^{\top}
\]
对于一个 MLP
\[
\mathrm{MLP}(x) = \sigma(XW_1)W_2
\]
我们可以将 \(W_1\) 做 Column-wise 拆分，\(W_2\) 做 Row-wise 拆分，这样子做完 \(XW_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{bmatrix}
O_1 & O_2 & \cdots & O_h
\end{bmatrix}W_o\\
O_t = \mathrm{Attention}\left(XW_q^{(t)}, XW_k^{(t)}, XW_v^{(t)}\right)
\]
我们将 \(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