Transformer 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.

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

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

Hands on LLM(1) Tokenizer

Sat, 24 May 2025 19:56:34 +0800

Introduction

在自然语言处理中，tokenizer的作用是将一个文本序列通过一个字典转化为一个token id的序列。我们回顾图片分类任务，模型在预测的时候，实际上预测的是类别对应的id，而不是类别本身。tokenizer做的事情就是提供一个类似于从类别到对应id的字典。

一般来说，一个tokenizer处理文本序列的过程有两步：

pre-tokenize，也就是预处理，我们需要将文本序列分割成合适大小的chunks (words)
tokenize，构建chunks (words)到token id的映射

实际上, huggingface的tokenizer包括四个步骤, 其中第二第三个步骤与上述一致. 在pre-tokenize之前, 我们有一个normalization过程, 该过程会对文本序列进行处理, 如将文本序列变为小写, 删掉声调符号等, 如下面例子所示：

1
2
3
4
5
6


from tokenizers import normalizers
from tokenizers.normalizers import NFD, StripAccents

normalizer = normalizers.Sequence([NFD(), StripAccents()])
normalizer.normalize_str("Héllò hôw are ü?")
# "Hello how are u?"

在tokenize之后, 我们会有一个post-processing过程, 比如BERT会在生成的token系列前后加入 [CLS] token 和 [SEP] token, 例子如下:

1
2
3
4
5
6
7


from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
token_ids = tokenizer.encode("I love NLP.")
print(token_ids)
# [101, 146, 1567, 21239, 2101, 119, 102]
# represents [[CLS], "I", "love", "NL", "##P", ".", [SEP]]

其完整流程如下图所示 (图源: huggingface llm-course)

构建好tokenizer之后, 我们还要保证tokenizer提供两个接口：

encoding, 给定文本序列, 将其映射到字典中去得到token id序列
decoding, 给定token id序列, 将其解码成文本序列

接下来, 我们将简单介绍一下word tokenizer, character tokenizer以及byte tokenizer, 并分析它们各自的不足。然后, 我们介绍现代大语言模型中使用最多的BPE tokenizer。最后, 我们介绍一些sub-word tokenizer。

Training-free tokenizer

本节我们将要介绍word tokenizer, character tokenizer以及byte tokenizer, 它们的特点就是简单易懂, 不需要额外的规则和学习.

Word tokenizer

给定一个文本序列, 我们现在需要将其转化为一个token序列。一个比较自然的想法是, 我们按照空格将序列拆分成若干个单词, 这样每个单词的语义都能比较好地保留。下面是一个例子：

1
2
3
4
5
6


import tiktoken

tokenizer = tiktoken.get_encoding("gpt2")
indices = tokenizer.encode("hello world")
# indices = [31373, 995]
# decode = ["hello", " world"]

接下来我们基于一个预定义好的词典, 将其转化为一个token id的序列。

word tokenizer的优点是能够保留语义信息，且压缩率比较高（每个token包含的bytes数），其问题是不能处理预定义好的词典之外的词 (out of vocabulary, OOV)。现有的处理方法是使用 <UNK> token来表示这些OOV的词。但这样显然会丢失语义信息, 因为我们编码成 <UNK> token之后, 就没办法再解码回原有的语义信息了。

word tokenizer的缺点为：

单词数量很大, 很多罕见单词的出现频率很低, 降低了tokenizer的利用率
对于不在词典内的单词只能用<UNK> token表示, 损害了语义信息

既然基于word的tokenizer有OOV的问题，我们能否想办法解决这个问题呢？答案是可以的, 我们可以使用 character tokenizer。

Character tokenizer

Character tokenizer的基本思想是使用字符而不是单词来编码文本序列。其实现方式如下：

1
2
3
4
5
6


class CharacterTokenizer:
 def encode(self, s: str) -> list[int]:
 return list(ord(c) for c in s)

 def decode(self, token_ids: list[int]) -> str:
 return "".join(chr(token_id) for token_id in token_ids)

character tokenizer的词表大小取决于我们的编码方式，UTF-8的编码大概有110K code points。character tokenizer的缺点总结如下：

character tokenizer会导致我们的词表非常大
和word tokenizer一样, 很多character非常罕见, 会降低词表的利用率
token序列的上下文语义信息较差

Byte tokenizer

我们发现, character tokenizer和word tokenizer的词表都很大, 我们能否想办法降低词表大小, 提升每个token的利用率呢？答案是使用Byte tokenizer.

Byte tokenizer的基本思想是, 所有的字符(character)都是由byte组成的, 比如对于UTF-8编码来说, 每个字符由1-4个byte组成。因此, 所有满足UTF-8编码的文本, 我们都可以将它们转换为基于byte的token序列。由于现在的大部分文本都是基于UTF-8的, 因此, 我们只讨论UTF-8编码的文本。

Byte tokenizer的实现如下：

1
2
3
4
5
6


