Distributed Training on Jiang Yi(姜祎)'s Homepage

LLM Parameter Computation

Tue, 22 Jul 2025 10:50:47 +0800

本文中，我们介绍一下如何计算 LLM 的参数量。我们将基于 Qwen3 模型架构出发，对模型架构进行拆解，然后给出 LLM 参数量计算公式。

Dense Model

我们首先来看一下 Qwen3 的架构，如下图所示

这里，Qwen3ForCausalLM 就是我们的的 LLM, 其关键代码如下

1
2
3
4


class Qwen3ForCausalLM:
 def __init__(self, config):
 self.model = Qwen3Model(config)
 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

我们假设 vocab_size, 也就是词表大小为 $|V|$ (我们用 $V$ 表示词表), hidden_size 为 $d$, 则总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3ForCausalLM}) =d|V| + \mathrm{parameter}(\texttt{Qwen3Model}) $$

Qwen3Model 包含三个（含参数的）模块，分别是 nn.Embedding, Qwen3DecodeLayer 以及 Qwen3RMSNorm, 分别代表了输入 token 的 embedding layer, Transformer block 和对输出的 normalization. 其关键代码如下：

1
2
3
4
5
6
7


class Qwen3Model:
 def __init__(self, config: Qwen3Config):
 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
 self.layers = nn.ModuleList(
 [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
 )
 self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

其中，nn.Embedding 参数量与 lm_head 一样，都是 $d|V|$.

对于 normalization, 现在大部分 LLM 用的都是 RMSNorm, 其定义如下：

$$ \mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\|x\|_2^2+\epsilon}}\odot \gamma $$

其参数量为：$d$.

如果说我们使用的是 LayerNorm, 则其定义如下：

$$ \mathrm{LayerNorm}(x) = \frac{x-\mathbb{E}[x]}{\sqrt{\mathrm{var}[x]+\epsilon}}\odot \beta + \gamma $$

其参数量为： $2d$.

因此，Qwen3Model 的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3Model})=d|V| + N*\mathrm{parameter}(\texttt{Qwen3DecoderLayer}) + d $$

这里第一项为 nn.Embedding, 第三项为 Qwen3RMSNorm，第二项里，$N$ 代表 decode layer 的个数，也就是 config.num_hidden_layers.

Qwen3DecoderLayer 包含了四个模块，其关键代码如下

1
2
3
4
5
6


class Qwen3DecoderLayer:
 def __init__(self, config: Qwen3Config):
 self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
 self.mlp = Qwen3MLP(config)
 self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

因此，Qwen3DecoderLayer 的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3DecoderLayer}) = 2d + \mathrm{parameter}(\texttt{Qwen3MLP}) + \mathrm{parameter}(\texttt{Qwen3Attention}) $$

其中，第一项是两个 Qwen3RMSNorm 的参数。

对于 Qwen3MLP, 其定义如下：

$$ y = W_2(W_3x\odot \mathrm{SwiGLU}(W_1x)) $$

这里 $W_3,W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $d_{ff}$ 是 MLP 的 hidden size, 代码中用 intermediate_size 来表示，因此

$$ \mathrm{parameter}(\texttt{Qwen3MLP}) = dd_{ff} + dd_{ff} + d_{ff}d=3dd_{ff} $$

这里三项分别代表 $W_1,W_3,W_2$ 的参数量。

如果说，我们使用原始 transformer 的 MLP, 也就是

$$ y = W_2\max(0,W_1x+b_1)+b_2 $$

其中 $W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $b_1\in\mathbb{R}^{d_{ff}}$, $b_2\in\mathbb{R}^d$, 则总参数为

$$ \mathrm{parameter}(\texttt{TransformerMLP}) = d_{ff}d + dd_{ff} + d_{ff} +d = 2d_{ff}d + d_{ff} + d $$

这里的四项分别代表了 $W_1,W_2,b_1,b_2$.

Qwen3Attention

接下来，就是 Attention 部分的参数，Qwen3Attention 的关键代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class Qwen3Attention:
 def __init__(self, config: Qwen3Config, layer_idx: int):
 self.q_proj = nn.Linear(
 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
 )
 self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
 self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias)
 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias)
 self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)
 self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps)

这里，我们先定义几个量：

