基础第 2 讲 / 共 17 讲

第 2 讲：PyTorch 与计算资源核算

1. PyTorch 核心概念

PyTorch 是当前深度学习研究和工程中最主流的框架之一。理解它的核心抽象—— Tensor、autograd、Module——是构建 LLM 的基础。

1.1 Tensor

Tensor 是 PyTorch 中的基本数据结构，本质上是一个多维数组。它与 NumPy 的 ndarray 类似，但有两个关键区别：

GPU 加速：Tensor 可以驻留在 GPU 显存中，利用 GPU 的并行计算能力
自动微分：当 requires_grad=True 时，PyTorch 会追踪对该 Tensor 的所有操作，自动构建计算图

import torch

# Create tensors
x = torch.randn(3, 4, device='cuda', requires_grad=True)
W = torch.randn(4, 5, device='cuda', requires_grad=True)

# Matrix multiplication
y = x @ W  # shape: (3, 5)

# y.grad_fn records the operation (MatMul) and its inputs
print(y.grad_fn)  #

在 LLM 训练中，Tensor 的 dtype（数据类型）至关重要：

dtype	每个元素字节数	用途
`float32` (FP32)	4 bytes	默认精度，用于参数更新
`float16` (FP16)	2 bytes	混合精度训练的前向/反向计算
`bfloat16` (BF16)	2 bytes	LLM 训练首选，范围更大
`float8` (FP8)	1 byte	H100+ 上的下一代训练格式

BF16 vs FP16

BFloat16 保留了与 FP32 相同的指数位数（8 位），因此它的数值范围与 FP32 一致，不容易溢出。FP16 的指数位数只有 5 位，范围较小，在 LLM 训练中更容易遇到 overflow/underflow 问题。这就是为什么现代 LLM 训练普遍使用 BF16 而非 FP16。

1.2 Autograd：自动微分

PyTorch 的 autograd 系统通过动态计算图（dynamic computational graph）实现自动微分。每次前向传播时，PyTorch 会自动构建一个有向无环图（DAG），记录所有操作及其输入输出关系。反向传播时，沿着这个图从输出到输入，利用链式法则（chain rule）逐层计算梯度。

# Forward pass: build computational graph
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x + 1  # y = x^2 + 3x + 1

# Backward pass: compute gradients
y.backward()

# dy/dx = 2x + 3 = 2*2 + 3 = 7
print(x.grad)  # tensor([7.])

在 LLM 训练中，计算图的结构如下：

Forward：输入 token IDs → Embedding → Transformer Layers → Logits → Loss
Backward：Loss → 反向遍历计算图 → 计算每个参数的梯度
Update：根据梯度更新参数（由 Optimizer 完成）

内存陷阱

计算图会保留前向传播中所有中间结果（激活值），因为反向传播需要它们来计算梯度。这就是为什么训练比推理消耗更多内存——不仅要存模型参数，还要存所有层的激活值。

1.3 nn.Module

nn.Module 是 PyTorch 中组织模型的基本单元。每个 Module 可以包含参数（nn.Parameter）和子 Module，形成树状层次结构。

