pytorch 与资源核算

资源核算的原因

场景一：时间估算

问题：假设你是一个 AI 工程师，老板问你："在 1024 张 H100 显卡上，训练一个 70B（700亿参数）的模型，数据量是 15T（15万亿 tokens），大概要多久？"

直接跑去写代码测试，那可能几天甚至几个月都出不来结果。因此，我们需要学会"Napkin math"（餐巾纸计算，即快速估算）。

第一步：计算总工作量:

FLOPs（Floating Point Operations，浮点运算次数）是衡量计算复杂度的常用指标，表示一个算法或模型在执行过程中所需的浮点加减乘除等运算的总次数。而训练模型本质上是进行浮点运算（FLOPs）。有一个经验公式：总计算量≈6×参数量×Token数量。

为什么是 6 倍？

bash 复制代码

前向传播 (Forward Pass): 计算一次约为 2× 参数量（乘法+加法）。
反向传播 (Backward Pass): 计算梯度约为前向传播的 2 倍工作量，即 4×参数量。

好的，这是您需要的 Markdown 格式内容，直接复制粘贴即可使用。

为什么训练 Transformer 的 FLOPs 公式是 6 × 参数量 × Token 数？

这个问题的答案非常经典，它源自于 Transformer 模型（基于 Decoder-only 架构，如 GPT） 的训练过程。下面详细拆解 2倍和 4倍的来源，以及为什么忽略了一些"零头"。

第一部分：前向传播 ------ 为什么是 2倍？

我们先看一个最基本的全连接层（线性层）：y = x * W（忽略偏置 bias 以简化计算）。假设权重矩阵 W W W 的形状是 $d i n , d o u t$ $d_{in}, d_{out}$ $din,dout$ ，参数量 P = d i n × d o u t P = d_{in} \times d_{out} P=din×dout。当输入一个 Token 的向量 x x x （长度为 d i n d_{in} din）时，要计算得到 y y y （长度为 d o u t d_{out} dout）：

乘法次数 ：输出层每一个位置都需要 d i n d_{in} din 次乘法。
加法次数 ：把 d i n d_{in} din 个乘积加起来，需要 d i n − 1 d_{in} - 1 din−1 次加法。

总运算量 ≈ 2 × d i n × d o u t = 2 P \approx 2 \times d_{in} \times d_{out} = 2P ≈2×din×dout=2P。

结论 1 ：对于矩阵乘法，计算量 (FLOPs) ≈ 2 × 参数量。

Transformer 前向传播的视角

尽管 Transformer 里还有 Attention（注意力机制）、LayerNorm（层归一化）、激活函数这些操作，但它们的计算复杂度远低于巨大的矩阵乘法层（MLP 和 QKV 投影）。在大模型训练（Token 数量巨大、维度数千维）的尺度下，矩阵乘法的 FLOPs 占据绝对主导地位（>99%）。

因此，前向传播 ≈ 2 × 参数量 FLOPs/Token。

第二部分：反向传播 ------ 为什么是前向的 2 倍（即 4倍参数量）？

这是理解 6 这个数字的关键。反向传播由两部分组成，正好是前向计算量的 2倍：

1. 计算对输入的梯度（后一层往前传，约 2倍参数量）

在反向传播中，我们需要计算损失 L L L 对输入 x x x 的梯度 ∂ L ∂ x \frac{\partial L}{\partial x} ∂x∂L，以便继续往更浅的网络层传递。

公式： ∂ L ∂ x = ∂ L ∂ y × W T \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \times W^T ∂x∂L=∂y∂L×WT。
这是一个 新的矩阵乘法（梯度向量乘以权重的转置）。
FLOPs 开销 ：和前向传播几乎一模一样，即 2倍参数量。

2. 计算对权重的梯度（更新参数用，约 2倍参数量）

同时，我们需要计算损失 L L L 对权重 W W W 的梯度 ∂ L ∂ W \frac{\partial L}{\partial W} ∂W∂L。

公式： ∂ L ∂ W = x T × ∂ L ∂ y \frac{\partial L}{\partial W} = x^T \times \frac{\partial L}{\partial y} ∂W∂L=xT×∂y∂L。
这也是一个 矩阵乘法（输入向量转置乘以输出梯度向量）。
FLOPs 开销 ：同样是 2倍参数量。

结论 2 ：反向传播的总计算量 = 计算输入梯度 (2P) + 计算权重梯度 (2P) = 4倍参数量。

第三部分：加总得出"6"与忽略的细节

把上述过程加起来：
总计 = 前向 ( 2 P ) + 反向 ( 4 P ) = 6 × 参数量 \text{总计} = \text{前向} (2P) + \text{反向} (4P) = 6 \times \text{参数量} 总计=前向(2P)+反向(4P)=6×参数量

乘以训练数据集的 Token 总数 T T T，就得到了著名的 6ND 公式：
Total FLOPs ≈ 6 × P × T \text{Total FLOPs} \approx 6 \times P \times T Total FLOPs≈6×P×T

Adam 优化器更新参数不需要计算吗？激活函数的计算不算 FLOPs 吗？

激活函数（GeLU, SiLU, ReLU）：
- 计算量是 O(d_model) （线性复杂度），而矩阵乘法是 O(d_model²)（平方复杂度）。
- 当模型维度 d_model 在 4096 以上时，矩阵乘法占了 99.9% 的计算时间，激活函数在 FLOPs 计数中可忽略不计。
Attention 中的 Softmax 和 Mask：
- 虽然 Q ⋅ K T Q \cdot K^T Q⋅KT 是矩阵乘法算过了，但 Softmax 涉及指数和对数运算，这在 理论 FLOPs 中占比极小，但在 实际推理延迟 中占比不小（内存受限）。在训练 FLOPs 估算中，同样被省略。
Adam 优化器的状态更新：
- 注意：FLOPs 定义的是模型计算梯度。
- Adam 步骤是：m = beta1*m + (1-beta1)*grad，v = ...，param -= lr * m / (sqrt(v) + eps)。
- 这确实需要计算，且计算量大约是 另外的 4倍 ~ 12倍参数量（取决于具体实现）。
- 但是：学术论文中提到的 Training FLOPs 通常 不包含 优化器更新步骤。它特指 前向+反向传播。如果算上优化器，可能是 10ND 甚至 18ND。

总结图示

阶段	核心操作	近似 FLOPs (每 Token)
前向传播	`x * W` (一次乘加)	2 × P
反向传播 (梯度回传)	`grad * W^T` (一次乘加)	2 × P
反向传播 (权重梯度)	`x^T * grad` (一次乘加)	2 × P
合计 (训练一个Token)		6 × P

第二步：计算硬件算力

查阅NVIDIA H100 的白皮书，其 FP16/BF16 的峰值算力约为 1979 TFLOPS（每秒万亿次浮点运算）。

但是要注意，这个值是 NVIDIA H100 GPU 在使用 FP16 或 BF16 数据类型、且启用结构化稀疏（Structured Sparsity）时可达到的理论最大计算吞吐量。对于我们训练普通的稠密模型（如标准的 LLaMA-1到LLaMA-3系列、Qwen3-0.6B/1.7B/4B/8B/14B/32B等），理论峰值减半，约 990 TFLOPS。

结构化稀疏，是一种模型压缩的方法，通常是对稠密模型按照50%的稀疏度进行剪枝**（稀疏形式为n:m，即m个连续权重里必须剪掉n个，类型有 2:4、4:8、8:16)。**还有一种灵活度更高的非结构化剪枝方法，可以按照百分比进行剪枝，但通常相对于结构化剪枝性能更差。

上述 990 TFLOPS 是 H100 的理论峰值，但实际运行模型时，由于各种软硬件开销，你几乎不可能达到 100% 模型算力利用率 (MFU, Model FLOPs Utilization)，通常按 30%--60% 的利用率估算更现实。这里取 50% 用作后续估计。

