Infra on Jiang Yi(姜祎)'s Homepage

MultiHeadSelfAttention

Wed, 08 Apr 2026 17:44:26 +0800

我们在本文中介绍 MultiHeadSelfAttention, MHA的实现

MultiHeadSelfAttention

 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


import torch
import torch.nn as nn
import math

class MultiHeadSelfAttention(nn.Module):
 def __init__(self,nums_head=8,hidden_dim=512):
 super().__init__()
 self.hidden_dim = hidden_dim
 self.nums_head = nums_head
 self.head_dim = hidden_dim//nums_head
 self.q_proj = nn.Linear(hidden_dim,hidden_dim)
 self.k_proj = nn.Linear(hidden_dim,hidden_dim)
 self.v_proj = nn.Linear(hidden_dim,hidden_dim)
 self.o_proj = nn.Linear(hidden_dim,hidden_dim)
 self.dropout=nn.Dropout(0.1)

 def forward(self,x,attn_mask=None):
 batch_size,seq_len,_ = x.size()
 q=self.q_proj(x)
 k=self.k_proj(x)
 v=self.v_proj(x)
 # split
 q_state = q.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)
 k_state = k.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)
 v_state = v.view(batch_size,seq_len,self.nums_head,self.head_dim).transpose(1,2)

 # q@k^t/sqrt(d)
 attn_score = q_state@k_state.transpose(-1,-2)/math.sqrt(self.head_dim)

 if attn_mask is not None:
 attn_score = attn_score.masked_fill(attn_mask ==0,-1e9)
 attn_weight = torch.softmax(attn_score,dim=-1)
 attn_weight = self.dropout(attn_weight)
 output_mid = attn_weight@v_state
 # concat
 output_mid = output_mid.transpose(-1,-2).contiguous()
 output_mid = output_mid.view(batch_size,seq_len,-1)
 output= self.o_proj(output_mid)
 return output

x= torch.rand(3,2,512)
attn_mask = torch.tensor([
 [0,1],
 [0,0],
 [1,0],
]).unsqueeze(1).unsqueeze(2).expand(3,8,2,2)

net = MultiHeadSelfAttention(8,512)
net(x,attn_mask).shape

利用hugo和github搭建个人网址

Wed, 01 Apr 2026 17:44:26 +0800

安装go和hugo

https://gentleostrich.github.io/posts/windows-11-%E5%88%9B%E5%BB%BA-hugo-%E7%BD%91%E7%AB%99/

官方教程不允许使用 CMD、Windows PowerShell 执行各个命令，建议使用 PowerShell 执行命令。Anaconda Powershell Prompt 应该是基于 PowerShell 开发的，因此，本文的命令均在 Anaconda Powershell Prompt 中执行。

安装 Hugo

安装 Hugo 前需要安装 Git、Go、Dart Sass。

安装 Git

无论是克隆 Hugo 的 GitHub 仓库，还是基于 Github Page 发布网站，都需要使用 Git。
安装 Go

Hugo 是基于 Go 语言开发，Go 作为一种高效的编程语言，使得 Hugo 在速度上远超 Hexo 等静态网站搭建技术。
安装 Dart Sass

Dart Sass 是 Hugo 开发的网页渲染插件，可以通过 Scoop 和 Chocolatey 两个 Windows 安装器安装，这里选择使用 Scoop 安装。
- 安装 Scoop
  
  首先 cd 到 C:\ 路径，在 C:\ 路径下执行下述命令。
  1 2
  
  Set-ExecutionPolicy -ExecutionPolicy RemoteSigned -Scope CurrentUser Invoke-RestMethod -Uri https://get.scoop.sh | Invoke-Expression
  按照上述官方命令安装时出现了 PowerShell 禁止运行脚本的报错，使用以下命令替代上面的第一条命令解决该问题。
  1 2 3
  
  Set-ExecutionPolicy RemoteSigned -Scope Process Set-ExecutionPolicy RemoteSigned -Scope CurrentUser Set-ExecutionPolicy RemoteSigned -Scope LocalMachine
安装 Scoop 后，可以执行下述命令安装 Dart Sass。
1

scoop install sass

安装 Git、Go、Dart Sass 后

在 Hugo GitHub 发布页下载并解压 hugo_0.154.5_windows-amd64.zip 。
将解压得到的文件 hugo.exe 移动到安装路径，这里选择的安装路径是 D:\Program Files\hugo，并将该路径添加到系统路径变量中。

此时，重新打开 PowerShell 并执行 hugo version，若输出 Hugo 的版本号等信息，则表明 Hugo 安装成功。

网站创建

成功安装 Hugo 后，就可以创建第一个网站了。

1
2
3
4
5
6


hugo new site quickstart
cd quickstart
git init
git submodule add https://github.com/theNewDynamic/gohugo-theme-ananke.git themes/ananke
echo "theme = 'ananke'" >> hugo.toml
hugo server

报错 1：

1

ERROR command error: failed to load config: "D:\Program Files\hugo\quickstart\hugo.toml:4:2": unmarshal failed: toml: expected character =

直接将 echo "theme = 'ananke'" >> hugo.toml 复制到 PowerShell 中执行会出现格式错误，修改 hugo.toml

报错 2：

1

ERROR error building site: render: [en v1.0.0 guest] failed to render pages: render of "/tags" failed: "D:\Program Files\hugo\quickstart\themes\ananke\layouts\baseof.html:26:15": execute of template failed: template: taxonomy.html:26:15: executing "taxonomy.html" at <partials.Include>: error calling Include: "D:\Program Files\hugo\quickstart\themes\ananke\layouts\_partials\site-style.html:2:32": execute of template failed: template: _partials/site-style.html:2:32: executing "_partials/site-style.html" at <.RelPermalink>: error calling RelPermalink: TOCSS: failed to transform "/ananke/css/main.css" (text/css). Check your Hugo installation; you need the extended version to build SCSS/SASS with transpiler set to 'libsass'.: this feature is not available in your current Hugo version, see https://goo.gl/YMrWcn for more information

这是因为通过 hugo_0.154.5_windows-amd64.zip 安装的不是 extended version，不能对 anake 主题进行渲染。执行命令 scoop install hugo-extended 直接从 scoop 中安装 hugo-extended 并删除系统路径中的 hugo 路径即可。

网站配置

通过目录下的 hugo.toml 配置网站，其中的 baseURL 设置为生产网站。

1
2
3
4


baseURL = 'https://example.org/'
languageCode = 'en-us'
title = 'My New Hugo Site'
theme = 'ananke'

内容添加

向网站中添加一个新文章：

1

hugo new content content/posts/my-first-post.md

在 content/posts 目录下出现了 my-first-post.md 文件，下述是该文件中的内容。其中 draft=true 表示这个文章不会被发表到网站上。

1
2
3
4
5


+++
title = 'My First Post'
date = 2024-01-14T07:07:07+01:00
draft = true
+++

通过下述命令查看含有 draft 文章的网站，当决定发表该文章后，将 draft 改为 false。

1
2


hugo server --buildDrafts
hugo server -D

网站发布

网站发布指的是 Hugo 在 public 目录下创建静态网站所需的全部文件（包括 HTML 文件、CSS 文件、Javascript 文件、图片得等等），命令如下：

1

hugo

网站部署

基于 GitHub Page 部署网站。

新建项目，项目名称为：<username>.github.io。这意味着这个项目是一个网站，每次 push 后都会进行 action 操作进行网站部署。

将本地 Hugo 仓库与 GitHub 项目仓库进行 remote 关联，代码如下：

1
2
3


git remote add origin https://github.com/GentleOstrich/gentleostrich.github.io.git
git branch -M main
git push -u origin main

报错：

1
2


git push -u origin mainerror: src refspec main does not match any
error: failed to push some refs to 'https://github.com/GentleOstrich/gentleostrich.github.io.git'

原因是最开始本地项目从未 commit 过，所以没有不存在分支 main。执行下述命令即可解决：

1
2


git commit -m "Initial commit"
git push -u origin main

设置图片缓存位置

在 hugo.toml 文件中添加下述内容：

1
2
3


[caches]
 [caches.images]
 dir = ':cacheDir/images'

设置 action 文件

action 文件的作用是指导 GitHub 如何部署网站的，首先创建相关 hugo.yaml 文件：

1
2


mkdir -p .github/workflows
touch .github/workflows/hugo.yaml

接下来在 hugo.yaml 文件中加入下述内容：

 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
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106


name: Build and deploy
on:
 push:
 branches:
 - main
 workflow_dispatch:
permissions:
 contents: read
 pages: write
 id-token: write
concurrency:
 group: pages
 cancel-in-progress: false
defaults:
 run:
 shell: bash
jobs:
 build:
 runs-on: ubuntu-latest
 env:
 DART_SASS_VERSION: 1.97.2
 GO_VERSION: 1.25.5
 HUGO_VERSION: 0.154.4
 NODE_VERSION: 24.12.0
 TZ: Europe/Oslo
 steps:
 - name: Checkout
 uses: actions/checkout@v6
 with:
 submodules: recursive
 fetch-depth: 0
 - name: Setup Go
 uses: actions/setup-go@v6
 with:
 go-version: ${{ env.GO_VERSION }}
 cache: false
 - name: Setup Node.js
 uses: actions/setup-node@v6
 with:
 node-version: ${{ env.NODE_VERSION }}
 - name: Setup Pages
 id: pages
 uses: actions/configure-pages@v5
 - name: Create directory for user-specific executable files
 run: |
 mkdir -p "${HOME}/.local"
 - name: Install Dart Sass
 run: |
 curl -sLJO "https://github.com/sass/dart-sass/releases/download/${DART_SASS_VERSION}/dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz"
 tar -C "${HOME}/.local" -xf "dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz"
 rm "dart-sass-${DART_SASS_VERSION}-linux-x64.tar.gz"
 echo "${HOME}/.local/dart-sass" >> "${GITHUB_PATH}"
 - name: Install Hugo
 run: |
 curl -sLJO "https://github.com/gohugoio/hugo/releases/download/v${HUGO_VERSION}/hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz"
 mkdir "${HOME}/.local/hugo"
 tar -C "${HOME}/.local/hugo" -xf "hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz"
 rm "hugo_extended_${HUGO_VERSION}_linux-amd64.tar.gz"
 echo "${HOME}/.local/hugo" >> "${GITHUB_PATH}"
 - name: Verify installations
 run: |
 echo "Dart Sass: $(sass --version)"
 echo "Go: $(go version)"
 echo "Hugo: $(hugo version)"
 echo "Node.js: $(node --version)"
 - name: Install Node.js dependencies
 run: |
 [[ -f package-lock.json || -f npm-shrinkwrap.json ]] && npm ci || true
 - name: Configure Git
 run: |
 git config core.quotepath false
 - name: Cache restore
 id: cache-restore
 uses: actions/cache/restore@v5
 with:
 path: ${{ runner.temp }}/hugo_cache
 key: hugo-${{ github.run_id }}
 restore-keys:
 hugo-
 - name: Build the site
 run: |
 hugo \
 --gc \
 --minify \
 --baseURL "${{ steps.pages.outputs.base_url }}/" \
 --cacheDir "${{ runner.temp }}/hugo_cache"
 - name: Cache save
 id: cache-save
 uses: actions/cache/save@v5
 with:
 path: ${{ runner.temp }}/hugo_cache
 key: ${{ steps.cache-restore.outputs.cache-primary-key }}
 - name: Upload artifact
 uses: actions/upload-pages-artifact@v4
 with:
 path: ./public
 deploy:
 environment:
 name: github-pages
 url: ${{ steps.deployment.outputs.page_url }}
 runs-on: ubuntu-latest
 needs: build
 steps:
 - name: Deploy to GitHub Pages
 id: deployment
 uses: actions/deploy-pages@v4

