Tutorial on Jiang Yi(姜祎)'s Homepage

LLM Memory Computation

Sat, 17 Jan 2026 10:04:32 +0800

Motivation

我们从一个简单的问题开始

假如我有一张 80GB 的显卡，我想训练/推理一个 4B 的模型，我应该设置多大的 batch size 和 sequence length?

在这个 tutorial 中，我们将基于这个问题来进行思考和分析。我们将考虑更一般的问题形式：

Motivation: 在训练和推理时 LLM 所需要的内存是多少？如何进行优化内存占用？

为了回答以上问题，我们先介绍训练/推理阶段的内存计算，再针对可优化部分进行分析并介绍相应优化算法。

Background

Transformer Architecture

以 Qwen3 为例，现代 LLM 的架构包含多层 Transformer Block，其中具体的模块不同的模型可能有改动。下图是对应的模型架构

Notation

与参数量、FLOPs 计算所用记号一致；参数量 $P$ 的推导见 LLM parameter analysis。

变量	含义
$P$	number of parameters
$L$	layers
$V$	vocabulary size
$d$	hidden size
$d_{\text{ff}}$	FFN hidden size
$s$	sequence length
$b$	batch size
$h$	number of attention heads
$d_h$	attention head dimension

Assumptions

若无特别说明，使用 BF16/FP16，每个参数 2 byte。
不使用 dropout（与现代大模型设定一致）。
Attention 基于原始 multi-head attention.
FFN 基于 SwiGLU.

Training Memory Analysis

Training Memory Components

训练部分的内存占用由四部分组成：

$$ \text{training\_memory} = \text{weight} + \text{activation} + \text{optimizer} + \text{gradient} $$

Weights: $\boxed{2P}$
Gradients（与权重同精度）: $\boxed{2P}$

Optimizer States

AdamW 优化器需要维护两个动量状态：

一阶动量 $m_t$：$2P$
二阶动量 $v_t$：$2P$
合计：$2 \times 2P = \boxed{4P}$

AdamW [2] 的更新规则如下：

$$ \begin{aligned} m_t &\leftarrow \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \\ v_t &\leftarrow \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \\ \hat{m}_t &\leftarrow \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t \leftarrow \frac{v_t}{1 - \beta_2^t} \\ \theta_t &\leftarrow \theta_{t-1} - \alpha \left( \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_{t-1} \right) \end{aligned} $$

Activation

激活值是前向传播过程中计算得到的中间结果，用于在反向传播时计算梯度。

我们仅针对 linear layer 进行推导：

$$ \begin{aligned} \text{forward:} \quad & \mathbf{z}_\ell = W_\ell \mathbf{a}_{\ell-1} + b_\ell, \quad \mathbf{a}_\ell = \phi(\mathbf{z}_\ell) \\ \text{backward:} \quad & \frac{\partial \mathcal{L}}{\partial W_\ell} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}_\ell} \cdot \frac{\partial \mathbf{z}_\ell}{\partial W_\ell} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}_\ell} \cdot \boxed{\mathbf{a}_{\ell-1}} \end{aligned} $$

可以看到，计算第 $\ell$ 层关于 $W_\ell$ 的梯度时需要其输入 $\mathbf{a}_{\ell-1}$，因此训练时需保存每个模块对应的输入，也就是激活值 (activation)。

Activation — Attention

按计算图（无优化）可得需保存的激活：

Q/K/V 投影：共享输入 → $2bsd$
$Q^\top K$：Q, K 均需保存 → $2 \times 2bsd = 4bsd$
softmax 输入：$2bhs^2$
weighted sum 输入：$2bhs^2 + 2bsd$
output projection 输入：$2bsd$

Attention 合计： $\boxed{10bsd + 4bhs^2}$

Activation — FFN & LayerNorm

FFN（SwiGLU，assume $d_{\text{ff}} = 4d$）：

第一层输入：$2bsd$
SwiGLU 输入：$2 \times d_{\text{ff}} \times d = 8bsd$
第二层输入：$2 \times d_{\text{ff}} \times d = 8bsd$
合计：$\boxed{18bsd}$

LayerNorm：保存输入 → $\boxed{2bsd}$

Activation — Output

Output 包含以下组成部分：

FinalNorm 输入：$2bsd$
lm_head 输入：$2bsd$
Loss 输入：$2bsV$

合计：$\boxed{4bsd + 2bsV}$

Activation — Total

将上面的结果汇总在一起，得到：

$$ \begin{aligned} \text{activation} &= L \cdot \text{transformer\_block} + \text{output} \\ &= L \cdot (\text{Pre\_Norm} + \textcolor{red}{\text{Attention}} + \text{Post\_Norm} + \text{FFN}) + \text{output} \\ &= \boxed{bs(32dL + \textcolor{red}{4hsL} + 4d + 2V)} \\ &\approx bsL(32d + \textcolor{red}{4hs}) \\ &\approx \textcolor{red}{4bs^2hL} \end{aligned} $$

注：在 Qwen3 中，$2V / 32dL \approx 6\%$，$32dL / 4hsL \approx 2.5\%$。

可以看到，未优化的情况下，$\text{activation} \propto bs^2$。这里 $s^2$ 主要由 attention 部分产生，后续 Flash Attention 就针对这一点进行了优化。

Total Training Memory

将上面的结果进行汇总：

$$ \begin{aligned} \text{training\_memory} &= \text{weight} + \text{activation} + \text{optimizer} + \text{gradient} \\ &= 2P + bs(32dL + 4hsL + 4d + 2V) + 4P + 2P \\ &= 8P + bs(32dL + 4hsL + 4d + 2V) \quad \text{(exact)} \\ &\approx 8P + 4bs^2hL \end{aligned} $$

可以看到，训练阶段的内存占用分为固定部分 ($8P$) 和动态部分 ($4bs^2hL$)，动态部分主要是 attention 的缓存。

Experiments

我们分别针对 80GB 的显卡计算 Qwen3 系列模型的最高配置：

Model	$P$	$L$	$h$	$s$	predicted $b$	actual $b$
Qwen3-0.6B	0.6	28	16	512	68	34
Qwen3-1.7B	1.7	28	16	512	42	28
Qwen3-4B	4	36	32	512	16	12
Qwen3-8B	8.1	36	32	512	4	2

其中 predicted $b$ 基于前面的准确公式计算得到；actual $b$ 通过实验验证得到。注意我们这里的 prediction 没有考虑任何优化手段与其他内存开销，因此与实际值有出入。

Case Study

我们分别使用 Qwen3-4B 和 Qwen3-8B 来进行实验（$b=1$, $s=512$）。参考 PyTorch 显存可视化与 Snapshot 数据分析。

Inference Memory Analysis

Inference Components

Inference 阶段内存占用主要与模型参数、KV cache 两部分相关：

$$ \text{Inference\_Memory} = \text{weight} + \text{activation} + \text{KV cache} $$

Weights: $\boxed{2P}$
Activations（经验值，batch size=1）: $\approx 0.4P$（参见 transformer-math）
KV cache：与序列长度相关，见后续分析。

KV Cache Mechanism

LLM 推理中为避免重复计算历史 token 的 key/value 而使用的空间换时间的缓存机制。

自回归时逐 token 生成，每步 attention 形式为（$\mathbf{q}_t$ 当前 query，$\mathbf{k}_{:,t}$ / $\mathbf{v}_{:,t}$ 历史 K/V）：

$$ \mathbf{q}_t = W_Q \mathbf{x}_t, \quad \mathbf{k}_{:,t} = W_K[\mathbf{x}_1, \ldots, \mathbf{x}_t], \quad \mathbf{v}_{:,t} = W_V[\mathbf{x}_1, \ldots, \mathbf{x}_t] $$

处理下一 token $\mathbf{x}_{t+1}$ 时只需在已有结果后追加当前步：

$$ \mathbf{k}_{:,t+1} = [\mathbf{k}_{:,t},\, W_K \mathbf{x}_{t+1}], \quad \mathbf{v}_{:,t+1} = [\mathbf{v}_{:,t},\, W_V \mathbf{x}_{t+1}] $$

缓存前：每生成一个 token 都重新计算 → 总计算量 $\sum_{t=1}^{s} \mathcal{O}(t) = \mathcal{O}(s^2)$
缓存后：每步只算当前 token $W_K \mathbf{x}_{t+1}$、$W_V \mathbf{x}_{t+1}$ → 计算量 $\mathcal{O}(s)$、空间占用 $\mathcal{O}(s)$ —— 以空间换时间

KV Cache Memory

对于 multi-head attention，KV cache 的显存占用为：

$$ \text{Memory}(\text{KV cache}) = s \times 2 \times 2 \times L \times h \times d_h = \boxed{4sLhd_h} $$

因子含义： $s$ 序列长，第一个 $2$ 为 K+V，第二个 $2$ 为 BF16 的 2 bytes，$L$ 层、$h$ 头、$d_h$ 头维度。

Remark:

KV 占用与模型配置（$L, h, d_h$）和序列长度 $s$ 都有关，token 越多占用越高。
实际中因 page granularity、padding、fragmentation 往往略高于理论值。
长输出时 KV 占比会超过权重，成为推理瓶颈 → 见后续「KV Cache 优化」。

Total Inference Memory

综合前面分析，推理阶段的总内存为：

$$ \boxed{\text{Inference\_Memory} \approx 2.4P + 4sLhd_h} $$

可以看到，推理阶段也由固定部分（参数量，activation）以及动态部分（KV cache）组成。

Dynamic vs. Static

由于 Qwen3 的 KV cache 计算为 $4sLh_{kv}d_h$，而不同模型只有 $L$ 不一样，因此对于更大的模型，KV cache 显存占用超过模型权重的上下文长度更高。

Optimization

Overview

阶段	核心方法	典型技术
Training	显存与效率提升	Activation Checkpointing, Mixed Precision Training, Flash Attention, ZeRO, Pipeline/Model/Data Parallelism
Inference	长序列与速度优化	KV Cache Optimization, Paged/Radix Attention, Faster Attention, Quantization

Mixed Precision Training

计算量大的部分用低精度，计算量小的部分用高精度。低精度参与运算，高精度避免 Overflow/Underflow。

下表是 DeepSeek-V3 [3] 使用的混合精度训练框架的显存分析：

Precision	BF16	FP32	BF16
AMP	No	Yes	Yes
Weights	BF16 (2)	FP32 (4)	BF16 (2)
Master weights	-	-	FP32 (4)
Gradients	BF16 (2)	FP32 (4)	BF16 (2)
Adam m	BF16 (2)	FP32 (4)	FP32 (4)
Adam v	BF16 (2)	FP32 (4)	FP32 (4)
Static total (bytes/param)	8	16	16

Remark:

BF16 (w/ AMP) 与 FP32 (w/ AMP) 的静态显存占用相同，但 BF16 (w/ AMP) 的动态显存占用更低。
主流框架基本都使用了 BF16/FP8 (w/ AMP) 的训练方式。

ZeRO

将 optimizer states / gradients / weights 按不同 GPU 切片存储，需要参与计算时再 all-gather 整合成完整参数。这样每张卡只需维护自己负责的一部分，大幅降低单卡显存需求 [4]。

ZeRO Stages:

ZeRO-1：shard optimizer states