单卡实际算力 ≈ ( 990 × 10 12 ) × 0.5 ≈ 5 × 10 14 \approx (990 \times 10^{12}) \times 0.5 \approx 5 \times 10^{14} ≈(990×1012)×0.5≈5×1014 FLOPS

1024 张卡总算力 ≈ 5 × 10 17 \approx 5 \times 10^{17} ≈5×1017 FLOPS。

第三步：得出结果

时间 = 总工作量总算力 = 6.3 × 10 24 5 × 10 17 ≈ 1.26 × 10 7 = \frac{总工作量}{总算力} = \frac{6.3 \times 10^{24}}{5 \times 10^{17}} \approx 1.26 \times 10^7 =总算力总工作量=5×10176.3×1024≈1.26×107 秒 ≈ 146 \approx 146 ≈146 天

场景二：内存估算

问题：在8张H100 GPU上，使用AdamW优化器（朴素实现，即所有数据都用float32存储，没有混合精度或压缩），你能训练的最大模型参数量是多少？

很多初学者认为：8张H100（每张80GB显存），显存总共640GB。参数采用float32存储，一个参数4字节，所以640GB/4字节=1600亿参数（160B）？错！

训练时，显存里存的不只是参数(Parameters)，还有：

梯度(Gradients):和参数一样大，需要4字节
优化器状态(Optimizer States):比如常用的AdamW优化器，需要存储每个参数的一阶矩估计（梯度的指数移动平均）和二阶矩估计（梯度平方的指数移动平均）
因此，每个参数在训练时大约占用16字节：
参数(FP32):4bytes
梯度(FP32):4bytes
优化器状态(FP32):8bytes(2个变量×4bytes)

最大参数量=640×10^9 bytes/16 bytes/param≈400亿(40B)

当然，这个估算较为粗略，因为它未考虑依赖于批次大小和序列长度的激活值内存占用，所以如果不考虑混合精度或压缩，实际能训练的模型可能会更小。

张量 (Tensors)

张量基础

在深度学习中，张量（tensor）是存储一切的基础数据结构：模型参数、梯度、优化器状态、输入数据、中间激活值等都以张量形式存在。更多关于张量的知识可见 PyTorch docs on tensors。
PyTorch提供了多种创建张量的方法：

bash 复制代码

import  torch
from torch import nn
x = torch.tensor([[1., 2, 3], [4, 5, 6]])  # 从Python列表创建
print(x)
x = torch.zeros(4, 8)  # 4x8的零矩阵
print("# 4x8的零矩阵: ",x)
x = torch.ones(4, 8)   # 4x8的全1矩阵
print("# 4x8的全1矩阵: ",x)
x = torch.randn(4, 8)  # 4x8的正态分布随机数
print("# 4x8的正态分布随机数: ",x)
 分配但不初始化值（用于自定义初始化的值）
x = torch.empty(4, 8)  # 未初始化的4x8矩阵
print("# 未初始化的4x8矩阵",nn.init.trunc_normal_(x, mean=0, std=1, a=-2, b=2) ) # 截断正态分布初始化，分布均值为0，标准差为1，截断范围，只保留[-2, 2]区间内的值，超出此范围的值会被重新采样

结果为：

张量的操作

许多操作只是提供张量的一个不同"视图"。这不会创建副本，因此在一个张量中的修改会影响另一个。

bash 复制代码

import torch

def same_storage(x, y):
    return x.storage().data_ptr() == y.storage().data_ptr()

x = torch.tensor([[1., 2, 3], [4, 5, 6]])
print("原始的x数值",x)
# 这些操作不复制数据
y = x[0]           # 获取第0行
print("获取第0行",y)
y = x[:, 1]        # 获取第1列
print("获取第1列",y)
y = x.view(3, 2)   # 重塑为3x2
print(" 重塑为3x2",y)
y = x.transpose(1, 0)  # 转置
print("转置后的数据y:",y)
print("原始的x数值",x)
# 验证是否共享存储
assert same_storage(x, y)  #
# 注意：修改x会影响y
x[0][0] = 100
print(y[0][0])  # 输出 100.0

输出结果：

注意：并非所有视图都是"连续"的。当一个视图的数据在内存中不是按顺序存储时，它就是非连续的（Non-contiguous）。
转置后的张量是非连续的：

bash 复制代码

x = torch.tensor([[1., 2, 3], [4, 5, 6]])
y = x.transpose(1, 0)
assert y.is_contiguous()  # 是否为连续张量

结果为：

不能对非连续张量直接执行某些操作，比如再次 view:

bash 复制代码

import torch
x = torch.tensor([[1., 2, 3], [4, 5, 6]])
y = x.transpose(1, 0)
# 尝试重塑会失败
try:
    y.view(2, 3)
except RuntimeError as e:
    print("view size is not compatible with input tensor's size and stride" in str(e))

输出结果：

如果需要对一个非连续张量进行进一步操作，可以先调用 .contiguous() 方法：

bash 复制代码

# 解决方案：先使张量连续
import torch
x = torch.tensor([[1., 2, 3], [4, 5, 6]])
y = x.transpose(1, 0)
y = x.transpose(1, 0).contiguous().view(2, 3)
def same_storage(x, y):
    return x.storage().data_ptr() == y.storage().data_ptr()
assert not same_storage(x, y)  # 现在创建了新存储

.contiguous() 会创建一个新的张量，并将数据按顺序复制到新的连续内存块中，这样后续操作就不会出错。

逐元素操作（Element-wise Operations）

逐元素操作是对每个元素独立应用函数，并返回相同形状的新张量。这意味着，当你对一个张量执行 x.pow(2) 或 x + x 时，操作会独立地作用于张量中的每一个数字，而不会像矩阵乘法那样考虑元素之间的行列关系。

下面代码演示了几个常用的逐元素数学函数：

bash 复制代码

x = torch.tensor([1, 4, 9])
assert torch.equal(x.pow(2), torch.tensor([1, 16, 81]))    # 幂运算 (pow)，每个元素平方
assert torch.equal(x.sqrt(), torch.tensor([1, 2, 3]))      # 对每个元素开平方根
assert torch.equal(x.rsqrt(), torch.tensor([1, 1/2, 1/3])) # rsqrt 是 reciprocal of sqrt 的缩写，即对每个元素先开方再取倒数

基本的逐元素算术运算：

bash 复制代码

assert torch.equal(x + x, torch.tensor([2, 8, 18])) # 每个元素自加
assert torch.equal(x * 2, torch.tensor([2, 8, 18])) # 每个元素乘以标量2
assert torch.equal(x / 0.5, torch.tensor([2, 8, 18])) # 每个元素除以标量0.5

一个非常实用的工具函数 triu，这在计算因果注意力掩码（causal attention mask）时非常有用。在语言模型中，为了确保模型在预测第 j 个词时只能看到第 j 个词之前的词（即不能"偷看"未来信息），就需要使用这种上三角矩阵作为掩码。其中 M $i, j$ 表示位置 i 对位置 j 的贡献，当 i > j 时（即 i 在 j 之后），贡献应为0。

bash 复制代码

import torch
# 实用操作：上三角矩阵（用于因果注意力掩码）
x = torch.ones(3, 3).triu()  # 创建一个3x3全1矩阵，然后取其上三角部分
# 结果：
# [[1, 1, 1],
#  [0, 1, 1],
#  [0, 0, 1]]
print(x)

矩阵乘法（Matrix Multiplication）

矩阵乘法是深度学习的基础。矩阵乘法是神经网络中最核心、最频繁的计算操作，无论是全连接层、卷积层还是注意力机制，其底层都离不开矩阵运算。

一个标准的矩阵乘法形式为：一个 M×K 的矩阵乘以一个 K×N 的矩阵，得到一个 M×N 的结果矩阵。示例如下：

bash 复制代码

# 基本矩阵乘法
x = torch.ones(16, 32)  # 16x32矩阵
w = torch.ones(32, 2)   # 32x2权重矩阵
y = x @ w               # 结果：16x2矩阵
assert y.size() == torch.Size([16, 2])