至此，等待 GitHub 项目主页中 Action 显示绿色对勾后，就可以通过网址 <username>.github.io 访问网站啦（可能显示绿色对勾后立刻访问网站会出现 403 错误，只需要再等待一会儿就可以正常访问了）

shortcode

可用于内嵌视频、图片等元素，分为 embedded custom inline 三类。

插入一个 Youtube 视频（是一个 embedded）：

插入日期（是一个 inline）：今天是。

插入二维码：

设置github actions部署 workflow

https://github.com/jiangyigithub/ai.github.io/actions action workflow .github\workflows\deploy.yml

设置github page的

hugo github pages https://github.com/jiangyigithub/ai.github.io/settings/pages

更新文章

content\post

个人网址配置

github.io website baseurl = “https://github.com/jiangyigithub/ai.github.io" config_default\config.toml

Performance and Scalability

Thu, 26 Mar 2026 17:44:26 +0800

我们在本文中介绍 strong scaling 和 weak scaling

Parallel Performance

给定一个算法，Parallel speedup 定义为在一个处理器上的运行时间 $T_S$ 和使用 $N$ 个处理器的运行时间 $T_P$ 之比，即

$$ \mathrm{speedup} = \frac{T_S}{T_P} $$

直觉上，使用的处理器越多，我们所花费的时间越少，并且理想情况下，我们的时间应该随 $N$ 增加而线性减少。但实际上，我们并不能实现 $N$ 倍加速，因此，我们使用 efficiency 来评估我们的并行效率：

$$ \mathrm{efficiency} = \frac{\mathrm{speedup}}{N} = \frac{T_S}{T_P\cdot N} $$

我们可以将 efficiency 理解为处理器做有用功的时间比例。

Amdahl’s Law

我们可以将一个算法或者程度的运行时间分为可并行部分以及不可并行部分，可并行部分可以运行在多个处理器上，而不可并行部分只能在一个处理器上运行。

我们将算法运行的时间分为并行部分 $P$ 和串行部分 $S$, 我们有

$$ T_s = T_SS+T_SP,\quad S+P=1 $$

当我们使用 $N$ 个处理器时，我们只能减少可并行部分的运行时间 $T_SP$，并且由于我们使用了并行算法，我们还可能会产生额外的并行开销，记为 $K$, 则我们有

$$ T_p=T_SS+T_S\frac{P}{N}+K $$

如果我们假设 $K< $$ \mathrm{speedup} = \frac{T_S}{T_P}=\frac{1}{S+\frac{P}{N}} $$

这个等式称为 Amdahl’s law. 它说明一个算法通过并行的 speedup 取决于能够被并行的那部分，即 $P$. 比如硕，如果 $P=0.5$, 则不论我们使用多少处理器，我们最多只能得到 $2$ 倍的加速。

注意到 Amdahl’s law 假设我们的算法要处理的任务不随处理器数量而变化，这种 scaling 被称为 strong scaling.

Gustafson’s Law

但是在某些情况下，我们希望使用更多的处理器来处理更大的问题，这种 scaling 被称为 weak scaling, 其示意图如下所示

在这种场景下，随着 $N$ 增大，问题的规模增长，使得并行运行时间保持不变，此时在单个处理器场景下我们并行与串行部分的比例仍然不变，如果我们要处理 $N$ 倍工作量的问题，则我们的串行总时间为

$$ T_S=S+N\cdot P $$

由于并行运行时间不变，我们可以将运行时间表示为

$$ T_P=S+P=1 $$

此时我们的 speedup 为

$$ \mathrm{speedup} = \frac{T_S}{T_P} = S+N\cdot P $$

References

Performance and Scalability

Notes on roofline model

Thu, 26 Feb 2026 17:23:40 +0800

Introduction

当我们在硬件上运行算法时，我们的算法效率取决于三个因素：

计算机的计算效率，也就是每秒可执行的操作数 (OPs/s)
数据移动效率，也就是每秒可传输的 bytes (bytes/s)
需要的内存空间 (bytes)

通过分析这三个因素，我们可以给出一个计算的上下界

Arithmetic Intensity

首先，我们要回答的第一个问题是，运行时间意味着什么？为什么实际运行时间比理想时间要更多？

容易知道，实际的运行时间由三部分组成

Computation 首先，运行时间与我们的算法和 GPU 设备相关，在深度学习中，我们的算法通常由 FLOPs 决定，GPU 决定了算法的运行时间：

$$ T_{\mathrm{math}} =\frac{\text{Computation FLOPs}}{\text{Accelerator FLOPs/s}} $$

Communication within a chip 其次，运行时间还受加载数据的速度的影响，在执行计算之前，我们需要先把数据从 HBM 加载到对应的 core 上，因此，加载数据的速度也会影响最终运行时间。我们将加载数据的速度称之为 HBM bandwidth.

Communication between chips 最后，当我们在多个 GPU 上运行模型时，我们还需要在 GPU 之间进行通信，相关的通信手段如下表所示 (intra node 代表同一台服务器，inter node 代表不同服务器)，我们将 GPU 的通信速度称之为 Network bandwidth.

	Method	Data Path	BandWith	Use cases
Intra Node	NVLink	Direct GPU-to-GPU	1.8 Tb/s (B200)	TP
	NVSwith	all-to-all	900 GB/s-1.8T b/s	full mesh sync in DGX
	GPUDirect P2P	Direct over PCIe Bus	32-128 GB/s	small clusters withou NVLink
	Shared Memory	GPU -> CPU RAM -> GPU	20-50 GB/s	Fallback when P2P is disabled
Inter Node	InfiniBand	GPU -> NIC -> IB Switch	200-400 Gbps	Large scale LLM training
	GPU Direct RDMA	GPU -> NIC -> Network	Same as NIC speed	Bypassing CPU for network sync
	RoCE (v2)	RAME over Ethernet	100-400 Gbps	Cost-effective Ethernet clusters
	TCP/IP sockets	GPU -> CPU -> Kernel -> Net	< 100 Gbps	Standard networking

我们将 GPU 的通信 (intra-node 和 inter-node) 效率定义如下

$$ T_{\mathrm{comm}} =\frac{\text{Communication Bytes}}{\text{Network/Memory Bandwith Bytes/s}} $$

通常来说，我们可以通过计算 - 通信重叠来降低整体的运行时间，比如说 DeepSeek-V3 等模型提出的优化方法, 此时理论上的最小运行时间就是 $T_{\mathrm{math}}$ 和 $T_{\mathrm{comm}}$ 之间的最大值。另一方面，如果我们不做任何优化，则理论上最大运行时间就是 $T_{\mathrm{math}}$ 和 $T_{\mathrm{comm}}$ 之和。我们定义对应的 bound 如下

$$ \begin{aligned} T_{lower} &= \max(T_{\mathrm{math}} ,T_{\mathrm{comm}} )\\ T_{upper} &= T_{\mathrm{math}}+T_{\mathrm{comm}} \end{aligned} $$

注意到

$$ T_{upper} = T_{\mathrm{math}}+T_{\mathrm{comm}} \leq 2\max(T_{\mathrm{math}} ,T_{\mathrm{comm}} )= 2T_{lower} $$

也就是最大运行时间最多为最小运行时间的两倍。因此，我们一般直接优化 $T_{lower}$.

如果我们假设我们可以完全达到计算 - 通信重叠，则：

当 $T_{\mathrm{math}}>T_{\mathrm{comm}}$ 时, 我们完全利用了 GPU 的资源，因为此时运行时间受硬件的计算效率约束，我们将其称之为 compute bound
当 $T_{\mathrm{comm}}> T_{\mathrm{math}}$ 时, 此时运行时间受硬件的通信效率影响，GPU 的（计算）利用率不完全，我们将其称为 communication bound

为了评估一个算法时 compute bound 还是 communication bound, 我们可以使用 arithmetic intensity 来进行评估。我们先分析 single GPU 的情况，再分析 multi GPU 的情况

Single GPU

Arithmetic Intensity 一个算法的 arithmetic intensity 定义为算法通信所需要的 FLOPs 和所需要的 bytes 之比：

$$ \text{Arithmetic Intensity} = \frac{\text{Computation FLOPs}}{\text{Communication Bytes}} $$

arithmetic intensity 衡量了一个算法的 FLOPs per bytes, 当 arithmetic intensity 比较高时，说明 $T_{\mathrm{math}}>T_{\mathrm{comm}}$, 此时 GPU 的利用效率较高，反之则说明 GPU 出现了空置的情况。这两种情况的交叉点由 GPU 的 Peak arithmetic intensity 决定：

$$ \begin{aligned} T_{\mathrm{math}}=T_{\mathrm{comm}} &\Leftrightarrow \frac{\text{Computation FLOPs}}{\text{Accelerator FLOPs/s}} = \frac{\text{Communication Bytes}}{\text{Bandwith Bytes/s}}\\ &\Leftrightarrow \frac{\text{Computation FLOPs}}{\text{Communication Bytes}} = \frac{\text{Accelerator FLOPs/s}}{\text{Bandwith Bytes/s}}\\ &\Leftrightarrow \text{Intensity}(\text{Computation}) = \text{Intensity}(\text{GPU}) \end{aligned} $$

这里 $\text{Intensity}(\text{GPU})$ 就是 GPU 达到 peak FLOPs 时对应的 arithmetic intensity.

不同 GPU Tensor Core 对应的 intensity 如下所示

GPU	Generation	bandwidth	peak BF16 Tensor Core	intensity
V100	Volta	300 GB/s	125 TFLOPs/s	427
A100	Amper	1,935 GB/s	156 TFLOPs/s	83
H100	Hopper	3.35TB/s	495 TFLOPs/s	147
H200	Hopper	4.8TB/s	495 TFLOPs/s	147
B200	Blackwell	576 TB/s	2250 TFLOPs/s	3.9

