CS336: Language Models From Scratch (Spring 2025)

1. Memory

float32，32位浮点数，包括1位符号位，8位指数位，23位尾数位。也称为fp32，单精度。是tensor的默认存储精度。

内存使用的估算：

x = torch.zeros(4, 8)  # 建立矩阵
assert x.dtype == torch.float32  # tensor默认精度为fp32
assert x.numel() == 4 * 8
assert x.element_size() == 4  # Float is 4 bytes
assert get_memory_usage(x) == 4 * 8 * 4  # 128 bytes

text("One matrix in the feedforward layer of GPT-3:")
assert get_memory_usage(torch.empty(12288 * 4, 12288)) == 2304 * 1024 * 1024  # 2.3 GB

float16，符号1位，指数5位，尾数10位，也叫半精度。相较于fp32内存可以减半。但是fp16的范围比较小，会出现上溢下溢的问题，影响模型。

bfloat16，bf16，符号1位，指数8位，尾数7位，在和fp16保持相同存储的同时和fp32有相同的动态范围，牺牲了部分精度但可以接受。

fp8，8位，有E4M3、E5M2两种形式。

训练时用fp32效果最好，内存开销也最高；fp16、bf16、fp8内存开销小，但不稳定；一种折中的办法是使用混合精度，只在关键的层使用高精度。


2. Compute

2.1 tensor

Pytorch中的tensor（张量）是一个多维数组，可以是1D的向量，2D的矩阵，3D的cube等。tensor是一个指向具体内存的指针+各种元数据。元数据包括shape和stride，shape告诉我们有几个维度，每个维度有多少个元素，stride告诉我们在内存中跳多少步才能访问下一个维度的元素。


tensor默认存储在cpu上，需要显式将其移动到gpu。

memory_allocated = torch.cuda.memory_allocated()

x = torch.zeros(32, 32)
assert x.device == torch.device("cpu")

text("为了利用GPU的并行计算能力，将tensor迁移到GPU")

text("Move the tensor to GPU memory (device 0).")
y = x.to("cuda:0")
assert y.device == torch.device("cuda", 0)

text("Or create a tensor directly on the GPU:")
z = torch.zeros(32, 32, device="cuda:0")

new_memory_allocated = torch.cuda.memory_allocated()

memory_used = new_memory_allocated - memory_allocated

assert memory_used == 2 * (32 * 32 * 4)  # 2 32x32 matrices of 4-byte floats

某些操作，如切片、转置、改变形状，并不会产生一个新的tensor，只是改变了tensor的元数据的值，例如：

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
y = x.view(3, 2) 
x[0][0] = 100
assert y[0][0] == 100

x y 实际上是指向同一片内存的指针。


对于改变tensor值的操作，会产生一个新的tensor。


有时候我们想对成batch的数据进行处理，Pytorch支持这种方法。

x = torch.ones(4, 8, 16, 32)  
w = torch.ones(32, 2)
y = x @ w
assert y.size() == torch.Size([4, 8, 16, 2])  # 末尾的两个维度进行相乘

2.2 contiguous

当x转置（transpose）得到y时，x和y实际上共享内存数据。但是此时x是一个contiguous（连续的）的矩阵，y是一个非contiguous的矩阵。

x=[[1,2,3],
	[4,5,6]]
	
tensor刚创立的时候是contiguous的。

pytorch中，默认按照行优先存储数据。实际上x在内存中是一维数组[1,2,3,4,5,6]，但是x的stride（步长）元数据为[3,1]，表示第 0 维（行）的步长为 3：从第 0 行到第 1 行（沿行维度移动 1 步），需要跳过 3 个元素（因为每行有 3 个元素）。
第 1 维（列）的步长为 1：从第 0 列到第 1 列（沿列维度移动 1 步），只需跳过 1 个元素（同一行内连续存储）。

然后将x transpose得到y，此时xy指针的值是一样的，仅仅stride的值进行了交换。y的stride为[1,3]，表示按第0维（行）每次跳过一个元素，按第一维（列）每次跳过3个元素。所以y输出以后看起来是3*2的矩阵：
y=[[1,4],
	[2,5],
	[3,6]]