使用 Einops 库对张量操作进行优化

在 pytorch 中，张量维度通常是 $batch, sequence, hidden$ 。用 PyTorch 原生的 .view() 和 .transpose() 操作维度时，需要时刻记住张量的维度顺序，很容易搞晕。并且，如果你未来修改了张量形状（比如在 transformers 模型中加了 heads 维度），这个代码就可能出错。

bash 复制代码

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)

用 jaxtyping 命名维度

bash 复制代码

# 传统方式（容易写错维度顺序）
x = torch.ones(2, 2, 1, 3)  # batch seq heads hidden

# Jaxtyping方式（在类型注解中命名维度）
from jaxtyping import Float
x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2, 2, 1, 3)

jaxtyping 提供的维度命名（如 "batch seq hidden"）在当前 PyTorch 生态中主要是"文档性质"的，不会在运行时自动强制检查维度是否真的匹配名称，但它能极大地提升代码的清晰度和可靠性。也就是说PyTorch 在运行时不会检查 x 的形状是否真的是 (batch=2, seq=3, hidden=4)，即使你写成：

bash 复制代码

y: Float[torch.Tensor, "batch seq hidden"] = torch.randn(100, 5)  # 形状与注解严重不符

代码依然会正常运行，不会抛出任何错误。Python 的类型注解（Type Hints）本身是可选的、非强制的，主要用于工具（如 IDE、mypy）做静态分析，而不是运行时校验。尽管不强制校验，jaxtyping 仍带来巨大工程价值，当你将来修改模型（比如增加 heads 维度）：

类型注解会提醒你哪里需要同步更新。

IDE（如 VS Code、PyCharm）能基于注解提供：
自动补全重构支持（rename dimension）错误高亮（如果你在 einsum 中拼错了维度名）

注意：另外需要注意的一点是，jaxtyping也提供了可选地运行时检验，第三方库如 beartype 也提供了类型注解做运行时校验的功能，读者感兴趣可自行探索。

用 einops.einsum 替代矩阵乘法 + 转置

广义矩阵乘法，具有清晰的维度跟踪：

bash 复制代码

import torch
from torchtyping import TensorType as Float
from einops import einsum

# 定义两个张量
x: Float["batch seq1 hidden"] = torch.ones(2, 3, 4)
y: Float["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")
print(z.shape)

输出结果：

这种方法会将没在输出中出现的 hidden 维度，自动对该维度做求和。这本质上是爱因斯坦求和约定（Einstein summation）的直观实现。更通用的写法是（支持任意前导维度）：

bash 复制代码

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

... 表示"任意数量的前导维度"，比如可能是 (device, batch) 或 (ensemble, batch, time)，代码依然适用。好处是维度逻辑显式、通用、不易出错，且天然支持广播。

用 einops.reduce 替代 mean(dim=...)

bash 复制代码

import torch
from torchtyping import TensorType as Float
from einops import einsum, reduce

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

# 传统写法
y = x.mean(dim=-1)
print("# 对最后一个维度 hidden 上求平均", y)

# Einops 写法
y = reduce(x, "... hidden -> ...", "mean")
print("# Einops 写法", y)

结果为：

从 ... hidden 变成 ...，说明 hidden 维度被"聚合"了，聚合方式是 "mean"（也可用 "sum", "max" 等）。同样... 表示"任意数量的前导维度"。这种写法的好处是明确表达了"我打算把哪个维度压缩掉"，语义清晰。

用 einops.rearrange 拆分/合并维度

这是 Einops 最强大的功能之一，用于处理扁平化与重组。

假设你有一个维度 total_hidden = 8，它实际上是 heads=2 和 hidden1=4 的乘积（即 2 * 4 = 8），也就是 total_hidden 维度实际上是 heads×hidden1 的扁平化表示。

bash 复制代码

from einops import rearrange

# 情景：hidden维度实际上是heads×hidden1的扁平化表示
x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)
w: Float[torch.Tensor, "hidden1 hidden2"] = torch.ones(4, 4)

内存（Memory）

张量的内存占用由两个因素决定：

元素数量：张量的形状（如4×8矩阵有32个元素）
数据类型：每个元素占用的字节数

浮点类型

在 LLM 训练中，选择数据类型本质上是在做显存占用、计算速度与数值稳定性之间的权衡（Trade-off）。我们关注的大多数张量（参数、梯度、激活值、优化器状态）都以浮点数形式进行存储，让我们深入理解不同的浮点数据类型：

Float32 (FP32 / 单精度浮点数)

规格：占用 4 字节 (32 bits)。结构为：1位符号位 + 8位指数位 + 23位尾数位。
地位：PyTorch 的默认数据类型，也是科学计算领域的"黄金标准"。
优点：数值精度高，动态范围大，训练最稳定，几乎不会出现数值溢出问题。
缺点：对于大模型而言太"奢侈"。它占用的显存是 16位格式的两倍，且在现代 GPU（如 H100）上的计算吞吐量远低于低精度格式。

炼丹用途：通常用于存储参数的主副本 (Master Weights) 和优化器状态 (Optimizer States)，以确保在梯度累积和参数更新时不会因为精度丢失而导致模型无法收敛。

Float16 (FP16 / 半精度浮点数)

规格：占用 2 字节 (16 bits)。结构为：1位符号位 + 5位指数位 + 10位尾数位。
优点：显存占用比 FP32 减少一半，计算速度快。
致命缺陷：动态范围（Dynamic Range）太窄。
由于指数位只有 5 位，它无法表示非常小的数（会发生下溢 Underflow，直接变成 0）或非常大的数（会发生上溢 Overflow，变成 Infinity）。

例如：在 FP16 中，1e-8 这样的小数会被直接当作 0 处理，导致梯度消失。
用途：这是上一代 GPU（如 V100）混合精度训练的主流。为了解决溢出问题，必须使用复杂的损失缩放 (Loss Scaling) 技术。目前在 LLM 训练中正逐渐被 BF16 取代。

BFloat16 (BF16 / Brain Floating Point)

规格：占用 2 字节 (16 bits)。结构为：1位符号位 + 8位指数位 + 7位尾数位。
来源：由 Google Brain 专为深度学习设计。
设计逻辑："要范围，不要精度"。

深度学习模型（尤其是神经网络）对小数点的后几位精度不敏感，但对数值的范围非常敏感。

BF16 直接截断了 FP32 的尾数，但保留了和 FP32 相同的 8 位指数位。

优点：

拥有和 FP32 一样宽广的动态范围，不需要 Loss Scaling 也能稳定训练。
显存占用和 FP16 一样少。
在 A100/H100 等新硬件上计算速度极快。

用途：当前 LLM 训练的绝对主流选择。通常用于存储激活值 (Activations) 以及进行前向和反向传播的矩阵乘法计算。

FP8 (8位浮点数)

规格：占用 1 字节 (8 bits)。
变体：
E4M3：4位指数，3位尾数（精度稍高，范围稍小）。
E5M2：5位指数，2位尾数（范围稍大，精度更低）。
优点：极致的显存压缩和计算吞吐量。
硬件限制：仅在 NVIDIA H100 及更新的架构（配合 Transformer Engine）上才被原生支持。
炼丹用途：目前主要用于推理阶段的量化 (Quantization)。

虽然理论上可以用于训练（H100 支持），但由于精度极低，训练极不稳定，属于比较前沿的研究领域。

🪜 为什么 BF16 优于 FP16？

在深度学习中，我们通常不关心小数点后第10位的精确度（精度），但非常关心能否表示非常大或非常小的数（动态范围）。FP16 的指数位太少，导致训练中经常出现 NaN 或 0（下溢）。BF16 截断了 FP32 的尾数，保留了指数位，因此它能表示的数值范围与 FP32 一样大，极大地提升了训练稳定性。

张量在内存中的存储机制