$$ \text{Training\_Memory} = \text{weight} + \text{activation} + \frac{\text{optimizer}}{\#\text{GPUs}} + \text{gradient} $$

ZeRO-2：shard optimizer states + gradients

$$ \text{Training\_Memory} = \text{weight} + \text{activation} + \frac{\text{optimizer} + \text{gradient}}{\#\text{GPUs}} $$

ZeRO-3：shard all

$$ \text{Training\_Memory} = \text{activation} + \frac{\text{weight} + \text{optimizer} + \text{gradient}}{\#\text{GPUs}} $$

ZeRO-3 可极大降低单卡显存上限，但通信量也会提高。

Model Parallelism

将模型切分到不同的 GPU 上，计算时，先 dispatch，再执行计算，最后通过 all-gather 等操作得到最终结果 [5]。切分方式包括 PP (Pipeline Parallelism)、TP (Tensor Parallelism)、EP (Expert Parallelism) 等。

$$ \text{Training\_Memory} = \frac{\text{Memory}(\text{weight})}{\text{PP degree} \times \text{TP degree}} $$

结合 ZeRO-1 与 Model Parallelism 时（activation 中与 TP 相关的部分按 TP degree 缩减）：

$$ \text{Memory}_{\text{train}} \approx \frac{\text{weight}}{\text{PP} \times \text{TP}} + \frac{\text{activation}}{\text{TP}} + \frac{\text{optimizer}}{\#\text{GPUs}} + \frac{\text{gradient}}{\text{PP}} $$

Activation Checkpointing

在反向传播时，重新计算所需的输入，来达到以时间换空间的目的 [6]。

	No ckpt	Selective ckpt	Full ckpt
memory	很高	中等	很低 $\sim 2bsd$
extra compute	无	中等	很高 $\sim 2Pbs$

实践中常结合 model parallelism 与 selective checkpointing 来实现 trade-off。

Flash Attention

通过将 Attention 的计算进行分块，来提高内存访问效率以及降低反向传播时所需要的 activation 大小 [7]。

Flash Attention 通过 tiling 与 online-softmax 降低该部分显存并提升效率（详见 notes on Flash Attention）。这样 attention 部分的显存就由 $\text{activation} \propto bs^2$ 降低到了 $\text{activation} \propto bs$。

Theorem: Flash Attention 输出 $O = \text{softmax}(QK^T)V$（correctness）。其时间复杂度为 $\mathcal{O}(s^2 d)$，空间复杂度为 $\mathcal{O}(s)$（memory savings）。

KV Cache Optimization

$$ \text{Memory}(\text{KV cache}) = s \times 2 \times 2 \times L \times h \times d_h $$

针对公式中各因子的优化方向 [8]：

$s$: KV cache compression, eviction, selection
$2$ (bytes): KV cache quantization
$2$ (K+V): key-value sharing, MLA [9]
$h \times d_h$: MQA [10], GQA [11], MLA

Weight Quantization

使用低精度来表示高精度数值的方法，来减少内存占用/提高计算效率。

量化时机	代表性工作
训练后量化 (PTQ)	GPTQ [12], AWQ [13], SmoothQuant [14], GGUF [15]
量化感知训练 (QAT)	LLM-QAT [16], PEQA [17]

Activation Offloading

将一部分参数/优化器状态/激活值等存储到 CPU 上，需要的时候再加载到 GPU 上。

Offloading 场景	代表性工作
训练阶段 Offloading	ZeRO-Offload [18] / ZeRO-Infinity, FSDP [19] CPU Offload
推理阶段 Offloading	FlexGen [20], vLLM [21] KV Cache Offload
MoE Offloading	KTransformers [22], DeepSpeed-MoE [23]

Real Systems

Training Frameworks

Framework	Memory Optimizations
Megatron-LM [5]	TP, SP, Activation Checkpointing
DeepSpeed [24]	ZeRO-1/2/3, CPU/NVMe Offload, Activation Checkpointing
FSDP [19]	Full parameter sharding, Gradient sharding, CPU Offload
Colossal-AI [25]	ZeRO, TP, PP, Activation Checkpointing

Inference Frameworks

Framework	Key Techniques	Memory Optimizations
vLLM [21]	Paged Attention	KV cache paging, Continuous batching
SGLang [26]	Radix Attention	KV cache reuse, Efficient scheduling
TensorRT-LLM [27]	Kernel fusion	Weight quantization, KV cache optimization

Conclusion

Takeaway

Components	Training	Inference	Optimization
weights	$2P$	$2P$	quantization
optimizer states	$4P$	0	ZeRO, offloading
gradients	$2P$	0	ZeRO
activations	$\sim 4Lhs^2b$	$\sim 0.4P$	ckpt, offloading, flash attention
KV cache	0	$4sLhd_h$	KV optimization, attention
TOTAL	$8P + 4Lhs^2b$	$2.4P + 4sLhd_h$

训练：瓶颈主要在激活值（随 batch size / 序列长度线性增长）。
推理：瓶颈主要在 KV Cache（随序列长度增长）。

Future Directions

More efficient architecture (attention, MoE).
Scalable training/inference framework.
Software-hardware co-design algorithms.

References

An Yang et al., “Qwen3 Technical Report,” arXiv:2505.09388, 2025.
Ilya Loshchilov and Frank Hutter, “Decoupled Weight Decay Regularization,” arXiv:1711.05101, 2019.
DeepSeek-AI, “DeepSeek-V3 Technical Report,” arXiv:2412.19437, 2025.
Samyam Rajbhandari et al., “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,” arXiv:1910.02054, 2020.
Mohammad Shoeybi et al., “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,” arXiv:1909.08053, 2020.
Vijay Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models,” arXiv:2205.05198, 2022.
Tri Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” arXiv:2205.14135, 2022.
Haoyang Li et al., “A Survey on Large Language Model Acceleration based on KV Cache Management,” arXiv:2412.19442, 2025.
DeepSeek-AI, “DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model,” arXiv:2405.04434, 2024.
Noam Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” arXiv:1911.02150, 2019.
Joshua Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” arXiv:2305.13245, 2023.
Elias Frantar et al., “GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers,” arXiv:2210.17323, 2023.
Ji Lin et al., “AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration,” arXiv:2306.00978, 2024.
Guangxuan Xiao et al., “SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models,” arXiv:2211.10438, 2024.
Georgi Gerganov, “ggml: Tensor library for machine learning,” GitHub, 2023.
Zechun Liu et al., “LLM-QAT: Data-Free Quantization Aware Training for Large Language Models,” arXiv:2305.17888, 2023.
Jeonghoon Kim et al., “Memory-Efficient Fine-Tuning of Compressed Large Language Models via sub-4-bit Integer Quantization,” arXiv:2305.14152, 2023.
Jie Ren et al., “ZeRO-Offload: Democratizing Billion-Scale Model Training,” arXiv:2101.06840, 2021.
Yanli Zhao et al., “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel,” arXiv:2304.11277, 2023.
Ying Sheng et al., “FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU,” arXiv:2303.06865, 2023.
Woosuk Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” SOSP, 2023.
Hongtao Chen et al., “KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models,” SOSP, 2025.
Samyam Rajbhandari et al., “DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale,” arXiv:2201.05596, 2022.
Jeff Rasley et al., “DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters,” KDD, 2020.
Shenggui Li et al., “Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training,” ICPP, 2023.
Lianmin Zheng et al., “SGLang: Efficient Execution of Structured Language Model Programs,” NeurIPS, 2024.
NVIDIA Corporation, “TensorRT-LLM: A TensorRT Toolset for Optimizing LLM Inference,” GitHub, 2023.

MoE tutorial

Sat, 13 Dec 2025 16:04:04 +0800

本 blog 详细介绍了 MoE 模型的一些关键设计与相关实验结果，为 MoE 模型的学习提供基础。

Introduction

Motivation

现有大部分大语言模型均是基于 Transformer 架构，Kaplan scaling law 通过实验说明，大语言模型的表现与算力，数据，模型参数量息息相关。但是，对于 dense 模型来说，我们提高模型参数量时，必须同时提高所使用的算力。这就限制了大模型的 scaling law.

而 MoE 模型的解决方法为在计算时只激活部分参数，这样，我们就可以在同等激活参数量/算力下训练更大参数量的模型，从而达到更好地表现。

因此，MoE 模型的核心思想在于

使用相同的激活参数量/算力，提高模型总参数量，从而达到更好的表现。

Definition

MoE 模型和 dense 模型的示意图如下，图源 olmoe

一个 MoE layer 包括两个模块：

Router：Router 负责为 token 指定合适的专家
Expert：Expert 负责处理 token

对于输入 $x\in\mathbb{R}^d$, 我们假设有 $N$ 个 Expert，router 一般是一个 linear layer 再加上一个 gating function (softmax 或者 sigmoid，我们本文中使用 softmax), 其构建了 $\mathbb{R}^d\to\mathbb{R}^N$ 的映射，定义为：

$$ G(x) =[G_1(x),\dots,G_N(x)] = \mathrm{softmax}(W_gx + b)\in\mathbb{R}^N $$

其中 $W_g\in\mathbb{R}^{N\times d}$, $b\in\mathbb{R}^N$ 是可学习的参数。$G_{i}(x)$ 代表了当前 token $x$ 选择第 $i$ 个 Expert 的概率。

一般来说，Expert 会使用和 dense 模型一样的 MLP, 我们记为

$$ E_i(x) = \mathrm{FFN}(x), \quad i = 1,\dots,N $$

接下来，基于 $G(x)$ 和 $E(x)$, 我们会使用合适的方法来挑选 $K $$ e_i \in \{i\mid G_i(x)\in \mathrm{TopK}(G_i(x), K)\},\ i=1,\dots,K $$

最终 MoE layer 的输出为

$$ y = \sum_{i=1}^K\mathrm{Normalize}(G_{e_i}(x)) E_{e_i}(x) $$

这里 $\mathrm{Normalize}(\cdot)$ 代表我们对于输出进行归一化，即

$$ \mathrm{Normalize}(G_{e_i}) = \frac{\exp(G_{e_i})}{\sum_{i=1}^K \exp(G_{e_i})},\quad i=1,\dots,K $$

Why MoE

选择 MoE 的原因有三点：效率, scaling law 以及表现。

Efficiency

MoE 训练更加高效，如下图所示

Switch Transformer 的实验结果说明，MoE model 的训练效率比 dense model 快 7 倍左右。其他模型也有类似结论。总的来说，MoE 模型相比于 dense 模型，训练效率更高。

Scaling Law

MoE 模型可以突破传统 scaling law 的限制，在算力固定的情况下，我们可以通过提高 MoE 模型的稀疏度来进一步提高模型的表现

如上图所示，在 FLOPs 给定的情况下，随着模型稀疏度的提高，模型的表现和效率都有提升

Performance

MoE 模型的表现更强，如下图所示，MoE 模型的训练，验证损失以及在下游任务上的表现均超过了 dense 模型

Timeline

激活参数比例变化图如下图所示

可以看到，当前大部分模型的激活比例都在 $5\%$ 左右。

另一方面，从专家的激活比例变化图如下图所示

可以看到，现在大部分模型总专家数都在 200-400 左右，Kimi-k2 认为提高专家个数可以提高模型表现，而 LongCat 则是使用了 phantom expert 机制

MoE Design

Experts Design

Number of Experts

一般来说，专家个数越多，模型越稀疏，模型表现越好。扩展专家个数有两个方式：

直接增加专家个数，这会导致模型参数量上升，如 Switch Transformer
对已有的专家进行切分，将大专家切分为小专家，如 DeepSeekMoE

Switch Transformer 也通过实验发现，增加专家个数可以显著提高模型的训练效率和表现，结果如下图所示

可以看到，当我们增加专家个数的时候，模型的表现是持续提升的。并且当我们增加专家个数之后，模型的训练效率也有所提升。

DeepSeekMoE 提出了 fine-granularity expert 的概念，其做法是通过减少 expert 的大小在相同参数量的场景下使用更多的专家。实验结果如下图所示

可以看到，在稀疏度 (激活专家个数占总专家个数比例) 不变的情况下，提高专家的粒度，可以提高模型的表现。

olmoe 对 DeepSeekMoE 的这个观点进行了验证，结果如下图所示，

结果显示，当专家粒度从 8E-1A 扩展到 32E-4A 时，模型在 HellaSwag 上的表现提升了 $10\%$, 但是进一步扩展到 64E-8A 时，模型的表现提升不到 $2\%$, 这说明了无限制提升专家粒度对模型的提升越来越有限。

Kimi-k2 探究了针对 MoE 模型 sparsity 的 scaling law, 结果也说明，提升 sparsity 可以提高模型的表现。因此，其相对于 DeepSeek-V3 使用了 $50\%$ 额外的的专家数。Ling-mini-beta 进一步验证了这个观点。

Shared Experts

Shared Expert 由 DeepSeekMoE 提出，其基本思想为，固定某几个专家，响应所有的 token，这样可以让某些专家学习到共有的知识，而让其他的专家学习到特定的知识。这个方法随后被 Qwen1.5, Qwen2 , Qwen2.5 以及 DeepSeek-V3 所采用。

DeepSeekMoE 给出的实验结果如下

作者发现，当使用 shared experts 之后，模型在大部分 benchmark 上的表现都有所提升。

olmoe 在 32 个专家下进行了实验，比较了 4 个激活专家和 3 个激活专家 +1 个共享专家两种设置的表现，结果如下：

作者认为，加入 shared experts 之后，组合的可能性有所减少，这会降低模型的泛化性。因此，在 olmoe 中，作者没有使用 shared experts.

虽然 Qwen1.5, Qwen2 和 Qwen2.5 都使用了 shared experts, 但是后续的 Qwen3 中却并没有使用。

Ling-mini-beta 通过实验得出的结论为，shared expert 应该是一个非零的尽可能小的值，作者认为将 shared expert 设置为 1 是一个比较合理的选择。

Activation Function

一般来说，在选取 top-K 专家时，我们会对 gating layer 的输出进行归一化，通常我们会使用 softmax function:

$$ G(x) = \mathrm{softmax}(W_gx + b)\in\mathbb{R}^N $$

但是，在 Loss-Free Balancing 中，作者通过实验发现，使用 sigmoid 作为激活函数效果更好，即

$$ G(x) = \mathrm{sigmoid}(W_gx + b)\in\mathbb{R}^N $$

其对应的实验结果如下图所示

实验结果显示，sigmoid function 对于超参数更加 robust, 且表现也更好一些。下面是一些使用不同激活函数的模型例子

Activation function	Models
softmax	Step 3, Kimi-K2, gpt-oss-120B
sigmoid	GLM-4.5, dots.llm1, DeepSeek-V3

Routing Z-loss

Routing Z-loss 由 ST-MoE 提出， Switch Transformer 发现在 gating layer 中使用 float32 精度可以提高训练稳定性，但是这还不够，因此 ST-MoE 使用了如下的 router Z-loss:

$$ \mathcal{L}_z(x) = \frac1B\sum_{i=1}^B\left(\log\sum_{j=1}^N e^{x_j^{(i)}}\right)^2 $$

其中 $B$ 是 batch size ，$x_j^{(i)}=[W_gx_i+b]_j$ 代表了第 $j$ 个专家对 $i$ 个 token 的激活 logits. olmoe 实验验证结果如下

可以看到，加入 router Z-loss 之后，模型训练的稳定性有所提升，因此 Olmoe 采取了这个改进，但是后续的 MoE 模型使用 Z-loss 较少，个人猜测原因是 Loss-Free Balancing 中提出的加入额外的 loss 会影响 nex-token prediction loss

Routing Strategy

routing 策略直接决定了 MoE 模型的有效性。在为专家分配 token 的时候，我们有如下方式：

为每个 token 选取若干个专家
为每个专家选取若个个 token
动态分配 token 与专家之间的关系

三种选择方式如下图所示，图源 MoE survey

Expert Choice

每个专家选取 top-k 的 token，此时每个专家处理的 token 个数是相同的，这个方法的好处是自带 load balance。缺点是自回归生成的方式没有完整序列长度的信息，从而导致 token dropping，也就是某些 token 不会被任何专家处理，某些 token 会被多个专家处理。

目前采用这个策略的有 OpenMoE-2，核心思想是 dLLM 的输出长度固定，expert choice 策略更有效

Token Choice

每个 token 选取 top-k 的专家，好处是每个 token 都会被处理，缺点是容易导致负载不均衡。因此，一般需要加上负载均衡或者 token dropping 策略来提高负载均衡

Capacity Factor 由 Switch Transformer 提出，其定义为

$$ \text{expert capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) * \text{capacity factor} $$

设置 capacity factor 之后，当某个专家处理的 token 个数超过 capacity 之后，概专家的计算就会直接跳过，退化为 residual connection. 后续 DeepSeek-V2 也采用了这种策略，但是 DeepSeek-V3 弃用

Load balancing Loss 在训练目标中加入负载均衡损失，要求每个专家处理的 token 个数的分布尽可能均匀。

这部分具体见 Load Balancing loss

Global Choice

全局分配决定 token 和专家的匹配关系，后续 Qwen 提出了 Global-batch load balancing 使用了这种方式来提高专家的特化程度

Dynamic Routing

根据输入 token 的难度动态决定激活专家的个数。LongCat 使用了一个 Phantom expert 的方法来实现根据 token 的难度动态分配专家。具体来说，除了 $N$ 个专家之外，MoE 还包括 $Z$ 个 zero-computation expert (现在一共有 $N+Z$ 个专家参与计算), 其计算方式如下

$$ \begin{aligned} y &= \sum_{i=1}^{K}\mathrm{Normalize}(G_{e_i}) E_{e_i}(x), e_i \in \{i\mid G_i(x)\in \mathrm{TopK}(G_i(x), K)\}\\ E_{e_i}(x) &= \begin{cases} FFN_{e_i}(x), & \text{ if }1\leq i\leq N\\ x, & \text{ otherwise }\\ \end{cases} \end{aligned} $$

注：LongCat 还是用了 Loss-Free Balancing, 我们这里省略掉了。

Overview

现在几乎所有的模型都选择方式 1，即每个 token 选取 top-k 的专家。 olmoe 对比了以下方式 1 和方式 2 的表现，如下图所示

可以看到，加入 load balancing loss 之后，相比于 Expert Choice, Token Choice 的表现更好。但是，expert choice 更加高效，作者认为 expert choice 更适用于多模态，因为丢掉 noise image tokens 对 text token 影响会比较小。因此，在 olmoe 中，作者使用 token choice 作为 routing 策略

Upcycling

upsampling 是一个将 dense model 转化为 MoEmodel 的方法，具体做法就是我们复制 dense model 中的 FFN layer 得到对应 MoE layer 中的 Expert，然后我们再结合 router 训练，这样可以提高整体的训练效率。相关模型有 MiniCPM, Qwen1.5 和 Mixtral MoE (疑似)

实验结果如下图所示

从已有的结果来看，MoE 模型会被 dense 模型的一些超参数所限制，且训练不是很稳定。因此，现在一般不采用这种方法

Analysis on MoE

Specialization of Experts

OpenMoE 分析了 MoE 模型的特化程度，其结论如下

作者发现，大部分专家对于不同的 domain 没有出现 specialization 情况，对于 math domain, specialization 现象比较明显，作者认为这是因为 math domain 包含更多的 special tokens
对于不同的语言，有部分专家出现 specialization 现象
部分专家对于 position ID 有 specialization 现象，并且连续的 token 更偏好相同的专家
作者还发现，部分专家对于 Token ID 有 specialization 现象，作者将这种现象称为 Context-independent Specialization.
专家还会对语义相似的 token 进行聚类，并且这种聚类在训练早期就已经发生，作者认为其原因在于重新分配 token 会增加最终的 loss
对于 token dropping, 作者发现越靠后的 token, 其被 drop 的概率比例也越高。并且对于指令跟随数据，更多的 token 都会被丢掉，因此作者认为指令跟随数据是 MoE 模型的一种 OOD 数据

Saturation of Experts

olmoe 探究了训练过程中激活的专家和训练结束后激活的专家的匹配程度，结果如下图所示

实验结果说明，训练 $1\%$ 的数据之后，就有 $40\%$ 的 routing 和训练完毕的 routing 一致，当训练 $40\%$ 的数据之后，这个比例提升到了 $80\%$. 作者认为，这是专家特化的结果，初始的 routing 如果改变的话会带来表现下降，因此模型倾向于使用固定的专家处理特定的 token

作者还发现，later layers 比 early layers 饱和更快，early layer, 特别是 layer 0, 饱和的非常慢。作者认为，这是 DeepSeekMoE 放弃在第一层使用 MoE layer 的原因，因为 load balancing loss 收敛更慢。后续 DeepSeek-V2 和 DeepSeek-V3 均在 early layer 上使用 dense layer 替换掉了 MoE layer

Optimization

MoE 模型的优势在于表现好，但是模型参数往往非常大，为了方便使用，我们需要对训练好的 MoE 模型进行优化，目前主要有蒸馏，专家剪枝/合并以及量化等优化方法

蒸馏是一个将大模型能力传递给小模型的做法，目前已有的包括：

Switch Transformer 通过蒸馏，在仅使用 $1/20$ 参数的情况下，保留了稀疏教师模型 $30\%$ 的表现
Gemini2.5 通过蒸馏 Gemini2.5 Pro 得到 Gemini2.5 Flash
DeepSeek-R1 通过蒸馏来提升小语言模型的 reasoning 能力
Qwen3 对于小语言模型的训练使用了 off-policy distillation 和 on-policy distillation 来训练小语言模型

Code

我们这里展示基于 olmoe 的代码，代码如下

 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
43
44
45
46
47
48
49


class OlmoeSparseMoeBlock(nn.Module):
 def __init__(self, config):
 super().__init__()
 self.num_experts = config.num_experts
 self.top_k = config.num_experts_per_tok
 self.norm_topk_prob = config.norm_topk_prob
 self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
 self.experts = nn.ModuleList([OlmoeMLP(config) for _ in range(self.num_experts)])

 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 batch_size, sequence_length, hidden_dim = hidden_states.shape
 # hidden_states: (batch * sequence_length, hidden_dim)
 hidden_states = hidden_states.view(-1, hidden_dim)
 # router_logits: (batch * sequence_length, n_experts)
 router_logits = self.gate(hidden_states)

 routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
 # routing_weights: (batch * sequence_length, top_k)
 # selected_experts: indices of top_k experts
 routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
 if self.norm_topk_prob:
 routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
 # we cast back to the input dtype
 routing_weights = routing_weights.to(hidden_states.dtype)

 final_hidden_states = torch.zeros(
 (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
 )

 # One hot encode the selected experts to create an expert mask
 # this will be used to easily index which expert is going to be selected
 expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

 # Loop over all available experts in the model and perform the computation on each expert
 for expert_idx in range(self.num_experts):
 expert_layer = self.experts[expert_idx]
 idx, top_x = torch.where(expert_mask[expert_idx])

 # Index the correct hidden states and compute the expert hidden state for
 # the current expert. We need to make sure to multiply the output hidden
 # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
 current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
 current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]

 # However `index_add_` only support torch tensors for indexing so we'll use
 # the `top_x` tensor here.
 final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
 final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
 return final_hidden_states, router_logits

Infra

针对 MoE 模型的 infra 主要涉及 expert parallelism (EP), EP 将 MoE layer 的计算分为了三个阶段：

all-to-all dispatch: 基于 gating layer 的结果，将 token 通信传输到对应专家所在的 GPU 上
computation: 执行计算，即 $E_i(x)$.
all-to-all combine: 收集专家计算的结果，即 $y = \sum_{i=1}^K\mathrm{Normalize}(G_{e_i}(x)) E_{e_i}(x)$.

其框架图如下所示

Challenges

构建针对 MoE 模型的 infra 比较困难，相关工作有 DeepSeek 提出的 DeepEP.
训练不稳定，特别是负载均衡。负载均衡做不好容易导致模型崩塌。

Conclusion

在本文中，我们系统性回顾了 MoE 的相关概念，MoE 模型已经是现在大语言模型的主流架构，比如商业模型 Gemini2.5, 开源领先的模型 DeepSeek-V3 , LLaMA4 以及 Qwen3 等都采用了 MoE 的架构，如何进一步优化 MoE 的训练方式是当前研究的一个重点方向。

Appendix

MoE model information

Model	Year	total parameters	activated parameters	shared expert	Routed experts	activated experts
ST-MoE	2022/4	269B	32B	0	64	2
Mistral	2024/1	47B	13B	0	8	2
DeepSeek-MoE	2024/1	145B	22B	4	128	12
DeepSeek-V2	2024/5	236B	21B	2	160	6
LLaMA4	2025/4	400B	17B	1	128	1
DeepSeek-V3	2024/12	671B	37B	1	256	8
Qwen3	2025/5	235B	22B	0	128	8
dots.llm1	2025/6	142B	14B	2	128	6
Step 3	2025/7	316B	38B	1	48	3
Kimi-K2	2025/7	1043B	32B	8	384	1
GLM-4.5	2025/8	355B	32B	1	160	8
gpt-oss	2025/8	116B	5B	0	128	4
LongCat	2025/9	560B	27B*	0	768*	12
Ling-1T	2025/10	1000B	51B	1	256	8

Remark LongCat 采用了动态激活的方式，因此其结果有一个浮动范围。

References

Load Balancing tutorial

Thu, 11 Dec 2025 16:10:08 +0800

我们在本文中探讨关于 load balancing loss 的定义，性质和推广

Introduction

我们在 MoE tutorial 中已经介绍了 MoE 模块，MoE 尽管可以在相同的算力下扩大模型的 size, 但是其问题在于训练时容易出现负载不均衡，也就是说只有少数几个专家被激活，其他专家处于闲置状态，从而导致模型性能下降

为了解决这个问题，一个通用的做法是使用 load balancing loss. load balancing loss 可以有效实现负载均衡，让各个专家被激活的概率差不多。

Load balancing loss 一共经历了如下几个阶段，发展路线如下图所示

1
2
3
4
5
6


flowchart TD;
 C[Bengio et.at 2015] --> A[Shazeer et al., 2017];
 A --> B[GShard];
 B --> F[Switch Transformer];
 F --> D[Loss-Free Balancing];
 F --> E[Global-batch load balancing];

Method

Notation

我们假设输入的 batch size 为 $B$, 总专家个数为 $N$, 激活专家个数为 $K$, $G_i(x_j)$ 代表 gating layer 预测的第 $i$ 个专家对于 token $x_j$ 的重要性程度。

Coefficient of Variation

第一个阶段的 load balancing loss 针对广义的 MoE 模型。

Bengio et.at 2015 作者提出了一个让 batch 里每个专家平均激活概率分布更均匀的损失函数，其定义如下所示

$$ L_b = \sum_{i=1}^N\left\Vert \frac1B\sum_{j=1}^BG_i(x_j) - \frac1N\right\Vert_2 $$

其主要目标是让每个专家的平均激活概率接近 $1/N$, 即均匀分布.

Noam et.al 2017 作者对 (Bengio et.at 2015) 进行了改进，提出了变异系数 $\mathrm{CV}(\cdot)$ 来保证负载均衡，其定义如下：

$$ L_{importance} = \mathrm{CV}(\mathrm{Importance}(X))^2 $$

其中

$$ \mathrm{Importance}_i(X) = \sum_{j=1}^B G_i(x_j),\ \mathrm{CV}(X)= \frac{\mathrm{var}[X]}{\mathbb{E}[X]},\ i=1,\dots,N $$

注意到 $\mathrm{var}[X]\geq0$, 当且仅当 $X$ 为均匀分布时 $\mathrm{var}[X]=0$, 因此 important loss 可以让每个专家在一个 batch 中的激活的平均概率尽可能一致。

但是平均概率一致不代表各个专家处理的 token 个数一致，比如一个专家可能处理比较少的 token, 但是每个 token 的权重都很大。因此，作者额外加入了 load balancing loss

$$ \mathcal{L}_{\mathrm{load}} = \mathrm{CV}(\mathrm{Load}(X))^2 $$

其中

$$ \begin{aligned} \mathrm{Load}_i(X) &= \mathrm{soft\_estimation}\left(\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i)\right)\\ &=\mathrm{soft\_estimation}\left(\sum_{j=1}^B\mathbb{1}\{i\in\mathrm{TopK}(\{G_i(x_j)\}_{i=1}^{E}, K)\}\right),\ i=1,\dots,N \end{aligned} $$

这里 $\mathrm{soft\_estimation}$ 是作者针对离散变量进行的一个光滑化处理，以方便反向传播。

GShard

GShard 进一步对 (Noam et.al 2017) 提出的 loss 进行了简化，首先注意到 $\sum_{i=1}^N\mathrm{Load}_i(X)=BK$, 因此我们有

$$ \begin{aligned} \mathcal{L}_{\mathrm{load}} &= \mathrm{CV}(\mathrm{Load}(X))^2 \\ &= \left(\frac{\mathrm{var}[\mathrm{Load}(X)]}{\mathbb{E}[\mathrm{Load}(X)]}\right)^2\\ &= \left(C_1[\mathrm{Load}(X)]\right)^2\\ &= C_1\mathbb{E}[\mathrm{Load}(X)^2]-C_2 \end{aligned} $$

这里 $C_1,C_2$ 均为常数。GShard 并没有使用 soft estimation, 因此这里的 $\mathrm{Load}_i(X)$ 定义为

$$ \mathrm{Load}_i(X)=\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i),\ i=1,\dots,N $$

其代表了一个 batch 里专家 $i$ 被激活的次数，理想情况下，所有专家被激活的次数应该是一致的，从而我们就实现了负载均衡。

但是现在的问题是，$\mathrm{Load}_i(X)$ 是一个离散变量，为了解决这个问题，作者使用了重要性来进行近似，即

$$ \begin{aligned} \mathcal{L}_{\mathrm{load}} &= \mathbb{E}[\mathrm{Load}(X)^2]\\ &\approx \mathbb{E}[\mathrm{Load}(X)\cdot\mathrm{Importance}(X)] \end{aligned} $$

Switch Transformer

Switch Transformer 进一步对 GShard 提出 load balancing loss 进行了规范化，也是现在大部分 load balancing loss 使用的形式，其定义如下：

$$ \mathcal{L}_{\mathrm{load}} = \sum_{i=1}^N f_i P_i $$

其中，

$$ f_i = \frac{1}{B}\mathrm{Load}_i(X)=\frac{1}{B}\sum_{j=1}^B\mathbb{1}(\text{token }j \text{ selects expert }i),\ i=1,\dots,N $$

代表了分配给专家 $i$ 的 token 比例，

$$ P_i = \frac{1}{B}\mathrm{Importance}_i(X) = \frac{1}{B}\sum_{j=1}^B G_i(x_j) $$

代表了这个 batch 里专家 $i$ 的平均激活概率。

从而，这个形式和 GShard 的形式实际上是等价的：

$$ \mathcal{L}_{\mathrm{load}} = \sum_{i=1}^N f_i P_i = \mathbb{E}[\mathbf{f}\cdot \mathbf{P}] = N \mathbb{E}[\mathrm{Load}(X)\cdot\mathrm{Importance}(X)] $$

Analysis

目前关于 load balancing loss 可以实现负载均衡的数学分析是比较少的。Switch Transformer 中进行了如下分析：

[!quote] The auxiliary loss of Equation 4 encourages uniform routing since it is minimized under a uniform distribution.

但实际上，这个分析是错误的，Switch transformer 中定义的 load balancing loss 并不是在均匀分布时取最小值，而是取最大值。考虑如下情形，

$$ \mathbf{f}=[1,\dots,0],\ \mathbf{P}=[0,\dots,1] $$

此时，我们有

$$ \mathbf{f}\cdot \mathbf{P} = 0\leq [1/N,\dots,1/N]\cdot [1/N,\dots,1/N] = 1/N. $$

我找了很多资料，最后发现苏剑林老师在 MoE环游记：2、不患寡而患不均这篇 blog 中进行了分析。我这里基于苏剑林老师的 blog 进行了一下总结，推荐大家看苏剑林老师的原文。

分析的核心思想就是，尽管 load balancing loss 目标函数没有意义，但是我们通过 straight-through estimator (STE)，就可以得到一个有意义的目标函数，其梯度与 load balancing loss 目标函数的梯度是一致的，这样 load balancing loss 目标函数就完成了其负载均衡的目标。

注意到 load balancing loss 的最终目标是让每个专家处理的 token 个数尽可能一致，也就是

$$ \mathcal{L}_{\mathrm{load}} = \sum_{j=1}^B \left(f_i-\frac1B\right)^2 $$

但是 $f_i=\mathrm{Load}_i(X)$ 是一个离散变量不可导，因此，基于 STE, 我们可以将其改写为

$$ \mathcal{L}_{\mathrm{load}} = \sum_{j=1}^B \left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)^2 $$

其中 $\mathrm{sg}[\cdot]$ 满足：

$$ \mathrm{sg}[\cdot]=\begin{cases} x, &\text{forward pass}\\ 0, &\text{backward pass} \end{cases} $$

此时我们目标函数的梯度就变成了

$$ \begin{aligned} \nabla_\theta\mathcal{L}_{\mathrm{load}} &= \sum_{j=1}^B 2\nabla_\theta\left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)\left(P_i+\mathrm{sg}[f_i-P_i]-\frac1B\right)\\ &= 2\sum_{j=1}^B\nabla_\theta P_i\left(f_i-\frac1B\right)\\ &= 2\sum_{j=1}^B\nabla_\theta P_i\left(f_i-\frac1B\right)\\ &= 2\nabla_\theta \left(\sum_{j=1}^Bf_iP_i\right) \end{aligned} $$

