Lecture 2 （一）

几个表格，快速回顾知识

缩写	全称	中文含义	常见解释
B	Batch size	批大小	一次并行送进模型的样本数
H	Heads	注意力头数	Multi-Head Attention 里的头数
T	Time steps / Sequence length	时间步 / 序列长度	序列中 token 的数量
D	Dimension / Hidden size	特征维度 / 隐藏维度	向量长度、embedding 维度
V	Vocabulary size	词表大小	不同 token 的总数
d	Head dim	每个 head 的维度	D / H

tensor shape：

名称	Shape	解释
tokens	`(B, T)`	原始输入的 Token ID 序列（整数）。
embeddings	`(B, T, D)`	嵌入后的向量，包含语义及位置信息。
Q / K / V	`(B, H, T, d)`	线性变换后的 Query, Key, Value；其中 $D = H \times d$
attn score	`(B, H, T, T)`	注意力权重矩阵
logits	`(B, T, V)`	未经归一化的预测概率，对应词表（Vocabulary）的大小。

reshape / transpose：

操作	做什么	理解
`view / reshape`	改 shape	不动数据
`transpose(i, j)`	换轴	第 i、j 维互换位置
典型用法	`(B,T,D) → (B,H,T,d)`	拆多头

broadcasting：

规则	说明
从右往左对齐	维度对齐规则
相同或为 1	可以广播
常见例子	`(B,T,D) + (T,D)`

课程笔记

1. tensor memory

数据类型

fp32

fp16

bfloat16

fp8

表格整理

数据类型	别名	占用字节 (Byte)	动态范围	精度 (分辨率)	评价
float32	fp32 / 单精度	4	极大	极高	默认格式，最稳，但最占空间
float16	fp16 / 半精度	2	较小	中	省空间，但容易产生“下溢” (Underflow)
bfloat16	bf16 / Brain	2	等同fp32	较低	省空间且不容易崩溃
fp8	8位浮点数	1	极小	极低	极致压缩，仅限 H100 等新型显卡

例：GPT-3 中的一个前馈层矩阵，大小为 $(12288 \times 4) \times 12288$，而 $内存 = 元素个数 \times 字节数$。

使用类型	每个元素大小	总内存占用	显存压力
float32	4 字节	约 2.3 GB	大
float16 / bf16	2 字节	约 1.15 GB	减半
fp8	1 字节	约 0.57 GB	极小

用 1e-8 测试了不同格式的极限。

数据类型	输入数值	计算机存储的结果	结论
float32	1e-8	1e-8	正常
float16	1e-8	0	下溢 (Underflow)：数值太小，被强行抹除，会导致训练中断
bfloat16	1e-8	1e-8 (或近似)	正常：虽然它精度低，但它能“容纳”这么小的数

方案：

方案	优点	缺点	建议
纯float32	最稳定，不会出错	耗显存，速度最慢	仅在模型极小或显存极多时使用
纯低精度 (fp16/fp8)	速度极快，省空间	危险	容易训练失败

解决方案：混合精度 (mixed precision training)，兼顾速度与稳定。

2. 计算与硬件

tensor on GPU

要把 tensor 从 CPU 移到 GPU 上。

import torch
x = torch.zeros(32, 32)

# `x.device` 查看 tensor 在哪
assert x.device == torch.device('cpu')

# `torch.cuda.is_available()` 查看是否有 GPU
assert torch.cuda.is_available()

# `torch.cuda.device(i)` 查看第 i 个 GPU 属性
num_gpus = torch.cuda.device_count()
for i in range(num_gpus):
    properties = torch.cuda.get_device_properties(i)

# 查看 GPU 上，所有张量实际占用的字节数
memory_allocated = torch.cuda.memory_allocated()

# 把 tensor 移到 GPU 上
y = x.to("cuda:0")
print("x.device: ", x.device)
print("y.device: ", y.device)
# x.device:  cpu
# y.device:  cuda:0

# 在 GPU 上创建 tensor
z = torch.zeros(32, 32, device="cuda:0")

操作

存储：

lecture_02.py 讲义中的重要笔记：