我们将 head 的个数记为 $h$, 即 num_attention_heads
我们将每个 head 的 hidden size 记为 $h_d$, 即 head_dim
我们将 key 和 value head 的个数记为 $h_{kv}$ , 即 num_key_value_heads

Qwen3Attention 的参数由以下几个部分组成：

Query projection: $W_{Q}\in\mathbb{R}^{hh_d\times d}$, $b_Q\in\mathbb{R}^{hh_d}$
Key projection: $W_K\in\mathbb{R}^{h_{kv}h_d\times d}$, $b_K\in\mathbb{R}^{h_{kv}h_d}$
Value projection: $W_V\in\mathbb{R}^{h_{kv}h_d\times d}$, $b_V\in\mathbb{R}^{h_{kv}h_d}$
output projection: $W_O\in\mathbb{R}^{d\times hh_d}$
RMSNorm：前文已经提到过，两个 normalization (query norm 以及 key norm) 的总参数量为 $2h_d$.

因此， Qwen3Attention 部分的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3Attention}) = hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d = 2hh_dd + 2h_{kv}h_dd + 2h_d $$

分别代表 $W_Q, W_O, W_K,W_V$ 和两个 normalization layer 的参数量。

注意，这里我们没有加入 bias, 这是因为 QKV bias 在 Qwen3 中被取消，取而代之的是两个 normalization.

如果我们查看 Qwen2Attention 的代码，我们可以得到 Qwen2Attention 的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen2Attention}) = 2hh_{d}d + 2h_{kv}h_dd +h_{kv}h_d+hh_d+h_{kv}h_d $$

分别代表 $W_Q, W_K, W_V,W_O$ 和 $b_Q,b_K, b_V$ 的参数量。

我们将计算结果汇总在一起就得到：

$$ \begin{aligned} \mathrm{parameter}(\texttt{Qwen3ForCausalLM}) &=d|V| + \mathrm{parameter}(\texttt{Qwen3Model})\\ &=2d|V| + N*\mathrm{parameter}(\texttt{Qwen3DecoderLayer}) + d\\ &= N*(2d + \mathrm{parameter}(\texttt{Qwen3MLP}) + \mathrm{parameter}(\texttt{Qwen3Attention})) + d(2|V|+1)\\ &= N*(2d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd + 2h_d) + d(2|V|+1)\\ \end{aligned} $$

这是针对 Qwen3ForCausalLM 的参数量计算。这里的变量定义如下：

Math Variable	Code Variable	Description
$N$	`num_hidden_layers`	Transformer block 个数
$\vert V\vert$	`vocab_size`	词表大小
$d$	`hidden_size`	token embedding 的维度
$d_{ff}$	`intermediate_size`	MLP 的中间层的维度
$h_d$	`head_dim`	每个 head 的维度
$h$	`num_attention_heads`	query head 的个数
$h_{kv}$	`num_key_value_heads`	key 和 value head 的个数

Verification

接下来，我们就可以基于 Qwen3 的模型来验证了，比如，Qwen3-32B 的配置如下:

Field	Value
$N$	64
$\vert V\vert$	151936
$d$	5120
$d_{ff}$	25600
$h_d$	128
$h$	64
$h_{kv}$	8

根据上式，最终的参数量为

$$ \mathrm{parameter}(\texttt{Qwen3-32B}) = 32.762123264*10^9\approx 32.8B $$

我们使用 Qwen3-32B 的 index.json 可以得到其真实的参数量为

$$ \texttt{total\_size}/\texttt{precision} = 65524246528/2 = 32762123264 $$

与我们计算的结果一致（这里除以 2 的原因是其表示模型权重文件的总大小，以 bytes 为单位，一般模型都是 bfloat16, 大小为 2 个 bytes, 因此总参数量为总大小除以权重的精度）

MoE Model

MoE model 与 Dense model 不同的地方在于每一层的 FFN, 因此，其总参数计算方式为：

$$ \mathrm{parameter}(\texttt{Qwen3MoeForCausalLM})=N*(2d+\mathrm{parameter}(\texttt{Qwen3MoE})+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2d) + d(2|V|+1) $$

对于 MoE layer, 其关键代码为：

1
2
3
4
5
6