class ByteTokenizer:
 def encode(self, s: str) -> list[int]:
 return list(s.encode("utf-8"))

 def decode(self, token_ids: list[int]) -> str:
 return bytes(token_ids).decode("utf-8")

byte tokenizer的词表很小, 其词表大小为 256, 这是因为一个byte可以有256中可能的值.

尽管byte tokenizer实现简单，并且词表也很小，可以说byte tokenizer解决了character tokenizer和word tokenizer的问题。但是，byte tokenizer的问题在于，其encode的到的token序列可能会非常长！我们知道，transformer计算量与token序列的长度是平方级关系的，也就是说token序列长度增加10倍，整体的计算量就会增加100倍，因此我们势必需要考虑token序列的长度。

总之，byte tokenizer的问题为：

产生的token序列过长, 增加了transformer的计算量
没有上下文语义信息

总结

我们总结一下word tokenizer, character tokenizer以及byte tokenizer三者各自的特点:

Feature	Word Tokenizer	Character Tokenizer	Byte Tokenizer
Granularity	Coarse	Medium	Fine
Vocabulary	Yes	No	No
Support OOV	Bad	Good	Best
#Tokens	Small	Large	Very Large
Chinese	Yes	Yes	Yes
Support Spell Error	Bad	Yes	Yes
Context	Good	Bad	Worst

因此, 这三种tokenizer尽管实现起来很简单, 但是其都有各自的问题. 为了解决这些问题, 我们的做法就是折衷, 使用sub-word tokenizer, 也就是介于word tokenizer和byte tokenizer之间的方法.

BPE

基本原理与实现

实际生活中，对于出现频率比较高的词，我们会有一个简写的方式，也就是我们使用一个新的单词来表示这个词。比如在英语中，我们会使用plz 来代替 please 以及使用how r u 来代替how are you。 BPE，即byte pair tokenizer的原理也是类似的，对于出现频率比较高的byte pair或者character pair, 我们会使用一个新的token来表示这个pair，这样就压缩了sequence的长度。

BPE算法包括以下几个步骤:

对文本序列进行pre-tokenize, 分割成不同的单词
当len(vocab)<vocab_size时, 重复以下步骤:
1. 对所有单词, 统计其相邻character或者byte pair的频率
2. 计算出现频率最高的pair, 使用一个新的token来表示这个pair
3. 将新的token和其对应的token_id加入到vocab中, 并更新单词的分割表示

算法如下图所示 (参考文献2)

注意：实际上，我们实现的是BBPE (byte BPE算法)，BBPE与BPE的区别在于我们的最小单元是character还是bytes。本质上原理是一致的

实现代码见 Github naive BPE

高效实现

BPE的原理很简单, 我们也实现了其naive版本, 但是naive版本的问题是太慢了。因此我们将要优化naive版本的效率。

首先我们发现, 我们不需要遍历所有的word, 只有含有best_pair的word我们才会进行处理, 因此, 我们的第一个改进就是使用 pair_to_word 来记录每个pair的来源, 比如：

1
2
3
4


pair_to_word = {
 (b' ', b't'): [b' the', b' it'],
 (b't', b'h'): [b'the']
}

这样, 我们在merge的时候, 直接使用 pair_to_word[best_pair] 来获取需要被更新的token序列就可以了。

其次, 注意到每次merge之后, 我们都需要重新计算一次 pair_freq, 而实际上, 只有被merge的token序列才需要被重新计数, 其他大部分token序列都是不需要重新计数的。因此, 一个改进点就是我们在merge的过程中就更新 pair_freq, 而不是重新计算。为了达到这个目标, 我们其实只需要两个操作。我们用(b'x', b'a', b'b', b'y') 和 best_pair=(b'a', b'b')来说明, merge之前, 这个序列贡献的pair_freq为：

1
2
3
4
5


{
 (b'x', b'a'): 1,
 (b'a', b'b'): 1,
 (b'b', b'y'): 1,
}

merge之后, token序列变成了(b'x', b'z', b'y') (假设best_pair对应的新的token为b'z'), 这时候的计数为:

1
2
3
4
5
6
7


{
 (b'x', b'a'): 0,
 (b'a', b'b'): 0,
 (b'b', b'y'): 0,
 (b'x', b'z'): 1,
 (b'z', b'y'): 1,
}

也就是说, merge之后, 三个pair的计数减少了1, 分别是(token_seq[i-1], merge_pair[0]),merge_pair 和 (merge_pair[1], token_seq[i+2])。两个pair的个数增加了1, 分别是 (token_seq[i-1], new_token)和(new_token, token_seq[i+2]) (这里我们假设merge_pair=(token_seq[i], token_seq[i+1]))。

基于这个结论，我们就可以优化BPE算法了，具体逻辑就是：

pretokenize, 将 text 切分为若干个 word
计算word_count, pair_freq, pair_to_word, 使用splits记录每个word对应的token分布
重复以下过程：
1. 挑选频率最高的pair将其merge为一个新的token, 基于pair_to_words更新对应的pair_freq
2. 对每个split, 按照上述方式更新pair_freq和split