在早期的 pytorch 版本中，张量在一台机器上的存储，占用了两部分内存，一个内存储存了这个张量的形状（size）、步长（stride）、数据类型（dtype）等元数据信息，我们把这一部分称之为头信息区（Tensor）；另一个内存储存的就是真正的数据，我们称为存储区（Storage）。但是将"存储区"作为一个独立的、可公开访问的对象，其带来的复杂性和潜在风险，远大于它为用户提供的一些灵活性。

因此，在现在的 pytorch 版本中，PyTorch 的设计者们选择将这一层实现细节进行封装，让 Tensor 对象本身成为一个集成了"指针+元数据"的单一实体。这意味着一个张量本身并不直接"包含"数据，而是指向一块连续的内存区域，并附带一套规则（元数据），告诉程序如何根据你请求的索引去这块内存里找到对应的数据。让我们定义一个4x4的二维张量作为例子：

bash 复制代码

x = torch.tensor([
    [0., 1, 2, 3],
    [4, 5, 6, 7],
    [8, 9, 10, 11],
    [12, 13, 14, 15],
])

张量在底层可以是按行存储也可以是按列存储。Numpy 和 Pytorch 都采用了按行存储的方式，任何维度的张量在底层存储都占据着内存中连续的空间，那么问题来了，我们如何访问到我们想要的位置的数据？

答案就是步长（Strides）！步长是一个元组，定义了在每个维度上移动一个单位时，在底层存储中需要"跳过"的元素数量。

bash 复制代码

# 步长（stride）：在每个维度上移动时跳过的元素数
assert x.stride(0) == 4  # 行维度：移动到下一行需跳过4个元素
assert x.stride(1) == 1  # 列维度：移动到下一列需跳过1个元素

# 计算元素位置：行索引r，列索引c
r, c = 1, 2
index = r * x.stride(0) + c * x.stride(1)  # 位置6

用一句话来说，PyTorch通过"步长"这一元数据，巧妙地将多维张量的逻辑结构映射到一维的物理内存上。

关键点：很多操作（如 view, transpose, slice）不会复制数据，只是修改了 stride。这叫 Zero-copy，非常高效。

bash 复制代码

x = torch.randn(32, 32)
y = x.view(1024) # 只是改变了看待数据的方式，内存地址没变
z = x.transpose(0, 1) # 修改了 stride，没有复制数据

坑点：如果你对转置后的张量调用 .view()，可能会报错。因为转置后内存不再连续（Contiguous）。你需要先调用 .contiguous()，这会触发数据复制（消耗内存和时间）。

计算效率

浮点运算总数（FLOPs）

概述：一个浮点运算（FLOP,Floating Point Operation）是一个基本的算术操作，比如加法 x + y x+y x+y或乘法 x ∗ y x*y x∗y，理解FLOPs对性能分析至关重要。

首先，我们先辨析几个常见概念：

FLOPs（小写s）：执行的浮点运算总数，这是一个衡量计算量的指标，
FLOP/s（每秒浮点运算）：指硬件每秒能执行的浮点运算次数，也写作FLOPS（大写S）。这是一个衡量"速度"的指标。例如，H100 GPU峰值性能为990 TFLOP/s。

重要提示：这两个缩写写法相似，但含义完全不同！本教程中会明确区分它们，避免混淆。

接下来，让我们通过几个例子让你对FLOPs和FLOP/s有一个直观认识：

GPT-3(2020年)：训练耗时约 3.14 × 10 23 3.14\times10^{23} 3.14×1023FLOPs文章
GPT-4(2023年)：据推测训练耗时约 2 × 10 25 2\times10^{25} 2×1025FLOPs文章
政策背景：美国曾有一项行政命令，要求任何训练FLOPs超过 1 × 10 26 1\times10^{26} 1×1026的基础模型必须向政府报告（该命令已于2025年被撤销）
NVIDIA A100：峰值性能为312 TFLOP/s(即 3.12 × 10 14 3.12\times10^{14} 3.12×1014FLOP/s)官方手册
NVIDIA H100：峰值性能为1979 TFLOP/s，但这通常是在启用"稀疏性"（sparsity）的情况下。对于密集矩阵乘法，其性能约为一半，即989.5 TFLOP/s官方手册。

线性模型的计算量

为了具体化FLOPs的计算方法，我们引入了一个简单的线性模型作为例子。这个线性模型有 B 个数据点，每个点是 D 维的，模型将其映射到 K 维输出。我们要做的核心操作是矩阵乘法 y = x @ w，目标是计算这个操作总共需要多少次浮点运算（FLOPs）？

bash 复制代码

if torch.cuda.is_available():
    B = 16384  # Number of points
    D = 32768  # Dimension
    K = 8192   # Number of outputs
else:
    B = 1024
    D = 256
    K = 64
device = get_device()
x = torch.ones(B, D, device=device)
w = torch.randn(D, K, device=device)
y = x @ w

接下来让我们进行分析：

x 是输入数据，形状为 (B, D)
B：batch size（数据点数量）
D：每个数据点的维度（输入特征数）
w 是权重矩阵，形状为 (D, K)
K：输出维度（比如分类任务的类别数）
y 是输出，形状为 (B, K) 考虑 y = x @ w 中的一个输出元素 y $i,k$ ：

bash 复制代码

y[i, k] = x[i, 0] * w[0, k] +
           x[i, 1] * w[1, k] +
           ...
           x[i, D-1] * w[D-1, k]

这个求和过程包含D次乘法（x $i,j$ *w $j,k$ ）和D-1次加法，总共≈2D次FLOPs（因为D通常很大，D-1≈D）因此，矩阵乘法的计算量：总FLOPs≈2×B×D×K

因为D×K正好是这个线性层的参数数量！所以我们可以重写为：
F L O P s = 2 × 数据量 × 模型参数量 FLOPs=2\times数据量\times模型参数量 FLOPs=2×数据量×模型参数量

或者在语言模型中，常用token代替数据点： F L O P s = 2 × t o k e n 数量 × 模型参数量 FLOPs=2\times token数量\times模型参数量 FLOPs=2×token数量×模型参数量。虽然相较于线性模型，Transformer还有注意力、LayerNorm等额外操作，但矩阵乘法占主导，所以这个公式对Transformer也基本成立。
其他操作的FLOPs

逐元素操作（Element-wise Operations）：对张量中每个元素独立应用某个函数（如x+1、x.sqrt()、torch.relu(x)等）。若张量形状为m×n，则逐元素操作的FLOPs为O(m×n)。
张量加法（Tensor Addition）：两个相同形状张量的对应元素相加，如z=x+y。若x和y均为m×n，则总FLOPs=m×n（每个元素1次加法）。

对比矩阵乘法A(m×k)@B(k×n)的FLOPs=2×m×k×n。当k较大时（如k=1024），矩阵乘法的FLOPs是加法的上千倍。对于现代深度学习模型（尤其是Transformer、MLP等），90%+的FLOPs来自矩阵乘法（线性层、QKV投影、FFN等）。其他操作（LayerNorm、Softmax、激活函数、加法残差连接等）虽然存在，但FLOPs可忽略不计。因此，在进行粗略但有效的资源估算时，只需计算所有矩阵乘法的FLOPs总和即可。

理论上的FLOPs只是一个衡量计算量的数字，如何将FLOPs转化为实际运行时间呢？

bash 复制代码

actual_time = time_matmul(x, w)  # 实际执行矩阵乘法所需的时间（秒）
actual_flop_per_sec = actual_num_flops / actual_time  # 实际每秒能完成多少次浮点运算

另外，FLOP/s 的性能很大程度上取决于数据类型！

bash 复制代码

promised_flop_per_sec = get_promised_flop_per_sec(device, x.dtype)  # 获取硬件的理论峰值性能

MFU (Model FLOPs Utilization)

MFU = (实际FLOP/s) / (硬件峰值FLOP/s)

MFU >= 0.5：被认为是相当不错的性能
MFU 接近 1.0：非常难达到，因为硬件不可能100%满负荷运转，总会有内存访问、数据传输等开销。