class Qwen3MoeSparseMoeBlock:
 def __init__(self, config):
 self.num_experts = config.num_experts
 self.top_k = config.num_experts_per_tok
 self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
 self.experts = nn.ModuleList([Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)])

我们记 $n$ 为总专家个数，即 num_experts，记 $k$ 为激活专家个数，即 num_experts_per_tok 或者 top_k,

首先 gate 的参数量为 $dn$, 接下来每个 expert 都是一个 Qwen3MLP, 因此 experts 总参数量为 $n * 3dd_{ff}$. 这样 MoE layer 的总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3MoE}) = nd + 3ndd_{ff} $$

在推理时，只有一部分专家，也就是 $k$ 个专家会参与计算，此时激活参数量为

$$ \mathrm{parameter\_activated}(\texttt{Qwen3MoE}) = nd + 3kdd_{ff} $$

我们带入到 Qwen3MoeForCausalLM 中就得到：

模型总参数量为：

$$ \mathrm{parameter}(\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3ndd_{ff}+hh_{d}d + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$

模型激活参数量为：

$$ \mathrm{parameter\_activated}(\texttt{Qwen3MoeForCausalLM})=N*(2d+nd + 3kdd_{ff} + 2h_{kv}h_dd +dhh_d + 2h_d) + d(2|V|+1) $$

这里的变量定义如下：

Math Variable	Code Variable	Description
$N$	`num_hidden_layers`	Transformer block 个数
$\vert V\vert$	`vocab_size`	词表大小
$d$	`hidden_size`	token embedding 的维度
$d_{ff}$	`intermediate_size`	MLP 的中间层的维度
$h_d$	`head_dim`	每个 head 的维度
$h$	`num_attention_heads`	query head 的个数
$h_{kv}$	`num_key_value_heads`	key 和 value head 的个数
$n$	`num_experts`	总专家个数
$k$	`top_k` (`num_experts_per_tok`)	激活专家个数

MoE Verification

我们用 Qwen3-235B-A22B 来验证，其配置如下：

Field	Value
$N$	94
$\vert V\vert$	151936
$d$	4096
$d_{ff}$	1536
$h_d$	128
$h$	64
$h_{kv}$	4
$n$	128
$k$	8

计算得到总参数量为

$$ \mathrm{parameter}(\texttt{Qwen3-235B-A22B}) = 235.09363456*10^9\approx 235B $$

计算得到激活参数量为

$$ \mathrm{parameter\_activated}(\texttt{Qwen3-235B-A22B}) = 22.14456064 * 10^9 $$

实际总参数量为

$$ 470187269120/2 = 235093634560.0\approx 235B $$

可以看到，计算结果与实际相符。

Extension

通过以上计算过程，我们可以很轻松将上述公式扩展到混合架构或者是 MLA 上，对于混合架构，我们分别计算不同 attention 的 layer 个数，然后分别计算。对于 MLA, 我们可以替换 Attention 的计算逻辑。

对于小语言模型系列，比如 0.6B 和 1.7B, 模型的大部分参数集中在 embedding 上，因此 Qwen3 采取了 tie embedding 的方式来减少参数量，具体做法就是 nn.Embedding 和 lm_head 共享参数。

Visualization

接下来，我们来可视化一下不同大小模型不同模块的参数量占比。计算的代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


def compute_param_distribution( N: int, V: int, d: int, d_ff: int, h_d: int, h: int, h_kv: int, tie_word_embeddings: bool) -> dict:
 """
 Compute parameter distribution across different components of the model.

 Args:
 N: Number of layers
 V: Vocabulary size
 d: Hidden dimension
 d_ff: Feed-forward dimension
 h_d: Head dimension
 h: Number of attention heads
 h_kv: Number of key/value heads
 tie_word_embeddings: Whether input and output embeddings are tied

 Returns:
 dict: Parameter counts for each component
 """
 # Embedding parameters
 embedding = V * d
 if not tie_word_embeddings:
 embedding *= 2

 # Attention parameters per layer
 attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d
 attention_total = attention_per_layer * N

 # FFN parameters per layer
 ffn_per_layer = 3 * d * d_ff
 ffn_total = ffn_per_layer * N

 # RMSNorm (2d), QK-Norm (2d), output normalization (d)
 others = N * (4 * d) + d

 return {
 "Embedding": embedding,
 "Attention": attention_total,
 "FFN": ffn_total,
 "Others": others,
 }

对于 dense 模型，每部分的参数量可视化如下：

可以看到，随着模型 size 增加，模型大部分参数量都集中在 FFN 上。

接下来，我们可视化一下 MoE 模型的参数分布，核心计算代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42


def compute_moe_param_distribution(N, V, d, d_ff, h_d, h, h_kv, n, k, activate=False) -> dict:
 """
 Compute parameter distribution across different components of the MoE model.

 Args:
 N: Number of layers
 V: Vocabulary size
 d: Hidden dimension
 d_ff: Feed-forward dimension
 h_d: Head dimension
 h: Number of attention heads
 h_kv: Number of key/value heads
 n: Number of experts
 k: Number of experts per token
 activate: Whether to compute activated parameters

 Returns:
 dict: Parameter counts for each component
 """
 # Embedding parameters
 embedding = 2 * V * d

 # Attention parameters per layer
 attention_per_layer = h * h_d * d + 2 * h_kv * h_d * d + d * h * h_d + 2 * h_d
 attention_total = attention_per_layer * N

 # FFN parameters per layer
 if activate:
 moe_per_layer = n * d + 3 * k * d * d_ff
 else:
 moe_per_layer = n * d + 3 * n * d * d_ff
 ffn_total = moe_per_layer * N

 # RMSNorm (2d), QK-Norm (2d), output normalization (d)
 others = N * (2 * d + 2 * h_d) + d

 return {
 "Embedding": embedding,
 "Attention": attention_total,
 "MoE": ffn_total,
 "Others": others,
 }

结果如下

可以看到，MoE 模型的大部分参数还是集中在 MoE 模块上，但是由于其稀疏机制，在激活的参数里，MoE 占比从 95% 以上降低到了 60% 左右。

Conclusion

在本文中，我们基于 Qwen3 大语言模型系列，介绍了如何计算 dense 模型和 MoE 模型的参数量。模型的参数量计算为后面的显存占用以及优化提供了基础。

References

Qwen3 Collection

分布式训练：参数量与计算量分析

Tue, 13 May 2025 11:26:36 +0800

在本文中，我们将要分析与大语言模型相关的参数量和计算量。在计算之前，我们会首先回顾一下大语言模型的架构

大语言模型架构

大语言模型参数计算

计算量估计

checkpointing

KV cache

参考文献

回旋托马斯x 文章

分布式训练：如何训练一个模型

Tue, 13 May 2025 11:26:36 +0800

本节中，我们将介绍模型训练的基本数学原理，以及在分布式训练中我们需要考虑的精度，优化器等问题。

训练的数学原理

在最优化里面，我们需要解决的问题一般有如下形式：

$$ \min_x\ f(x) $$

这里 $f$是我们的目标函数, $x$是我们的变量。一个比较简单的例子就是求一个给定函数的最小值。

如果说，我们想要基于数据来训练一个模型，这个时候，我们目标函数的输入就包括两部分，一部分是模型参数，另一部分是数据，为了方便起见，我们使用$\theta$来代表模型的参数，用 $\{x_i,y_i\}_{i=1}^N$ 来表示模型的训练集。上述的优化问题改写如下：

$$ \min_{\theta}\ \frac{1}{N}\sum_{i=1}^Nf(x_i,y_i\;\theta) $$

比如我们训练 resnet 作为分类器，那么 resnet 的模型参数就是这里的 $\theta$, 训练集就是我们的图片和对应的标签，比如ImageNet等，对应的$f$可以设定为 cross entropy loss.

现在有了优化问题之后，我们就需要设计算法求解这个优化问题。一个最简单的优化算法就是梯度下降算法：

$$ \theta^{k+1} = \theta^k - \alpha_k\frac{1}{N}\sum_{i=1}^N\nabla_{\theta}f(x_i,y_i;\theta^k) $$

这里 $\nabla_{\theta}f(x;\theta^k)$ 是 $f$ 相对于 $\theta$ 在 $\theta^k$ 处的梯度。

但是，当我们模型过于复杂的时候，梯度往往计算起来非常复杂。为了简化模型的训练，现在的框架如tensorflow和pytorch都支持自动微分。因此，我们只需要定义如何从输入 $(x_i,y_i)$计算得到 $f(x_i,y_i;\theta)$ 就可以了，框架会帮我们计算参数的梯度。

自动微分

自动微分的目的是将求导的过程交给框架，从而让用户专注于模型的开发（也就是设计forward函数）。

自动微分的核心思想就是链式法则.

假设我们有一个复合函数 $y= f(g(h(x)))$, 一般来说，我们求 $y$ 关于 $x$的导数过程为

$$ \frac{dy}{dx} = \frac{dy}{df}\frac{df}{dg}\frac{dg}{dh}\frac{dh}{dx} $$

如果我们的中间函数 $g$, $h$非常复杂的话，那么整个求导过程就会非常复杂。而链式法则则是将这样一个全局过程给分解成了若干个局部过程。我们将 $y$ 表示为：

$$ \begin{aligned} y &= f(y_1)\\ y_1&=g(y_2)\\ y_2&=h(y_3)\\ y_3&=x \end{aligned} $$

接下来，

总结

在本文中，我们简单介绍了一下如何训练一个模型，我们使用pytorch作为例子展示了现在训练框架的工作方式。在下一篇博客中，我们将会探究训练精度和优化器。训练精度和优化器是模型在训练过程中需要考虑的重点之一。

训练

优化器

SGD

Adam

AdamW

精度

Distributed training--Basic

Mon, 12 May 2025 10:15:17 +0800

说明：本文参考了 nanotron/ultrascale-playbook 和 Colossal-AI Concepts

什么是分布式系统

分布式系统允许一个软件的多个组件运行在不同的机器上。与传统集中式系统不一样，分布式系统可以有效提高系统的稳健性。一个比较比较经典的分布式就是Git，Git允许我们把代码保存在多个remote上。这样当一个remote宕机时，其他remote也能提供服务。

评估一个分布式系统的重要标准就是规模效益(scalablity)，也就是说，我们希望使用8台设备应该要比4台设备快2倍。但是，由于通信带宽等原因，实际上加速比并不是和设备数量成线性关系。因此，我们需要设计分布式算法，来有效提高分布式系统的效率。

为什么需要分布式训练

我们需要分布式训练的原因主要是以下几点：

模型越来越大。当下（2025）领先模型如Qwen，LLaMA系列的最大模型都超过了100B [2][3]。LLaMA系列最大的模型甚至超过了1000B。Scaling law告诉我们模型表现与参数量，算力，数据量成正相关关系。
数据集越来越大。现在领先的模型需要的数据量基本都需要100M以上，而大语言模型训练需要的token数量也都超过了10T的量级 [2][3].
算力越来越强。现有最强的GPU H100其显存为80GB，拥有3.35TB/s 的带宽 (PcIe)，这让训练大规模模型成为可能。

超大的模型使得我们很难在一张GPU上进行训练，甚至我们都很难使用单张GPU进行部署。而10T级的数据也也需要几个月的时间才能训练完毕。因此，如何高效利用多张GPU在大规模数据上训练超大模型就是我们需要解决的问题。

基本概念

我们先来熟悉一下分布式训练中的一些基本概念：

Host: host (master address)是分布式训练中通信网络的主设备(main device). 一般我们需要对其进行初始化
Node: 一个物理或虚拟的计算单元，可以是一台机器，一个容器或者一个虚拟机
Port: port (master port)主要是用于通信的master port
Rank: rank是通信网络中每个设备唯一的ID
world size: world size是通信网络中设备的数量
process group: 一个process group是通信网络中所有设备集合的一个子集。通过process group, 我们可以限制device只在group内部进行通信

我们以下图为例：上图中一共包含2个node (2台机器)，每台机器包含4个GPU (device)，当我们初始化分布式环境时，我们一共启动了8个进程（每台机器4个进程），每个进程绑定一个GPU。

在初始化分布式环境之间，我们需要指定host和port。假设我们指定host为node 0和port为 29500，接下来，所有的进程都会基于这个host和port来与其他进程连接。默认的process group（包含所有device）的world size 为8. 其细节展示如下

process ID	rank	Node index	GPU index
0	0	0	0
1	1	0	1
2	2	0	2
3	3	0	3
4	4	1	0
5	5	1	1
6	6	1	2
7	7	1	3

我们可以创建一个新的process group，使其仅包含ID为偶数的process：

process ID	rank	Node index	GPU index
0	0	0	0
2	1	0	2
4	2	1	0
6	3	1	2

Remark: 注意，rank与process group相关，一个process在不同的process group里可能会有不同的rank.

通信方式

接下来，我们需要介绍一下设备间的通信方式，这是我们后面分布式训练算法的基础。根据设备数量的不同，我们可以将设备间通信分为：

one-to-one: 两个device之间互相进行通信
one-to-many: 一个device与多个device进行通信
many-to-one: 多个device与一个device之间进行通信
many-to-many: 多个device之间互相进行通信

One-to-one

One-to-one的情况很简单，一个process与另一个process进行通信，通信通过 send 和 recv 完成。还有对应的 immediate版本，即 isend 和 irecv，示意图如下所示

测试代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


# send_recv.py
import os

import torch
import torch.distributed as dist


def init_process():
 dist.init_process_group(backend='nccl')
 torch.cuda.set_device(dist.get_rank())

def example_send():
 if dist.get_rank() == 0:
 tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda()
 dist.send(tensor, dst=1)
 elif dist.get_rank() == 1:
 tensor = torch.zeros(5, dtype=torch.float32).cuda()
 print(f"Before send on rank {dist.get_rank()}: {tensor}")
 dist.recv(tensor, src=0)
 print(f"After send on rank {dist.get_rank()}: {tensor}")


init_process()
example_send()

# run with
# torchrun --nproc_per_node=2 send_recv.py

结果输出

1
2


Before send on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
After send on rank 1: tensor([1., 2., 3., 4., 5.], device='cuda:1')

注：为了方便，后续代码仅定义函数和运行方式，init_process()和import部分省略

send/recv的特点是在完成通信之前，两个process是锁住的。与之相反，isend/irecv 则不会加锁，代码会继续执行然后返回Work对象，为了让通信顺利进行，我们可以在返回之前加入wait()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


# isend_irecv.py
def example_isend():
 req = None
 if dist.get_rank() == 0:
 tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda()
 req = dist.isend(tensor, dst=1)
 print("Rank 0 is sending")

 elif dist.get_rank() == 1:
 tensor = torch.zeros(5, dtype=torch.float32).cuda()
 print(f"Before irecv on rank {dist.get_rank()}: {tensor}")
 req = dist.irecv(tensor, src=0)
 print("Rank 1 is receiving")
 req.wait()
 print(f"After isend on rank {dist.get_rank()}: {tensor}")

init_process()
example_isend()

# run with
# torchrun --nproc_per_node=2 isend_irecv.py

结果输出

1
2
3
4


Before irecv on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
Rank 0 is sending
Rank 1 is receiving
After isend on rank 1: tensor([1., 2., 3., 4., 5.], device='cuda:1')

由于isend/irecv这种不锁的特性，我们不应该

在dist.isend()之前修改发送的内容tensor
在dist.irecv()之后读取接受的内容tensor

req.wait() 可以保证这次通信顺利完成，因此我们可以在req.wait()之后再进行修改和读取。

One-to-many

One-to-many 情形下，可以分为两种：scatter 和 broadcast

scatter的作用是将一个process的数据均分并散布到其他process。broadcast的作用是将一个process的数据广播到其他process。两者不同的地方在于其他process获取到的是全量数据(copy)还是部分数据(slice)，其示意图如下所示

scatter 测试代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21


# scatter.py
def example_scatter():
 if dist.get_rank() == 0:
 scatter_list = [
 torch.tensor([i+1] * 5, dtype=torch.float32).cuda()
 for i in range(dist.get_world_size())
 ]
 print(f"Rank 0 scatter list: {scatter_list}")
 else:
 scatter_list = None

 tensor = torch.zeros(5, dtype=torch.float32).cuda()
 print(f"Before scatter on rank {dist.get_rank()}: {tensor}")
 dist.scatter(tensor, scatter_list, src=0)
 print(f"After scatter on rank {dist.get_rank()}: {tensor}")

init_process()
example_broadcast()

# run with
# torchrun --nproc_per_node=4 broadcast.py

结果输出以下内容（输出内容有优化，后续不再说明）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


Rank 0 scatter list: [
 tensor([1., 1., 1., 1., 1.], device='cuda:0'),
 tensor([2., 2., 2., 2., 2.], device='cuda:0'),
 tensor([3., 3., 3., 3., 3.], device='cuda:0'),
 tensor([4., 4., 4., 4., 4.], device='cuda:0')
]
Before scatter on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2')
Before scatter on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
Before scatter on rank 3: tensor([0., 0., 0., 0., 0.], device='cuda:3')
Before scatter on rank 0: tensor([0., 0., 0., 0., 0.], device='cuda:0')

After scatter on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
After scatter on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')
After scatter on rank 3: tensor([4., 4., 4., 4., 4.], device='cuda:3')
After scatter on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')

broadcast 测试代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


# broadcast.py
def example_broadcast():
 if dist.get_rank() == 0:
 tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda()
 else:
 tensor = torch.zeros(5, dtype=torch.float32).cuda()
 print(f"Before broadcast on rank {dist.get_rank()}: {tensor}")
 dist.broadcast(tensor, src=0)
 print(f"After broadcast on rank {dist.get_rank()}: {tensor}")


init_process()
example_broadcast()

# run with
# torchrun --nproc_per_node=3 broadcast.py

结果输出：

1
2
3
4
5
6
7


Before broadcast on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1')
Before broadcast on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2')
Before broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0')

After broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0')
After broadcast on rank 1: tensor([1., 2., 3., 4., 5.], device='cuda:1')
After broadcast on rank 2: tensor([1., 2., 3., 4., 5.], device='cuda:2')

Many-to-one

Many-to-one 情形下，也可以分为两种：gather 和 reduce, Gather对应one-to-many的scatter操作，负责将多个process的内容汇聚到一起，形成一个完整的向量。而reduce的操作则是通过一个函数 $f(\cdot)$ 来把数据进行汇总，常见的函数有求和以及求平均，示意图如下所示

gather 测试代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


# gather.py
def example_gather():
 tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda()
 if dist.get_rank() == 0:
 gather_list = [
 torch.zeros(5, dtype=torch.float32).cuda()
 for _ in range(dist.get_world_size())
 ]
 print(f"Rank 0 gather list: {gather_list}")
 else:
 gather_list = None

 print(f"Before gather on rank {dist.get_rank()}: {tensor}")
 dist.gather(tensor, gather_list, dst=0)
 if dist.get_rank() == 0:
 print(f"After gather on rank {dist.get_rank()}: {gather_list}")

init_process()
example_gather()

# run with
# torchrun --nproc_per_node=4 gather.py

结果输出：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


Before gather on rank 3: tensor([4., 4., 4., 4., 4.], device='cuda:3')
Before gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')
Before gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')

After gather on rank 0: [
 tensor([1., 1., 1., 1., 1.], device='cuda:0'),
 tensor([2., 2., 2., 2., 2.], device='cuda:0'),
 tensor([3., 3., 3., 3., 3.], device='cuda:0'),
 tensor([4., 4., 4., 4., 4.], device='cuda:0')
]

reduce 测试代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


# example_reduce.py
def example_reduce():
 tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda()
 print(f"Before reduce on rank {dist.get_rank()}: {tensor}")

 dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM)
 if dist.get_rank() == 0:
 print(f"After reduce on rank {dist.get_rank()}: {tensor}")