这样就对应上了我们之前的目标函数。总之，目标函数尽管本身没有意义，但是其对应的梯度却可以让模型实现负载均衡的目标。

Extension

Loss Free Load Balancing Strategy

Loss-Free load balancing 是 DeepSeek 提出来的一个无需 load balancing loss 实现负载均衡的方法，其核心思想是 load balancing loss 会影响语言建模性能，因此在 $\mathrm{TopK}$ 操作时，作者加入一些扰动，从而让负载高的专家被选择的概率降低，让负载低的专家被选择的概率提高。

具体细节见 Loss-Free Balancing.

Global-batch Load Balancing Loss Strategy

Global-batch Load Balancing 是 Qwen 提出来的一个提高负载均衡的方法。其核心思想是，现有的 token choice 只是在 batch 层面达到负载均衡，这种局部均衡的性质可能会对模型的全局均衡性产生影响。因此，作者就预先在 global batch 层面对专家进行分配，从而提高专家在不同 domain 上的特化程度。

具体细节见 Global-batch load balancing

Experiments

olmoe 对比了加入 load balancing loss 之后模型的表现变化情况，结果如下图所示

可以看到，加入 load balancing loss 之后，模型的表现均超过了不加时的表现。

作者进一步分析了不同专家在加/不加 load balancing loss 时的激活情况，结果如下图所示