Example 假设我们使用 BF16 精度计算两个长度为 $N$ 的向量的内积，计算时的流程如下：

从 HBM 中加载两个向量，通信量为 $2*2N$
计算两个向量内积，需要 $N$ 次乘法和 $N-1$ 次加法
将结果写回 HBM, 通信量为 $2$.

从而内积的 intensity 为

$$ \text{Intensity}(\text{dot product}) = \frac{\text{Total FLOPs}}{\text{Total Bytes}} = \frac{N+N-1}{2N+2N+2} = \frac{2N-1}{4N+2}\to \frac12 $$

也就是说，内积的 arithmetic intensity 最大为 $1/2$, 说明这是一个 memory-bound operation.

接下来，我们可以对 communication-bound 和 compute-bound 进行可视化，如下图所示

这里我们展示了两个算法 Algo 1 和 Algo 2 的 arithmetic intensity, 当 intensity 较小（小于硬件的 intensity）是，算法为 communication bound, 随着 intensity 增加，算法逐渐由 communication bound 过渡到 compute bound. 上面这个模型被称为 Roofline model, 其展示了算法计算效率与硬件之间的关系。

为了提高算法计算效率，我们可以通过提升算法的 arithmetic intensity 或者提高 memory bandwith.

Example

接下来我们来看一下矩阵计算的 intensity, 我们仍然假设在 BF16 精度下进行计算，假设输入矩阵为 $X\in\mathbb{R}^{b\times d}$ 以及 $Y\in\mathbb{R}^{d\times f}$ , 则输出为 $Z=XY\in\mathbb{R}^{b\times f}$. 计算过程为：

从 HBM 中加载 $X$ 和 $Y$, 通信量为 $2df+2bd$ bytes
计算矩阵乘法，计算量为 $2bdf$ FLOPs
将结果写回到 HBM, 通信量为 $2bf$.

因此对应的 intensity 为

$$ \text{Intensity}(\text{MatMul}) = \frac{2bdf}{2bd+2df+2bf} = \frac{bfd}{bd+df+bf} $$

假设 $b << \min(d,f)$ (在 LLM 中这个假设一般成立)，我们有

$$ \text{Intensity}(\text{MatMul}) = \frac{bfd}{bd+df+bf}\approx \frac{bdf}{df}=b $$

因此我们就可以通过调整 $b$ (batch size) 来将算法从 memory-bound 转换为 compute bound.

Multi GPU

在 multi GPU 的场景下，我们需要关注 GPU 之间通信的效率，一般来说，GPU 之间效率要远小于 GPU 内部通信效率，下面是不同 GPU 的内部通信（HBM bandwidth）和 GPU 之间通信效率对比

GPU	Generation	HBM Type	HBM bandwidth	NVLink Generation	NVLink Bandwidth
V100	Volta	HBM2	300 GB/s	2.0	300 GB/s
A100	Amper	HBM2e	1,935 GB/s	3.0	600 GB/s
H100	Hopper	HBM3	3.35TB/s	4.0	900 GB/s
H200	Hopper	HBM3e	4.8TB/s	4.0	900 GB/s
B200	Blackwell	HBM3e	576 TB/s	5.0	1800 GB/s

Example 我们考虑一个分布式矩阵乘法，假设我们有两个 GPU, 计算精度为 BF16, 矩阵 $X\in\mathbb{R}^{b\times d}$ 以及 $Y\in\mathbb{R}^{d\times f}$ 分别被切分为相同大小存储在这两个 GPU 上，此时，计算流程为：

在 GPU0/GPU1 上，分别从对应的 HBM 中加载 $X_1, Y_1$ 以及 $X_2, Y_2$
在 GPU0/GPU1 上，分别计算对应的矩阵乘法 $A=X_1Y_1\in\mathbb{R}^{b\times f}$ 和 $B=X_2Y_2\in\mathbb{R}^{b\times f}$, 计算量为 $2bdf$.
将 GPU1 的结果 $B$ 通信传输到 GPU0 上，通信量为 $2bf$.
对结果进行相加，$Z=A+B$. 计算量为 $bf$.

此时，我们的 intensity 为（注意这里计算时间除以 2 代表我们使用两张 GPU 因而效率翻倍了）

$$ \text{Intensity}(\text{MatMul (2-GPU)}) = \frac{2bfd/2}{2bf}= \frac{d}{2} $$

可以看到，此时决定性因素在于 $d$ 而不是 $b$, 这对于我们实现并行算法至关重要

megatron-lm

Wed, 21 Jan 2026 18:04:12 +0800

Introduction

随着模型参数变大，现有的 GPU 已经很难使用单一 GPU 来训练模型。对于多 GPU 训练场景，目前主要采用了 pipeline parallelism, 比如 GPipe 等，但是，这些策略需要我们对代码进行比较大的改动，这提高了开发成本。

为了解决多 GPU 训练大规模 LLM 的效率，降低开发成本，目前主要使用了 model parallelism 策略，即对模型进行切分部署在多个 GPU 上。model parallelism 有两种范式：

pipeline parallelism (PP): 将模型按照 layer 进行切分，如 GPipe 等，这种方法的问题是需要额外的逻辑来处理通信以及存在 pipeline bubbles
tensor parallelism (TP): 将模型的按照权重进行切分，部署在不同的 GPU 上。

作者在本文中基于 TP 策略来对 attention, FFN layer 进行简单改动来实现训练效率的提升。

作者通过实现验证了 tensor parallelism 的有效性和高效率，结果发现在 512 张 GPU 的场景下，TP 可以达到 $76\%$ 的 scaling efficiency (相比于 1 张 GPU 带来的性能提升)

Method

作者使用的 transformer 架构如下图所示

本文中，作者探究了 BERT 和 GPT-2 两种架构。

首先，我们假设 transformer layer 输入为 $X\in\mathbb{R}^{bs\times d}$, 这里 $b, s$ 分别为 batch size, sequence length, 接下来我们介绍如何针对 FFN, attention 以及 embedding 构建 TP 策略

FFN

论文中使用的 FFN 为 Linear-GeLU-Linear 的结构，对应第一层权重为 $W_1\in\mathbb{R}^{d\times d_{ff}}$, 第二层权重为 $W_2\in\mathbb{R}^{d_{ff}\times d}$, 对应数学表达式为

$$ Y = \mathrm{GeLU}(XW_1)W_2\in\mathbb{R}^{bs\times d} $$

我们首先对 $W_1$ 按照 column 进行切分，得到

$$ W_1 = [W_{11}, W_{12}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{11}\in\mathbb{R}^{d\times d_1}, W_{12}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff} $$

这里 $d_1, d_2$ 与我们并行的 GPU 数 (x-way TP) 相关，这样，我们就有

$$ \mathrm{GeLU}(XW_1) = \mathrm{GeLU}(X[W_{11}, W_{12}]) = \mathrm{GeLU}([XW_{11}, XW_{22}]) = [\mathrm{GeLU}(XW_{11}), \mathrm{GeLU}(XW_{12})] $$

从而我们可以分别将 $W_{11}$ 和 $W_{12}$ 部署在两个 GPU 上，然后并行计算。

论文中还介绍如果我们对 $W_1$ 按照 row 进行切分，则最终由于 $\mathrm{GeLU}(A+B)\neq \mathrm{GeLU}(A)+\mathrm{GeLU}(B)$ 计算时会产生一次额外的同步。

接下来，对于 $W_2$, 我们按照 row 进行切分得到

$$ W_2 = \begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix}\in\mathbb{R}^{d_{ff}\times d}, \text{ where }W_{21}\in\mathbb{R}^{d_1\times d}, W_{22}\in\mathbb{R}^{d_2\times d}, d_1+d_2=d_{ff} $$

计算时，我们有

$$ \mathrm{GeLU}(XW_1)W_2 = [\mathrm{GeLU}(XW_{11}), \mathrm{GeLU}(XW_{12})]\begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix} = \mathrm{GeLU}(XW_{11})W_{21} + \mathrm{GeLU}(XW_{12})W_{22} $$

可以看到，通过按照 row 进行切分，我们可以将 $W_{11}, W_{21}$ 部署在一个 GPU 上，将 $W_{12}, W_{22}$ 部署在另一个 GPU 上，分别计算出 $\mathrm{GeLU}(XW_{11})W_{21}$ 和 $\mathrm{GeLU}(XW_{12})W_{22}$ 之后，再通过一此 all-reduce 操作得到最终的输出结果。计算图如下所示

这里 $f$ 和 $g$ 是两个对偶算子，代表了 TP 产生的额外通信开销

operator	forward	backward
$f$	identity	`all-reduce`
$g$	`all-reduce`	identity

如果说我们使用的是 SwiGLU FFN, 即

$$ Y = (XW_3\odot \mathrm{Swish}(XW_1))W_2 $$

我们按照 column 对 $W_1, W_3$ 进行切分，按照 row 对 $W_2$ 进行切分（假设我们有 2 个 GPU），得到

$$ \begin{aligned} W_1 &= [W_{11}, W_{12}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{11}\in\mathbb{R}^{d\times d_1}, W_{12}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff}\\ W_3 &= [W_{31}, W_{32}]\in\mathbb{R}^{d\times d_{ff}}, \text{ where } W_{31}\in\mathbb{R}^{d\times d_1}, W_{32}\in\mathbb{R}^{d\times d_2}, d_1+d_2=d_{ff}\\ W_2 &= \begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix}\in\mathbb{R}^{d_{ff}\times d}, \text{ where }W_{21}\in\mathbb{R}^{d_1\times d}, W_{22}\in\mathbb{R}^{d_2\times d}, d_1+d_2=d_{ff} \end{aligned} $$

然后我们将 $W_{11}, W_{31}, W_{21}$ 放在第一个 GPU 上，将 $W_{12}, W_{32}, W_{22}$ 放在第二个 GPU 上，此时，