其具体实现见 Github

Other subword tokenizers

WordPiece

WordPiece是Google在预训练BERT时采用的tokenizer，WordPiece的基本思想和BPE差不多，都是从一个较小的vocab开始的。

首先，WordPiece会通过加上prefix ##来把单词进行切分，比如 word 会被拆分为：

1

['w', '##o', '##r', '##d']

接下来，对于pair $(a, b)$, WordPiece定义了merge pair的规则如下：

$$ \mathrm{score}((a, b)) = \frac{\#(a, b)}{\#a \times \#b} $$

其中 $\#(a, b)$, $\#a$, $\#b$ 分别代表 pair $(a, b)$, 元素 $a$ 和元素 $b$ 的频率。通过这个方式，我们会给予包含元素出现频率较低的pair更高的优先级。通过这个方式，我们选取score最高的pair，然后将其用一个新的token表示，然后和BPE算法一样，继续这一过程直到我们的vocab size达到指定大小。

在tokenize的时候，WordPiece会找出现在vocab中的最长的subword, 比如对于'hugs', 假设从左向右在词典中的最长subword是'hug', 那么'hugs' 就会被拆分为 ['hug', '##s']。如果我们在词表中找不到对应的subword，这个时候我们就会使用'[UNK]'来表示。

具体实现见 Github wordpiece (基于huggingface llm course)。代码实现除了选择最优pair的方式不同之外，和BPE基本一致。

Unigram

Unigram也是由Google提出来的tokenizer，与BPE和wordpiece不同，unigram从一个非常大的vocab开始，然后merge token来降低vocab的size，直到达到指定大小。初始的vocab可以基于BPE算法或者使用prefix subword来构建。并且，初始vocab还包含所有的base characters来保证所有的word都可以被tokenize。

算法的描述如下:

我们来看一下算法的细节, 首先对于一个word, 我们有多种切割方式, 比如'bug'可以被切分为如下三种形式:

1

[['b', 'u', 'g'], ['b', 'ug'], ['bu', 'g']]

unigram 假设每个 word 出现的概率是其 subword 出现概率的乘积, 即对于包含 $n$个subword的单词 $\bm{x}=(x_1,\dots,x_n)$, 我们有:

$$ p(\bm{x}) = \prod_{i=1}^n p(x_i) $$

其中，对于给定的vocab $\mathcal{V}$, 我们有：

$$\sum_{v\in\mathcal{V}} p(x)=1$$

unigram的目的就是选择合适的切分 $\bm{x}\in S(\bf{x})$ (这里我们用 $\bf{x}$ 表示单词本身, 用 $\bm{x}$ 表示 $\bf{x}$ 的一个切分), 使得 $p(\bm{x})$的概率最大. 这样我们就可以写出unigram的损失函数了:

$$ \mathcal{L} = \sum_{i=1}^{N} \log\left(\sum_{\bm{x}\in S(\bf{x})}p(\bm{x})\right) $$

其本质就是: 我们希望对每个单词找到一种合适的切分, 切分得到的subword的概率分布满足其求和为1, 并且使得每个单词的概率最大.

但是直接对上面概率最大化的问题就是我们每个subword的概率是未知的, unigram的做法是使用EM算法求解这个问题.

当我们求解完成之后, 对每个subword, 我们都尝试将其从 $\mathcal{V}$中移除, 然后计算移除后的损失 $loss_i$, 我们依照$loss_i$对subword进行排序, 然后我们去掉 $\eta \%$ 比例的subword.

unigram的伪代码逻辑如下:

1
2
3
4
5
6
7
8
9


while len(model) > vocab_size:
 scores = compute_scores(model)
 sorted_scores = sorted(scores.items(), key=lambda x: x[1])
 # Remove percent_to_remove tokens with the lowest scores.
 for i in range(int(len(model) * percent_to_remove)):
 _ = token_freqs.pop(sorted_scores[i][0])

 total_sum = sum([freq for token, freq in token_freqs.items()])
 model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

其中 compute_scores 用于计算最优分割以及从model中去掉每个token之后的loss.

具体实现见 Github wordpiece (基于huggingface llm course)。代码实现的关键在于为每个word选取最优分割，huggingface是采取了动态规划的方法，也就是我们使用 dp[i] 来表示 word[:i] 的最优score，这样我们有：

$$ dp[i] = \max_{0 \leq j < i} dp[j]* p(word[j:i]),\quad \mathrm{s.t.}\ word[j:i]\in \mathcal{V} $$

这里的乘法代表 $p(\bm{x}) = \prod_{i=1}^n p(x_i)$, 在实现的时候我们会取log变成加法，然后概率会由频率来代替。

Subword tokenizer总结

sub-word tokenizer的对比 (来自huggingface llm course)：

Method	BPE	WordPiece	Unigram
Start Vocabulary	Small	Small	Large
Train	Merge tokens	Merge tokens	Remove tokens
Training Step	Merge with most frequent pair	Merge with best score	Remove all tokens minimized the loss
Learns	Merge rules and a vocab	A vocab	A vocab with a score for each token
Encoding	Splits into words and applies merge rules	Find the longest subword from the beginning that is in the vocab	Finds the most likely split into tokens with learned scores
Model	GPT	BERT	T5

实践

tiktoken

tiktoken是openAI提出来的一个BPE tokenizer, openAI的模型都基于这个tokenizer, 其主要用于调用GPT系列模型是对token进行计数, 我们可以在tokenizer 这个网站查看其分词情况.

SentencePiece

SentencePiece是google开源的一个无监督的text tokenizer，其实现了BPE和unigram两种算法，SentencePiece还是一个语言无关的tokenizer，使其更适合多语种大语言模型的开发。

Tokenizer

tokenizer 是huggingface推出的为基于transformer服务的tokenizer库, 其支持BPE, wordpiece和unigram等分词算法, 使用简便. 并且, huggingface的tokenizer包括两种:

fast tokenizer, 即Tokenizer库, 这个库是基于Rust开发的
slow tokenizer, 这个是transformer库里模型自带的, 比如ChatGLM就有自己开发的tokenizer

huggingface比较了并行处理时两者的区别:

Setting	Fast Tokenizer	Slow Tokenizer
`batched=True`	10.8s	4min41s
`batched=False`	59.2s	5min3s

huggingface提供的tokenizer库已经非常齐全了, 如果我们要训练新的基于transformer的模型的话，建议直接使用Huggingface的AutoTokenizer。

总结

特性	SentencePiece	Tokenizer	tiktoken
是否适合中文	√	√	×
是否适合英文	√	√	√
是否适合训练	√	√	×
是否快速	√	√	fast
是否用于 GPT 系列	×	×	√
是否可解码	√	√	√
是否支持多语言	√	√	×

结论

本文中, 我们介绍了大语言模型中的tokenizer, 我们从byte level, word level到sub-word level, 再到现代大语言模型最常使用的BPE tokenizer, 并给出了其（高效版本）实现。最后, 我们介绍了一下tokenizer-free的大语言模型和huggingface的tokenizer库。在未来, 我们将继续深入了解大语言模型的基本原理和实现细节。

参考文献

Notes on attention bias

Thu, 22 May 2025 15:25:07 +0800

Introduction

我们知道，transformer使用position encoding的一个原因就是，attention layer具有置换不变性，也就是说，我们随机打乱输入token的顺序，并不影响其最终结果 (我们后面会证明，实际上只对key和value具有置换不变性，对query具有置换等变性，也就是改变query的顺序之后，结果的顺序也相应改变)。因此为了让模型学习到正确的上下文知识，我们需要加上position encoding。

已有的工作大部分都在讨论如何构建更好的position encoding，但是鲜有工作探究为什么attention layer具有置换不变性. 因此，本文将从这一点出发，抽丝剥茧探究其内在原因，最后通过数学公式证明原始transformer是如何具有置换不变性的。

attention layer介绍

原始transformer layer的架构比较简单，其结构具有attention-LayerNorm-FFN-LayerNorm的形式。给定输入 $X\in\mathbb{R}^{d\times m}$ 和上下文 $Y\in\mathbb{R}^{d\times n}$. 其中，attention的定义为

$$ \mathrm{Attn}(X, Y, Y) = V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)\in\mathbb{R}^{d\times m} $$