结果显示，load balancing loss 确实可以让不同专家的激活概率大致相当。

Conclusion

在本文中，我们介绍了 load balancing loss 的演化，定义，分析以及相关的实验结果。

Reference

LLM FLOPs Computation

Wed, 15 Oct 2025 16:33:39 +0800

本文中，我们介绍如何计算基于 transformer 架构的 LLM 的 FLOPs, 计算完成之后，我们可以推导出算力 $C$ 与模型参数量 $N$，数据集大小 $D$ 之间的关系，即 $C\approx 6ND$.

Background

FLOPs

FLOPs，floating point operations，表示浮点数运算次数，一般计算 FLOPs 仅考虑加法和乘法。

假设 $A\in\mathbb{R}^{m\times p}$, $B\in\mathbb{R}^{p\times n}$, 则计算 $C=AB$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp$ 次加法运算，因此总的 FLOPs 为 $2mnp$.

Assumption

假设 $A\in\mathbb{R}^{m\times p}$, $B\in\mathbb{R}^{p\times n}$, $C\in\mathbb{R}^{m\times n}$ 则计算 $D=AB+C$ 的过程中一共需要进行 $mnp$ 次乘法运算和 $mnp+mn$ 次加法运算，因此总的 FLOPs 为 $2mnp+mn\approx 2mnp$.

基于上述结论，我们计算 FLOPs 时，忽略 element-wise 的操作，即我们做如下假设：

忽略 normalization 中的小常数项运算
忽略 residual connection 和 bias term 的加法
忽略注意力的 MASK 和 softmax，因为这两者都是 element-wise operation.
使用 look-up 计算 embedding layer

注，基于以上假设，我们的结果与 Chinchilla Scaling law 的结果稍有不同，但结论不变。

Notation

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$	`num_attention_heads`	query head 的个数
$s$	`seq_len`	length of token sequence

注：为了避免混淆，我们使用 $n$ 来表示 decode layer 的个数。

Computation

我们计算训练阶段的总 FLOPs, 记为 $C$, Kaplan scaling law 用 PF-days 作为单位，$1\text{ PF-Days}=10^{15}\times 24\times 3600\ FLOPs$. 训练阶段包括前向阶段 (forward pass) 和反向传播阶段 (backward pass). 因此

$$ C = FLOPs(\text{forward}) + FLOPs(\text{backward}) $$

Forward

decoder-only transformer 的模型架构包含三个模块：

1 层 embedding layer
$n$ 层 decoder layer
1 层 lm head layer

因此模型总的 FLOPs 为

