CS336 Lecture_02
本节课介绍了如何利用资源，包括内存（GB）和计算（FLOPs），以及模型中的张量，参数，优化器和混合精度训练等内容
Memory accounting（内存成本）
tensors_memory
张量基本都是以浮点数形式存储
float32（FP32，多用于存储优化器状态和参数）：32位（4字节），1位符号，8位指数（动态范围），23位尾数（数值精度）
memory_usage = 元素个数 * 元素大小
float16（FP16）：16位（2字节），1位符号，5位指数，10位尾数，内存减半，动态范围小
bfloat16（brain，BF16，多用于计算）：16位，1位符号，8位指数，7位字节，动态范围和fp32一样，但是精度稍低
float8（FP8）：用的比较少，H100显卡支持
根据流水线具体环节决定精度，称为混合精度训练，比如自注意力用FP32保证准确率，涉及矩阵乘法的简单前向传播可以用BF16
对于需要长时间积累的东西，希望的精度越高（FP32），BF16可以看作是更具有临时性的东西，基本上就是取出参数转为BF16然后推进模型运行
Compute accounting（计算成本）
tensors_on_gpus
如果要估计计算和数据移动开销：时刻注意tensor的位置（torch.cuda.is_available()）
# 将 CPU 数据移到 GPU
y = x.to("cuda:0")
assert y.device == torch.device("cuda", 0)
# 在 GPU 上生成数据
z = torch.zeros(32, 32, device="cuda:0")tensor_operations
tensor_storage
tensor在pytorch中其实就是一些指向已分配内存的指针，假如有一个4 x 4矩阵，在内存中看起来像是一个长数组，张量本身拥有的是元数据（stride[0]，stride[1]），元数据指定了如何在数组中找到特定地址
stride[0]：移动到下一行元素，需要跳过的元素，这里是4
stride[1]：移动到下一列元素，需要跳过的元素，这里是1
（行r，列c）计算index = r * stride[0] + c * stride[1]
tensor_slicing
对张量slice、view、transpose，其实只是创建了一个视图，而不是复制，改动视图同时会改变原张量的值
需要注意的一点是，有些视图（view）是连续的，就像只是平滑地遍历存储中的这个数组一样（连续要求各维度步幅从左到右严格递减，比如（3,1）连续，（1,3）非连续，访问同一行的步幅需要为1）），transpose会破坏张量连续性，而要创建视图则张量必须连续
y = x.transpose(1, 0).contiguous().view(2, 3)在这种情况下，x和y不具有相同的存储空间，因为使它连续的操作（contiguous）会创建副本（copy）（reshape同样是创建副本）
tensor_matmul
输入张量的前两位一般为 batch（批次）、sequence（序列长度，即token数量），进行乘法时，遍历 x 前 2 维与 y 相乘
x = torch.ones(4, 8, 16, 32)
w = torch.ones(32, 2)
y = x @ w
assert y.size() == torch.Size([4, 8, 16, 2])tensor_einops（张量操作库）
为什么引入einops（Einstein Operations）：传统 pytorch 的维度混乱
我们只需要命名所有的维度，而不是本质上只依赖于索引
x = torch.ones(2, 2, 3) # batch, sequence, hidden
y = torch.ones(2, 2, 3)
z = x @ y.transpose(-2, -1) # 容易搞混 -2, -1 是什么jaxtyping_basics
使用这个库，我们可以为各个维度添加语义标注
# 旧方式
x = torch.ones(2, 2, 1, 3) # 注释: batch seq heads hidden
# 新方式
x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2, 2, 1, 3)einops_einsum
利用这个方式我们只需要关注输入和输出，并给出标注，而不关注过程
输出中未命名的维度（如 hidden）会被求和，任何被命名的维度都会被遍历
"..." 表示广播任意数量的维度（在矩阵乘法中不会变化的维度）
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
# 新方式
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")
# 另一种写法
z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2") einops_reduce
聚合指定的维度（如 hidden），且支持 sum、mean、max、min 等操作
# 旧方式
y = x.mean(dim=-1)
# Einops 方式
y = reduce(x, "... hidden -> ...", "sum")einops_rearrange
拆分某个展平的维度为多个维度，应用于多头注意力等场景
# 这里的 8 维向量是一个 flattened representation, 表示为 total_hidden = heads * hidden1
x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)
# 拆分维度
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2) # (2, 3, 2, 4)
# 变换后合并
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")tensor_operations_flops
张量操作的计算成本
FLOPs：浮点操作次数（衡量完成的计算量），矩阵乘法在计算中占主导地位
FLOP/s：每秒浮点运算次数（衡量硬件速度），取决于硬件和数据类型
Linear model 前向传播的矩阵乘法中（（B，D）@（D，K））：FLOPs = 2 * B * D * K（包括乘法，加法），B 是数据点数量，D 是输入维度，K 是输出维度，D x K 是这个矩阵的参数数量
因此前向传播需要的 FLOPs = 2 * (data points) * (parameters)
MFU（Model FLOPs utilization）模型浮点数运算利用率：实际每秒浮点运算次数 / 承诺的每秒浮点运算次数，通常 MFU >= 0.5 是好的
gradients_flops
还有一部分浮点运算次数来自计算梯度
还是在 Linear model 中，经过反向传播需要的 FLOPs = 4 * (data points) * (parameters)
所以总的运算量为：6 * (data points) * (parameters)
Models
Parameter initialization
输入 -> 输出之间的数值增长的量级本质上与隐藏维度（hidden）的平方根成正比
当模型很大时，将会爆炸（数值发散），并且训练方差不稳定，所以需要一种方式进行初始化，这种方式对隐藏层具有“不变性”
Xavier 初始化：让输出值在某个值附近保持稳定（正态分布）
# 除以 sqrt(input_dim) 进行缩放
w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
# 现在 output 的量级是常数
output = x @ w 但是正态分布的尾部是无界的，对它进行截断到 [-3，-3] 的范围，这样就不会得到异常大的值
custom_model
这部分设计了一个简单的线性模型，一共三层参数（D * D，D * D，D），最后的 D 为 head / final layer，模型参数为 D^2 + D，并记得将模型移动 GPU 上
然后生成了一些随机数据输入模型，经过前向传播就是依次通过各层，最后应用头部层
note_about_randomness
设定随机种子（seed）是复现模型训练的必要手段，比如有时候需要重现某个 bug
训练模型中的很多地方都具有随机性（参数初始化，随机失活(dropout)，数据排序）
建议总是传递一个固定的随机种子，可以尽可能重现模型
# Torch
seed = 0
torch.manual_seed(seed)
# NumPy
import numpy as np
np.random.seed(seed)
# Python
import random
random.seed(seed)data_loading
文本经过分词器得到的数据是一个整数序列，转化为 numpy 数值会很方便
我们一般不希望一次性加载全部数据到内存中，使用 memmap 函数可以假装加载，本质上提供了一个映射到文件的变量，它会按需加载文件
orig_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=np.int32)
# 将 NumPy 数组保存为二进制文件
orig_data.tofile("data.npy")
# 使用 memmap 函数假装加载
data = np.memmap("data.npy", dtype=np.int32)
assert np.array_equal(data, orig_data)还可以利用数据加载器 DataLoader，从每个 batch 中获取数据
B = 2 # Batch size（批次大小）
L = 4 # Sequence length（序列长度）
x = get_batch(data, batch_size=B, sequence_length=L, device=get_device())Optimizer（优化器）
先定义几个通用符号
：第 t 步的模型参数（想要优化的对象）
：第 t 步的梯度（gradient）
：学习率（Learning Rate）
：防止分母为 0 而加的一个极小值（通常为 1e-8）
SGD
梯度下降算法：直接沿着梯度反方向走
特点：没有记忆，每次更新独立（无状态）
class SGD(torch.optim.Optimizer):
def __init__(self, params, lr=0.01):
super(SGD, self).__init__(params, dict(lr=lr))
def step(self):
for group in self.param_groups:
lr = group["lr"]
for p in group["params"]:
grad = p.grad.data
p.data -= lr * grad # θ = θ - η·∇L
Momentum
动量法：SGD + 梯度的指数加权平均
是动量系数（通常取 0.9），
不仅取决于当前梯度
，还保留了 90% 的上一步的方向
AdaGrad
引入了自适应学习率：SGD + 累积历史梯度的平方和
这个方法有个问题：分母会随时间不断变大，最后趋近 0（后期“学不动”）
特点：有记忆，记住所有历史梯度的平方和，调整学习率（有状态）
def step(self):
for p in params:
state = self.state[p] # 🔑 这就是"状态"
grad = p.grad.data
# 从状态中取出历史信息
g2 = state.get("g2", torch.zeros_like(grad))
# 更新状态（累积梯度平方）
g2 += grad ** 2
state["g2"] = g2 # 保存回状态
# 使用状态调整学习率
p.data -= lr * grad / sqrt(g2 + 1e-5)模型训练中什么是状态？g2 一直在累积，这就是状态
grad = 3.0
g2 = 0 + 3^2 = 9
w = 2.0 - 0.1 * 3.0 / sqrt(9) = 2.0 - 0.1 = 1.9
状态：{"g2": 9}
grad = 4.0
g2 = 9 + 4^2 = 25 # 累积之前的 9
w = 1.9 - 0.1 * 4.0 / sqrt(25) = 1.9 - 0.08 = 1.82
状态：{"g2": 25}RMSProp
解决了 AdaGrad 分母无限变大的问题：AdaGrad 梯度的平方和从"直接累加" -> "指数加权平均"
这里的 是衰减率（通常取 0.99），使得很久以前的梯度平方会被忘记，
不会一直变大
Adam
Adam = RMSProp + Momentum
步骤 1：指数加权
步骤 2：偏差修正
步骤 3：更新参数（分母提供了动量方向，分子提供了自适应步长）
不同优化器的状态对比
| 优化器 | 状态内容 | 每个参数的状态大小 |
|---|---|---|
| SGD | 无 | 0 |
| Momentum | v（动量） | 1× 参数大小 |
| AdaGrad | g2（梯度平方和） | 1× 参数大小 |
| RMSProp | v（梯度平方指数平均） | 1× 参数大小 |
| Adam | m（一阶动量）+ v（二阶动量） | 2× 参数大小 |
Memory
内存构成：
- 参数：模型权重本身
# 模型权重本身
W1 = torch.randn(4, 4) # 16 个数
W2 = torch.randn(4, 1) # 4 个数
# 总共：20 个数- 激活值：神经网络每一层的输出，反向传播要用（batch_size × 维度 × 层数）
- 对每一层、每一个数据点、每一个维度，都必须存储激活值
- 但其实不需要存储所有激活值，可以重新计算它们，这种技术为激活检查点（activation checkpointing）
# 输入
x = torch.randn(2, 4) # (batch=2, dim=4)
# 第 1 层
W1 = torch.randn(4, 4)
h1 = x @ W1 # (2, 4) ← 保存！
# 第 2 层
W2 = torch.randn(4, 1)
h2 = h1 @ W2 # (2, 1) ← 保存！
# 计算损失
loss = (h2 - target).pow(2).mean()
# 计算 W2 的梯度
dL/dW2 = h1.T @ (dL/dh2) # 🔑 需要 h1！
# 计算 W1 的梯度
dL/dW1 = x.T @ (dL/dh1) # 🔑 需要 x！- 梯度：每个参数的导数
# 每个参数的梯度
W1.grad = torch.randn(4, 4) # 16 个数
W2.grad = torch.randn(4, 1) # 4 个数
# 总共：20 个数（与参数量相同）- 优化器状态：历史梯度信息（SGD 无，Adam 2倍参数量）
state = {
W1: {
"m": torch.zeros(4, 4), # 一阶动量：16 个数
"v": torch.zeros(4, 4) # 二阶动量：16 个数
},
W2: {
"m": torch.zeros(4, 1), # 4 个数
"v": torch.zeros(4, 1) # 4 个数
}
}
# 总共：40 个数（参数量的 2 倍）train_loop
模型训练的基本步骤：
def train(name: str, get_batch,
D: int, num_layers: int,
B: int, num_train_steps: int, lr: float):
# 模型和优化器
model = Cruncher(dim=D, num_layers=0).to(get_device())
optimizer = SGD(model.parameters(), lr=0.01)
for t in range(num_train_steps):
# 获取数据
x, y = get_batch(B=B)
# 前向传播（计算 loss）
pred_y = model(x)
loss = F.mse_loss(pred_y, y)
# 反向传播（计算梯度）
loss.backward()
# 更新参数，梯度清零
optimizer.step()
optimizer.zero_grad(set_to_none=True)checkpointing
检查点：训练模型时间很长可能会遇到程序崩溃，所以如果不希望丢失训练进度，需要定期将模型（及优化器）保存到磁盘，之后就可以直接加载它们了
def checkpointing():
model = Cruncher(dim=64, num_layers=3).to(get_device())
optimizer = AdaGrad(model.parameters(), lr=0.01)
# 保存检查点
checkpoint = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "model_checkpoint.pt")
# 加载检查点
loaded_checkpoint = torch.load("model_checkpoint.pt")mixed_precision_training
训练中的混合精度：默认推荐使用 FP32，可以在前向传播用低精度如 BF16 / FP8
pytorch 中提供了一些工具可以自动进行混合精度训练 AMP（automatic mixed precision）
link("https://pytorch.org/docs/stable/amp.html")
link("https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/")
虽然训练时需要高精度，但是到了推理阶段就可以大胆尝试了
模型量化：降低神经网络模型中的数值精度，来减小模型体积、提升推理速度
更多推荐
13
0- 0
已为社区贡献5条内容
所有评论(0)