其中 $d$是模型的hidden_size, $Q=W_QX\in\mathbb{R}^{d\times m}$, $K=W_KY\in\mathbb{R}^{d\times n}$, $V=W_VY\in\mathbb{R}^{d\times n}$, $W_Q, W_K, W_V\in\mathbb{R}^{d\times d}$ 分别是QKV projection layer的参数.

LayerNorm的定义为：

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

其中 $\epsilon>0$是一个超参数， $\gamma, \beta\in\mathbb{R}^d$ 是可学习的参数.

FFN的定义为：

$$ \mathrm{FFN}(x) = 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$ 是可学习的参数。

最后，一个attention layer的的结构可以表达为：

$$ X = X + \mathrm{LayerNorm}(\mathrm{Attn}(X, Y, Y))\\ X = X + \mathrm{LayerNorm}(\mathrm{FFN}(X))\\ $$

置换不变性的定义

置换不变性(permutation invariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

$$ f(\sigma(\bm{x})) = (f(\bm{x})) $$

则我们说 $f$是置换不变的. 这里 $\sigma:\mathbb{R}^n\to\mathbb{R}^n$ 是一个置换函数 (permutation function). 当输入的是一个矩阵时，我们默认置换其列，即对 $X=[X_1,\dots,X_n]\in\mathbb{R}^{d\times n}$, 我们有 $\sigma(X)=[X_{\sigma_1},\dots, X_{\sigma_n}]=Y\Pi $, 其中 $\Pi\in\mathbb{R}^{n\times n}\in \{0,1\}^{n\times n}$ 是一个置换矩阵 (permutation matrix)。

置换等变性 (permutation equivariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

$$ f(\sigma(\bm{x})) = \sigma(f(\bm{x})) $$

则我们说 $f$是置换等变的.

attention的置换不变性与置换等变性

我们首先证明attention 对于key和value是置换不变的，即

$$ \boxed{\mathrm{Attn}(X, \sigma(Y),\sigma(Y)) = \mathrm{Attn}(X, Y, Y)} $$

证明: 我们直接计算即可得到：

$$ \begin{aligned} \mathrm{Attn}(X, \sigma(Y),\sigma(Y)) &= V\Pi\mathrm{softmax}\left(\frac{(K\Pi)^TQ}{\sqrt{d}}\right)\\ &=V\Pi\mathrm{softmax}\left(\frac{\Pi^TK^TQ}{\sqrt{d}}\right)\\ \end{aligned} $$

由于softmax是按列计算的，置换只是改变了元素的顺序，因此我们自然有

$$ \mathrm{Attn}(X, \sigma(Y),\sigma(Y)) = V\Pi\Pi^T\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)=V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)=\mathrm{Attn}(X, Y, Y) $$