init_process()
example_reduce()

# run with
# torchrun --nproc_per_node=3 example_reduce.py

这里我们使用求和dist.ReduceOp.SUM作为我们的汇总操作，Pytorch还支持其他的reduce operations. 结果输出以下内容：

1
2
3
4
5
6


Before reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')
Before reduce on rank 3: tensor([4., 4., 4., 4., 4.], device='cuda:3')
Before reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')

After reduce on rank 0: tensor([10., 10., 10., 10., 10.], device='cuda:0')

Many-to-many

Many-to-many 情形下的两种通信方式为：All-Reduce 和 All-Gather，分别是reduce和gather的升级版，all-reduce对所有process都执行一次reduce操作，而all-gather则对所有process执行一次gather操作，其示意图如下所示

all-gather 测试代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


# example_all_gather.py
def example_all_gather():
 tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda()
 gather_list = [
 torch.zeros(5, dtype=torch.float32).cuda()
 for _ in range(dist.get_world_size())
 ]
 print(f"Before all gather on rank {dist.get_rank()}: {tensor}")
 dist.all_gather(gather_list, tensor)
 print(f"After all gather on rank {dist.get_rank()}: {gather_list}")


init_process()
example_all_gather()

# run with
# torchrun --nproc_per_node=3 example_all_gather.py