$$ FLOPs(\text{forward}) = FLOPs(\text{embedding}) + n*FLOPs(\mathrm{decode\_layer})+FLOPs(\mathrm{lm\_head}) $$

Embedding & Lm Head

首先，对于 embedding layer, embedding layer 本质上是一个 look up table, 计算过程中不涉及浮点数运算，因此 $\boxed{FLOPs(\text{embedding})=0}$.

接下来，对于 lm_head, 这是一个 linear layer, 其权重大小为 $W\in\mathbb{R}^{d\times |V|}$, 输入为 $x\in\mathbb{R}^{s\times d}$, 因此 $\boxed{FLOPs(\mathrm{lm\_head})=2sd|V|}$.

因此，我们有

$$ FLOPs(\text{forward}) = n*FLOPs(\mathrm{decode\_layer})+ 2sd|V| $$

Decode Layer

对于 decode_layer, 其又包含了四个模块：

pre-normalization
attention
post-normalization
FFN

pre-normalization 和 post-normalization 一般是一样的，因此

$$ \begin{aligned} FLOPs(\mathrm{decode\_layer}) &= FLOPs(\mathrm{pre\_normoalization}) + FLOPs(\mathrm{Attention}) + FLOPs(\mathrm{post\_normoalization}) +FLOPs(\mathrm{FFN})\\ &= 2*FLOPs(\mathrm{normoalization}) + FLOPs(\mathrm{Attention})+FLOPs(\mathrm{FFN}) \end{aligned} $$

Normalization

现在有两种常见的 normalization, 也就是 LayerNorm 和 RMSNorm

LayerNorm 定义如下

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

其中 $\beta,\gamma\in\mathbb{R}^d$ 是可学习的参数。

对输入 $x\in\mathbb{R}^{s\times d}$, 均值需要约 $sd$ 次 FLOPs，方差 $\mathrm{var}[x]$ 只需要约 $3sd$ 次 FLOPs，这两者的计算都可以忽略。接下来就是 scaling 和 shift, 这两者都是 element-wise 操作，我们这里，因此总的 FLOPs 为

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单

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

其中 $\gamma\in\mathbb{R}^d$ 是可学习的参数

对于 RMSNorm，其分析方式与 LayerNorm 基本一致，因此总的 FLOPs 为

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

总之，不管使用哪种 normalization，其 FLOPs 都是

$$ \boxed{FLOPs(\mathrm{normoalization}) = 4sd} $$

Attention

Attention 定义如下

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

其中 $X\in\mathbb{R}^{s\times d}$, $W_Q,W_K,W_V\in\mathbb{R}^{d\times d}$

$$ Q = W_QX\in\mathbb{R}^{s\times d},\quad K =W_KX\in\mathbb{R}^{s\times d},\quad V = W_VX\in\mathbb{R}^{s\times d} $$

$Q,K,V$ 计算的 FLOPs 为 $6*sd^2$. $QK^T$ 的 FlOPs 为 $2s^2d$, $\mathrm{softmax}(\cdot)V$ 的 FLOPs 为 $2s^2d$, 最后对于 multi-head attention 还有一个 output projection layer, 其权重为 $W_O\in\mathbb{R}^{d\times d}$, 因此 FLOPs 为 $2sd^2$. 故 attention 最终的 FLOPs 为

$$ FLOPs(\mathrm{Attention})=6sd^2+2s^2d+2s^2d+2sd^2=\boxed{8sd^2+4s^2d} $$

FFN

对于 FFN，有两种常见的形式，一种基于 ReLU 激活函数，其定义如下

$$ y = \max(xW_1+b_1, 0)W_2 + 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}$.

对输入 $x\in\mathbb{R}^{s\times d}$, 其 FLOPs 为

$$ FLOPs(\mathrm{FFN_{ReLU}}) = 2sdd_{ff} + 2sd_{ff}d =\boxed{4sdd_{ff}} $$

其中第一项和第二项分别为为 $xW_1$ 与 $\max(xW_1+b_1, 0)W_2$ 的 FLOPs.

另一种基于 SwiGLU 激活函数，其定义为

$$ \mathrm{SwiGLU}(x) = x\odot \sigma(x) $$

其中 $\sigma(\cdot)$ 是 sigmoid 函数

FFN 的定义为

$$ 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}}$.

对输入 $x\in\mathbb{R}^{s\times d}$, 其 FLOPs 为

$$ FLOPs(\mathrm{FFN_{SwiGLU}}) = 2sdd_{ff} + 2sdd_{ff} + 2sd_{ff}d = \boxed{6sdd_{ff}} $$

Summary

最终，decoder-only transformer 的 FLOPs 计算量为（我们假设使用 multi-head attention, 基于 SwiGLU 的 MLP）

$$ \begin{aligned} FLOPs(\text{forward}) &= FLOPs(\text{embedding}) + n*FLOPs(\mathrm{decode\_layer})+FLOPs(\mathrm{lm\_head})\\ &= n*FLOPs(\mathrm{decode\_layer})+2sd|V|\\ &= n*(2*FLOPs(\mathrm{normoalization}) + FLOPs(\mathrm{Attention})+FLOPs(\mathrm{FFN}))+2sd|V|\\ &= n*(8sd+8sd^2+4s^2d+6sdd_{ff})+2sd|V|\\ &= nsd^2\left(\frac8d + 8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right)\\ &\approx \boxed{nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right)} \end{aligned} $$

这里我们丢弃了 normalization 项，因为 normalization 的 FLOPs 是一个低阶项

Backward

首先，我们有如下结论：

神经网络 backward 过程的计算量（FLOPs）为 forward 过程的两倍

我们使用一个简单的例子来证明这个结论，考虑线性层 $h=Wx$, 其中 $W\in\mathbb{R}^{m\times d}$, 对于输入 $x\in\mathbb{R}^{d\times 1}$ 其 forward 过程的计算量为 $2md$.

对于反向过程，我们需要分别计算损失 $L$ 对权重和输入的梯度，即

$$ \frac{\partial L}{\partial x} = W^T\frac{\partial L}{\partial h}\in\mathbb{R}^{d\times 1}, \quad\frac{\partial L}{\partial W} = \frac{\partial L}{\partial h}\otimes x^T\in{m\times d}, $$

这里 $\frac{\partial L}{\partial h}\in\mathbb{R}^{m\times 1}$ 为损失对输出 $h$ 的梯度。因此反向传播的总计算量为

$$ 2dm + 2md = 4md $$

这里的两项分别是对 $x$ 和 $W$ 求梯度的 FLOPs.

Overall

将前向传播和反向传播的计算量汇总我们就得到一次前向传播和一次反向传播过程中，对于长度为 $s$ 的 token, 其 FLOPs 为 (multi-head attention, SwiGLU-FFN)

$$ \begin{aligned} C &= FLOPs(\text{forward}) + FLOPs(\text{backward})\\ &= 3FLOPs(\mathrm{forward}) \\ &\approx \boxed{3nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2V}{nd}\right)} \end{aligned} $$

Extension

GQA

GQA 与 MHA 不同的地方在于通过共享 key 和 value 来降低 KV cache 的占用，我们假设 group number 为 $g$, 则 key 和 value 的 FLOPs 现在变成了

$$ 2sdd_h\frac{h}{g}+2sdd_h\frac{h}{g}=4sdd_h\frac{h}{g} $$

因此 attention 部分总的 FLOPs 变成了

$$ FLOPs(\mathrm{Attention})=4sd^2+2s^2d+2s^2d+4sdd_h\frac{h}{g}=4sd^2+4s^2d+4sdd_h\frac{h}{g} $$

当 $g=h$ 时，GQA 就变成了 MHA, 此时的 FLOPs 也一致。

MoE

MoE 是针对 Dense FFN 的一个改进，介绍见 MoE, 我们假设一共有 $e$ 个路由专家，其中激活 $k$ 个。

Gate layer 一般是一个 linear layer, 其权重矩阵大小为 $W_{G}\in\mathbb{R}^{d\times e}$, 因此 $FLOPs(\text{router})= 2sde$.

Expert layer 和前面提到的 FFN 一致，我们每次挑选出 $k$ 个专家进行计算，因此 expert 部分 $FLOPs(\text{expert})=6ksdd_{ff}$.

从而对于 MoE 来说，FFN 部分的 FLOPs 为

$$ FLOPs(\text{MoE}) = FLOPs(\text{router})+FLOPs(\text{expert})= \boxed{2sde+6ksdd_{ff}} $$

Simplification

我们已经得到了 transformer 的 FLOPs 计算表达式，但是其表达式比较繁琐，因此，在研究 scaling law 时，一般会进行简化。

首先，在 LLM parameter analysis 中，我们已经给出了 LLM 参数量 $N$ （基于 Qwen3）的计算结果

$$ N=n*(4d+3dd_{ff}+2hh_{d}d + 2h_{kv}h_dd) + d(2|V|+1) $$

我们这里对其进行简化，一般来说 $d_{ff}=8/3d$, $h_{kv}=h$, $h_d=d/h$, 带入就得到

$$ N = n(4d+8d^2+2d^2+2d^2) + d(2|V|+1)=n(12d^2+4d)+ d(2|V|+1) $$

我们忽略关于 $d$ 的一阶项，并且我们假设 $|V| << 12nd$, 则最终模型参数量可以近似为

$$ \boxed{N \approx 12nd^2} $$

接下来，我们基于上面的配置简化 FLOPs 表达式

$$ \begin{aligned} C &= 3nsd^2\left(8+\frac{4s}{d}+\frac{6d_{ff}}{d}+\frac{2|V|}{nd}\right) \\ &= 3nsd^2\left(24+\frac{4s}{d}+\frac{2|V|}{nd}\right)\\ &\approx 72nsd^2 \\ &= 6sN \end{aligned} $$

这里我们利用了前面的 $|V| << 12nd$ 假设，为了简便我们还舍弃了 $4s/d$.

注意到 $s$ 代表 token 序列长度，如果训练集的总 token 个数为 $D$, 则最终对于包含 $D$ tokens 的数据集和包含 $N$ 参数量的 LLM, 其训练总 FLOPs 可以近似估计为

$$ \boxed{C\approx 6ND} $$

Experiments

Setting

接下来我们定量分析一些模型的 FLOPs. 我们基于 Chinchilla scaling law 给出的实验配置 (Table A9), 我们筛掉 kv_size * n_heads != d_model 的配置，$|V|=32,000$.

各部分的 FLOPs 计算代码如下

1
2
3
4
5


def compute_flops(n, V, d, d_ff, s):
 lm_head_flops = 2 * V * d * s
 attention_flops = 8 * s * d * d + 4 * s * s * d
 feed_forward_flops = 4 * s * d * d_ff
 return lm_head_flops, attention_flops, feed_forward_flops

首先我们看一下不同大小模型的 FLOPs 分布情况

可以看到，当模型越来越大，FFN 层的算力占比越来越高，这也是为什么后来采用 MoE 架构的一个原因。

接下来，我们看一下模型 FLOPs 分布情况随上下文长度变化的情况

可以看到，随 context length 增加，attention 的算力占比逐渐上升，这符合 attention 是一个平方复杂度的算法的结论。并且，当上下文足够长之后，计算还出现了 overflow (图像右端的突然下降)。

Conclusion

在本文中，我们介绍了如何计算基于 decoder-only transformer LLM 的 FLOPs, 并推导除了 Kaplan scaling law 中使用的公式 $C=6ND$, 这为后面的 infra 和 scaling law 的学习提供了基础。

References

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

Hands on LLM(2) Transformer

Sun, 29 Jun 2025 11:40:39 +0800

Transformer 实现

我们采用 top-down 的形式构建 transformer 的代码

架构

我们以 Qwen3 的代码为例子讲解 Assignment1 的代码实现

我们通过在 transformer 架构上加上一个 linear layer 就可以完成不同的下游任务，比如：

Qwen3ForQuestionAnswering
Qwen3ForCausalLM
Qwen3ForSequenceClassification

因此，大语言模型是 transformer 的一个附加产物

CausalLM

编写大语言模型的第一步为定义 Qwen3ForCausalLM

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


class CausalLM(nn.Module):
 def __init__(self, config):
 self.model = Transformer(config)
 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

 def forward(self, ...):
 outputs = self.model(***)
 logits = self.lm_head(outputs)

 return logits