这里我们使用了性质 $\Pi\Pi^T=\mathbf{I}$.

接下来我们证明，attention对于query是置换等变的，即

$$ \boxed{\mathrm{Attn}(\sigma(X), Y, Y) = \sigma(\mathrm{Attn}(X,Y,Y))} $$

证明:

$$ \begin{aligned} \mathrm{Attn}(\sigma(X), Y, Y) &= V\mathrm{softmax}\left(\frac{K^TQ\Pi}{\sqrt{d}}\right)\\ &= V\mathrm{softmax}\left(\frac{K^TQ\Pi}{\sqrt{d}}\right)\\ &= V\mathrm{softmax}\left(\frac{K^TQ}{\sqrt{d}}\right)\Pi\\ &= \mathrm{Attn}(X,Y,Y)\Pi\\ &= \sigma(\mathrm{Attn}(X,Y,Y)) \end{aligned} $$

从以上的证明可以看到，attention layer对于key和value具有置换不变性，也就是说，我们改变文字顺序不影响最终的输出结果。但是，我们发现，尽管我们证明了attention具有置换不变性，我们却忽略了一件事：那就是我们计算query, key和value的时候，没有加上bias! 为什么bias如此重要呢？这是因为，$W\sigma(x) = \sigma(Wx)$, 但是 $W\sigma(X)\neq \sigma(Wx+b)$. 因此，我们就会思考，难道是transformer实际上可以通过增加bias的方式来让模型学习到上下文知识？事实上并非如此，我们将要通过分析表明，我们计算query, key和value时，增加的query bias和key bias会被softmax操作给消除掉，而key bias则会被LayerNorm消除掉。因此，我们加与加bias，对attention的置换不变性没有任何影响。

Bias对attention layer的影响

接下来，我们考虑在计算query, key和value时加入bias。为了简化，我们只考虑query为一个向量的情况，即 $X=\bm{x}\in\mathbb{R}^d$, 我们计算query, key和value如下：

$$ \bm{q} = W_Q\bm{x}+\bm{b}_Q\in\mathbb{R}^{d}\\ K = W_KY + \bm{b}_K\mathbf{1}^T\in\mathbb{R}^{d\times n}\\ V = W_VY + \bm{b}_V\mathbf{1}^T\in\mathbb{R}^{d\times n} $$

这里 $\mathbf{1}^T\in\mathbb{R}^{n}$. 我们这里简化了scaling的操作，因为其不对结果产生影响。

注：以下证明参考了【参考文献2】

我们首先展开attention中的 $V$:

$$ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= V\mathrm{softmax}\left(K^T\bm{q}\right)\\ &= \left(W_VY + \bm{b}_V\mathbb{1}^T\right)\mathrm{softmax}\left(K^T\bm{q}\right)\\ &= W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V\mathbb{1}^T \mathrm{softmax}\left(K^T\bm{q}\right) \end{aligned} $$

由于 $\mathrm{softmax}\left(K^T\bm{q}\right)\in\mathbb{R}^{n}$的列求和为$1$, 因此，$\mathbb{1}^T\mathrm{softmax}\left(K^T\bm{q}\right)=1$, 我们有

$$ \mathrm{Attn}(\bm{x}, Y, Y) = W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V $$

接下来，我们展开 $K$:

$$ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= W_VY\mathrm{softmax}\left(K^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left((W_KY + \bm{b}_K\mathbf{1}^T)^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left(Y^TW_K^Tq + \mathbf{1}\bm{b}_K^T\bm{q}\right) + \bm{b}_V\\ \end{aligned} $$

这里，我们需要用到softmax函数的平移不变性，即 $\mathrm{softmax}(\bm{x}+\delta\mathbf{1})=\mathrm{softmax}(\bm{x})$, 这里 $\delta\in\mathbb{R}$ 是一个常数，证明起来很简单：

$$ \mathrm{softmax}(\bm{x}+\delta)_i = \frac{e^{x_i+\delta}}{\sum_{j}e^{x_j+\delta}} = \frac{e^{x_i} * e^{\delta}}{\sum_{j}e^{x_j} * e^{\delta}} = \mathrm{softmax}(\bm{x})_i $$

而这里 $\bm{b}_K^T\bm{q}\in\mathbb{R}$，因此我们可以将这一项给去掉，我们得到：

$$ \mathrm{Attn}(\bm{x}, Y, Y) = W_VY\mathrm{softmax}\left(Y^TW_K^T\bm{q}\right) + \bm{b}_V $$

接下来，我们展开 $\bm{q}$,

$$ \boxed{ \begin{aligned} \mathrm{Attn}(\bm{x}, Y, Y) &= W_VY\mathrm{softmax}\left(Y^TW_K^T\bm{q}\right) + \bm{b}_V\\ &= W_VY\mathrm{softmax}\left(Y^TW_K^T(W_Q\bm{x}+\bm{b}_Q)\right) + \bm{b}_V\\ \end{aligned}} $$

因此，我们最终的结论为： key bias对attention输出没有任何贡献，query bias和key bias会影响结果。

到这里，看了参考文献3，我本以为可以进一步简化。但实际上并不行。参考文献3关于“transformer block is equivariant"的结果是错的，因为在attention layer之后还有一个LayerNorm，而LayerNorm不是置换不变的，这也是LayerNorm和BatchNorm之间的区别。也就是如果我们在nn.Linear后加一个BatchNorm，那么nn.Linear的bias是无效的，反之如果是LayerNorm的话，则bias是有效的.

为什么没有bias

实际上这个问题并没有定论。特别是加入position encoding之后，就更难探究bias对最终结果的影响了。但是，我认为一个原因就是bias其实就是某种先验知识，假设输入满足高斯分布，那么我们有

$$ \mathbb{E}[W\bm{x}+b] = b $$

加上先验知识后，当训练数据出现distribution shift之后，模型在训练过程中可能就会不稳定(PaLM). 而后来将LayerNorm替换为RMSNorm，使用RoPE而不是其他的additive position encoding, 我认为也是避免模型学习到先验知识，从而影响其泛化性。在未来，我认为transformer里应该是没有bias的，尽管这样效果可能会差一些，但是其稳定性更好，泛化性应该也会更好。

结论

在本文中，我们分析了attention的性质，我们发现，在原始transformer架构中，attention对于key和value有置换不变性，对于query有置换等变性。然后，我们给出了一些猜测，也就是bias会让模型产生先验知识，而这种先验知识很可能会影响训练的稳定性和模型的泛化性。

参考文献

附录

下面是测试上面结论的python代码

 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
50
51
52
53
54
55
56
57
58
59
60
61
62
63


import torch
import torch.nn as nn
import torch.nn.functional as F

# 设置随机种子，确保可复现性
torch.manual_seed(42)

# 输入参数
batch_size = 1
seq_len = 16
embed_dim = 1024 # 嵌入维度
num_heads = 32 # 多头注意力头数
head_dim = embed_dim // num_heads

# 输入张量 (batch_size, seq_len, embed_dim)
x = torch.randn(batch_size, seq_len, embed_dim)


# 有 bias 的 QKV 线性层
class Attention(nn.Module):
 def __init__(self, embed_dim, q_bias=False, k_bias=False, v_bias=False):
 super().__init__()
 self.q = nn.Linear(embed_dim, embed_dim, bias=q_bias)
 self.k = nn.Linear(embed_dim, embed_dim, bias=k_bias)
 self.v = nn.Linear(embed_dim, embed_dim, bias=v_bias)

 def forward(self, x):
 B, N, C = x.shape
 q = self.q(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
 k = self.k(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
 v = self.v(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)

 attn = (q @ k.transpose(-2, -1)) * (1.0 / (head_dim**0.5))
 attn = attn.softmax(dim=-1)
 x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 return x, attn


# 初始化模型
model_no_bias = Attention(embed_dim, q_bias=False, k_bias=False, v_bias=False)
model_with_bias = Attention(embed_dim, q_bias=False, k_bias=True, v_bias=False)

model_with_bias.q.weight.data = model_no_bias.q.weight.data
model_with_bias.k.weight.data = model_no_bias.k.weight.data
model_with_bias.v.weight.data = model_no_bias.v.weight.data

# 推理
out_no_bias, attn_no_bias = model_no_bias(x)
out_with_bias, attn_with_bias = model_with_bias(x)


# 比较差异
diff_output = torch.abs(out_no_bias - out_with_bias).mean()
diff_variance = torch.abs(out_no_bias - out_with_bias).var()
diff_attn = torch.abs(attn_no_bias - attn_with_bias).mean()

print("\nMean difference in output:", diff_output.item())
print("Mean difference in variance:", diff_variance.item())
print("Mean difference in attention weights:", diff_attn.item())

# Mean difference in output: 1.2734082233123445e-08
# Mean difference in variance: 1.7173628739783402e-16
# Mean difference in attention weights: 3.949708116124384e-09

Formal Algorithms for Transformer

Thu, 02 May 2024 13:13:12 +0800

This post is a notes on understanding how transformer works in an algorithm perspective.

Introduction

Transformer is a neural architecture that is used for neural language processing. Transformer receives an embedding matrix, which represents a sentence as input, and outputs a matrix of the same size as the embedding matrix, then the output can be used for downstream tasks.

Notation

We denote $V=[N_V]:=\{1,\dots,N_V\}$ as the vocabulary of tokens or words or characters.
We denote $\bm{x}=x[1...n]:=x[1]...x[n]\in V^n$ be a sequence of tokens, for example, a sentence or a paragraph.
Given a matrix $M\in\mathbb{R}^{m\times n}$, $M[i,:]\in\mathbb{R}^n$ is the $i$-th row of $M$, $M[:, j]\in\mathbb{R}^m$ is the $j$-th column of $M$.

Tokenization

Tokenization determines how the text are represented. Given a piece of text, for example, "I have an apple", we seek to find a proper way to represent this sentence.

Character level tokenization. In this setting, $V$ is the English alphabet plus punctuation. This tends to yield very long sequences (depends on the character contained in the raw text).
Word level tokenization. In this setting, $V$ is the set of all English words plus punctuation. Word level tokenization is straightforward, but it tends to required a very large vocabulary and cannot handle new vocabulary at test time.
Subword tokenization. This is the most common way used in nowadays, $V$ is the set containing the commonly used word segments like “ing”, “est”. This can be computed via Byte-Pair Encoding (BPE) algorithm.
We suppose the length of the input text is $L$, if the length input text exceeds $L$, we chunk it.

After tokenization, each element in the vocabulary is assigned to a unique index $i\in\{1,\dots,N_V-3\}$, and a number of special tokens are added to the vocabulary. For example:

mask_token$=N_V-2$, used in masked language modeling
bos_token$=N_V-1$ and eos_token$=N_V$, these two tokens are used to represent the beginning and the end of the sequence.

Finally, a piece of raw text is represented as a sequence of indices, often called token IDs corresponding to its subwords, preceded by bos_token and followed by eos_token.

Embedding

The embedding layer is used to represent each token as a vector that contains richer semantic information. The embedding contains two parts:

token embedding, where each token is embedded into a vector space
positional embedding, where embeds the position information of the tokens.

Token embedding

Given a sequence of token ID, we now need to represent each token as a vector in $\mathbb{R}^d$.

The simplest way is to use one-hot embedding, where each token $i$ is represented a vector $[0,\dots,1,\dots,0]\in\mathbb{R}^{N_V}$ whose elements are all $0$ excepts that $i$-th position is equal to $1$. However, the problem is that the vocabulary size $N_V$ is two large.

To solve this problem, we can train a learnable embedding model, of which parameter is a matrix $W_e\in\mathbb{R}^{d\times N_V}$, its $i$-th row corresponds to vector representation of the token $i$:

$$ \bm{e} = W_{e}[:, i]\in\mathbb{R}^d $$

Position embedding

There is a problem in token embedding, that is, it doesn’t contain consider the order of tokens. In latter, we show that the self-attention mechanism is equivariant to a permutation matrix $X\Pi$ of data $X$, where $\Pi$ is a permutation matrix, that is,

$$ \mathrm{Sa}(X\Pi) = \mathrm{Sa}(X)\Pi $$

the above equation indicates that the self-attention layer learn no position information at all!

To solve this problem, we add a positional embedding to token embedding. There are two kinds of embeddings:

Absolute positional embeddings. In this setting, a matrix $W_P\in\mathbb{R}^{d\times N}$ is learned or design to indicate the position of tokens. Mathematically, we have

$$ \bm{e}_p = W_p[:, \mathrm{index}(i)]\in\mathbb{R}^{d} $$

where $\mathrm{index}(i)$ is the index of token $i$ in the input sequence. 2. Relative positional embeddings. We leave this in latter notes. Compared to absolute positional embeddings, relative positional embeddings uses offset information, which performs well when the input sequence is too long.

The final embedding of a token $i$ is given by

$$ \bm{e} = W_e[:, i] + W_p[:, \mathrm{index}(i)]\in\mathbb{R}^{d} $$

Attention

The idea of attention mechanism is: Given a sequence of token, to predict the current token, which token should I pay attention to? For example, I opened the door with my ___, we may answer key, password or fingerprint etc. This is because we notice that we opened the door, so to predict the next token, we should make use of the information. What attention mechanism does is to quantify this process and make them parallel and learnable.

Single query attention

We first consider a simple example. Given the embedding of the current token $\bm{e}\in\mathbb{R}^d$ and the list of context tokens $[\bm{e}_1,\dots,\bm{e}_N]\in\mathbb{R}^{d\times N}$, the attention is given as follows:

compute query vector: $\bm{q}=W_q\bm{e}+b_q\in\mathbb{R}^{d}$
compute key vectors: for $i=1,\dots,L$, $\bm{k}_i=W_k\bm{e}_i+b_k\in\mathbb{R}^{d}$
compute value vectors: for $i=1,\dots,L$, $\bm{v}_i=W_v\bm{e}_i+b_v\in\mathbb{R}^{d}$
compute attention weights: let $\bm{s}=[\bm{q}^T\bm{k}_1,\dots,\bm{q}^T\bm{k}_L]\in\mathbb{R}^{N}$, then:

$$ \bm\alpha = \mathrm{softmax}\left(\frac{\bm{s}}{\sqrt{d}}\right)\in\mathbb{R}^{N}$$

5. compute vector representation of the token and context combined:

$$ \bm{v}'= \sum_{i=1}^N\alpha_i\bm{v}_i\in\mathbb{R}^{d} $$

where $W_q,W_k,W_v\in\mathbb{R}^{d\times d}$, $b_q,b_k,b_v\in\mathbb{R}$.

General attention

To extend the single query attention to general form, we consider the embedding matrix $X\in\mathbb{R}^{D\times N}$, the context matrix $Z\in\mathbb{R}^{d\times C}$ and a mask matrix $M\in\mathbb{R}^{D\times D}$, then the attention is computed as follows:

compute query matrix:

$$ Q=W_qX+\bm{b}_q\in\mathbb{R}^{D\times N}$$

2. compute key matrix:

$$ K=W_kZ+\bm{b}_k\in\mathbb{R}^{D\times C}$$

3. compute value vectors:

$$ V=W_vZ+\bm{b}_v\in\mathbb{R}^{D\times C}$$

4. compute attention weights:

$$\mathrm{Sa}(X) = \mathrm{softmax}\left(M\odot \frac{K^TQ}{\sqrt{D}}\right) \in\mathbb{R}^{C\times N} $$

where $\odot$ is the element-wise product. 5. output the updated representations of tokens in $X$, folding the information from tokens in $Z$

$$ \tilde{V} = V\odot \mathrm{Sa}(X)\in\mathbb{R}^{D\times N} $$

There are two kinds of mask matrices depending on which attention we are using:

Bidirectional attention, in this case, $M=\bm{1}\bm{1}^T\in\mathbb{R}^{C\times N}$.
Undirectional attention, in this case, $M[i,j] = \bm{1}_{i\leq j}$, where $\bm{1}_{i\leq j}$ is the indicator function.

Multi-head Attention

The previous describes the operation of a single head. In practice, transformers run multiple attention heads in parallel and combine their outputs, this is called multi-head attention. The idea behind multi-head attention can be summarized as follows:

In high dimensional space, two vectors are usually far from each other, with multiple attention heads, we can reduce the dimension of the representation.
With multiple attention heads, each heads may focus on specific semantics of the representation. For example, one head focuses on positiveness, and one another head focuses on noun/verb semantics.

For simplicity, we denote the single-head attention as $\mathrm{attention}(X, Z\mid M)$, suppose we have $H$ heads, then we compute $\tilde{V}_i$ for each heads:

$$ \tilde{V}_i = \mathrm{attention}(X, Z\mid M)\in\mathbb{R}^{D\times N},i=1,\dots,H $$

then attention representation are concatenated together:

$$ V = [\tilde{V}_1^T, \dots, \tilde{V}_H^T]^T\in\mathbb{R}^{HD\times N} $$

combined via an output matrix $W_o\in\mathbb{R}^{D\times HD}$:

$$ \tilde{V} = W_oV+\bm{b}_o\in\mathbb{R}^{D\times N} $$

We denote the multi head attention as $\mathrm{MhSa}(X, Z\mid M)$.

Transformer layer

After computing the multi head attention, we can now construct a transformer layer, which can also be stacked as convolution neural networks. A transformer layer can be constructed by the following operations:

Multi head attention (residual), $X\gets X + \mathrm{MhSa}(X, Z\mid M)$.
Layer norm, $X \gets \mathrm{LayerNorm}(X)$
Multi layer perception $\bm{x}_i\gets \bm{x}_i + \mathrm{mlp}(\bm{x}_i), i=1,\dots,N$
Layer norm, $X \gets \mathrm{LayerNorm}(X)$

where $\mathrm{LayerNorm}$ is the layer norm operation. $\mathrm{mlp}$ is a multi layer perception, usually it consists of one hidden layer of size $4D$, that is, then umber of neurons in three layers are $D, 4D, D$.

Usually, a large language model consists of multiple transformer layers.

Unembedding

The unembedding learns to convert a vector representation of a token and its context $\bm{e}$ into a distribution over the vocabulary elements.

$$ \bm{p} = \mathrm{softmax}(W_u\bm{e})\in \Delta(V)\subseteq \mathbb{R}^d $$

where $\Delta(V)$ is a simplex over the set $V$.

变量	含义
\(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\)