测试输出结果：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


Before all gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')
Before all gather on rank 3: tensor([4., 4., 4., 4., 4.], device='cuda:3')
Before all gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before all gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')

After all gather on rank 0: [
 tensor([1., 1., 1., 1., 1.], device='cuda:0'),
 tensor([2., 2., 2., 2., 2.], device='cuda:0'),
 tensor([3., 3., 3., 3., 3.], device='cuda:0'),
 tensor([4., 4., 4., 4., 4.], device='cuda:0')]
After all gather on rank 2: [
 tensor([1., 1., 1., 1., 1.], device='cuda:2'),
 tensor([2., 2., 2., 2., 2.], device='cuda:2'),
 tensor([3., 3., 3., 3., 3.], device='cuda:2'),
 tensor([4., 4., 4., 4., 4.], device='cuda:2')
]
After all gather on rank 3: [
 tensor([1., 1., 1., 1., 1.], device='cuda:3'),
 tensor([2., 2., 2., 2., 2.], device='cuda:3'),
 tensor([3., 3., 3., 3., 3.], device='cuda:3'),
 tensor([4., 4., 4., 4., 4.], device='cuda:3')
]
After all gather on rank 1: [
 tensor([1., 1., 1., 1., 1.], device='cuda:1'),
 tensor([2., 2., 2., 2., 2.], device='cuda:1'),
 tensor([3., 3., 3., 3., 3.], device='cuda:1'),
 tensor([4., 4., 4., 4., 4.], device='cuda:1')
]