这里 lm_head 的作用就是构建 embedding space 到 vocabulary 的映射，即 $\mathbb{R}^d\to\mathbb{R}^{|V|}$

Transformer

transformer 部分包括四个部分：

Embedding Layer：将 token 映射到 embedding space
layers：Transformer 的主体部分，由 $n$ 个 DecodeLayer 组成
Norm：在输出之前，进行一次 Normalization
Position Embedding：由于输入的 sequence 长度是固定的，因此我们提前计算好每一层的 position embedding

Transformer 部分的代码

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


class Transformer(nn.Module):
 def __init__(self, config):
 self.embedding = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
 self.layers = nn.ModuleList(
 [DecodeLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
 )
 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.rotary_emb = RotaryEmbedding(config)

 def forward(self, input_ids,...):
 input_embeds = self.embedding(input_ids)

 hidden_states = input_embeds
 position_embeddings = self.rotary_emb(hidden_states, position_ids)

 for decode_layer in self.layers:
 layer_outputs = decode_layer(hidden_states, position_ids, position_embeddings)
 hidden_states = layer_outputs[0]

 hidden_states = self.norm(hidden_states)

 return logits

DecodeLayer

DecodeLayer 就是 transformer 的核心部分，里面包含四个模块：

Pre-Normalization：一般是 RMSNorm 或者 LayerNorm
Attention：self-attention
Post-Normalization：与 Pre-Normalization 一致
MLP：FFN，SwiGLU 或者 MoE

DecodeLayer 还会使用 residual connection 来防止梯度消失

DecodeLayer 部分的代码

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


class DecodeLayer(nn.Module):
 def __init__(self, config, layer_idx):
 self.attn = Attention(config, layer_idx)
 self.mlp = MLP(config)
 self.pre_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 self.post_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

 def forward(self, hidden_states, position_ids, position_embeddings):
 residual = hidden_states
 hidden_states = self.pre_norm(hidden_states)
 hidden_states = self.attn(hidden_states, position_ids, position_embeddings)
 # residual
 hidden_states = hidden_states + residual

 residual = hidden_states
 hidden_states = self.post_norm(hidden_states)
 hidden_states = self.mlp(hidden_states)
 # residual
 hidden_states = hidden_states + residual

 return hidden_states

我们接下来按照

Normalization
MLP
Attention
Position embedding

的顺序来介绍

RMSNorm

RMSNorm 的作用和 LayerNorm 是一样的，但是实现上更简单

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

其中 $\beta,\gamma\in\mathbb{R}^d$ 是可学习的参数

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

其中 $\gamma\in\mathbb{R}^d$ 是可学习的参数

RMSNorm 代码实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


class RMSNorm(nn.Module):
 def __init__(self, d, eps):
 self.weight = nn.Parameter(torch.ones(d))
 self.eps = eps

 def forward(self, x):
 input_dtype = x.dtype
 x = x.to(torch.float32)
 variance = x.pow(2).mean(-1, keepdim=True)
 x = x * torch.rsqrt(variance + self.eps)
 return self.weight * x.to(input_dtype)

MLP

现在大语言模型的 MLP 使用的激活函数一般都是 SwiGLU, 其定义为

$$ \mathrm{SwiGLU}(x) = x\odot \sigma(x) $$

其中 $\sigma(\cdot)$ 是 sigmoid 函数

MLP 的定义为

$$ 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}}$

一般地，由于 FFN 只有两个权重矩阵，且 $d_{ff}=4d$, 在 SwiGLU 中，为了保证参数量一致，其隐藏层大小设置为 $d_{ff}'=\frac23d_{ff}=\frac83 d$.

MLP 的代码如下所示

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


def SwiGLU(x):
 return x * torch.sigmoid(x)

class MLP(nn.Module):
 def __init__(self, d, d_ff):
 self.gate_proj = nn.Linear(d, d_ff, bias=False)
 self.up_proj = nn.Linear(d, d_ff, bias=False)
 self.down_proj = nn.Linear(d_ff, d, bias=False)

 def forward(self, x):
 return self.down_proj(SwiGLU(self.gate_proj(x)) * self.up_proj(x))

Attention

我们先不考虑 position embedding，直接看 attention，attention 定义为

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

其中 $X\in\mathbb{R}^{m\times d}$,

$$ Q = W_QX\in\mathbb{R}^{m\times d},\quad K =W_KX\in\mathbb{R}^{n\times d},\quad V = W_VX\in\mathbb{R}^{n\times d} $$

在自回归模型里，我们还会加上 mask, 让每个 token 只能看见前面的 token 的信息

$$ \mathrm{Attention}(X) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\odot M\right)V $$

其中

$$ M = [M_{ij}] = \begin{cases} 1, &\text{ if } i < j\\ 0, &\text{ otherwise} \end{cases} $$

self-attention 的代码如下：

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


def scaled_dot_product_attention(Q, K, V, mask) -> torch.Tensor:
 d_k = Q.shape[-1] # d_k
 scaled_factor = 1 / d_k**0.5

 scores = torch.einsum("... s_q d_k, ... s_k d_k -> ... s_q s_k", Q, K)
 scores *= scaled_factor

 if mask is not None:
 scores = scores.masked_fill(mask == 0, float("-inf"))

 scores = scores.softmax(dim=-1)

 return torch.einsum("... s_q s_k, ... s_k d_v -> ... s_q d_v", scores, V)

Multi-Head Attention

Multi-Head Attention 定义如下

$$ \mathrm{MultiHeadAttention}(X) = [\mathrm{Attention}_1(X),\dots,\mathrm{Attention}_h(X)]W_o\in\mathbb{R}^{m\times d} $$

其中 $W_o\in\mathbb{R}^{d\times d}$, 且每一个 Attention heads 的维度会从 $d\to d/h$.

Multi-Head Attention 的主要作用为：

让不同的 head 关注不同的信息
并行计算，提高计算效率

MHA 代码

 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


class MultiHeadAttention(nn.Module):
 def __init__(self, d_model: int, num_heads: int) -> None:
 self.q_proj = Linear(d_model, d_model)
 self.k_proj = Linear(d_model, d_model)
 self.v_proj = Linear(d_model, d_model)
 self.output_proj = Linear(d_model, d_model,)

 def forward(self, x, position_embeddings, mask):
 Q = rearrange(self.q_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)
 K = rearrange( self.k_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)

 if mask is None:
 mask = torch.ones(Q.shape[-2], K.shape[-2])
 mask = torch.tril(mask)

 if position_embeddings is not None:
 sin, cos = position_embeddings
 Q = apply_rotary_pos_emb(Q, sin, cos)
 K = apply_rotary_pos_emb(K, sin, cos)

 V = rearrange(self.v_proj(x), "... seq_len (num_heads head_dim) -> ... num_heads seq_len head_dim",
 num_heads=self.num_heads, head_dim=self.head_dim)

 output = scaled_dot_product_attention(Q, K, V, mask=mask)
 output = rearrange(output, "... num_heads seq_len head_dim -> ... seq_len (num_heads head_dim)")
 return self.output_proj(output)

Position Encoding

Attention 对于输入的顺序是不敏感的，也就是

$$ \mathrm{Attention}(Q, \Pi K, \Pi V) = \mathrm{Attention}(Q, K, V) $$

这里 $\Pi\in \{0,1\}^{d\times d}$ 是一个置换矩阵 (permutation matrix)

Transformer 的解决方法是在 query 和 key 上加上位置信息：

$$ Q' = Q + PE(Q),\ K'=K + PE(K) $$

这样

$$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{(Q + PE(Q))(K + PE(K))^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

就包含了位置信息

绝对位置编码

Transformer 的使用的位置编码如下所示

$$ PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right) $$$$ PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right) $$

RoPE

苏剑林老师提出了 Position Encoding，现在已经被广泛使用

$$ q' = R_{\theta,m}^dq, k' = R_{\theta,n}^d k $$

这样 $\langle q, k\rangle$ 就仅包含两者的相对位置信息

$$ \langle q_m, k_n\rangle = x^TW_qR_{\theta, n-m}^d W_kx_n $$

RoPE 的矩阵定义如下

$$ R_{\theta,m}^d = \mathrm{diag}(M_1,\dots,M_{d/2}) $$

其中

$$ M_i = \begin{bmatrix} \cos m\theta_i & -\sin m\theta_i\\ \sin m\theta_i & \cos m\theta_i \end{bmatrix} $$

这里

$$ \theta_i = \frac{1}{10000^{2(i-1)/d}}, i\in\{1,2,\dots,d/2\} $$

简化后得到

$$ R_{\theta,m}^dq = \begin{bmatrix} \cos m\theta_0\\ \cos m\theta_0\\ \vdots\\ \cos m\theta_{d/2}\\ \cos m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} x1\\ x2\\ \vdots\\ x_{d-1}\\ x_d\\ \end{bmatrix} + \begin{bmatrix} \sin m\theta_0\\ \sin m\theta_0\\ \vdots\\ \sin m\theta_{d/2}\\ \sin m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} -\ x_2\\ x_1\\ \vdots\\ -x_d\\ x_{d-1}\\ \end{bmatrix} $$

RoPE 代码 Naive 实现

RotaryEmbedding 代码

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


class RotaryEmbedding(nn.Module):
 def __init__(self, dim):
 self.inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

 def forward(self, x, position_ids):
 freqs = einsum(torch.arange(self.max_seq_len),
 self.inv_freq, "seq_len, d_k_half -> seq_len d_k_half")[token_positions]

 sin = torch.sin(freqs)
 cos = torch.cos(freqs)

 return sin, cos

计算部分代码

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


def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
 x_even = x[..., ::2] # (seq_len, d_k_half)
 x_odd = x[..., 1::2] # (seq_len, d_k_half)
 odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half)
 evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half)
 stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half)
 stacked_trans = rearrange(
 stacked, "... seq_len double d_k_half -> ... seq_len d_k_half double"
 ) # (...,seq_len, d_k_half, 2)
 out = rearrange(
 stacked_trans, "... seq_len d_k_half double -> ... seq_len (d_k_half double)"
 ) # (..., seq_len, d_k)
 return out

RoPE 标准实现

RotaryEmbedding 代码 (LLaMA)

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


class LlamaRotaryEmbedding(nn.Module):
 def __init__(self, config: LlamaConfig, device=None):
 inv_freq = inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim)) # d_k_half

 def forward(self, x, position_ids):
 # (bsz, d_k_half, 1)
 inv_freq_expanded = self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
 # (bsz, 1, seq_len)
 position_ids_expanded = position_ids[:, None, :]
 # (bsz, seq_len, d_k_half)
 freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
 emb = torch.cat((freqs, freqs), dim=-1) # (..., seq_len, d_k)
 cos = emb.cos() # (..., seq_len, d_k)
 sin = emb.sin() # (..., seq_len, d_k)

 return cos, sin

计算部分代码

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


