# Infra: Tensor Parallel

Author: Fanyi Pu

Published: 2026-09-07

Canonical: <https://pufanyi.com/blog/ml/ml-revisit/infra/infra-tp>

Notes for Tensor Parallel

对于 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](https://pufanyi.com/blog/ml/ml-revisit/infra/infra-tp#bib-shoeybi2020megatronlmtrainingmultibillionparameter)) 里没有考虑 number of attention heads 大于 TP 的情况（[代码位置](https://github.com/NVIDIA/Megatron-LM/blob/b1fe7599e18a213292600177ad7ff290a1127160/megatron/core/transformer/transformer_config.py#L1586-L1590)），我们也只讨论这一点

```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](https://arxiv.org/abs/1909.08053 "https://arxiv.org/abs/1909.08053")