$$ \begin{aligned} \mathrm{Swish}(XW_1) &= \mathrm{Swish}(X[W_{11}, W_{12}]) = \mathrm{Swish}([XW_{11}, XW_{12}])=[\mathrm{Swish}(XW_{11}, \mathrm{Swish}(XW_{12}]\\ XW_3\odot \mathrm{Swish}(XW_1) &= [XW_{31}, XW_{32}]\mathrm{Swish}(XW_1) = [XW_{31}\mathrm{Swish}(XW_{11}), XW_{32}\mathrm{Swish}(XW_{12})]\\ Y = (XW_3\odot \mathrm{Swish}(XW_1))W_2&=(XW_3\odot \mathrm{Swish}(XW_1))\begin{bmatrix} W_{21}\\ W_{22} \end{bmatrix} = XW_{31}\mathrm{Swish}(XW_{11})W_{21}+ XW_{32}\mathrm{Swish}(XW_{12})W_{22} \end{aligned} $$

这样我们通过一次 all-reduce 也可以完成 SwiGLU FFN 的 tensor parallelism, 示意图如下所示

Attention

Attention 的处理与 MLP 非常相似，论文中的做法就是将不同 head 部署到不同 gpu 上分别进行计算，最后在计算 output projection 时再通过一次 all-reduce 来合并输出，这里我们假设有 $h$ 个 heads, 每个 head 的 dimension 为 $d_h$, 我们先对 query, key, value layer 的 weight $W_Q, W_K, W_V\in\mathbb{R}^{d\times hd_h}$ 进行切分

$$ W_Q = [W_{Q1}, \dots, W_{Qh}], W_K = [W_{K1}, \dots, W_{kh}], W_V = [W_{V1}, \dots, W_{Vh}] $$

其中 $W_{Qi}, W_{Ki}, W_{Vi}\in\mathbb{R}^{d\times d_h}$ 为每个 head 对应的 query, key, value weight. 我们将切分后的 $W_{Qi}, W_{Ki}, W_{Vi}$ 部署在一个 GPU 上（也可以将若干个 head 部署在一个 GPU 上），然后分别计算出每个 GPU 的 attention 结果，最后再进行汇总，如下所示

$$ \begin{aligned} o_i &= \mathrm{softmax}\left(\frac{(XW_{Qi})(XW_{Ki})^T}{d_h}\right) XW_{Vi}, i=1,\dots,h\\ O &= [o_1,\dots,o_h]W_O \end{aligned} $$

下面是 multi-head attention 对应的 TP 示意图

Embedding

对于 Input embedding, 作者将 embedding matrix $E\in\mathbb{E}^{V\times d}$ 按照 row 进行切分（论文中使用了转置，因此是按照 column 进行切分），得到 $E=[E_1,E_2]^T$, 这里 $E_i\in\mathbb{R}^{d\times V_i}$, $V_1+V_2=V$, 接下来我们把切分后的 embedding matrix 部署在不同的 GPU 上，由于每个 GPU 只有部分结果，因此我们还需要进行 all-reduce 来进行汇总。

而对于 output embedding, 我们也可以使用类似的做法进行切分，每个 GPU 上计算完结果之后我们还需要一个 all-gather 来汇总结果。

作者在这里还额外介绍了针对 output embedding 的优化方法，由于 embedding 的输出大小为 $[bs, V]$, 而 $V$ 通常比较大，因此，为了降低通信开销，作者将 cross-entropy-loss 与 output embedding kernel 进行融合，这样我们传输的数据量就减少到了 $bs$.

Experiments

作者首先对 GPT-2 模型进行了修正，首先将 vocab_size 从 50257 提升到 128 的倍数，即 51200. 对于 model+data parallelism, 作者固定 global batch size 为 512. (64-way DP)

配置如下表所示（head size 为 96）

Hidden size	attention heads	layers	parameters (B)	TP	TP+DP
1536	16	40	1.2	1	64
1920	20	54	2.5	2	128
2304	24	64	4.2	4	256
3072	32	72	8.3	8	512

对应的 scaling （使用多卡训练后，每个 GPU 相对于单卡训练的利用率）如下表所示

parallelism	TP-1	TP-2	TP-4	TP-8	TP-1+DP-64	TP-2+DP-64	TP-4+DP-64	TP-8+DP-64
scaling	100%	95%	82%	77%	96%	83%	79%	74%

Implementation

首先是 linear layer 的 TP 版本，如下所示

 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


import torch
import torch.nn as nn
import torch.distributed as dist


class ColumnParallelLinear(nn.Module):
 def __init__(self, in_features, out_features, bias=True):
 super().__init__()
 self.rank = dist.get_rank()
 self.world_size = dist.get_world_size()

 self.local_out = out_features // self.world_size
 self.weight = nn.Parameter(torch.empty(in_features, self.local_out))


 def forward(self, x):
 out = x @ self.weight

 gather_list = [torch.empty_like(out) for _ in range(self.world_size)]
 dist.all_gather(gather_list, out)
 return torch.cat(gather_list, dim=-1)


class RowParallelLinear(nn.Module):
 def __init__(self, in_features, out_features, bias=True):
 super().__init__()
 self.rank = dist.get_rank()
 self.world_size = dist.get_world_size()

 self.local_in = in_features // self.world_size
 self.weight = nn.Parameter(torch.empty(self.local_in, out_features))

 def forward(self, x):
 x_local = torch.chunk(x, self.world_size, dim=-1)[self.rank]

 out = x_local @ self.weight

 dist.all_reduce(out, op=dist.ReduceOp.SUM)
 return out

接下来是针对 LLM 中使用的 SwiGLU FFN 进行的优化，基于前面的介绍，我们不需要对基于 column linear 进行 all-reduce, 代码如下所示

SwiGLU

 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


import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist


world_size = 1
rank = 0

class ColumnParallelLinear(nn.Module):
 def __init__(self, in_features: int, out_features: int, dtype = None):
 assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})"
 self.part_out_features = out_features // world_size
 self.weight = nn.Parameter(torch.empty(part_out_features, part_in_features, dtype=dtype))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 y = x @ self.weight
 return y


class RowParallelLinear(nn.Module):
 def __init__(self, in_features: int, out_features: int, dtype = None):
 assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})"
 self.part_in_features = in_features // world_size
 self.weight = nn.Parameter(torch.empty(out_features, part_in_features, dtype=dtype))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 y = x @ self.weight
 if world_size > 1:
 dist.all_reduce(y)
 return y


class MLP(nn.Module):
 def __init__(self, dim: int, inter_dim: int):
 super().__init__()
 self.w1 = ColumnParallelLinear(dim, inter_dim)
 self.w2 = RowParallelLinear(inter_dim, dim)
 self.w3 = ColumnParallelLinear(dim, inter_dim)

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 return self.w2(F.silu(self.w1(x)) * self.w3(x))

attention

 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


class TPMultiHeadAttention(nn.Module):
 def __init__(self, d_model: int, num_heads: int, head_dim: int = None):
 self.d_model = d_model
 self.num_heads = num_heads
 self.head_dim = head_dim if head_dim is not None else d_model // num_heads

 assert self.num_heads % world_size == 0, "num_heads must be divisible by world size"
 assert self.d_model == self.num_heads * self.head_dim, "d_model must equal num_heads * head_dim"

 # heads of different GPU
 self.local_num_heads = self.num_heads // world_size
 self.local_qkv_dim = self.local_num_heads * self.head_dim

 self.qkv_proj = ColumnParallelLinear(in_features=d_model, out_features=3 * self.local_qkv_dim)
 self.out_proj = RowParallelLinear(in_features=d_model, out_features=d_model)

 self.scale = 1.0 / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))

 def _split_heads(self):
 batch_size, seq_len, _ = x.shape
 # [batch, seq_len, local_num_heads, head_dim]
 x = x.reshape(batch_size, seq_len, self.local_num_heads, self.head_dim)
 # [batch, local_num_heads, seq_len, head_dim]
 return x.transpose(1, 2)

 def forward(self, x: torch.Tensor, attn_mask: torch.Tensor = None):
 batch_size, seq_len, _ = x.shape
 # [batch, seq_len, local_qkv_dim * 3]
 qkv = self.qkv_proj(x)
 # [batch, seq_len, local_qkv_dim]
 q, k, v = torch.split(qkv, self.local_qkv_dim, dim=-1)
 # [batch, local_num_heads, seq_len, head_dim]
 q = self._split_heads(q)
 k = self._split_heads(k)
 v = self._split_heads(v)

 attn_scores = (q @ k.transpose(-2, -1)) * self.scale

 if attn_mask is not None:
 attn_scores = attn_scores + attn_mask

 attn_weights = F.softmax(attn_scores, dim=-1)
 # [batch, local_num_heads, seq_len, head_dim]
 attn_output = attn_weights @ v
 # [batch, seq_len, local_num_heads, head_dim]
 attn_output = attn_output.transpose(1, 2)
 # [batch, seq_len, local_qkv_dim]
 attn_output = attn_output.reshape(batch_size, seq_len, self.local_qkv_dim)

 gather_list = [torch.empty_like(attn_output) for _ in range(world_size)]
 dist.all_gather(gather_list, attn_output)
 # [batch, seq_len, d_model]
 full_attn_output = torch.cat(gather_list, dim=-1)

 out = self.out_proj(full_attn_output)

 return out

Embedding

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


class ParallelEmbedding(nn.Module):
 def __init__(self, vocab_size: int, dim: int):
 super().__init__()
 self.vocab_size = vocab_size
 self.dim = dim
 assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})"
 self.part_vocab_size = (vocab_size // world_size)
 self.vocab_start_idx = rank * self.part_vocab_size
 self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
 self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))

 def forward(self, x: torch.Tensor) -> torch.Tensor:
 if world_size > 1:
 mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)
 x = x - self.vocab_start_idx
 x[mask] = 0
 y = F.embedding(x, self.weight)
 if world_size > 1:
 y[mask] = 0
 dist.all_reduce(y)
 return y

Conclusion

作者提出了针对 transformer 架构的 tensor parallelism 策略来提高整体的训练效率，通过在训练过程加入四次 all-reduce 通信我们就可以训练更大规模的模型。

References

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.

Nvidia-GPU specs

Wed, 14 Jan 2026 11:09:19 +0800

V100

V100 关键改进

Volta architecture
SM architecture: 支持深度学习
2nd NVIDIA NVLink
HBM2 memory
Volta Multi-process Service

V100 技术规格

Tesla Product	Tesla K40	Tesla M40	Tesla P100	Tesla V100
GPU	GK180 (Kepler)	GM200 (Maxwell)	GP100 (Pascal)	GV100 (Volta)
SMs	15	24	56	80
TPCs	15	24	28	40
FP32 Cores / GPU	2880	3072	3584	5120
FP64 Cores / GPU	960	96	1792	2560
Tensor Cores / GPU	NA	NA	NA	640
GPU Boost Clock	810/875 MHz	1114 MHz	1480 MHz	1530 MHz
Peak FP32 TFLOPS²	5	6.8	10.6	15.7
Peak FP64 TFLOPS²	1.7	.21	5.3	7.8
Peak Tensor TFLOPS²	NA	NA	NA	125
Memory Size	Up to 12 GB	Up to 24 GB	16 GB	16 GB
Memory Interface	384-bit GDDR5	384-bit GDDR5	4096-bit HBM2	4096-bit HBM2
TDP	235 Watts	250 Watts	300 Watts	300 Watts
Manufacturing Process	28 nm	28 nm	16 nm FinFET+	12 nm FFN