def rotate_half(x):
 x1 = x[..., : x.shape[-1] // 2]
 x2 = x[..., x.shape[-1] // 2 :]
 return torch.cat((-x2, x1), dim=-1)



def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 return q_embed, k_embed

References

Notes on Position encoding

Mon, 19 May 2025 10:46:39 +0800

本文前半部分参考参考文献1，推荐大家看博客原文。

Position encoding总结

在上一篇blog 中, 我们介绍了 Attention 的两个性质，也就是在不加 position encoding 的情况下，Attention 对于 query 是 permutation equivariant 的，对于 key 和 value 是 permutation invariant 的。

但是“我爱你”和“你爱我”这两句话所表示的含义应该是不一样的，我们将这两句话作为 key 和 value 的时候，我们发现模型的输出是一致的，这显然是不能接受的。因此，我们就需要加入 position encoding，让模型学习到语序信息，从而明白不同的语序有不同的含义。

下面是测试代码（来自参考文献1）

 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


import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
 for param in mha.parameters():
 nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True

Position encoding 可以分为绝对位置编码 (absolute position encoding, APE)，相对位置编码 (relative position encoding, RPE) 以及可学习的位置编码。可学习位置编码主要是 BERT 类的模型在使用，其训练成本比较高，本文不做讨论。绝对位置编码是原始 transformer 里提出的编码模式，现在的大多数基于 transformer 模型使用的都是相对位置编码。

本文中，我们先介绍位置编码应该具有的性质，然后我们分别介绍绝对位置编码和相对位置编码，我们将着重关注苏剑林老师提出来的 RoPE。

位置编码

在介绍位置编码之前，我们首先应该关注位置编码的性质，位置编码的目标是为输入的 token embedding 增加位置信息，那么理想的位置编码应该是怎么样的呢？

我们这里直接引用参考文献1 中给定的性质：

性质 1: token sequence 中每个位置的位置编码都是唯一的。这个很好理解，如果不唯一的话，那么根据前面推导的性质，这两个位置的 attention 输出就完全一致了
性质 2: 线性相关性。也就是说，如果我们知道了位置 $p$ 处的位置编码，那么理想情况下，我们应该能比较简单地得到 $p+k$ 处的位置编码，理想情况下，我们应该有 $PE(p+k)=W_kPE(p)$.
性质 3: 泛化到长上下文中去。我们希望位置编码不仅在 8K 的上下文起作用，还希望位置编码能够泛化到 32K 的上下文
性质 4: 生成模式是固定的。固定的模式有助于模型更好地学习位置相关的信息
性质 5: 可以扩展到多维。我们希望位置编码可以从文本扩展到图片再到视频，也就是从 $1D$ 到 $nD$.

绝对位置编码

绝对位置编码依照其名称，其思想就是为每个位置的 token 分配一个固定的位置信息，也就是对于输入的 hidden states $\bm{x}=[\bm{x}_1,\dots,\ \bm{x}_m]\in\mathbb{R}^{m\times d}$, 我们有

$$ \bm{x}_i' = \bm{x}_i + p_i, i=1,\dots, m $$

这里，$p_i\in\mathbb{R}^d$. 我们的 attention 就变成了

$$ \mathrm{Attn}(X) = \mathrm{softmax}\left(\frac{(Q+P)(K+P)^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

这里

$$ P = [p_1,\dots,p_m]\in\mathbb{R}^{m\times d}， Q= W_QX\in\mathbb{R}^{m\times d}, K=W_KX, V=W_VX\in\mathbb{R}^{n\times d} $$

整数位置编码

一个最简单的想法就是我们使用正整数来标记 token 所在的位置，也就是

$$ PE(i) = [i, \dots, i]=i\mathbf{1}_{d\times 1}\in\mathbb{R}^d,\ i=1,\dots,m $$

可以看到，这个简单的设计满足性质 1，性质 2，性质 3，性质 4.

但是，注意到 attention 的输入 $X$ 通常是经过 Layer Normalization 处理过后的，因此其按列符合正态分布，并且均值和方差一般较小。当我们加上整数位置编码之后，其 token 本身的信息就会被污染，也就是信噪比非常低。一个解决方法就是我们对 $PE(i)$ 进行 normalization，即

$$ PE(i)' = \frac{1}{m}PE(i) = \frac{i}{m}\mathbf{1}_{d\times 1} $$

现在所有的位置编码的值都比较小，但是我们发现新的位置编码不满足性质 2 了，这是因为现在位置编码还和 sequence 长度有关，我们从位置 $p$ 到位置 $p+k$ 不仅取决于 $k$ 还取决于 sequence 长度 $m$

二进制位置编码

既然整数位置编码的主要问题是对输入影响太大，我们能否找一个不影响输入的整数位置编码方式呢？参考文献1 提出了二进制位置编码，因为每个 token 是 $d$ 维的，因此我们可以使用 $d$ 位二进制来表示 $i$. 比如说，当 $d=3$, $m=4$ 时，我们的位置编码分别为

$$ PE(0) =p_{(000)_2} = [0, 0, 0],\ PE(1) =p_{(001)_2}= [0, 0, 1],\ PE(2) =p_{(010)_2} = [0, 1, 0],\ PE(3) =p_{(011)_2} = [0, 1, 1] $$

现在，我们二进制位置编码满足性质 1，性质 2. 对于性质 3，由于 $d$ 位二进制的表示范围为 $[0, 2^d-1]$，因此其泛化性受到 $d$ 的影响。

参考文献1 画出了不同位置的值的变化情况。我们这里也模仿绘制出类似的曲线图

我们发现，二进制位置编码高位，也就是 $PE(i)_{d}$ 的变化很慢，而低位，也就是 $PE(i)_{0}$ 变化很快，

二进制位置编码解决了整数位置编码的信噪比过低和线性相关性。但是其问题是其对不同位置的 token embedding 产生的影响是不一样的。比如位置 1 和位置 2 的相同的 token embedding 之间的区别是：

$$ (\bm{x}_2 + PE(2)) - (\bm{x}_1 + PE(1)) = (\bm{x}_2-\bm{x}_1)+ [0, 1, -1] $$

一般来说, $\bm{x}_2-\bm{x}_1$ 比较小，因此使用二进制位置编码的问题是输入位置的微小变化（增加一个 token 或减少一个 token）都会对最终结果产生巨大影响。因此，我们需要想办法解决这个问题。

Sinusoidal

前面提到二进制位置编码的问题是相邻 token 之间变化太大，不够光滑。因此我们想要增加一个光滑性质，也就是说我们希望：

位置编码值在 $[-1, 1]$ 之间，防止对 token embedding 产生影响
相邻 token 的位置编码尽可能相近，即 $|PE(k+p)-PE(p)| \leq \delta |k|$, 其中 $\delta>0$ 是一个比较小的数。
与二进制一样，高位的变化比较慢，低位的变化比较快

一个想法就是利用三角函数 $\sin$ 或者 $\cos$，三角函数满足前两个性质，对于第三个性质，我们可以通过控制频率来满足。这样我们得到的位置编码就具有如下形式：

$$ PE(p, i) = \sin\left(\frac{p}{\theta^{i/d}}\right) $$

其中 $\theta$ 是我们的超参数。

我们现在来推导一下上面位置编码的线性相关性：

$$ PE(p+k) = \sin\left(\frac{p+k}{\theta^{i/d}}\right)=PE(p)\cos\left(\frac{k}{\theta^{i/d}}\right) + \cos\left(\frac{p}{\theta^{i/d}}\right)\sin\left(\frac{k}{\theta^{i/d}}\right) $$

我们发现，$\sin$ 位置编码不满足线性相关性。但是出现的 $\cos$ 给了我们启发，也就是我们可以同时使用 $\sin$ 和 $\cos$ 来完成位置编码，这也是原始 transformer 里提出来的 Sinusoidal 位置编码，其形式为：

$$ \begin{aligned} PE(p, 2i) &= \sin\left(\frac{p}{\theta^{2i/d}}\right)\\ PE(p, 2i+1) &= \cos\left(\frac{p}{\theta^{2i/d}}\right) \end{aligned} $$

现在，记 $\omega_i=1/\theta^{2i/d}$, 我们再推导一下线性相关性，就得到：

$$ \begin{aligned} \begin{bmatrix} PE(p+k, 2i)\\ PE(p+k, 2i+1)\\ \end{bmatrix}&=\begin{bmatrix} \sin \omega_i(p+k)\\ \cos \omega_i(p+k) \end{bmatrix}\\ &=\begin{bmatrix} \sin \omega_i(\omega_ip)\cos(\omega_ik)+\cos w_i(\omega_ip)\sin(\omega_ik)\\ \cos \omega_i(\omega_ip)\cos(\omega_ik)-\sin w_i(\omega_ip)\sin(\omega_ik) \end{bmatrix}\\ &= \begin{bmatrix} \cos(\omega_ik)& \sin(\omega_ik)\\ -\sin(\omega_ik)& \cos(\omega_ik) \end{bmatrix}\begin{bmatrix} PE(p, 2i)\\ PE(p, 2i+1)\\ \end{bmatrix} \end{aligned} $$

也就是说，Sinusoidal 位置编码满足线性相关性。对于 Sinusoidal 位置编码我们也可以进行可视化：

相对位置编码

前面介绍了绝对位置编码，每个位置的位置编码是固定的。但是绝对位置编码的问题是，模型比较难以学习相对位置关系。

举个例子，我们提到上下文时，通常会使用“上一节”，“上一章”这些表示相对位置关系的词。

因此，我们希望让模型学习相对位置关系而不是绝对位置关系，因为相对关系更符合我们的认知。

RoPE

RoPE 由苏剑林老师提出，最早应用于 LLaMA 架构（没有确认），后续被大多数模型所采用。

之前的 PE 大多数关注于加性位置编码，也就是假设位置编码的形式为 $\bm{x}+\bm{p}$, 基于这种假设，已有的工作基本都集中于优化下面的 Q 和 K 的内积

$$ \langle f_q(\bm{x}_q, m), f_k(\bm{x}_k, n) \rangle $$

这里 $f_q(\bm{x}_q, m)=W_q(\bm{x}_q+\bm{p}_m)$, $f_k(\bm{x}_k, n)=W_k(\bm{x}_k+ \bm{p}_n)$.

而 RoPE 里面，作者使用了一个不同的假设： 假设内积应该仅包含两者的相对信息，也就是

$$ \langle f_q(\bm{x}_q, m), f_q(\bm{x}_k, n)\rangle := g(\bm{x}_q,\bm{x}_k, m-n) $$

这里的 $f$ 和 $g$ 都是未知函数。我们的目标就是从这个公式中推导出一个合适的位置编码出来。

不失一般性，我们可以假设

$$ f_q(\bm{x}_m,0) = \bm{x}_q,\quad f_q(\bm{x}_n, 0) = \bm{x}_k $$

这个假设代表初始条件下，我们不对输入做任何改变，也就是不增加位置信息。

2D 推导

与 RoPE 一样，我们直接使用复平面来进行推导。

我们假设 $d=2$, 注意到二维平面上的每个点都可以表示为如下形式

$$ \bm{z} = (x,y) = re^{i\theta} $$

其中 ($\mathrm{atan2}$ 定义参考维基百科)

$$ r = \|\bm{z}\|_2 = \sqrt{x^2+y^2}\in\mathbb{R},\quad \theta = \mathrm{atan2}(y, x)\in\mathbb{R}, $$

现在，对于三个向量 $f_q(\bm{x}_q, m)$, $f_q(\bm{x}_k, n)$, $g(\bm{x}_q,\bm{x}_k, m-n)$ 我们可以写出其极坐标形式：

$$ \begin{aligned} f_q(\bm{x}_q,m) &:= r_q(\bm{x}_q,m)e^{i\theta_q(\bm{x}_q,m)}\\ f_k(\bm{x}_k, n) &:= r_k(\bm{x}_k, n)e^{i\theta_k(\bm{x}_k, n)}\\ g(\bm{x}_q,\bm{x}_k, m-n) &:= r_g(\bm{x}_q,\bm{x}_k, m-n)e^{i\theta_g(\bm{x}_q,\bm{x}_k, m-n)} \end{aligned} $$

我们计算内积并比较同类项得到：

$$ \begin{aligned} r_g(\bm{x}_q,\bm{x}_k, m-n) &:= r_q(\bm{x}_q,m)r_k(\bm{x}_k, n)\\ \theta_g(\bm{x}_q,\bm{x}_k, m-n) &:= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, n) \end{aligned}\tag{3} $$

我们接下来分别推导 $r_g(\bm{x}_q,\bm{x}_k, m-n)$ 和 $\theta_g(\bm{x}_q,\bm{x}_k, m-n)$ 的形式

$r_g(\bm{x}_q,\bm{x}_k, m-n)$

我们令 $m=n=0$ 可以得到初始条件

$$ r_g(\bm{x}_q,\bm{x}_k, 0) = r_q(\bm{x}_q,0)r_k(\bm{x}_k, 0)=\|\bm{q}\|_2\|\bm{k}\|_2 $$

我们再令 $n=0$,得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_q(\bm{x}_q,m)r_k(\bm{x}_k, 0)=r_q(\bm{x}_q,m)\|\bm{k}\|_2=\frac{r_g(\bm{x}_q,\bm{x}_k, m-n)}{r_k(\bm{x}_k, n)}\|\bm{k}\|_2 $$

这里最后一个等式带入了原始等式 (3)，注意到左侧与 $n$ 无关，因此右侧我们选取 $n=1$, 得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = \frac{r_g(\bm{x}_q,\bm{x}_k, m-1)}{r_k(\bm{x}_k, 1)}\|\bm{k}\|_2 =\cdots= r_g(\bm{x}_q,\bm{x}_k, 0)\left(\frac{\|\bm{k}\|_2 }{r_k(\bm{x}_k, 1)}\right)^{m+1} $$

令 $m=0$ 我们有

$$ r_k(\bm{x}_k, 1) = \|\bm{k}\|_2. $$

因此我们最终的表达式为：

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_g(\bm{x}_q,\bm{x}_k, 0) = \|\bm{q}\|_2\|\bm{k}\|_2. $$

并且，通过分别设置 $m=0$ 以及 $n=0$ 我们还可以得到

$$ r_q(\bm{x}_q,m) = \|\bm{q}\|_2,\quad r_k(\bm{x}_k, n) = \|\bm{k}\|_2 $$

$\theta_g(\bm{x}_q,\bm{x}_k, m-n)$

令 $m=n=0$, 我们得到初始条件

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_q(\bm{x}_q,0)-\theta_k(\bm{x}_k, 0)=\theta_q-\theta_k $$

令 $n=1$, 我们有

$$ \begin{aligned} \theta_g(\bm{x}_q,\bm{x}_k, m-1) &= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 1)\\ &=\theta_g(\bm{x}_q,\bm{x}_k, m-n) + \theta_k(\bm{x}_k, n)-\theta_k(\bm{x}_k, 1) \end{aligned} $$

这里我们带入了公式 (3)，注意到公式左边与 $n$ 无关，因此在公式右侧我们令 $n=0$, 得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m-1) = \theta_g(\bm{x}_q,\bm{x}_k, m)+ \theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1) $$