x = torch.tensor([[1., 2, 3], [4, 5, 6]])  # @inspect x
# Many operations simply provide a different view of the tensor.
# This does not make a copy, and therefore mutations in one tensor affects the other.
# Get row 0:
y = x[0]  # @inspect y
assert torch.equal(y, torch.tensor([1., 2, 3]))
assert same_storage(x, y)
# Get column 1:
y = x[:, 1]  # @inspect y
assert torch.equal(y, torch.tensor([2, 5]))
assert same_storage(x, y)
# View 2x3 matrix as 3x2 matrix:
y = x.view(3, 2)  # @inspect y
assert torch.equal(y, torch.tensor([[1, 2], [3, 4], [5, 6]]))
assert same_storage(x, y)
# Transpose the matrix:
y = x.transpose(1, 0)  # @inspect y
assert torch.equal(y, torch.tensor([[1, 4], [2, 5], [3, 6]]))
assert same_storage(x, y)
# Check that mutating x also mutates y.
x[0][0] = 100  # @inspect x, @inspect y
assert y[0][0] == 100
# Note that some views are non-contiguous entries, which means that further views aren't possible.
x = torch.tensor([[1., 2, 3], [4, 5, 6]])  # @inspect x
y = x.transpose(1, 0)  # @inspect y
assert not y.is_contiguous()
try:
    y.view(2, 3)
    assert False
except RuntimeError as e:
    assert "view size is not compatible with input tensor's size and stride" in str(e)
# One can enforce a tensor to be contiguous first:
y = x.transpose(1, 0).contiguous().view(2, 3)  # @inspect y
assert not same_storage(x, y)
# Views are free, copying take both (additional) memory and compute.

特性	view	reshape
内存要求	必须是连续的 (Contiguous)	没有要求，连续不连续都能用
底层逻辑	永远不复制数据，只改变读取方式	如果内存连续就用 view；如果不连续，复制一份数据使其连续
安全性	容易报错（如果不连续的话）	更安全，基本不会报错

示例：

1 2	x = torch.randn(4, 4) # 16 个元素 y = x.reshape(2, -1) # 指定 2 行，-1 自动算出列数是 8

讲义中的其他重要知识点：einops

einops 非常高效。

einsum:

from einops import rearrange, einsum, reduce

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")

reduce:

# You can reduce a single tensor via some operation (e.g., sum, mean, max, min).
x: Float[torch.Tensor, "batch seq hidden"] = torch.ones(2, 3, 4)
# Old way:
y = x.mean(dim=-1)
# New (einops) way:
y = reduce(x, "... hidden -> ...", "sum")

rearrange:

# Sometimes, a dimension represents two dimensions
# ...and you want to operate on one of them.
x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)
# ...where total_hidden is a flattened representation of heads * hidden1
w: Float[torch.Tensor, "hidden1 hidden2"] = torch.ones(4, 4)
# Break up total_hidden into two dimensions (heads and hidden1):
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
# Perform the transformation by w:
x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")
# Combine heads and hidden2 back together:
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")

其他知识点：tensor_operations_flops, gradients_basics, gradients_flops,

Parameter Initialization

标准正态分布随机初始化参数，可能梯度爆炸：

x = nn.Parameter(torch.randn(input_dim))
output = x @ w
assert output.size() == torch.Size([output_dim])
# Note that each element of output scales as sqrt(input_dim): 18.919979095458984.
# Large values can cause gradients to blow up and cause training to be unstable.

Large values can cause gradients to blow up，即梯度爆炸。初始值太大，经过激活函数（比如 Sigmoid 或 Tanh）时，会落入饱和区；在反向传播时，梯度逐层累积可能变得无穷大，Loss 变成 NaN。

Xavier 初始化，即 Glorot 初始化

希望无论输入维度多大，输出值的方差保持一致。

由数学可知，标准差变大了 $\sqrt{\text{input_dim}}$ 倍。

所以初始化权重时，预先除以 $\sqrt{\text{input_dim}}$。

1	w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))

为了更安全，把正态分布两端罕见值切掉。例：[-3, 3]。

1	w = nn.Parameter(nn.init.trunc_normal_(torch.empty(input_dim, output_dim), std=1 / np.sqrt(input_dim), a=-3, b=3))

Palind's Blog

CS336 学习笔记（二）