Lecture 2 PyTorch && Resource Accounting
FP32 float32
float32 是default精度，也被叫做fp32
单精度，默认设置。每个数据用 32 个 bit 存储，也就是 4 个 byte
The float32 data type (also known as fp32 or single precision) is the default.
FP16 float16
float16 又称半精度，存储的 memory 更少，但是 dynamic range 太低，如果用 16 位表示 1e-8，结果为 0（下溢），使用该精度会导致模型的训练不稳定
bfloat16
bfloat16 和 float16 一样的 memory，和 float32 一样的 dynamic range
fp8
H100s support two variants of FP8: E4M3 (range [-448, 448]) and E5M2 ([-57344, 57344]).
Storing
By default, tensors are stored in CPU memory.
如何把数据转移到GPU
y = x.to("cuda:0")
PyTorch维度转换技巧
import torch
from torchtyping import TensorType
# 基本张量操作
x = torch.tensor([1, 2, 3])
y = x.pow(2)
print(y.sqrt()) # 输出: tensor([1.0000, 1.4142, 1.7321])
# 张量乘法示例
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
print(z)
print(z.shape) # 输出: torch.Size([2, 2, 2])
# 使用纯字母作为einsum的维度标识（严格遵循a-z/A-Z规则）
x: TensorType["batch seq1 hidden", torch.float32] = torch.ones(2, 3, 4)
y: TensorType["batch seq2 hidden", torch.float32] = torch.ones(2, 3, 4)
# 用纯字母表示维度：b=batch, i=seq1, j=seq2, h=hidden
z = torch.einsum("bih,bjh->bij", x, y)
print(z)
print(z.shape) # 输出: torch.Size([2, 3, 3])
计算FLOPs
Training GPT-3 (2020) took 3.14e23 FLOPs.
Training GPT-4 (2023) is speculated to take 2e25 FLOPs
FLOP/s 是每秒的（floating-point operations），用来评估硬件的运算速度。那其中训练 GPT-3 花费了 3.14e23 FLOPs，训练 GPT-4 推测花费了近 2e25 FLOPs。
矩阵乘法：
We have one multiplication (x[i][j] * w[j][k]) and one addition per (i, j, k) triple.
Model: x –w1–> h1 –w2–> h2 -> loss
符号含义
- B ：batch size（一次前向/反向里并行处理的 数据点个数 ）。
- D ：中间特征维度/隐藏维度。在这段代码里，
h1的形状是[B, D]，w1的形状是[D, D]，w2的形状是[D, K]，所以输入与隐藏层都用同一个维度记作 D。 - K ：输出维度（比如分类问题中的 类别数 ）。
h2的形状是[B, K]。
下面计算 1. 训练 Transformer 模型所需的计算量（ C）2. 其参数数量或模型大小（ *N ） 3. 模型训练的标记数量（ D ）*
我们来看看 Transformer FLOPs 公式 ： C ≈ 6ND
该方程的一个稍微复杂一些的版本将计算 C 表示为集群吞吐量𝜏和训练时间 T 的乘积： 𝜏T = 6ND.
我们一步一步来分析一下是如何计算的：
前向传播 (Forward pass)
在神经网络中，前向传播就是把输入数据传到网络里，依次经过各层计算，得到输出。
在这里，计算复杂度（FLOPs）被估算为： Forward pass: 2 × (数据点个数) × (参数个数)
也就是说，前向传播每个参数大概需要 2 次浮点运算（一次乘法 + 一次加法）
反向传播 (Backward pass)
反向传播需要计算梯度，复杂度会更高。
计算w2 的梯度： w2.grad[j,k] = sum_i h1[i,j] * h2.grad[i,k] 这里每个三元组 (i,j,k) 都要做一次乘法加法 → 2 次 FLOPs 总计2 * B * D * K
计算h1 的梯度： h1.grad[i,j] = sum_k w2[j,k] * h2.grad[i,k] 同样，每个 (i,j,k) 要做乘加运算 → 2 次 FLOPs 总计2 * B * D * K
计算w1 的梯度 (输入层权重)：这里不需要 x.grad，但是 w1 也有 D×D 个参数，计算梯度时依然要做乘加运算。
所以：(2 + 2) * B * D * D
综合来看就是说，训练一次神经网络（前向 + 反向）大约需要 6 倍的数据量 × 参数量的浮点运算