注意：这个公式忽略了通信和系统开销，只关注纯粹的计算效率。
接下来在 bfloat16 上计算一下 MFU：

bash 复制代码

import torch
import time


# ------------------------------------------------------------
# 1. 定义辅助函数
# ------------------------------------------------------------
def time_matmul(x, w, num_warmup=10, num_iter=50):
    """
    测量矩阵乘法 x @ w.t() 的平均执行时间（秒）
    使用 CUDA Events 进行精确计时，包含预热阶段。
    """
    # 确保张量在 GPU 上且为连续内存布局
    x = x.contiguous()
    w = w.contiguous()

    # 创建 CUDA 事件用于计时
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    # 预热 GPU
    for _ in range(num_warmup):
        _ = x @ w.t()
    torch.cuda.synchronize()

    # 正式测量
    start_event.record()
    for _ in range(num_iter):
        _ = x @ w.t()
    end_event.record()
    torch.cuda.synchronize()

    elapsed_time_ms = start_event.elapsed_time(end_event)  # 毫秒
    avg_time_sec = (elapsed_time_ms / 1000.0) / num_iter
    return avg_time_sec


def get_promised_flop_per_sec(device, dtype):
    device_name = torch.cuda.get_device_name(device)
    flops_table = {
        "A100": {torch.bfloat16: 312e12, torch.float16: 312e12, torch.float32: 77.6e12},
        "H100": {torch.bfloat16: 989e12, torch.float16: 989e12, torch.float32: 67e12},
        "RTX 4090": {torch.bfloat16: 82.6e12, torch.float16: 82.6e12, torch.float32: 35.6e12},
    }
    for key, vals in flops_table.items():
        if key in device_name:
            return vals.get(dtype, 100e12)
    return 100e12


# ------------------------------------------------------------
# 2. 主测试逻辑
# ------------------------------------------------------------
def main():
    # 设置设备和矩阵尺寸
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cpu":
        print("This example requires a CUDA GPU. Exiting.")
        return

    # 定义矩阵大小 (M, K) 和 (K, N)
    M, K, N = 8192, 8192, 8192  # 可调整以观察不同 MFU

    # 创建随机张量（在 GPU 上）
    x = torch.randn(M, K, device=device)
    w = torch.randn(N, K, device=device)  # 注意形状以便计算 x @ w.t()

    # ------------------------------------------------------------
    # 将张量转换为 bfloat16
    x = x.to(torch.bfloat16)
    w = w.to(torch.bfloat16)

    # 计算一次矩阵乘法的理论浮点运算次数（乘加计为2次操作）
    # FLOPs = 2 * M * N * K
    actual_num_flops = 2 * M * N * K

    # 测量实际性能
    bf16_actual_time = time_matmul(x, w)
    bf16_actual_flop_per_sec = actual_num_flops / bf16_actual_time

    # 获取 bfloat16 的理论峰值
    bf16_promised_flop_per_sec = get_promised_flop_per_sec(device, x.dtype)

    # 计算 MFU (Model FLOPs Utilization)
    bf16_mfu = bf16_actual_flop_per_sec / bf16_promised_flop_per_sec

    # 打印结果
    print(f"Device: {torch.cuda.get_device_name(device)}")
    print(f"Matrix size: {M}x{K} * {K}x{N}")
    print(f"Data type: {x.dtype}")
    print(f"Actual time per matmul: {bf16_actual_time * 1000:.3f} ms")
    print(f"Actual FLOPS: {bf16_actual_flop_per_sec:.3e} ({bf16_actual_flop_per_sec / 1e12:.2f} TFLOPS)")
    print(f"Promised FLOPS (BF16): {bf16_promised_flop_per_sec:.3e} ({bf16_promised_flop_per_sec / 1e12:.2f} TFLOPS)")
    print(f"MFU: {bf16_mfu * 100:.2f}%")

if __name__ == "__main__":
    main()

PyTorch 中的梯度计算

到目前为止，我们已经构建了张量（对应于参数或数据），并将它们传递给操作（前向传播）。现在，我们将计算梯度（反向传播）。

梯度反向传播

我们继续以一个简单的线性模型为例：

bash 复制代码

y = 0.5 * (x * w - 5)²

这里，x 是输入数据，w 是模型参数，y 是损失值。前向传播代码:

bash 复制代码

x = torch.tensor([1., 2, 3]) # 输入数据，不需要计算梯度
w = torch.tensor([1., 1, 1], requires_grad=True) # 模型参数，需要计算梯度
pred_y = x @ w # 预测值（矩阵乘法）
loss = 0.5 * (pred_y - 5).pow(2) # 计算损失值

requires_grad=True 这个参数它告诉 PyTorch "请为这个张量 w 构建计算图，并在反向传播时计算它的梯度。" 对于输入数据 x，我们通常不需要它的梯度，所以不设置此标志。

反向传播代码：

bash 复制代码

loss.backward() # 触发反向传播
assert loss.grad is None # 损失值 loss 本身是一个标量，它没有"上游"的梯度，所以它的 .grad 属性是 None
assert pred_y.grad is None # 中间变量 pred_y 默认情况下也不会保存梯度，除非你显式地调用 pred_y.retain_grad()
assert x.grad is None # 因为 x 在创建时没有设置 requires_grad=True，所以它不会被追踪，其 .grad 也为 None
assert torch.equal(w.grad, torch.tensor([1, 2, 3])) # 验证模型参数 w 的梯度是否计算正确

调用 loss.backward() 后，PyTorch 会从 loss 开始，沿着计算图回溯，自动计算出每个需要梯度的张量的梯度。

计算反向传播（梯度）所需的浮点运算次数

首先定义的线性模型，并将其扩展为一个简单的两层的线性模型：

bash 复制代码

if torch.cuda.is_available():
    B = 16384  # Number of points
    D = 32768  # Dimension
    K = 8192   # Number of outputs
else:
    B = 1024
    D = 256
    K = 64
device = get_device()
x = torch.ones(B, D, device=device)
w1 = torch.randn(D, D, device=device, requires_grad=True)
w2 = torch.randn(D, K, device=device, requires_grad=True)

Model: x --w1--> h1 --w2--> h2 -> loss
h1 = x @ w1 # 第一层：输入x乘以权重w1得到隐状态h1
h2 = h1 @ w2 # 第二层：h1乘以权重w2得到输出h2
loss = h2.pow(2).mean()  # 计算损失（均方误差）

回顾一下前向传播 FLOPs 的计算，第一层 x @ w1 需要 2 * B * D * D FLOPs；第二层 h1 @ w2 需要 2 * B * D * K FLOPs，总计为 num_forward_flops = (2 * B * D * D) + (2 * B * D * K)

接下来让我们计算反向传播 FLOPs，调用 loss.backward() 后，PyTorch 需要计算以下四个梯度：

h1.grad = d(loss) / d(h1) （中间激活值的梯度）

h2.grad = d(loss) / d(h2) （最后一层输出的梯度）

w1.grad = d(loss) / d(w1) （第一层权重的梯度）

w2.grad = d(loss) / d(w2) （第二层权重的梯度）

重点分析 w2.grad 的计算，根据链式法则（ h 2 = h 1 W 2 h_{2} = h_{1} W_{2} h2=h1W2，故 ∂ L ∂ W 2 = h 1 T ∂ L ∂ h 2 \frac{\partial L}{\partial W_2} = h_1^{\mathsf{T}}\frac{\partial L}{\partial h_2} ∂W2∂L=h1T∂h2∂L）：

∂ L ∂ W 2 $j , k$ = ∑ i h 1 $i , j$ ∂ L ∂ h 2 $i , k$ \frac{\partial L}{\partial W_{2} $j,k$ } = \sum_{i} h_{1} $i,j$ \,\frac{\partial L}{\partial h_{2} $i,k$ } ∂W2 $j,k$ ∂L=i∑h1 $i,j$ ∂h2 $i,k$ ∂L