内存规格

GPU	Kepler GK180	Maxwell GM200	Pascal GP100	Volta GV100
Compute Capability	3.5	5.2	6.0	7.0
Threads / Warp	32	32	32	32
Max Warps / SM	64	64	64	64
Max Threads / SM	2048	2048	2048	2048
Max Thread Blocks / SM	32	32	32	32
Max 32-bit Registers / SM	65536	65536	65536	65536
Max Registers / Block	65536	65536	65536	65536
Max Registers / Thread	255	255	255	255
Max Thread Block Size	1024	1024	1024	1024
FP32 Cores / SM	192	128	64	64
Ratio of SM Registers to FP32 Cores	341	512	1024	1024
Shared Memory Size / SM	16 KB/32 KB/ 48 KB	96 KB	64 KB	Configurable up to 96 KB

系统规格

Specification	DGX-1 (Tesla P100)	DGX-1 (Tesla V100)
GPU	8x Tesla P100 GPUs	8x Tesla V100 GPUs
TFLOPS	170 (GPU FP16) + 3 (CPU FP32)	1 (GPU Tensor PFLOP)
GPU Memory	16 GB per GPU / 128 GB per DGX-1 Node	16 GB or 32 GB per GPU / 128-256 GB per DGX-1 Node
CPU	Dual 20-core Intel® Xeon® E5-2698 v4	Dual 20-core Intel® Xeon® E5-2698 v4
FP32 CUDA Cores	28,672 Cores	40,960 Cores
System Memory	Up to 512 GB 2133 MHz DDR4 LRDIMM	Up to 512 GB 2133 MHz DDR4 LRDIMM
Storage	4x 1.92 TB SSD RAID 0	4x 1.92 TB SSD RAID 0
Network Interconnect	Dual 10 GbE, 4 IB EDR	Dual 10 GbE, 4 IB EDR
System Dimensions	866 D x 444 W x 131 H (mm)	866 D x 444 W x 131 H (mm)
System Weight	80 lbs	80 lbs
Max Power TDP	3200 W	3200 W
Operating Temp	10 - 35°C	10 - 35°C

A100

A100 关键改进

Ampere 架构：使用 MIG 来将 A100 切分为更小的实例或者链接更多 GPU
Tensor Cores: 312 TFLOPs/s
NVLink: 更高的 throughput
MIG (multi-instance GPU): 一个 A100 可以切分为至多 7 个硬件层面隔离的实例
HBM2e: 更大的 HBM, 更快的 bandwidth, 更高的 DRAM 使用效率
structure sparsity: 稀疏运算可以带来 2 倍的算力提升

A100 技术规格

	A100 80GB PCIe	A100 80GB SXM
FP64	9.7 TFLOPS	9.7 TFLOPS
FP64 Tensor Core	19.5 TFLOPS	19.5 TFLOPS
FP32	19.5 TFLOPS	19.5 TFLOPS
Tensor Float 32 (TF32)	156 TFLOPS \| 312 TFLOPS	156 TFLOPS \| 312 TFLOPS*
BFLOAT16 Tensor Core	312 TFLOPS \| 624 TFLOPS*	312 TFLOPS \| 624 TFLOPS*
FP16 Tensor Core	312 TFLOPS \| 624 TFLOPS*	312 TFLOPS \| 624 TFLOPS*
INT8 Tensor Core	624 TOPS \| 1248 TOPS*	624 TOPS \| 1248 TOPS*
GPU Memory	80GB HBM2e	80GB HBM2e
GPU Memory Bandwidth	1,935 GB/s	2,039 GB/s
Max Thermal Design Power (TDP)	300W	400W ***
Multi-Instance GPU	Up to 7 MIGs @ 10GB	Up to 7 MIGs @ 10GB
Form Factor	PCIe Dual-slot air-cooled or single-slot liquid-cooled	SXM
Interconnect	NVIDIA® NVLink® Bridge for 2 GPUs: 600 GB/s ** PCIe Gen4: 64 GB/s	NVLink: 600 GB/s PCIe Gen4: 64 GB/s
Server Options	Partner and NVIDIA-Certified Systems™ with 1-8 GPUs	NVIDIA HGX™ A100-Partner and NVIDIA-Certified Systems with 4,8, or 16 GPUs NVIDIA DGX™ A100 with 8 GPUs

H100

H100 关键改进

Hopper 架构
Tensor Core: 更强的 tensor core
transformer engine: 加速基于 transformer 架构模型的训练
NVLink: 900GB/s 的 bandwidth
2nd MIG: 支持 multi-tenant, multi-user 使用
DPX: 基于 DPX 指令集加速动态规划算法

H100 技术规格

	H100 SXM	H100 NVL
FP64	34 teraFLOPS	30 teraFLOPs
FP64 Tensor Core	67 teraFLOPS	60 teraFLOPs
FP32	67 teraFLOPS	60 teraFLOPs
TF32 Tensor Core*	989 teraFLOPS	835 teraFLOPs
BFLOAT16 Tensor Core*	1,979 teraFLOPS	1,671 teraFLOPS
FP16 Tensor Core*	1,979 teraFLOPS	1,671 teraFLOPS
FP8 Tensor Core*	3,958 teraFLOPS	3,341 teraFLOPS
INT8 Tensor Core*	3,958 teraFLOPS	3,341 teraFLOPS
GPU Memory	80GB	94GB
GPU Memory Bandwidth	3.35TB/s	3.9TB/s
Decoders	7 NVDEC 7 JPEG	7 NVDEC 7 JPEG
Max Thermal Design Power (TDP)	Up to 700W (configurable)	350-400W (configurable)
Multi-Instance GPUs	Up to 7 MIGS @ 10GB each	Up to 7 MIGS @ 12GB each
Form Factor	SXM	PCIe dual-slot air-cooled
Interconnect	NVIDIA NVLink™: 900GB/s PCIe Gen5: 128GB/s	NVIDIA NVLink: 600GB/s PCIe Gen5: 128GB/s
Server Options	NVIDIA HGX H100 Partner and NVIDIA- Certified Systems™ with 4 or 8 GPUs NVIDIA DGX H100 with 8 GPUs	Partner and NVIDIA-Certified Systems with 1–8 GPUs
NVIDIA AI Enterprise	Add-on	Included

H200

H200 关键改进

更高的 HBM 内存和带宽
更高的 LLM inference 速度

H200 技术规格

	H200 SXM	H200 NVL
FP64	34 teraFLOPS	30 teraFLOPs
FP64 Tensor Core	67 teraFLOPS	60 teraFLOPs
FP32	67 teraFLOPS	60 teraFLOPs
TF32 Tensor Core*	989 teraFLOPS	835 teraFLOPs
BFLOAT16 Tensor Core*	1,979 teraFLOPS	1,671 teraFLOPS
FP16 Tensor Core*	1,979 teraFLOPS	1,671 teraFLOPS
FP8 Tensor Core*	3,958 teraFLOPS	3,341 teraFLOPS
INT8 Tensor Core*	3,958 teraFLOPS	3,341 teraFLOPS
GPU Memory	141GB	141GB
GPU Memory Bandwidth	4.8TB/s	4.8TB/s
Decoders	7 NVDEC 7 JPEG	7 NVDEC 7 JPEG
Confidential Computing	Supported	Supported
Max Thermal Design Power (TDP)	Up to 700W (configurable)	Up to 600W (configurable)
Multi-Instance GPUs	Up to 7 MIGS @ 18GB each	Up to 7 MIGS @ 18GB each
Form Factor	SXM	PCIe dual-slot air-cooled
Interconnect	NVIDIA NVLink™: 900GB/s PCIe Gen5: 128GB/s	2- or 4-way NVIDIA NVLink bridge: 900GB/s per GPU PCIe Gen5: 128GB/s
Server Options	NVIDIA HGX H200 Partner and NVIDIA- Certified Systems™ with 4 or 8 GPUs	NVIDIA MGX™ H200 NVL partner and NVIDIA-Certified Systems with up to 8 GPUs
NVIDIA AI Enterprise	Add-on	Included

相比于 H100, H200 升级了 HBM 和 bandwidth

B200

B200 关键改进

blackwell 架构： GPU 之间的通信效率大幅度提升
Grace CPU: GPU 可以与 Grace CPu 之间达到 900GB/s 的 bidirectional bandwidth
5th NVIDIA NVLink: 可以链接 576 块 GPU 来支持计算，NVlink 的带宽可以达到 130TB/s
RAS engine: 自动识别故障来提高效率
NVIDIA networking

B2100 技术规格

system specification 如下

Specification	GB200 NVL72	GB200 NVL4	HGX B200
NVIDIA Blackwell GPUs \| Grace CPUs	72 \| 36	4 \| 2	8 \| 0
CPU Cores	2,592 Arm® Neoverse V2 Cores	144 Arm Neoverse V2 Cores	-
Total NVFP4 Tensor Core²	1,440 \| 720 PFLOPS	80 \| 40 PFLOPS	144 \| 72 PFLOPS
Total FP8/FP6 Tensor Core²	720 PFLOPS	40 PFLOPS	72 PFLOPS
Total Fast Memory	31 TB	1.8 TB	1.4 TB
Total Memory Bandwidth	576 TB/s	32 TB/s	62 TB/s
Total NVLink Bandwidth	130 TB/s	7.2 TB/s	14.4 TB/s

individual specification 如下

Specification	GB200 NVL72	GB200 NVL4	HGX B200
FP4 Tensor Core	20 PFLOPS	20 PFLOPS	18 PFLOPS
FP8/FP6 Tensor Core²	10 PFLOPS	10 PFLOPS	9 PFLOPS
INT8 Tensor Core²	10 POPS	10 POPS	9 POPS
FP16/BF16 Tensor Core²	5 PFLOPS	5 PFLOPS	4.5 PFLOPS
TF32 Tensor Core²	2.5 PFLOPS	2.5 PFLOPS	2.2 PFLOPS
FP32	80 TFLOPS	80 TFLOPS	75 TFLOPS
FP64 / FP64 Tensor Core	40 TFLOPS	40 TFLOPS	37 TFLOPS
GPU Memory Bandwidth	186 GB HBM3E 8 TB/s	186 GB HBM3E 8 TB/s	180 GB HBM3E 7.7 TB/s
Multi-Instance GPU (MIG)	-	7	-
Decompression Engine	-	Yes	-
Decoders	-	7 NVDEC³ 7 nvJPEG	-
Max Thermal Design Power (TDP)	Configurable up to 1,200 W	Configurable up to 1,200 W	Configurable up to 1,000 W
Interconnect	-	Fifth-generation NVLink: 1.8 TB/s PCIe Gen5: 128 GB/s	-
Server Options	NVIDIA GB200 NVL72 partner and NVIDIA-Certified Systems™ with 72 GPUs	NVIDIA MGX partner and NVIDIA-Certified Systems	NVIDIA HGX B200 partner and NVIDIA-Certified Systems with 8 GPUs