分别令 $m=1,2,\dots$ 并相加这些等式，我们得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_g(\bm{x}_q,\bm{x}_k, m) + m(\theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1)) $$

即

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+(\theta_q-\theta_k)\tag{4} $$

注意到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 0)=\theta_q(\bm{x}_q,m)-\theta_k $$

带入上式我们就得到

$$ \theta_q(\bm{x}_q,m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_q $$

在 (4) 式中再令 $m=m-n$，并带入 $\theta_q(\bm{x}_q,m)$ 就有

$$ \theta_k(\bm{x}_k,n) = n(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_k $$

汇总

最后，我们将以上结果放在一起，就得到

$$ f_q(\bm{x}_q,m) = \bm{q}e^{im\theta}, f_v(\bm{x}_k,m) = \bm{k}e^{in\theta} $$

这里 $\theta=\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0)$ 是一个超参数，用于控制频率。

我们记

$$ R_{\theta,m} = \begin{bmatrix} \cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta \end{bmatrix} $$

则我们有：

$$ f_q(\bm{x}_q,m) = R_{\theta,m}\bm{q}, f_v(\bm{x}_k,m) = R_{\theta,n}\bm{k}. $$

并且

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} \tag{5} $$

多维扩展

上面是 2D 的情况，对于多维情况，苏剑林老师通过将两个元素组对，然后分别进行处理，得到了多维的情形：

$$ R_{\theta,m}^d = \begin{bmatrix} R_{\theta_1,m} & & && & \\ & & R_{\theta_2,m} & & & \\ &&&& \ddots & \\ &&&& & R_{\theta_{d/2},m} \end{bmatrix}\in\mathbb{R}^{d\times d} $$

我们可以验证公式 (5) 仍然是成立的。

RoPE 的远程衰减性质

我们接下来看一下结果与相对距离 $m-n$ 之间的关系，注意到

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} = \sum_{i=1}^{d/2} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i $$

这里 $\bm{q}_i=[q_{2i},q_{2i+1}]^T$, $\bm{k}_i=[k_{2i},k_{2i+1}]^T$ 分别是对应的 pair，我们考虑其中一个分量，不妨假设 $\|\bm{q}\|_2=\|\bm{k}\|_2=1$, 我们有

$$ \begin{aligned} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i &\leq \bm{q}_i^TR_{\theta, m-n}\bm{q}_i\\ &= \bm{q}_i^T\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right)\bm{q}_i\\ &\leq \lambda_{\max}\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right) \\ &= \cos (m-n)\theta_i \end{aligned} $$

其中，第一个不等式是因为两个向量相等时其内积最大，第二个不等式是由于二次型最大值为矩阵的特征值。

这样我们就有

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle \leq \sum_{i=1}^{d/2}\cos (m-n)\theta_i. $$

我们可以简单画出对应的曲线：

这里针对不同的配置，结果也会略微不同，具体分析可以参加知乎回答 RoPE的远距离衰减

RoPE 代码实现与理解

Naive 实现

我们接下来看一下如何实现 RoPE

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0 & -\sin m\theta_0 & &&\cdots &\cdots &\cdots \\ \sin m\theta_0 & \cos m\theta_0 & &&\cdots &\cdots &\cdots\\ & & \cos m\theta_1 & -\sin m\theta_1 & \cdots &\cdots &\cdots\\ & & \sin m\theta_1 & \cos m\theta_1 & \cdots &\cdots &\cdots \\ &&&& \ddots &\vdots &\vdots\\ &&&& & \cos m\theta_{d/2} & -\sin m\theta_{d/2}\\ &&&& & \sin m\theta_{d/2} & \cos m\theta_{d/2} \end{bmatrix}\begin{bmatrix} x_1\\ x_2\\ \vdots \\ x_d \end{bmatrix} $$

在实现的时候，我们一般根据 $\sin$ 和 $\cos$ 进行分组，也就是

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0\\ \cos m\theta_0\\ \vdots\\ \cos m\theta_{d/2}\\ \cos m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} x1\\ x2\\ \vdots\\ x_{d-1}\\ x_d\\ \end{bmatrix} + \begin{bmatrix} \sin m\theta_0\\ \sin m\theta_0\\ \vdots\\ \sin m\theta_{d/2}\\ \sin m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} -\ x_2\\ x_1\\ \vdots\\ -x_d\\ x_{d-1}\\ \end{bmatrix} $$

我们通常按照奇偶 index 来分别计算，然后通过重排序来得到最终的结果，实现代码如下：

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


def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
 x_even = x[..., ::2] # (seq_len, d_k_half)
 x_odd = x[..., 1::2] # (seq_len, d_k_half)
 odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half)
 evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half)
 stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half)
 stacked_trans = rearrange(
 stacked, "... seq_len double d_k_half -> ... seq_len d_k_half double"
 ) # (...,seq_len, d_k_half, 2)
 out = rearrange(
 stacked_trans, "... seq_len d_k_half double -> ... seq_len (d_k_half double)"
 ) # (..., seq_len, d_k)
 return out

LLaMA 实现

 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


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
 t = torch.arange(end, device=freqs.device) # type: ignore
 freqs = torch.outer(t, freqs).float() # type: ignore
 freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
 return freqs_cis

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
 ndim = x.ndim
 assert 0 <= 1 < ndim
 assert freqs_cis.shape == (x.shape[1], x.shape[-1])
 shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
 return freqs_cis.view(*shape)


def apply_rotary_emb(
 xq: torch.Tensor,
 xk: torch.Tensor,
 freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
 xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
 xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
 freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
 xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
 xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
 return xq_out.type_as(xq), xk_out.type_as(xk)

在 LLaMA 中，我们首先还是计算 $\theta_i$, 然后在计算的过程中，我们将 $(x_i,x_{i+1})$ 视作一个复数，然后乘以 $\exp(im\theta)$, 最后再取实部得到最终的结果

通用实现

实际上，naive 版本的实现与现在大语言模型所采用的实现并不一致，我们先看一下现有的大语言模型的 RoPE 实现，这里我们将 LLaMA的transformer代码放在下面，

 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


def rotate_half(x):
 """Rotates half the hidden dims of the input."""
 x1 = x[..., : x.shape[-1] // 2]
 x2 = x[..., x.shape[-1] // 2 :]
 return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 return q_embed, k_embed

class LlamaRotaryEmbedding(torch.nn.Module):
 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
 super().__init__()

 self.dim = dim
 self.max_position_embeddings = max_position_embeddings
 self.base = base
 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
 self.register_buffer("inv_freq", inv_freq)

 self.max_seq_len_cached = seq_len
 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

 freqs = torch.einsum("i,j->ij", t, self.inv_freq)
 # Different from paper, but it uses a different permutation in order to obtain the same calculation
 emb = torch.cat((freqs, freqs), dim=-1)
 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)

 def forward(self, x, seq_len=None):
 # x: [bs, num_attention_heads, seq_len, head_size]
 if seq_len > self.max_seq_len_cached:
 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

 return (
 self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 )

我们将上述代码翻译成公式，现在我们的 $\Theta$ 变成了 (对应 emb = torch.cat((freqs, freqs), dim=-1))

$$ \Theta = [\theta_0,\dots,\theta_{d/2},\theta_0,\dots,\theta_{d/2}]^T $$

实际上 $\sin$ 部分对应的向量现在变成了

$$ [-x_{d/2+1}, -x_{d/2+2}, \dots, -x_{d}, x_1, \dots, x_{d/2}]^T $$

我们带回到原始公式，可以得到对应的 RoPE 操作变成了

$$ R_{\theta,m}^d=\begin{bmatrix} \cos m\theta_0 & & & -\sin m\theta_0 & \cdots &\cdots &\cdots \\ & & \cos m\theta_1 & &-\sin m\theta_1 & \cdots &\cdots \\ & & & \cos m\theta_2 & &-\sin m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ & & & &\cos m\theta_{d/2 - 1} & & -\sin m\theta_{d/2 - 1} \\ \sin m\theta_0 && & \cos m\theta_0 &&\cdots &\cdots \\ & & \sin m\theta_1 & &\cos m\theta_1 & \cdots &\cdots \\ & & & \sin m\theta_2 & &\cos m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ &&&& \sin m\theta_{d/2 - 1}& & \cos m\theta_{d/2 - 1} \end{bmatrix} $$

这列每一行的 $\cos$ 和 $\sin$ 都相差了 $d/2$ 列.

因此，这里的区别在于，原始 RoPE 计算的 pair 为 $(x_{i}, x_{i+1})$, 而 LLaMA 里的 RoPE 计算的 pair 为 $(x_{i}, x_{i+d/2})$. transformers library 使用这种方式，可以减少计算量，提高整体的计算效率。

为了适应使用 LLaMA 中实现的 RoPE 的，Huggingface 对权重进行了转换，使得基于原始 RoPE 实现的模型也可以获得加速.

假设 $d=8$，原始 RoPE 的 pair 为 [(q_0, q_1), (q_2, q_3), (q_4, q_5), (q_6, q_7)], 新的 pair 为 [(q_0, q_4), (q_1, q_5), (q_2, q_6), (q_3, q_7)]. 我们希望对 index 进行 remap，我们发现一个满足条件的 permutation 为 [0, 2, 4, 6, 1, 3, 5, 7], 也就是 q_0->q_0, q_2->q_1, …, q_7->q_7.

但是，如果我们在推理时这样做，就会降低整体速度，因此 Huggingface 的做法是改变 $W_Q$ 和 $W_K$ 的权重，具体来说，就是 $\Pi q=(\Pi W_Q)x$，左边是在线转换，右侧离线转换。转换好 $W_Q$ 之后，正常计算就可以了。具体代码为

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


# permute for sliced rotary
def permute(w, n_heads=n_heads, dim1=dim, dim2=dim):
 return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)

state_dict = {
 f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wq.weight"]
 ),
 f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wk.weight"]
 ),
 ...
}

结论

本文中，我们回顾了位置编码，包括绝对位置编码和相对位置编码，我们着重介绍了 RoPE 的原理，推导以及代码实现。

参考文献

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的时候再一并讲解。

变量	含义
\(P\)	number of parameters
\(L\)	layers
\(V\)	vocabulary size
\(d\)	hidden size
\(d_{\text{ff}}\)	FFN hidden size
\(s\)	sequence length
\(b\)	batch size
\(h\)	number of attention heads
\(d_h\)	attention head dimension

Components	Training	Inference	Optimization
weights	\(2P\)	\(2P\)	quantization
optimizer states	\(4P\)	0	ZeRO, offloading
gradients	\(2P\)	0	ZeRO
activations	\(\sim 4Lhs^2b\)	\(\sim 0.4P\)	ckpt, offloading, flash attention
KV cache	0	\(4sLhd_h\)	KV optimization, attention
TOTAL	\(8P + 4Lhs^2b\)	\(2.4P + 4sLhd_h\)