其中 i i i 为 batch 维下标（与上文 h 1 h_1 h1 的行对应）。与 PyTorch 一致地记 h2.grad[i,k] = ∂L/∂h2[i,k]，即 w2.grad[j,k] = ∑i h1[i,j] * h2.grad[i,k]，等价于矩阵形式 w2.grad = h1.T @ h2.grad。

这本质上是一个矩阵乘法： h 1 h_{1} h1 的形状是 ( B , D ) (B, D) (B,D)，h2.grad 的形状是 ( B , K ) (B, K) (B,K)， h 1 T @ h2.grad h_{1}^{T} @ \text{h2.grad} h1T@h2.grad 的形状是 ( D , K ) (D, K) (D,K)，正好是 w2.grad 的形状。因此，计算 w2.grad 的 FLOPs 为：2 × B × D × K。

为了将梯度继续传回第一层（随后才能算 w1.grad），需要先求 ∂ L / ∂ h 1 \partial L/\partial h_1 ∂L/∂h1。由 h 2 = h 1 W 2 h_2 = h_1 W_2 h2=h1W2 对 h 1 h_1 h1 求导得 ∂ L ∂ h 1 = ∂ L ∂ h 2 W 2 T \frac{\partial L}{\partial h_1} = \frac{\partial L}{\partial h_2} W_2^{\mathsf{T}} ∂h1∂L=∂h2∂LW2T：

∂ L ∂ h 1 $i , j$ = ∑ k ∂ L ∂ h 2 $i , k$ W 2 $j , k$ \frac{\partial L}{\partial h_{1} $i,j$ } = \sum_{k} \frac{\partial L}{\partial h_{2} $i,k$ }\, W_{2} $j,k$ ∂h1 $i,j$ ∂L=k∑∂h2 $i,k$ ∂LW2 $j,k$

其中 k k k 为输出维下标。记 PyTorch 的 h1.grad、h2.grad 即上式左、右端的 ∂ L / ∂ h 1 \partial L/\partial h_1 ∂L/∂h1、 ∂ L / ∂ h 2 \partial L/\partial h_2 ∂L/∂h2，则元素形式为 h1.grad[i,j] = ∑k h2.grad[i,k] * w2[j,k]，矩阵形式为 h1.grad = h2.grad @ w2.T。

w2 的形状是 ( D , K ) (D, K) (D,K)，h2.grad 的形状是 ( B , K ) (B, K) (B,K)，故 h2.grad @ w2.T 为 ( B , D ) (B, D) (B,D)，与 h1.grad 一致。计算 h1.grad 的 FLOPs 也是：2 × B × D × K。

同理，第一层 h 1 = x W 1 h_1 = x\, W_1 h1=xW1 对 W 1 W_1 W1 的梯度为 ∂ L ∂ W 1 = x T ∂ L ∂ h 1 \frac{\partial L}{\partial W_1} = x^{\mathsf{T}}\frac{\partial L}{\partial h_1} ∂W1∂L=xT∂h1∂L：

∂ L ∂ W 1 $j , k$ = ∑ i x $i , j$ ∂ L ∂ h 1 $i , k$ \frac{\partial L}{\partial W_{1} $j,k$ } = \sum_{i} x $i,j$ \,\frac{\partial L}{\partial h_{1} $i,k$ } ∂W1 $j,k$ ∂L=i∑x $i,j$ ∂h1 $i,k$ ∂L

即 w1.grad[j,k] = ∑i x[i,j] * h1.grad[i,k]，矩阵形式 w1.grad = x.T @ h1.grad。其 FLOPs 为：2 × B × D × D。

将这个过程可视化：）

模型构建与训练基础

参数初始化 (Initialization)：

模型参数的存储：在 PyTorch 中，可训练的模型参数被封装为 nn.Parameter 对象：

bash 复制代码

w = nn.Parameter(torch.randn(input_dim, output_dim))

nn.Parameter 是 torch.Tensor 的子类，因此它"表现得像一个张量"，可以进行所有张量操作。它有一个 .data 属性，用于访问其底层的 torch.Tensor 数据。
参数初始化的重要性及常用方法

如果你直接用 torch.randn 进行标准高斯分布初始化参数，随着层数变深，数值会变得非常大（爆炸）或非常小（消失）。

bash 复制代码

x = nn.Parameter(torch.randn(input_dim)) # 输入向量
output = x @ w # 输出向量

当 input_dim = 16384 时，输出值的大小约为 18.9，这是一个非常大的数值。这种大数值会逐层放大，导致梯度爆炸（gradient explosion），使训练过程变得极不稳定，甚至无法收敛。

为了克服这个问题，我们需要一种对输入维度 input_dim 不敏感的初始化方法。解决方法是使用 Xavier初始化(paper stackexchange社区)。通过除以输入维度来缩放权重，保持数值稳定性。

bash 复制代码

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

经过缩放后，输出 output 的每个元素的值变得稳定在一个较小的范围内，不再随 input_dim 增长。即使使用了 Xavier 初始化，由于正态分布的尾部是无界的，仍然存在产生极端值（outliers）的可能性。解决方案是使用截断正态分布（truncated normal distribution），将生成的随机数限制在一个合理的范围内（如 $-3, 3$ ）。

bash 复制代码

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

使用 pytorch 自定义模型

这一节是关于如何在 PyTorch 中从零开始构建一个自定义的深度线性模型（deep linear model）。它展示了使用 nn.Parameter 来定义可学习参数，并通过组合这些参数来创建一个简单的神经网络。

定义模型结构。这里定义了一个名为 Cruncher 的自定义模型类，它是一个"深度线性模型"，包含 num_layers 个隐藏层和一个输出层。

Cruncher 类将多个 Linear 层组合在一起，形成一个更深的网络。在 init 方法中，使用 nn.ModuleList 创建一个包含 num_layers 个 Linear 层的列表。每个 Linear 层的输入和输出维度都是 dim，这意味着它们是"恒等"变换，但内部的权重是可学习的。

bash 复制代码

class Cruncher(nn.Module):
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([
            Linear(dim, dim)
            for i in range(num_layers)
        ])
        self.final = Linear(dim, 1)  # 创建一个最终的 Linear 层，它将 dim 维的特征映射到 1 维的标量输出
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Apply linear layers
        B, D = x.size()
        for layer in self.layers:
            x = layer(x)
        # Apply final head
        x = self.final(x)
        assert x.size() == torch.Size([B, 1])
        # Remove the last dimension
        x = x.squeeze(-1) # 移除最后一个维度（即大小为1的那个维度），使得最终输出是一个一维张量 (B,)，更符合我们对"预测值"的直观理解（例如，对每个样本输出一个分数）
        assert x.size() == torch.Size([B])
        return x

Linear 类实现了神经网络中最基本的构建块------线性层（也称为全连接层或密集层）。

bash 复制代码

class Linear(nn.Module):
    """Simple linear layer."""
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight

检查模型参数。使用 model.state_dict() 将返回一个字典，包含了模型中所有可学习参数（nn.Parameter）的名称和值

bash 复制代码

param_sizes = [
    (name, param.numel())  # param.numel()：返回该参数张量中的元素总数
    for name, param in model.state_dict().items()
]
assert param_sizes == [
    ("layers.0.weight", D * D),
    ("layers.1.weight", D * D),
    ("final.weight", D),
]

将模型移至 GPU。在实际训练中，为了利用 GPU 的并行计算能力，需要将模型和数据移动到 GPU 上。get_device() 函数会自动选择可用的 GPU 或 CPU。

bash 复制代码

device = get_device()
model = model.to(device)

在数据上运行模型。创建一个批次大小为 B=8 的随机输入数据 x，调用 model(x) 执行前向传播，得到输出 y。验证输出 y 的形状为 (B,)，即每个样本对应一个标量输出。

bash 复制代码

B = 8 # Batch size
x = torch.randn(B, D, device=device)
y = model(x)
assert y.size() == torch.Size([B])

如何管理随机性（Randomness）以确保结果的可重现性