B300

B300 关键改进

Blackwell 架构
AI reasoning inference: 支持 test-time scaling, 对 attention layer 和 FLOPs 都有加速
HBM3e: 支持更大的 batch size 和 throughput
ConnectX-8 SuperNIC, 一个 host2 个 ConnectX-8 设备，支持 800Gb/s 的 GPU 之间通信
Grace-CPU: 更强的表现和带宽
5th NVIDIA NVLink: 更高的通信效率

B3100 技术规格

system specification 如下

	GB300 NVL72	HGX B300
Blackwell Ultra GPUs\| Grace CPUs	72 \| 36	8 \| 0
CPU Cores	2,592 Arm Neoverse V2 Cores	-
Total FP4 Tensor Core	1 1,440 PFLOPS \| 1,080 PFLOPS	144 PFLOPS \| 108 PFLOPS
Total FP8/FP6 Tensor Core	2 720 PFLOPS	72 PFLOPS
Total Fast Memory	37 TB	2.1 TB
Total Memory Bandwidth	576 TB/s	62 TB/s
Total NVLink Switch Bandwidth	130 TB/s	14.4 TB/s

individual specification 如下

	GB300 NVL72	HGX B300
FP4 Tensor Core	20 PFLOPS \| 15 PFLOPS	18 PFLOPS \| 14 PFLOPS
FP8/FP6 Tensor Core2	10 PFLOPS	9 PFLOPS
INT8 Tensor Core2	330 TOPS	307 TOPS
FP16/BF16 Tensor Core	5 PFLOPS	4.5 PLFOPS
TF32 Tensor Core2	2.5 PFLOPS	2.2 PFLOPS
FP32	80 TFLOPS	75 TFLOPS
FP64/FP64 Tensor Core	1.3 TFLOPS	1.2 TFLOPS
GPU Memory \| Bandwidth	279 GB HBM3E \| 8 TB/s	270 GB HBM3E \| 7.7 TB/s
Multi-Instance GPU (MIG)	7	7
Decompression Engine	Yes	Yes
Decoders	7 NVDEC3 7 nvJPEG	7 NVDEC3 7 nvJPEG
Max Thermal Design Power (TDP)	Configurable up to 1,400 W	Configurable up to 1,100 W
Interconnect	Fifth-Generation NVLink: 1.8 TB/s PCIe Gen6: 256 GB/s	Fifth-Generation NVLink: 1.8 TB/s PCIe Gen6: 256 GB/s
Server Options	NVIDIA GB300 NVL72 partner and NVIDIA-Certified Systems™	NVIDIA HGX B300 partner and NVIDIA-Certified Systems

References

GPipe

Tue, 23 Dec 2025 16:49:25 +0800

google 在 2018 年提出了 GPipe, 一个使用 pipeline parallelism 来训练大规模神经网络的并行策略

Introduction

大规模神经网络已经在计算机视觉和自然语言处理等任务上取得了突破性进展。但是目前训练大规模神经网络存在的问题时，我们无法在单一 GPU 上训练我们的模型。基于多 GPU 训练模型需要考虑模型的切分以及通信优化。

为了解决这个问题，作者提出了 GPipe, 一个用于将大规模性模型分割部署到不同设备上的并行计算策略。

Method

Notation

作者首先定义 notation 如下表所示

notation	description
$L$	number of layers
$w_i$	weights of a layer
$f_i$	forward function of a layer
$F_k=f_j\circ\cdots\circ f_i$	forward of a partition
$B_k$	backward of a partition
$K$	number of partitions
$N$	batch size
$M$	micro batch size

GPipe

首先是 naive pipeline parallelism (naive PP), 我们的输入为一个 batch, 然后我们依次计算 $F_1$, 通信传输，计算 $F_2$, 计算完成之后，我们再进行反向传播，更新参数。最后继续下一个 batch 的计算。

总体的过程如下图所示

下面是一个按照时间轴给出的例子

naive PP 的问题在于，每个时刻只有一个 GPU 在工作，GPU 的利用效率很低。因此，GPipe 的做法在于将一个 batch 切分为 $M$ 个更小的 micro-batch, 下面是一个 $M=4$ 的例子

通过切分更小的 batch，我们可以提高 GPU 的利用率

Analysis

Bubble

接下来作者分析了 GPipe 的 bubble 情况，bubble 指的是 PP 过程中的 GPU idle time.

对于 naive PP 来说，一个 GPU 工作时，其余 GPU 都处于空闲状态，因此其 bubble 为

$$ T_{bubble} = (K-1)(F+B) $$

总的计算时间为

$$ T_{total} = K(F+B) $$

从而 bubble rate 为

$$ Bubble_{naive} = \frac{T_{bubble}}{T_{total}} = \frac{K-1}{K} $$

当 $K=8$ 时，我们有 $Bubble_{naive}=87.5\%$, 也就是说，当前训练的 GPU 空闲率为 $87.5\%$.

对于 GPipe 来说，由于我们将一个 batch 拆分为了更小的 batch, 我们可以提高 GPU 的利用率。

此时，我们的 bubble time 仍然是 $T_{bubble} = (K-1)(F+B)$. 但是，现在同一时刻工作的 GPU 变多了，从上面的示意图可以看到，前向过程所需要的时间为第一个 micro batch 运行的时间加上 $M-1$ 个 batch 运行所需要的时间，反向同理，因此，GPipe 的总计算时间为

$$ T_{total} = (M+K-1)(F+B) $$

从而 GPipe 的 bubble rate 为

$$ Bubble_{naive} = \frac{T_{bubble}}{T_{total}} = \frac{K-1}{M+K-1} $$

当我们令 $M=8, K=8$ 时，我们有 $Bubble_{naive}=46.7\%$, 可以看到，通过提高 micro batch 数量，我们可以显著降低 bubble rate.

Activation Memory

对于 naive PP 来说，我们需要缓存每一层的输入，因此 activation memory 为 $\mathcal{O}(N\times L/K)$, 而使用 activation checkpointing 之后，我们现在的 activation memory 为

$$ \mathcal{O}(N + \frac{L}{K}\times\frac{N}{M}) $$

其中第一项代表了 boundary activation, 第二项代表了 Internal activation.

Experiments

作者在 image classification, machine translation 任务上进行了实验。

作者还进一步分析了影响 GPipe 性能的因素，结果如下图所示

可以看到，activation checkpointing 是 GPipe 的主要开销来源。

Conclusion

作者在本文中提出了 GPipe, 一个针对大规模神经网络训练的并行策略。通过将模型切分部署在不同的设备上以及使用 micro batch, 我们可以显著提高硬件的利用效率以及训练稳定性。

References

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

Notes on flashattention

Thu, 21 Aug 2025 11:32:53 +0800

作者提出了 flashattention, 一个通过降低 multi head attention 内存访问开销来提高 attention 计算效率的方法

Introduction

Transformer 的 attention 是一个平方度复杂度的算法，这个平方复杂度既体现在时间复杂度上（矩阵乘法），也体现在空间复杂度上（需要存储中间结果）。因此，要降低 attention 的复杂度，我们有两种思路：

从时间复杂度上入手，比如使用稀疏 attention 机制或者线性注意力机制
从空间复杂度上入手，比如使用 GQA, MQA 等减少内存的访问开销

本文提出的 flashattention 就属于降低空间复杂度的一种做法。作者认为，我们应该设计一种 IO-aware 的 attention 算法，来减少 attention 计算式的内存访问开销，进而提高 attention 的计算效率。

作者首先提到，一个未解决的问题就是：

降低 attention 的内存访问开销是否可以提高 attention 的计算效率？

作者发现，已有的一些工作虽然在理论上降低了 attention 的计算效率，但是在实际中，他们的效果并没有提升太多。作者分析原因认为，已有工作主要关注于降低 FLOPs, 但是忽略了内存访问开销。

因此，作者在本文中就提出了 flashattention, 一个 IO-aware 的 attention 算法，作者通过尽可能降低内存访问开销来提高模型的计算效率。具体做法就是，避免从内存中读写 attention matrix, 作者认为这个目标有两个挑战：

计算 softmax 的时候不访问所有的输入
在反向传播时不存储中间的 attention matrix

作者提出了两个方法来分别解决这两个问题：

作者使用了 tiling 技巧，将 input 分成多个 block, 然后分别进行处理，进而降低 softmax 的内存访问开销
作者使用了 recompute 技巧，在反向传播时，重新计算 softmax normalization factor

通过这些改进，我们可以让 attention 运行更快，并且降低内存访问开销。

作者还从理论上分析了 flashattention 的复杂度，提供了理论基础。

作者通过实验验证了 flashattention 的有效性，主要是三点：

训练效率更高：相比于 Huggingface 和 Megatron, flashattention 的训练效率提升了 2-3 倍
模型的表现更好：相比于 GPT-2, 模型的 perplexity 提升了 0.7 个点左右
速度更快：flashattention 比标准的 attention 实现快 3 倍以上

Background

Hardware Performance

作者首先介绍了以下 GPU 的内存架构，如下图所示

可以看到，GPU 内存可以分为三个层级：

SRAM: GPU 的寄存器，容量小，但是访问速度极快
High bandwith memory (HBM): GPU 的高速内存，访问速度较快，容量中等
DRAM: CPU 内存，容量最大，但是访问速度较慢

接下来作者介绍了 Execution model 的概念，GPU 有多个线程来执行同一个操作（SPMD），这个操作也被称为 kernel, kernel 会从 HBM 中加载输入到 SRAM 中进行计算，然后写回 HBM.

对一个算法，我们可以将其归类为 compute-bound 和 memory-bound 两类，我们可以用 arithmetic intensity 来进行区分，arithmetic intensity 定义为 arithmetic operations 与 memory access 的比率。

compute bound: 算法的瓶颈在于算力，由于算力不足导致运行时间慢，比如矩阵乘法
memory-bound: 算法的瓶颈在于内存访问效率，比如 element-wise 操作或者是 reduction

为了提高 memory-bound 类型算法的效率，我们进行 kernel fusion, 即把多个访问同一片内存的操作放在一起处理，避免多次读写内存

Standard Attention Implementation

作者还回顾了一下标准化的 attention 实现。

Forward Pass

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

这里 $\mathrm{softmax}$ 是逐行计算的。

算法的执行过程如下

我们有第一个结论

Proposition 1 标准化 attention 前向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于 attention, 我们需要从 HBM 中加载 $Q,K,V\in\mathbb{R}^{N\times d}$, 然后输出 $O\in\mathbb{R}^{N\times d}$ 并保存到内存中。