import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, n_heads)
        self.ffn = FeedForward(d_model)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)

    def forward(self, x):
        # Pre-Norm architecture
        x = x + self.attention(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

# Count parameters
model = TransformerBlock(768, 12)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")

2. 计算图与反向传播深入

理解计算图对于优化训练性能至关重要。以下是几个关键概念：

2.1 梯度累积（Gradient Accumulation）

当 batch size 太大，无法放入 GPU 显存时，可以将一个大 batch 拆分为多个 micro-batch，分别前向和反向传播，累积梯度后再执行一次参数更新。

optimizer.zero_grad()

# Accumulate gradients over 4 micro-batches
accumulation_steps = 4
for i in range(accumulation_steps):
    micro_batch = get_micro_batch(i)
    loss = model(micro_batch) / accumulation_steps  # Scale loss
    loss.backward()  # Gradients accumulate in .grad

# Single optimizer step for the effective batch
optimizer.step()

2.2 混合精度训练

混合精度训练（Mixed Precision Training）使用低精度（FP16/BF16）进行前向和反向计算，但保留 FP32 的主参数副本（master weights）用于参数更新。这样既减少了显存占用和加速了计算，又不损失训练稳定性。

from torch.amp import autocast, GradScaler

scaler = GradScaler()  # Only needed for FP16, not BF16

for batch in dataloader:
    optimizer.zero_grad()

    # Forward pass in BF16
    with autocast(device_type='cuda', dtype=torch.bfloat16):
        logits = model(batch['input_ids'])
        loss = cross_entropy_loss(logits, batch['labels'])

    # Backward pass (still in BF16 where possible)
    loss.backward()

    # Optimizer step in FP32 (master weights)
    optimizer.step()

3. FLOPs 计算

FLOPs（Floating Point Operations）是衡量计算量的核心指标。准确估算 FLOPs 是计算资源规划和效率优化的基础。

3.1 矩阵乘法的 FLOPs

矩阵乘法是神经网络中最核心的运算。对于矩阵乘法 C = A × B：

A ∈ ℝ^m×k, B ∈ ℝ^k×n → C ∈ ℝ^m×n
FLOPs = 2 × m × k × n

这里的 2 是因为每计算一个输出元素需要 k 次乘法和 k-1 次加法，近似为 2k 次浮点运算。

3.2 Transformer 各层的 FLOPs

一个 Transformer layer 的主要计算量来自以下几个矩阵乘法。假设输入序列长度为 T，模型维度为 d， batch size 为 B：

操作	矩阵维度	FLOPs（前向）
QKV Projection	(BT, d) × (d, 3d)	2 × BT × d × 3d = 6BTd²
Attention Score	(BT, d) × (d, BT)	2 × BT² × d（per head, total）
Attention × V	(BT, BT) × (BT, d)	2 × BT² × d
Output Projection	(BT, d) × (d, d)	2BTd²
FFN (up + down)	(BT, d) × (d, 4d) × 2	16BTd²（if d_ff = 4d）

简化估算公式（忽略 attention score 的计算，因为当 T << d 时它占比较小）：

FLOPs per layer ≈ 24 × B × T × d² (forward only)
FLOPs per layer ≈ 72 × B × T × d² (forward + backward, backward ≈ 2× forward)

对于整个模型（L 层 + embedding + output projection）：

Total FLOPs ≈ 6 × N × B × T（前向+反向，其中 N 为参数量）

6N 的直觉

这个著名的 "6N" 近似来自：前向传播约 2N FLOPs（每个参数参与一次乘法和一次加法），反向传播约 4N FLOPs（计算梯度需要约 2 倍前向的计算量）。因此总计约 6N FLOPs per token。

4. 内存使用分析

GPU 显存是 LLM 训练中最紧张的资源之一。理解显存的组成对于选择合适的模型大小和训练策略至关重要。

4.1 显存组成

训练时的 GPU 显存使用可以分为四大部分：

组成部分	每参数字节数（BF16 + Adam）	说明
Model Parameters	2 bytes (BF16)	模型权重本身
Gradients	2 bytes (BF16)	每个参数的梯度
Optimizer States (Adam)	4 + 4 + 4 = 12 bytes	FP32 master weights + first moment + second moment
Activations	变化	取决于序列长度、batch size、是否使用 activation checkpointing

因此，对于一个 N 参数的模型（使用 BF16 混合精度 + Adam optimizer）：

固定显存 ≈ (2 + 2 + 12) × N = 16N bytes
例如：7B 参数模型需要 16 × 7 × 10⁹ = 112 GB（仅参数和优化器状态）

关键瓶颈

Adam optimizer 的状态占据了大部分显存！这也是为什么很多工作在研究更省内存的 optimizer（如 AdaFactor, 8-bit Adam）。使用 SGD 只需要存参数和梯度（4N bytes），但收敛性通常不如 Adam。

4.2 Activation 内存

激活值（Activations）是前向传播中各层的中间输出，需要保留到反向传播时使用。它的大小取决于：

Batch size (B)：线性增长
Sequence length (T)：线性增长（attention 矩阵是 T² 但可以通过 Flash Attention 优化）
Model dimension (d)：线性增长
Number of layers (L)：线性增长

# Approximate activation memory per Transformer layer
# (without activation checkpointing)
activation_per_layer = (
    B * T * d * 2 +     # Input to attention (BF16)
    B * T * 3 * d * 2 + # QKV projections
    B * T * T * 2 +     # Attention weights (per head, simplified)
    B * T * d * 2 +     # Attention output
    B * T * 4 * d * 2 + # FFN intermediate
    B * T * d * 2       # FFN output
)  # Rough estimate in bytes

4.3 Activation Checkpointing

Activation Checkpointing（也称 Gradient Checkpointing）是一种用时间换空间的技术：不保存所有层的激活值，而是只保存某些"检查点"层的激活值。反向传播到某一层时，从最近的检查点重新执行前向传播来恢复所需的激活值。

# PyTorch activation checkpointing
from torch.utils.checkpoint import checkpoint

class TransformerWithCheckpointing(nn.Module):
    def forward(self, x):
        for layer in self.layers:
            # Recompute this layer's forward during backward
            x = checkpoint(layer, x, use_reentrant=False)
        return x

这可以将 activation 内存从 O(L) 降低到 O(√L)（如果每隔 √L 层设一个检查点），但会增加约 33% 的计算量（因为需要重新计算前向传播）。

5. GPU 利用率与 MFU

MFU（Model FLOPs Utilization）是衡量训练效率的最重要指标之一。它表示实际用于模型计算的 FLOPs 占 GPU 理论峰值 FLOPs 的比例。

MFU = (Model FLOPs per second) / (GPU Peak FLOPs per second)

注意 MFU 计算的是模型本身的理论 FLOPs，不包括重新计算（如 activation checkpointing 引起的）或通信开销。以下是一些参考数值：

GPU	BF16 Peak FLOPS	典型 MFU
A100 (80GB)	312 TFLOPS	40-55%
H100 (SXM)	989 TFLOPS	30-45%

为什么 MFU 达不到 100%？

MFU 永远不可能达到 100%，因为除了矩阵乘法之外，训练还包括：数据加载、非矩阵运算（LayerNorm, Softmax 等）、GPU 间通信、内存读写延迟、kernel launch overhead 等。好的训练系统的目标是尽可能提高 MFU。

6. 训练时间估算

有了 FLOPs 和 MFU，我们就可以估算训练所需的时间和成本。

6.1 估算公式

Training Time = Total FLOPs / (num_GPUs × GPU_peak_FLOPS × MFU)

其中 Total FLOPs ≈ 6 × N × D（N 为参数量，D 为训练 token 总数）。

6.2 具体例子

# Example: Training a 7B parameter model on 2T tokens
N = 7e9           # 7 billion parameters
D = 2e12          # 2 trillion tokens
total_flops = 6 * N * D  # = 8.4e22 FLOPs

# Hardware: 256 × H100 GPUs
num_gpus = 256
peak_flops = 989e12  # H100 BF16 peak: 989 TFLOPS
mfu = 0.40           # 40% utilization

# Training time
effective_flops_per_sec = num_gpus * peak_flops * mfu
training_seconds = total_flops / effective_flops_per_sec

training_hours = training_seconds / 3600
training_days = training_hours / 24

print(f"Total FLOPs: {total_flops:.2e}")
print(f"Effective FLOPS: {effective_flops_per_sec:.2e}")
print(f"Training time: {training_days:.1f} days")
# Output:
# Total FLOPs: 8.40e+22
# Effective FLOPS: 1.01e+17
# Training time: ~9.6 days

6.3 成本估算

训练成本 = 训练时间 × GPU 数量 × 每 GPU 小时价格。以云上 H100 约 $2-3/GPU-hour 计算：

成本 ≈ 9.6 天 × 24 小时 × 256 GPUs × $2.5/GPU-hour ≈ $147,000

这还只是一次训练的成本。考虑到 hyperparameter tuning、失败重试、evaluation 等，实际成本可能是 2-5 倍。

Token 的"价格"

另一个有用的视角是计算每个训练 token 的成本：$147,000 / 2T tokens ≈ $0.000073/1K tokens。这个数字可以帮助你快速估算不同规模训练的成本。

7. 实践建议

7.1 Profile First

在优化之前，先用 profiling 工具了解时间花在哪里：

# PyTorch Profiler
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    with_stack=True,
) as prof:
    model(input_ids)

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

7.2 常用检查清单

确认 GPU 确实在被使用：nvidia-smi 查看 GPU 利用率
确认使用了混合精度：检查 Tensor dtype 是否为 BF16
确认 data loading 不是瓶颈：使用 num_workers > 0
监控显存使用：torch.cuda.max_memory_allocated()
检查梯度范数：确保训练稳定，无 NaN

8. 本讲总结

PyTorch 的三个核心抽象：Tensor（数据）、autograd（自动微分）、Module（模型组织）
BF16 是 LLM 训练的首选精度，兼顾性能和数值稳定性
矩阵乘法 FLOPs = 2mnk；整个模型每 token 约 6N FLOPs
训练内存 = 参数（2N）+ 梯度（2N）+ 优化器状态（12N）+ 激活值
MFU 是衡量训练效率的核心指标，好的系统可以达到 40-55%
训练时间 = Total FLOPs / (GPUs × Peak FLOPS × MFU)