随机性在深度学习的许多地方都会出现：

参数初始化：模型权重通常从随机分布（如正态分布）中采样。
Dropout：训练时随机"关闭"一部分神经元。
数据排序：数据加载器通常会打乱数据顺序。
其他：如数据增强、优化器中的动量等。

这些随机性会导致每次运行代码时得到不同的结果，这在调试和比较不同模型时是一个巨大的障碍。为了确保实验的可靠性和可比性，我们需要让程序的行为是确定性的（Deterministic）。
设置随机种子的方法

有三个主要的库需要分别设置随机种子，通常在程序开始时一次性全部设置。
PyTorch

下面是设置 PyTorch 自身随机数生成器的种子，影响所有使用 torch.randn、torch.randint 等函数的操作。

bash 复制代码

seed = 0
torch.manual_seed(seed)

NumPy

许多数据预处理操作（如数据加载、划分）都依赖于 NumPy 的随机函数，因此也需要单独设置其种子。

bash 复制代码

import numpy as np
np.random.seed(seed)

Python 标准库 (random)

Python 内置的 random 模块也常用于数据打乱等操作，同样需要设置种子。

bash 复制代码

import random
random.seed(seed)

数据加载

在语言建模中，输入数据通常是经过分词器（tokenizer）处理后的整数序列。例如，句子 "Hello world" 可能被编码为 $1, 2$ 。为了方便处理，这些整数序列通常会被保存为 NumPy 数组文件（.npy 格式）。

bash 复制代码

orig_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=np.int32)
orig_data.tofile("data.npy") # 将数组保存到文件

对于像 LLaMA 这样高达 2.8TB 的超大数据集，不可能将其全部加载到内存中。解决方案是使用 numpy.memmap 创建一个"内存映射"对象。

bash 复制代码

data = np.memmap("data.npy", dtype=np.int32) # 创建内存映射
assert np.array_equal(data, orig_data) # 验证加载的数据是否正确

data 对象本身不包含数据，而是一个指向磁盘上文件的"指针"。当你访问 data $0$ 或 data $1:10$ 时，系统才会按需将这部分数据从磁盘加载到内存中。这种方式极大地节省了内存占用，允许你在内存有限的情况下处理远超内存大小的数据集。

在实际训练中，我们不会直接操作整个数据集，而是通过一个"数据加载器"来生成训练所需的批次。

bash 复制代码

B = 2  # 批次大小，即一次训练处理多少个样本
L = 4  # 序列长度，即每个样本包含多少个词元（token）
x = get_batch(data, batch_size=B, sequence_length=L, device=get_device())
assert x.size() == torch.Size([B, L]) # 验证输出张量的形状

调用 get_batch 函数会从 data 中随机采样 B 个起始位置，然后截取每个位置后长度为 L 的序列，最终返回一个形状为 (B, L) 的张量 x。接下来，我们详细看下get_batch 函数的内部处理逻辑：

bash 复制代码

def get_batch(data: np.array, batch_size: int, sequence_length: int, device: str) -> torch.Tensor:
    # 随机采样起始位置
    start_indices = torch.randint(len(data) - sequence_length, (batch_size,)) # 使用 torch.randint 在 [0, len(data) - sequence_length] 范围内随机生成 batch_size 个起始索引。这样可以确保每个序列都能完整地从数据中截取出来，不会越界
    assert start_indices.size() == torch.Size([batch_size]) # 断言验证了生成的索引数量正确

    # 根据起始索引提取数据
    x = torch.tensor([data[start:start + sequence_length] for start in start_indices]) # 遍历 start_indices 中的每一个起始位置 start，从 data 中切片出长度为 sequence_length 的子序列。最后将这些子序列转换成一个 PyTorch 张量 x，其形状为 (batch_size, sequence_length)
    assert x.size() == torch.Size([batch_size, sequence_length])

    # 固定内存（Pinned Memory）优化
    if torch.cuda.is_available():
        x = x.pin_memory()

    # 异步数据传输
    x = x.to(device, non_blocking=True)

这个函数的核心功能是从一个巨大的一维数组 data 中随机抽取 batch_size 个序列，每个序列长度为 sequence_length。
参数说明：

data: 输入数据，是一个 NumPy 数组，通常是经过分词器处理后的整数序列。 batch_size: 要采样的批次大小。
sequence_length: 每个序列的长度。 device: 目标设备（如 "cuda:0" 或 "cpu"）。
固定内存（Pinned Memory）优化：默认情况下 CPU 上的张量存储在"分页内存"（paged memory）中。当需要将数据从CPU 传输到 GPU 时，操作系统必须先将数据复制到一个"非分页"的内存区域，然后再通过 PCIe 总线发送给 GPU。这个过程是同步的，会阻塞当前进程。通过调用 .pin_memory()，我们可以将 CPU张量显式地标记为"固定"，这意味着它的物理内存地址是固定的，不会被操作系统换出到磁盘。这使得 GPU驱动程序可以直接访问这块内存，无需额外的拷贝步骤。
异步传输：设置 non_blocking=True 参数，这告诉 PyTorch 数据传输可以在后台异步进行，不会阻塞当前的 Python线程通过结合"固定内存"和"异步传输"，我们可以实现以下高效的流水线操作：在 GPU 上处理当前批次的数据。在 CPU 上同时加载下一个批次的数据（例如，从硬盘读取或从内存映射文件中加载）。这种并行化可以显著减少 GPU 的空闲等待时间，从而大幅提升整体训练吞吐量
关于数据加载更多优化的知识，可以参考https://developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc/、Tricks to Speed Up Data Loading with PyTorch

优化器 (Optimizer)

常用优化器介绍

SGD (随机梯度下降)：最基础的优化器，直接用学习率乘以梯度更新参数。
Momentum (动量法)：在 SGD 基础上增加了一个"动量"项，即梯度的指数移动平均值，有助于加速收敛并减少震荡。
AdaGrad：根据历史梯度的平方值来调整每个参数的学习率，对稀疏特征更友好。
RMSProp：对 AdaGrad 的改进，使用梯度平方的指数加权平均代替简单累加，避免学习率过早衰减。
Adam：融合了 RMSProp 和 Momentum 的思想，是目前最流行的优化器。

我们以 AdaGrad 为例（虽然现在常用 AdamW，但原理类似）。优化器不仅要更新参数，还要记住每个参数的历史梯度信息（状态）。

bash 复制代码

class AdaGrad(torch.optim.Optimizer):
    def step(self):
        for group in self.param_groups:
            for p in group['params']:
                grad = p.grad.data
                # 获取状态（梯度平方和）
                state = self.state[p]
                if 'sum_squared_grad' not in state:
                    state['sum_squared_grad'] = torch.zeros_like(p.data)
                
                # 更新状态：累加梯度的平方
                state['sum_squared_grad'] += grad ** 2
                
                # 更新参数：除以 根号下(状态)
                std = state['sum_squared_grad'].sqrt() + 1e-10
                p.data -= group['lr'] * grad / std

优化器的使用

在 PyTorch 中实例化和使用一个 AdaGrad 优化器：

bash 复制代码

# 实例化优化器
optimizer = AdaGrad(model.parameters(), lr=0.01) # model.parameters()：将模型中所有可学习的参数传递给优化器
# 计算梯度
loss.backward() # 计算损失函数对所有参数的梯度
# 执行一步更新
optimizer.step() # 根据梯度和优化器内部状态，更新模型参数

释放内存（可选）

在每次迭代开始前，需要清空上一次计算得到的梯度，否则梯度会累积。优化器在调用 zero_grad 函数清空梯度时，设置 set_to_none=True 是一种更高效的内存管理方式，它会将梯度指针设为 None，而不是将其置零，可以节省内存。

bash 复制代码

optimizer.zero_grad(set_to_none=True)

优化器的底层实现

这部分展示了 SGD 和 AdaGrad 两个优化器的自定义实现，让你深入理解其内部工作原理。下面是最基础的梯度下降算法：