首先我们需要计算 $S = QK^T$, 这一步需要加载 $Q,K$ 并将 $S$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

接下来，我们需要计算 $P = \mathrm{softmax}(S)$, 这一步需要加载 $S$ 然后将 $P$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(N^2)$.

最后，我们需要计算 $O = PV$, 这一步需要加载 $P$ 和 $V$ 然后将 $O$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

总的来说，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Backward Pass

标准 attention 反向传播过程如下图所示

Proposition 2 标准化 attention 反向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于标准化 attention 的反向传播，我们需要从 HBM 中加载 $Q,K,V,dO\in\mathbb{R}^{N\times d}$ , 然后输出 $dQ,dK,dV$ 并保存到 HBM 中。

首先我们计算 $dV=P^TdO$, 这一步需要加载 $P,dO$ 并将 $dV$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd+N^2)$.

接下来我们计算 $dP=dOV^T$, 这一步需要加载 $dO, V$ 并将 $dP$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd)$.

然后我们计算 $dS$, 这一步需要加载 $P$ 并将 $dS$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(N^2)$.

对于 $dQ$ 和 $dK$ 的计算，内存访问开销都是 $\mathcal{O}(Nd+N^2)$.

因此，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Method

作者在本节首先介绍了 flashattention 算法，然后作者证明了 flashattention 的正确性以及分析了复杂度。最后作者对 flashattention 进行扩展得到了 Block-sparse Flashattention.

Flashattention

attention 模块的输入是 $Q,K,V\in\mathbb{R}^{N\times d}$, 输出是 $O\in\mathbb{R}^{N\times d}$, 作者的目标是减少计算过程中的 HBM 访问次数

作者分别使用了 tiling 和 recomputation 来解决 attention 前向传播和反向传播中的内存访问开销。flashattention 的核心思想是，我们将 $Q,K,V$ 分割成 block, 然后在 block 层面进行加载和计算。

Tiling

首先作者介绍了一下如何使用 tiling 来计算 softmax.

给定一个向量 $x\in\mathbb{R}^{B}$, 其 softmax 计算方式如下

$$ m(x) = \max_i x_i,\ f(x) = [e^{x_1-m(x)},\dots,e^{x_B-m(x)}], \ \ell(x)=\sum_if(x)_i, \ \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} $$

如果我们现在有两个向量 $x^{(1)}, x^{(2)}\in\mathbb{R}^{B}$, 记 $x=[x^{(1)}, x^{(2)}]^T\in\mathbb{R}^{2B}$, 我们可以将 $\mathrm{softmax}(x)$ 的计算分解为

$$ \begin{aligned} m(x) &= \max(m(x^{(1)}), m(x^{(2)}))， f(x) = [e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})]\\ \ell(x) &= e^{m(x^{(1)})-m(x)}\ell(x^{(1)}) + e^{m(x^{(2)})-m(x)}\ell(x^{(2)}), \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} \end{aligned} $$

因此，如果我们额外记录 $m(x)$ 以及 $\ell(x)$ 这两个量，那么我们可以每次仅计算 softmax 的一个 block. 具体细节见softmax.

Recomputation

在反向传播过程中，一般我们需要存储 $S,P\in\mathbb{R}^{N\times N}$, 需要的空间复杂度为 $\mathcal{O}(N^2)$. 但是，通过存储 $O\in\mathbb{R}^{N\times d}$ 以及 $(m,\ell)$, 我们可以避免重新计算 $S,P$,这可以看做是 gradient checkpointing. 但是与 checkpointing 相比，因为 flashattention 减少了内存访问开销，因此其反向过程并没有变得更慢。

Algorithm

最终，flashattention 的算法如下图所示

Analysis

Correctness

算法的正确性由定理 1 给出

Theorem 1 flashattention (即算法 1) 输出 $O=\mathrm{softmax}(QK^T)V$, 其时间复杂度为 $\mathcal{O}(N^2d)$, 空间复杂度为 $\mathcal{O}(N)$.

证明：时间复杂度主要由矩阵乘法决定。在计算 $S_{ij}=Q_iK_j^T$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 在计算 $\tilde{P}_{ij}V_j$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 循环一共执行了

$$ T_cT_r = \left\lceil\frac{N}{B_c}\right\rceil\left\lceil\frac{N}{B_r}\right\rceil $$

从而总的 FLOPS 为

$$ \mathcal{O}\left(\frac{N^2}{B_rB_c}B_rB_cd\right) = \mathcal{O}(N^2d) $$

在 flashattention 的计算过程中，我们只需要保存 $(\ell, m)$ 即可，因此需要的额外内存空间为 $\mathcal{O}(N)$.

接下来，我们可以证明 flashattention 的正确性，我们使用归纳法来证明。令 $j$ 满足 $0\leq j\leq T_c$, $K_{:j}\in\mathbb{R}^{jB_c\times d}$, $V_{:j}\in\mathbb{R}^{jB_c\times d}$ 分别为 $K$ 和 $V$ 的前 $jB_c$ 行。 $S_{:, :j}=QK_{:j}^T\in\mathbb{R}^{N\times jB_c}$, $P_{:,:j}=\mathrm{softmax}(S_{:,:j})\in\mathbb{R}^{N\times jB_c}$, $m^{(j)}, \ell^{(j)}, O^{(j)}$ 分别为 $m,\ell, O$ 的第 $j$ 个元素。我们证明经过第 $j$ 次迭代后，HBM 中保存的是

$$ m^{(j)}=\mathrm{rowmax}(S_{:,:j})\in\mathbb{R}^N, \ell^{(j)}=\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)}))\in\mathbb{R}^N, O^{(j)} = P_{:,:j}V_{:j}\in\mathbb{R}^{N\times d} $$

当 $j=0$ 时，上面的结果显然成立。现在我们假设对某个 $j=0,\dots, T_c-1$ 上面的结果成立，我们需要证明对 $j+1$ 也成立。

首先

$$ m^{(j+1)}=\max(m^{(j)}， \tilde{m}) = \max(\mathrm{rowmax}(S_{:,:j}), \mathrm{rowmax}(S_{:,j:j+1}))=\mathrm{rowmax}(S_{:,:j+1}) $$

接下来

$$ \begin{aligned} \ell^{(j+1)} &= \exp(m^{(j)}-m^{(j+1)})\ell^{(j)} + \exp(\tilde{m}-m^{(j+1)})\tilde{\ell}\\ &=\exp(m^{(j)}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)})) + \exp(\tilde{m}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,j:j+1}-\tilde{m}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j}-m^{(j+1)})) + \mathrm{rowsum}(\exp(S_{:,j:j+1}-m^{(j+1)}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j+1}-m^{(j+1)})) \end{aligned} $$

最后，我们计算 $O^{(j+1)}$ 得到：

$$ \begin{aligned} O^{(j+1)} &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})O^{(j)}+\exp(\tilde{m}-m^{(j+1)})\exp(S_{:,j:j+1}-\tilde{m})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})P_{:,:j}V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})\mathrm{diag}(\ell^{(j)})^{-1}\exp(S_{:,:j}-m^{(j)})V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\exp(-m^{(j+1)})\exp(S_{:,:j}))V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}( \begin{bmatrix} \exp(S_{:,:j}-m^{(j+1)}) & \exp(S_{:,:j}-m^{(j+1)}) \end{bmatrix}\begin{bmatrix} V_{:,:j} \\ V_{:,j:j+1} \end{bmatrix}\\ &= \mathrm{softmax}(S_{:,:j+1})V_{:,:j+1} \end{aligned} $$

因此上面的结果对 $j+1$ 也成立，从而 flashattention 的结果对 $j=0,\dots,T_c$ 都成立。

Forward Pass of flashattention

第一个问题是如何提高 softmax 计算的效率，作者的做法先先计算 normalization constant 然后再分别计算不同的 column.

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

我们有 $S_{ij}=q_i^Tk_j$, 这里 $q_i$ 和 $k_j$ 分别是 $Q$ 和 $K$ 的第 $i$ 列以及第 $j$ 列， normalization constant 定义为：

$$ L_i = \sum_{j=1}^N \exp\left(q_i^Tk_j\right) $$

对任意 $i$, 计算 $L_i$ 只需要 $\mathcal{O}(N)$ 的空间复杂度。

令 $v_j$ 是 $V$ 的第 $i$ 列，则输出 $O$ 的第 $i$ 列 $o_i$ 为

$$ o_i = P_{i:}V = \sum_{j=1}^N P_{ij}v_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j $$

这个过程中，对任意 $i$, 计算 $o_i$ 也只需要 $\mathcal{O}(N)$ 的空间复杂度。

因此，在 $L_i$ 已经计算好的情况下，我们可以在 $\mathcal{O}(N)$ 的空间复杂度下计算 $o_i$.

最终，flashattention 的 forward pass 过程如下图所示

接下来，作者分析了 flashattention 的内存访问开销。结论如下

Theorem 2 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 前向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

证明：由 Algorithm 1（或者 Algorithm 2）可以知道，$K$ 和 $V$ 的每一个元素都只需要从 HBM 中加载一次，而每一次外层循环都会从 HBM 中加载一次 $O$ 和 $Q$, 因此总的 HBM 访问次数为 $\mathcal{O}(Nd+NdT_c)=\mathcal{O}(NdT_c)$.

接下来，我们给出每一次内层循环的内存访问开销，这是由 SRAM 的大小决定的。由于我们需要 SRAM 可以存储 $K_j\in\mathbb{R}^{B_c\times d}$ 以及 $V_j\in\mathbb{R}^{B_c\times d}$ ，我们的 block size 需要满足

$$ B_cd = \mathcal{O}(M) \Rightarrow B_c = \mathcal{O}\left(\frac{M}{d}\right) $$

同理，对于 $O$ 和 $Q$, 我们有

$$ B_rd = \mathcal{O}(M) \Rightarrow B_r = \mathcal{O}\left(\frac{M}{d}\right) $$

最后，我们还需要 SRAM 可以存储 $S_{ij}\in\mathbb{R}^{B_r\times B_c}$, 因此

$$ B_rB_c=\mathcal{O}(M) $$

这样，

$$ B_c = \mathcal{O}\left(\frac{M}{d}\right), B_r=\mathcal{O}\left(\min\left(\frac{M}{d},\frac{M}{B_c}\right)\right)=\mathcal{O}\left(\min\left(\frac{M}{d},d\right)\right) $$

从而

$$ T_c = \frac{N}{B_c} = \mathcal{O}\left(\frac{Nd}{M}\right) $$

最终，总的内存访问开销为

$$ \mathcal{O}(NdT_c) = \mathcal{O}\left(\frac{N^2d^2}{M}\right) $$

一般来说, $d$ 的大小为 $64-128$, $M$ 的大小为 $100 KB$ 左右, $d^2« M, 因此 flashattention 的内存访问开销远小于标准化 attention 的内存访问开销。

作者还证明 flashattention 的内存访问开销是一个下界，即

Proposition 3 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则不存在一个对任意 $M\in[d,Nd]$ 都可以在内存访问开销为 $\Theta(N^2d^2M^{-1})$ 的条件下完成 attention 计算的算法。

证明可以用反证法，基本思想是加载 $Q,K,V$ 的 HBM 访问次数至少为 $\mathcal{O}(Nd)$.

Backward Pass of flashattention

第二个问题是能否在线性空间复杂度下计算 attention 的反向传播过程。

首先我们记损失函数为 $\phi$, 然后令 $\phi$ 对 $O,Q,K,V$ 的梯度分别为 $dO,dQ,dK, dV\in\mathbb{R}^{N\times d}$, 我们的目标是计算 $dQ, dK, dV$.

$dV$ 的计算是最容易的，我们有 $dV=P^TdO$, 因此

$$ dv_j = \sum_{i=1}^N P_{ij}do_i = \sum_{i=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i $$

由于我们已经计算了 $L_i$, 因此，$dv_j$ 只需要 $\mathcal{O}(d)$ 的空间复杂度。

接下来，注意到 $dP=dOV^T$, 因此我们有

$$ dP_{ij} = do_i^Tv_j $$

计算的空间复杂度也是要 $\mathcal{O}(N)$ 的

注意到 $P_{i:}=\mathrm{softmax}(s_{i:})$, 且 $y=\mathrm{softmax}(x)$ 的 Jacobian 是 $\mathrm{diag}(y)-yy^T$ (推导过程见 softmax), 我们有

$$ dS_{i:} = (\mathrm{diag}(P_{i:})-P_{i:}P_{i:}^T)dP_{i:} = P_{i:} \odot dP_{i:} - (P_{i:}^TdP_{i:})P_{i:} $$

我们定义

$$ D_i = P_{i:}^TdP_{i:}= \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i^Tv_j = do_i^T\sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j = do_i^To_i $$

$D_i$ 的空间复杂度也只需要 $\mathcal{O}(N)$.

则

$$ dS_{i:} =P_{i:} \odot dP_{i:} - D_iP_{i:} $$

我们有

$$ dS_{ij} = P_{ij}dP_{ij} - D_iP_{ij} = P_{ij}(dP_{ij}-D_i) $$

注意到 $S_{ij}=q_i^Tk_j$, 我们有

$$ dq_i = \sum_{j=1}^N dS_{ij}k_j = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)k_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)k_j $$