对于不连续的矩阵，不能再接着使用view、flatten，因为它们要重新解释张量的布局，需要假设数据在内存中是连续存储的。不连续的矩阵需要使用.contiguous()先变为连续的矩阵。

y = x.transpose(1, 0).contiguous().view(2, 3)

pytorch中可使用.is_contiguous()来判断矩阵是否连续。判断逻辑为：

沿第 i 个维度的步长（stride[i]）等于后一个维度的步长（stride[i+1]）乘以该维度的大小（shape[i+1]），即：

stride[i] = stride[i+1] * shape[i+1]  （对所有 0 ≤ i < N-1 成立）

且最后一个维度的步长必须为 1（stride[-1] = 1）。

比如：

转置前：

x = torch.tensor([[1, 2, 3], [4, 5, 6]])  # shape=(2, 3)
print(x.stride())  # (3, 1)
print(x.is_contiguous())  # True
  • 检查条件：stride[0] = 3，stride[1] * shape[1] = 1 * 3 = 3，满足 stride[0] = stride[1] * shape[1]；且最后一维步长 stride[1] = 1，故连续。

转置后：

x_t = x.transpose(0, 1)  # shape=(3, 2)
print(x_t.stride())  # (1, 3)
print(x_t.is_contiguous())  # False
  • 检查条件：stride[0] = 1，stride[1] * shape[1] = 3 * 2 = 6，显然 1 ≠ 6，不满足；故非连续。

3. Einops

Einops是一个用于操作张量的库，其中的维度都有明确的名称。

Einops教程

3.1 Einops的设计动机

传统的PyTorch代码：

x = torch.ones(2, 2, 3)  # 分别表示batch, sequence, hidden维度
y = torch.ones(2, 2, 3)  # 分别表示batch, sequence, hidden维度
z = x @ y.transpose(-2, -1)  # 结果维度为batch, sequence, sequence

这种写法很容易搞混维度（比如-2和-1到底指的是什么维度？）


3.2 Jaxtyping写法

如何记录张量各维度的意义？

旧方法，通过自己加注释：

x = torch.ones(2, 2, 1, 3)  # 代表batch seq heads hidden

新方法（使用jaxtyping）：

x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2, 2, 1, 3)

# 如果你想尝试einops，别忘了导入以下的库：
import torch
from jaxtyping import Float
from einops import einsum, reduce, rearrange

3.3 Einops的einsum函数

einsum是一种广义的矩阵乘法，带有良好的维度记录功能。

定义两个张量：

x: Float[torch.Tensor, "batch seq1 hidden"] = torch.ones(2, 3, 4)
y: Float[torch.Tensor, "batch seq2 hidden"] = torch.ones(2, 3, 4)

旧方法：

z = x @ y.transpose(-2, -1)  # 结果维度为batch, sequence, sequence

新方法（使用einops）：

z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")

在输出中没有命名的维度会被求和。

也可以使用...来表示对任意数量的维度进行广播：

z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

3.4 Einops的reduce函数

你可以通过某些操作（如sum、mean、max、min）对单个张量进行降维。

x: Float[torch.Tensor, "batch seq hidden"] = torch.ones(2, 3, 4)

旧方法：

y = x.mean(dim=-1)  # 对最后一个维度求平均值

新方法（使用einops）：

y = reduce(x, "... hidden -> ...", "sum")  # 对hidden维度求和

3.5 Einops的rearrange函数

有时，一个维度实际上代表了两个维度，而你想要对其中一个进行操作。

x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)

这里的total_hidden是heads * hidden1的扁平化表示。

将total_hidden拆分为两个维度（heads和hidden1）：

x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)

通过w执行转换：

w: Float[torch.Tensor, "hidden1 hidden2"] = torch.ones(4, 4)

x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")

将heads和hidden2重新组合在一起：

x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")

4.FLOPs&FLOP/s