all-reduce 测试代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


# example_all_reduce.py
def example_all_reduce():
 tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda()
 print(f"Before all reduce on rank {dist.get_rank()}: {tensor}")
 dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
 print(f"After all reduce on rank {dist.get_rank()}: {tensor}")

init_process()
example_all_reduce()

# run with
# torchrun --nproc_per_node=3 example_all_reduce.py

测试输出结果

1
2
3
4
5
6
7
8


Before all reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1')
Before all reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0')
Before all reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')
Before all reduce on rank 3: tensor([4., 4., 4., 4., 4.], device='cuda:3')
After all reduce on rank 0: tensor([10., 10., 10., 10., 10.], device='cuda:0')
After all reduce on rank 2: tensor([10., 10., 10., 10., 10.], device='cuda:2')
After all reduce on rank 3: tensor([10., 10., 10., 10., 10.], device='cuda:3')
After all reduce on rank 1: tensor([10., 10., 10., 10., 10.], device='cuda:1')

Barrier

除了之前这些传输数据的方式之外，我们还有Barrier，用于在所有process之间进行同步。Barrier会确保所有的process在同一时间点完成某些操作。其流程为，先让每个process完成各自的任务，然后当process到达barrier时，process会通知系统自己已到达。最后当所有process都到达barrier之后，阻塞会解除，所有process继续执行下一步操作。

barrier 测试代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


# example_barrier.py
def example_barrier():
 import time
 rank = dist.get_rank()
 t_start = time.time()
 print(f"Rank {rank} sleeps {rank} seconds")
 time.sleep(rank)
 dist.barrier()
 print(f"Rank {rank} is done at {time.time() - t_start:.4f} seconds")

init_process()
example_barrier()

# run with
# torchrun --nproc_per_node=3 example_barrier.py

结果输出

1
2
3
4
5
6
7
8
9


Rank 2 sleeps 2 seconds
Rank 0 sleeps 0 seconds
Rank 1 sleeps 1 seconds
Rank 3 sleeps 3 seconds

Rank 3 is done at 3.3046 seconds
Rank 1 is done at 3.3229 seconds
Rank 2 is done at 3.8437 seconds
Rank 0 is done at 3.6613 seconds

可以看到，四个process的到达时间都在3s左右，这是因为rank 3需要3s才能完成当前任务

Advanced

除了前面的通信方式之外，还有 Reduce-Scatter和Ring All-Reduce，这两个通信方式等我们学习ZeRO的时候再一并讲解。