因此计算 $dq_i$ 的空间复杂度为 $\mathcal{O}(d)$.

同样的，

$$ dk_j = \sum_{j=1}^N dS_{ij}q_i = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)q_i = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)q_i $$

其空间复杂度为 $\mathcal{O}(N)$.

总之，attention 的反向传播过程所需要的空间复杂度为 $\mathcal{O}(N)$.

作者发现有两点可以改进：

attention mask 不需要存储，我们只需要保存 forward pass 时的输入，然后在 backward pass 时重新生成即可，这样只需要 $\mathcal{O}(N)$ 的空间复杂度。
计算 softmax 的梯度是，如果使用公式 $D_i=P_{i:}^TdP_{i:}$ 来计算的话，由于 $P_{i:}\in\mathbb{R}^N$, 可能会导致超过 SRAM 的内存使用限制，因此，我们可以使用 $D_i=do_i^To_i$ 来避免这个问题，其中 $o_i\in\mathbb{R}^d$.

最终，flashattention 的 backward pass 过程如下图所示

经过前面的分析，flashattention 的反向传播的时间复杂度为 $\mathcal{O}(N^2)$, 空间复杂度为 $\mathcal{O}(N)$.

Theorem 5 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 反向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

定理的证明与 Theorem 2 基本一致，我们此处不再赘述。

Block-sparse Flashattention

当 attention 具有 block sparsity 的性质时，作者提出了 blck-sparse flashattention 来进一步提高 attention 的计算效率。

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 以及一个 mask $M\in\{0,1\}^{N\times N}$, 我们需要计算

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S\odot \mathbb{1}_{M})\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

其中当 $M_{kl}=1$ 时， $(S\odot \mathbb{1}_ {M})_ {kl}=S_ {kl}$, 否则 $(S\odot \mathbb{1}_ {M})_{kl}=0$.

Block-sparse attention 的算法如下所示

Proposition 4 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 block-sparse attention 的内存访问开销为 $\Theta(Nd+N^2d^2M^{-1}s)$, 其中 $s$ 是 block-sparse mask 中的非零 block 的比例

证明与 Theorem 2 的证明是类似的，总的内存访问开销为 $\mathcal{O}(Nd+NdT_c)$, 但是在计算的过程中，由于 mask 矩阵的 block-sparsity, 我们实际上只需要计算一小部分 $M_{ij}\neq0$ 的情况，因此最终的内存访问开销为

$$ \mathcal{O}\left(Nd+\frac{N^2d^2}{M}s\right) $$

可以看到，attention mask 的 sparsity 越高，block-sparse flashattention 的效率也就越高。当 $N$ 非常大时，$s 通常为 $1/\sqrt{N}$ 或者 $N^{-1}\log N$, 从而最终的内存访问开销为 $\mathcal{O}(N\sqrt{N})$ 或者 $\mathcal{O}(N\log N)$.

作者对比了以下 block-sparse flashattention 和 flashattention 的效率对比，结果如下图所示

Experiment

作者通过实验验证了 flashattention 的有效性，如下表所示

Attention	Standard	FlashAttention
GFLOPs	66.6	75.2
HBM R/W (GB)	40.3	4.4
Runtime (ms)	41.7	7.3

可以看到，尽管 flashattention 相比于标准化 attention 需要更多的算力，但是由于其内存访问开销更少，所以最终的运行时间大有了大幅度降低

作者还探究了 block size 对 flashattention 性能对的影响，实验结果如下图所示

可以看到，随着 block size 增加，循环次数降低，内存访问开销也逐渐降低。但是当 block size 充分大 ( $> 256$) 之后，运行时间就会被别的因素所限制，并且过大的 block size 可能会导致 SRAM 的内存溢出

作者首先在 BERT 和 GPT-2 上验证了 flashattention 的表现，BERT 的实验结果如下表所示

BERT Implementation	Training time (minutes)
Nvidia MLPerf 1.1	$20.0\pm1.5$
FlashAttention (ours)	$17.4\pm1.4$

GPT-2 的实验结果如下表所示

Model implementations	OpenWebText (ppl)	Training time (speedup)
GPT-2 small - Huggingface	18.2	9.5 days (1.0 )
GPT-2 small - Megatron-LM	18.2	4.7 days (2.0 )
GPT-2 small - FlashAttention	18.2	2.7 days (3.5 )
GPT-2 medium - Huggingface	14.2	21.0 days (1.0 )
GPT-2 medium - Megatron-LM	14.3	11.5 days (1.8 )
GPT-2 medium - FlashAttention	14.3	6.9 days (3.0 )

实验结果显示，flashattention 比 Huggingface 快 3 倍左右，比 Megatron 快 1.7 倍左右

训练速度：实验显示，flashattention 在 BERT 上，比 MLPerf 1.1 快 $15\%$, 在 GPT-2 上比 HuggingFace 快 3 倍，比 Megatron 快 1.8 倍
准确率：flashattention 是第一个在 Path-X 上比随机表现更好的 transformer 模型；block-sparse flashattention 是第一个在 Path-256 上比随机表现更好的的 sequence model

Conclusion

作者提出了 flashattention, 一个通过优化标准 attention 内存访问效率来提高 attention 计算效率的方法，作者详细介绍了算法设计的原理与证明，并通过实验证明了结果的有效性。

References

arxiv

LLM Parameter Computation

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

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

Dense Model

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

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

1
2
3
4


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

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

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

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

1
2
3
4
5
6
7


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

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

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

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

其参数量为：$d$.

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

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

其参数量为： $2d$.

因此，Qwen3Model 的参数量为

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

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

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

1
2
3
4
5
6


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

因此，Qwen3DecoderLayer 的参数量为

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

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

对于 Qwen3MLP, 其定义如下：

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

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

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

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

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

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

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

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

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

Qwen3Attention

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Verification

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

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

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

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

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

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

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

MoE Model

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

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

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

1
2
3
4
5
6


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

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

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

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

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

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

我们带入到 Qwen3MoeForCausalLM 中就得到：

模型总参数量为：

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

模型激活参数量为：

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

这里的变量定义如下：

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

MoE Verification

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

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

计算得到总参数量为

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

计算得到激活参数量为

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

实际总参数量为

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

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

Extension

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

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

Visualization

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

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


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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

结果如下

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

Conclusion

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

References

Qwen3 Collection

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

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

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

大语言模型架构

大语言模型参数计算

计算量估计

checkpointing

KV cache

参考文献

回旋托马斯x 文章

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

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

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

训练的数学原理

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

$$ \min_x\ f(x) $$

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

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

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

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

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

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

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

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

自动微分

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

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

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

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

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

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

接下来，

总结

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

训练

优化器

SGD

Adam

AdamW

精度

Distributed training--Basic

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

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

什么是分布式系统

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

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

为什么需要分布式训练

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

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

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

基本概念

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

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

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

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

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

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

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

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

通信方式

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

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

One-to-one

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

测试代码如下：

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


# send_recv.py
import os

import torch
import torch.distributed as dist


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

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


init_process()
example_send()

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

结果输出

1
2


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

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

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

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


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

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

init_process()
example_isend()

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

结果输出

1
2
3
4


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

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

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

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

One-to-many

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

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

scatter 测试代码：

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


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

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

init_process()
example_broadcast()

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

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

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


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

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

broadcast 测试代码:

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


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


init_process()
example_broadcast()

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

结果输出：

1
2
3
4
5
6
7


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

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

Many-to-one

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

gather 测试代码：

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


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

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

init_process()
example_gather()

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

结果输出：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


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

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

reduce 测试代码:

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


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

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


init_process()
example_reduce()

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

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

1
2
3
4
5
6


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

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

Many-to-many

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

all-gather 测试代码：

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


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


init_process()
example_all_gather()

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

测试输出结果：

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


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

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

all-reduce 测试代码：

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


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

init_process()
example_all_reduce()

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

测试输出结果

1
2
3
4
5
6
7
8


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

Barrier

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

barrier 测试代码

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


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

init_process()
example_barrier()

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

结果输出

1
2
3
4
5
6
7
8
9


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

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

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

Advanced

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

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