FLOPs是表示浮点运算次数的单位。一次基础操作（加、乘）记为一次浮点操作。

训练GPT-3 (2020) 需要 3.14e23 FLOPs。

FLOP/s或者FLOPS表示每秒浮点运算次数。

比如 A100 峰值运算 312 teraFLOP/s。

FLOP/s 取决于使用的硬件和数据类型。


Model FLOPs utilization (MFU)

mfu = 实际的 FLOP/s 除以理论上的 FLOP/s

通常当MFU ≥ 0.5时被认为较好的利用了硬件。


估算简单线性模型前向传播的FLOPS

假设有B个token，每个token为D维，要变换为K维，那么相当于[B, D] [D, K] 矩阵相乘。

B = 16384  # Number of points(batchsize)
D = 32768  # Dimension
K = 8192   # Number of outputs
x = torch.ones(B, D, device=device)
w = torch.randn(D, K, device=device)
y = x @ w

必要的操作有 (x[i][j] * w[j][k]) 的乘法运算以及每个 (i, j, k) 对进行一次加法运算，总共需要的flops为2 * B * D * K，即三个维度的乘积。把DK看作parameter，相当于2 * token * parameter。


估算简单线性模型总的的FLOPS

包括前向传播，反向传播。推导过程省略。

前向传播：2*数据点数*参数量

后向传播：4*数据点数*参数量

总计：6*数据点数*参数量


5.Model

5.1参数初始化

参数在pytorch中以nn.Parameter形式存储，是tensor。

假设 x 是一个形如 (input_dim,) 的向量，w 是一个形如 (input_dim, output_dim) 的矩阵，output = x @ w。如果随机初始化，会发现output大约随着sqrt(input_dim)变化。需要所以对参数乘以 1/sqrt(input_dim)。

如果 output 随 input_dim 膨胀，会导致两个严重问题：

  1. 梯度消失 / 爆炸

    深层网络中，若每一层的输出都随维度膨胀（或收缩），经过多层传递后，数值会变得极大（梯度爆炸）或极小（梯度消失），导致模型无法训练。

  2. 激活函数失效

    例如 ReLU 激活函数（f (x)=max (0,x)），若 input_dim 很大导致 output 数值过大，大部分神经元会处于 “激活状态”（x>0），丧失非线性表达能力；反之，数值过小则大部分神经元 “死亡”（x<0）。

可以对权重 w 乘以 1/√(input_dim) 进行缩放（Xavier initialization）


5.2 Randomness

保证实验的可复现是非常有用且必要的，因此要设置随机种子，一共有三个地方

# Torch
seed = 0
torch.manual_seed(seed)
# NumPy
import numpy as np
np.random.seed(seed)
# Python
import random
random.seed(seed)

5.3 pin_memory 和 non_blocking

默认的 CPU tensor 是 paged memory（分页内存）。分页内存是操作系统常规分配给进程的内存，用于正常运行没问题，但是 GPU 无法直接访问这类内存，从 CPU 拷贝到 GPU 较慢。

if torch.cuda.is_available():
x = x.pin_memory()

pin_memory() 会将张量固定到物理内存中，禁止其换页。这使得 GPU 可以使用 DMA（直接内存访问）异步拷贝数据，速度更快。

x = x.to(device, non_blocking=True)

加了 non_blocking=True，就表示这个 .to() 操作是 异步的（non-blocking），只要 x 是 pinned memory，就能真正做到异步。

你可以在同时干两件事：

  1. CPU 异步准备 + 传输下一批数据
  2. GPU 同时处理上一批数据

最终的效果就是提高 GPU 利用率，减少等待数据加载的空转时间

参考：详解Pytorch里的pin_memory 和 non_blocking - 知乎