bash 复制代码

class SGD(torch.optim.Optimizer):
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super(SGD, self).__init__(params, dict(lr=lr)) # 调用父类 Optimizer 的初始化方法，并将学习率 lr 存储在参数组的字典中

    def step(self):
        for group in self.param_groups: # 遍历所有的参数组（通常只有一个组）
            lr = group["lr"]
            for p in group["params"]: # 遍历该组中的每一个参数 p
                grad = p.grad.data  # 获取该参数的梯度 grad
                p.data -= lr * grad # 更新参数

AdaGrad 的核心思想是"自适应学习率"，它根据每个参数的历史梯度信息来动态调整其学习率，从而在处理稀疏数据时表现更好。

bash 复制代码

class AdaGrad(torch.optim.Optimizer):
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super(AdaGrad, self).__init__(params, dict(lr=lr))

    def step(self):
        for group in self.param_groups:
            lr = group["lr"]
            for p in group["params"]:
                state = self.state[p] # 获取优化器状态
                grad = p.grad.data

                g2 = state.get("g2", torch.zeros_like(grad)) # 获取或初始化梯度平方和
                g2 += torch.square(grad) # 累加当前梯度的平方
                state["g2"] = g2 # 更新状态

                p.data -= lr * grad / torch.sqrt(g2 + 1e-5) # 更新参数

state：这是一个字典，用于存储每个参数的优化器状态。对于 AdaGrad，状态是 g2，即历史梯度平方的累加和。
state.get("g2", torch.zeros_like(grad))：如果 g2
不存在，则初始化为全零张量；如果存在，则获取其值。 g2 += torch.square(grad)：将当前梯度的平方累加到 g2 中。
p.data -= lr * grad / torch.sqrt(g2 + 1e-5)：更新参数。这里的除法操作使得学习率会随着 g2
的增大而减小，从而对频繁更新的参数进行惩罚，对稀疏更新的参数给予更大的步长。1e-5 是为了防止除零错误。

资源核算

这部分旨在教你如何估算一个模型在训练过程中所需的内存和计算资源。
内存占用分析

对于一个深度线性模型，总内存需求由四部分组成：

参数 (Parameters)：模型中所有可学习权重的数量。
激活值 (Activations)：前向传播过程中产生的中间结果，需要保存下来用于反向传播。
梯度 (Gradients)：反向传播计算出的梯度，其数量与参数相同。
优化器状态 (Optimizer States)：优化器维护的额外状态信息（如 AdaGrad 的 g2），其数量也与参数相同。
假设所有数据都使用 float32 格式（每个元素占 4 字节），则总内存为：

bash 复制代码

total_memory = 4 * (num_parameters + num_activations + num_gradients + num_optimizer_states)

参数数量：(D * D * num_layers) + D
激活值数量：B * D * num_layers （每个批次、每个样本、每层都需要保存激活值）
梯度数量：num_parameters
优化器状态数量：num_parameters
计算量（FLOPs）分析
根据前面课程的内容，我们知道训练一次（前向+反向）的总 FLOPs 约等于 6 × (数据量) × (参数量)。因此，对于一个训练步骤，其计算量为：

bash 复制代码

flops = 6 * B * num_parameters

训练循环：整合所有组件

这部分是一个完整的、从零开始的深度学习训练循环（Training Loop）。它们以一个简单的线性回归任务为例，清晰地演示了如何将之前学到的所有组件------数据生成、模型构建、优化器、前向传播、反向传播和参数更新------整合在一起，形成一个可运行的训练流程。

bash 复制代码

def train_loop():
    # 数据生成
    D = 16
    true_w = torch.arange(D, dtype=torch.float32, device=get_device()) # 创建一个真实的权重向量 [0, 1, 2, ..., 15]

    # 数据加载器，用于生成训练批次
    def get_batch(B: int) -> tuple[torch.Tensor, torch.Tensor]:
        x = torch.randn(B, D).to(get_device()) # 从标准正态分布中随机采样 B 个 D 维的输入样本
        true_y = x @ true_w # 根据真实的权重 true_w 计算出对应的"真实标签"，即 y = x @ w_true
        return (x, true_y) # 作为模型的输入和目标输出

    # 执行训练
    train("simple", get_batch, D=D, num_layers=0, B=4, num_train_steps=10, lr=0.01)
    # 超参数调优（hyperparameter tuning）
    train("simple", get_batch, D=D, num_layers=0, B=4, num_train_steps=10, lr=0.1) # 通过改变学习率 lr 从 0.01 到 0.1，观察不同学习率对训练效果的影响

train() 函数详细实现了深度学习的标准训练步骤

bash 复制代码

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=lr)

    # 主训练循环
    for t in range(num_train_steps): # 循环执行 num_train_steps 次，每次迭代称为一个"训练步"
        # 获取数据
        x, y = get_batch(B=B) # 获取当前批次的输入 x 和目标输出 y

        # 前向传播（计算损失）
        pred_y = model(x)
        loss = F.mse_loss(pred_y, y) # 计算预测值 pred_y 和真实值 y 之间的均方误差（Mean Squared Error），作为损失值 loss

        # 反向传播（计算梯度）
        loss.backward() # 调用 loss.backward()，触发自动微分机制，计算损失函数对所有模型参数的梯度，并将这些梯度存储在 .grad 属性中
        
        # 更新参数
        optimizer.step() # 根据计算出的梯度和优化器的更新规则（这里是 SGD），更新模型参数
        optimizer.zero_grad(set_to_none=True) # 清空上一步计算得到的梯度，为下一次迭代做准备

train() 函数参数说明：

"simple": 实验名称。
get_batch: 数据生成函数。
D=16: 输入维度。
num_layers=0: 模型层数为0，意味着这是一个单层线性模型。
B=4: 批次大小。
num_train_steps=10:训练步数。
lr=0.01: 学习率。

检查点（Checkpointing）

大型语言模型的训练耗时极长，过程中几乎不可避免地会发生崩溃或中断，我们需要一种机制来保存进度，以便在中断后能够从中断点恢复，而不是从头开始。

为什么需要检查点？
训练时间长：训练一个大型语言模型可能需要数天、数周甚至数月。
系统不稳定：硬件故障、软件错误、电源中断、人为操作失误等都可能导致训练进程意外终止。
避免损失：如果没有任何备份，一次崩溃就意味着之前所有的计算和时间都付诸东流。因此，定期保存"检查点"是保证项目顺利进行的基本保障。
检查点包含哪些内容？
为了能够完整地恢复训练，检查点必须包含所有必要的状态信息。主要包括两个核心部分：模型参数：这是模型的核心，包含了所有可学习权重（nn.Parameter）的当前值。通过 model.state_dict() 可以获取一个字典，其中键是参数名称，值是对应的张量数据。优化器状态 (optimizer.state_dict())：这是很多人容易忽略但极其关键的部分。优化器（如 Adam、AdaGrad）不仅存储了当前的学习率，还维护着一些内部状态变量。例如：Adam：存储了动量（momentum）和方差（variance）的移动平均值。AdaGrad：存储了历史梯度平方的累加和。
如果只保存了模型参数而没有保存优化器状态，那么在恢复训练时，优化器会从零开始，这会导致训练过程不连续，性能下降，甚至可能无法收敛。
如何保存和加载检查点？
保存检查点：

bash 复制代码

# 1. 创建一个字典，包含模型和优化器的状态
checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}

# 2. 使用 torch.save 将字典序列化并保存到磁盘文件
torch.save(checkpoint, "model_checkpoint.pt") # 文件名通常使用 .pt 或 .pth 后缀

加载检查点

bash 复制代码

# 1. 从磁盘文件加载保存的字典
loaded_checkpoint = torch.load("model_checkpoint.pt")
# 2. （后续步骤）将状态加载回模型和优化器
# model.load_state_dict(loaded_checkpoint["model"])
# optimizer.load_state_dict(loaded_checkpoint["optimizer"])

参考

pytorch 与资源核算