5.4 Optimizer

  • SGD（随机梯度下降）：

    最基础的优化器。特点为每一步都朝着当前梯度方向走；对所有参数使用同一个固定学习率；收敛慢、容易卡在鞍点或局部最小值。$\theta_{t + 1} = \theta_t - \eta \cdot \nabla_{\theta} J(\theta)$ 其中，$\theta$为模型参数，$\eta$为学习率，$\nabla_{\theta} J(\theta)$为损失函数对参数的梯度。

  • Momentum(SGD + exponential averaging of grad)： 加了“惯性”，用指数加权移动平均平滑梯度，减少震荡、加快收敛。直觉上就好像在山谷中下坡，Momentum像是加了惯性的小球，可以跨过小波动，不轻易被困住。 $v_t = \beta v_{t - 1} + (1 - \beta)\nabla_{\theta} J(\theta)$ $\theta_{t + 1} = \theta_t - \eta \cdot v_t$ 其中，$v_t$是速度（梯度的指数平均），$\beta \in [0.9, 0.99]$控制历史记忆。

  • AdaGrad(SGD + averaging by grad^2)：

    对每个参数都使用不同的学习率，且梯度越大，学习率衰减得越快。优点是参数稀疏时非常有效；缺点是随着训练进行，$G_t$越来越大，学习率会衰减到趋近0，也就是学不动了。 $G_t = G_{t - 1} + \nabla_{\theta} J(\theta)^2$ $\theta_{t + 1} = \theta_t - \frac{\eta}{\sqrt{G_t + \epsilon}} \cdot \nabla_{\theta} J(\theta)$ 其中，$G_t$是累计的平方梯度（对每个参数单独计算），$\epsilon$防止除以0。

  • RMSProp(AdaGrad + exponentially averaging of grad^2)： 是对AdaGrad的修正，用指数移动平均替代了“累加所有历史梯度平方”。优点是避免了AdaGrad的“学习率消失”问题，在非凸优化问题上表现稳定. $s_t = \beta s_{t - 1} + (1 - \beta)\nabla_{\theta} J(\theta)^2$ $\theta_{t + 1} = \theta_t - \frac{\eta}{\sqrt{s_t + \epsilon}} \cdot \nabla_{\theta} J(\theta) $ 其中，$s_t$是对平方梯度的指数平均。

  • Adam(RMSProp + momentum)：

    • 综合了Momentum（动量）+ RMSProp（自适应学习率）的优点。优点是自动调整每个参数的学习率；在许多任务中无需调参效果也很好；收敛速度快，表现稳。
    • 一阶动量估计：$m_t = \beta_1 m_{t - 1} + (1 - \beta_1)\nabla_{\theta} J(\theta)$
    • 二阶动量估计：$v_t = \beta_2 v_{t - 1} + (1 - \beta_2)(\nabla_{\theta} J(\theta))^2$
    • 偏差修正：$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$，$\hat{v}_t = \frac{v_t}{1 - \beta_2^t}$
    • 最终更新规则：$\theta_{t + 1} = \theta_t - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$

Optimizer主要机制
SGD基础版本，直接减梯度
Momentum平滑梯度方向，加速下降
AdaGrad每个参数学习率按梯度平方缩放（历史平均）
RMSProp改进AdaGrad，用指数平均避免学习率变太小
AdamRMSProp + Momentum（双重指数平均）

  • SGD：走一步看一步
  • Momentum：考虑过去的速度，减少来回震荡
  • AdaGrad：走路时每次都踩相同的地方会越来越慢
  • RMSProp：记得最近的坑，别总踩同一个坑
  • Adam：不仅避坑，还顺着山坡惯性滑下去

5.5 Train

如果要使用混合精度进行训练，一个常规的做法是：

  • 在前向传播中使用bfloat16/bfloat8（保存/计算激活值）
  • 在其余部分使用float32（反向传播，保存参数、梯度等，保证精度）

5.6 Checkpoint

在保存的时候不光是要保存模型参数，如果需要继续训练，还需要保存优化器的参数

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "model_checkpoint.pt")

这一讲主要讲了内存和计算的估算，有助于对模型及其训练有一个宏观的把控。还讲到了tensor的一些操作。最后讲了模型训练的的流程，其中有很多有用的技巧，希望以后的训练中能用上，提高效率。