3.0 PyTorch框架入门

发表于 更新于 分类于 ， ，

PyTorch 是当前大模型训练和推理的事实标准框架。本文聚焦于 AI Infra 工程师需要掌握的核心能力——从 Tensor 操作到完整训练循环，再到用 profiler 定位性能瓶颈，为后续的 CUDA 编程、分布式训练和推理优化打下基础。

📑 目录

1. Tensor：PyTorch 的基本数据单元

如果把深度学习比作搭积木，Tensor 就是最基本的积木块。所有的模型参数、输入数据、中间计算结果，在 PyTorch 里都以 Tensor 的形式存在。

正式定义：Tensor（张量）是一个多维数组，可以看作 NumPy ndarray 的 GPU 加速版本，同时支持自动微分。

1.1 创建 Tensor 的常用方式

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch
import numpy as np

a = torch.tensor([1.0, 2.0, 3.0])          # 从 Python 列表创建
zeros = torch.zeros(3, 4)                    # 3x4 全零矩阵
ones = torch.ones(2, 3, 4)                   # 2x3x4 全一张量
rand_normal = torch.randn(3, 4)              # 标准正态分布
seq = torch.arange(0, 10, 2)                 # tensor([0, 2, 4, 6, 8])

x = torch.randn(3, 4)
y = torch.zeros_like(x)                      # 和 x 形状一样的全零 Tensor

t = torch.from_numpy(np.array([1.0, 2.0]))   # 从 NumPy 转换（共享内存）

1.2 形状操作：view, reshape, permute, squeeze, unsqueeze

把 Tensor 想象成一块可以任意捏形的橡皮泥——数据不变，只是换个排列方式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch

x = torch.arange(12)             # 一维，12 个元素
a = x.view(3, 4)                 # 变成 3x4（要求内存连续）
b = x.reshape(3, 4)              # 和 view 类似，但内存不连续时也能用
c = x.view(-1, 4)                # -1 自动推断为 3

# permute：交换维度，Transformer 中常用
t = torch.randn(2, 8, 12, 64)              # (batch, seq_len, heads, dim)
t_permuted = t.permute(0, 2, 1, 3)         # (batch, heads, seq_len, dim)

# squeeze/unsqueeze：去掉或插入大小为 1 的维度
e = torch.randn(1, 3, 1, 4)
print(e.squeeze().shape)                    # torch.Size([3, 4])
f = torch.randn(3, 4)
print(f.unsqueeze(0).shape)                 # torch.Size([1, 3, 4])

1.3 设备管理：CPU 与 GPU 数据搬运

1
2
3
4
5
6
7
8
9
10
11
import torch

print(torch.cuda.is_available())            # GPU 是否可用

x = torch.randn(3, 4)
x_gpu = x.to('cuda')                        # CPU → GPU（推荐写法）
x_gpu = x.cuda()                             # 等价简写
x_cpu = x_gpu.cpu()                          # GPU → CPU

# 直接在 GPU 上创建（避免多余的搬运）
x_gpu = torch.randn(3, 4, device='cuda')

AI Infra 视角：CPU-GPU 数据搬运走 PCIe 总线，带宽远低于 GPU 显存带宽。频繁的 .cpu() 和 .cuda() 调用往往是隐蔽的性能瓶颈，第 6 节会讲如何用 profiler 发现这类问题。

1.4 dtype 与内存：fp32, fp16, bf16

不同的数据精度直接决定显存占用和计算速度：

dtype 每元素字节 数值范围 使用场景
float32 (fp32) 4 大 默认精度，优化器状态
float16 (fp16) 2 小，易溢出 混合精度训练（需 Loss Scaling）
bfloat16 (bf16) 2 与 fp32 相同 混合精度训练（推荐，Ampere+ GPU）

白话理解：fp32 像高清照片，又大又清晰；fp16 像压缩照片，体积减半但偶尔”失真”（溢出）；bf16 是聪明的压缩，体积和 fp16 一样小，但”失真”概率大大降低。

1
2
3
4
5
6
7
8
9
10
import torch

x_fp32 = torch.randn(1000, 1000, dtype=torch.float32)   # 4 bytes/element
x_bf16 = torch.randn(1000, 1000, dtype=torch.bfloat16)  # 2 bytes/element
print(f"fp32: {x_fp32.nelement() * x_fp32.element_size() / 1024:.0f} KB")  # 3906 KB
print(f"bf16: {x_bf16.nelement() * x_bf16.element_size() / 1024:.0f} KB")  # 1953 KB

x = torch.randn(3, 4)                     # 默认 fp32
x_half = x.half()                          # 转 fp16
x_bf16 = x.to(torch.bfloat16)             # 转 bf16

2. 自动微分（autograd）

训练神经网络的核心是”根据损失调整参数”，而调整的依据就是梯度。PyTorch 的 autograd 引擎帮你自动完成这件事。

2.1 什么是计算图

比喻：想象你在做一道菜。你把食材 A 和 B 混合得到 C，再把 C 加热得到 D（成品）。如果 D 味道不对（loss 太大），你需要反推——是 C 的问题？还是 A、B 的比例不对？计算图就是 PyTorch 帮你记下的”食谱”：它记录每一步操作，这样就能从结果反向推导出每种食材对结果的影响（梯度）。

正式定义：计算图（Computational Graph）是一个有向无环图（DAG），节点代表 Tensor，边代表运算操作。PyTorch 在前向传播时动态构建图，反向传播时沿图计算梯度。PyTorch 采用动态计算图（Define-by-Run），支持 if/else、for 循环等 Python 控制流。

2.2 requires_grad 和 backward()

1
2
3
4
5
6
7
import torch

x = torch.tensor([2.0, 3.0], requires_grad=True)  # 追踪所有操作
y = x * 3
z = y.sum()        # z = 3*2 + 3*3 = 15
z.backward()       # 反向传播
print(x.grad)      # tensor([3., 3.])，dz/dx = 3

要点：backward() 只能对标量调用；叶节点才保存梯度；计算图用完即释放。

2.3 梯度累积与清零

白话理解：PyTorch 默认累加梯度，不自动清零。就像一个计数器，每次 backward() 都往上加。标准训练循环里，每个 step 开始前必须清零——否则上一轮梯度会混进来。

这个设计是有意为之的——梯度累积是显存不足时模拟大 batch 的常用技巧。

1
2
3
4
5
6
7
8
9
10
11
import torch

x = torch.tensor([1.0], requires_grad=True)
(x * 2).sum().backward()
print(x.grad)         # tensor([2.])
(x * 3).sum().backward()
print(x.grad)         # tensor([5.]) ← 2+3，梯度被累加！
x.grad.zero_()        # 手动清零
(x * 4).sum().backward()
print(x.grad)         # tensor([4.]) ← 清零后正确
# 实际训练中用 optimizer.zero_grad() 一次性清零所有参数的梯度

代码示例：手动梯度下降学习 y = 2x + 1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import torch

w = torch.tensor([0.0], requires_grad=True)
b = torch.tensor([0.0], requires_grad=True)
x_train = torch.tensor([1.0, 2.0, 3.0, 4.0])
y_train = torch.tensor([3.0, 5.0, 7.0, 9.0])

for epoch in range(100):
    loss = ((w * x_train + b - y_train) ** 2).mean()
    loss.backward()
    with torch.no_grad():
        w -= 0.01 * w.grad
        b -= 0.01 * b.grad
    w.grad.zero_()
    b.grad.zero_()

print(f"学到的模型: y = {w.item():.2f}x + {b.item():.2f}")
# 输出接近 y = 2.00x + 1.00

3. nn.Module：模型的组织方式

比喻：如果 Tensor 是积木块，nn.Module 就是积木的”说明书”——它定义了积木怎么拼接（前向传播），并帮你清点所有零件（参数管理）。

3.1 定义模型：继承 nn.Module，实现 __init__ 和 forward

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

model = SimpleModel(784, 256, 10)
output = model(torch.randn(32, 784))    # 用 model(x)，不要直接调用 forward
print(output.shape)                       # torch.Size([32, 10])

3.2 参数管理：parameters(), named_parameters(), state_dict()

1
2
3
4
5
6
7
8
9
10
11
12
import torch.nn as nn

model = nn.Linear(10, 5)
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")   # weight: [5, 10], bias: [5]

# 统计参数量（AI Infra 中非常常用）
total_params = sum(p.numel() for p in model.parameters())
print(f"参数量: {total_params}")       # 55

# state_dict 用于保存和加载模型
print(model.state_dict().keys())       # odict_keys(['weight', 'bias'])

3.3 常用层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import torch
import torch.nn as nn

# nn.Linear：大模型中用最多的层（Attention Q/K/V 投影、FFN）
linear = nn.Linear(512, 256)

# nn.Embedding：将 token id 映射为稠密向量（大模型第一层）
embedding = nn.Embedding(num_embeddings=50000, embedding_dim=768)
print(embedding(torch.tensor([0, 42])).shape)   # torch.Size([2, 768])

# nn.LayerNorm：Transformer 每个子层后都有，稳定训练
layer_norm = nn.LayerNorm(768)

# nn.Dropout：训练时随机丢弃神经元，model.eval() 后自动关闭
dropout = nn.Dropout(p=0.1)

代码示例：两层 MLP

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import torch
import torch.nn as nn

class TwoLayerMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        return self.fc2(self.dropout(self.relu(self.fc1(x))))

model = TwoLayerMLP(784, 256, 10)
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")  # 203,530
print(model(torch.randn(64, 784)).shape)  # torch.Size([64, 10])

4. 完整训练循环

训练循环是 AI Infra 工程师每天都要接触的代码——即使在分布式场景下，核心循环依然是同一个模式。

4.1 数据加载：Dataset 和 DataLoader

  • Dataset：定义”数据集里有什么”和”怎么取一条数据”
  • DataLoader：定义”怎么分 batch 喂给模型”（batching、shuffling、多进程加载）
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, data, labels):
        self.data, self.labels = data, labels
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

dataset = SimpleDataset(torch.randn(1000, 784), torch.randint(0, 10, (1000,)))
loader = DataLoader(dataset, batch_size=64, shuffle=True,
                    num_workers=4, pin_memory=True, drop_last=True)

AI Infra 视角：num_workers 过小会导致 GPU 等数据（data loading bottleneck）；pin_memory=True 使传输走异步 DMA，避免 CPU 拷贝开销。

4.2 标准训练流程

每个 step 的核心五步：forward → loss → backward → optimizer.step → zero_grad。

4.3 学习率调度

1
2
3
4
5
6
7
import torch

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
# 常用调度器
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
# 每个 epoch 结束后调用 scheduler.step()

4.4 Checkpoint 保存与恢复

1
2
3
4
5
6
7
8
9
10
11
12
import torch

# 保存（模型 + 优化器 + 训练进度，断点恢复需要全部保存）
torch.save({
    'epoch': epoch, 'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(), 'loss': avg_loss,
}, 'checkpoint.pt')

# 加载
ckpt = torch.load('checkpoint.pt', weights_only=False)
model.load_state_dict(ckpt['model_state_dict'])
optimizer.load_state_dict(ckpt['optimizer_state_dict'])

代码示例：完整 MNIST 训练脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

batch_size, lr, num_epochs = 128, 1e-3, 5
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('./data', train=False, transform=transform)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=batch_size, num_workers=2, pin_memory=True)

# 模型
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Flatten(), nn.Linear(784, 256),
                                 nn.ReLU(), nn.Dropout(0.1), nn.Linear(256, 10))
    def forward(self, x):
        return self.net(x)

model = MLP().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

# 训练
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        loss = criterion(model(images), labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        total_loss += loss.item()
    scheduler.step()

    # 验证
    model.eval()
    correct = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            correct += (model(images).argmax(1) == labels).sum().item()
    print(f"Epoch {epoch+1}/{num_epochs} | Loss: {total_loss/len(train_loader):.4f} "
          f"| Acc: {correct/len(test_set)*100:.2f}%")

    # 保存 checkpoint
    torch.save({'epoch': epoch, 'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict()}, f'ckpt_ep{epoch+1}.pt')

5. GPU 训练基础

5.1 将模型和数据搬到 GPU

1
2
3
4
5
6
7
import torch
import torch.nn as nn

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = nn.Linear(100, 10).to(device)
x = torch.randn(32, 100).to(device)     # 模型和数据必须在同一设备上
y = model(x)

5.2 混合精度训练入门

核心思路：前向和反向用低精度（fp16/bf16）加速计算，参数更新用高精度（fp32）保证精度。

白话理解：算草稿用铅笔（低精度，快），写定稿用钢笔（高精度，准）。关键的参数更新步骤用高精度保证不丢信息，其他大量计算用低精度跑得更快。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import torch
import torch.nn as nn

device = torch.device('cuda')
model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# BF16 混合精度（推荐，Ampere+ GPU，不需要 GradScaler）
for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        loss = nn.functional.cross_entropy(model(images), labels)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# FP16 混合精度（需要 GradScaler 防止梯度下溢）
scaler = torch.cuda.amp.GradScaler()
for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        loss = nn.functional.cross_entropy(model(images), labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

5.3 显存查看

1
2
3
4
5
6
7
8
9
10
import torch

print(f"已分配: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"峰值:   {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
print(torch.cuda.memory_summary(abbreviated=True))

# 分析某段代码的显存
torch.cuda.reset_peak_memory_stats()
# ... 运行目标代码 ...
print(f"峰值: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

代码示例：混合精度 vs fp32 显存对比

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import torch
import torch.nn as nn

assert torch.cuda.is_available(), "需要 GPU"
device = torch.device('cuda')
criterion = nn.CrossEntropyLoss()
data = torch.randn(256, 1024, device=device)
labels = torch.randint(0, 10, (256,), device=device)

def make_model():
    return nn.Sequential(nn.Linear(1024, 2048), nn.ReLU(),
                         nn.Linear(2048, 1024), nn.ReLU(),
                         nn.Linear(1024, 10)).to(device)

# fp32 训练
model_fp32, opt_fp32 = make_model(), None
opt_fp32 = torch.optim.AdamW(model_fp32.parameters(), lr=1e-3)
torch.cuda.reset_peak_memory_stats()
criterion(model_fp32(data), labels).backward()
opt_fp32.step(); opt_fp32.zero_grad()
fp32_peak = torch.cuda.max_memory_allocated() / 1024**2

# bf16 混合精度训练
model_bf16, opt_bf16 = make_model(), None
opt_bf16 = torch.optim.AdamW(model_bf16.parameters(), lr=1e-3)
torch.cuda.reset_peak_memory_stats()
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    loss = criterion(model_bf16(data), labels)
loss.backward()
opt_bf16.step(); opt_bf16.zero_grad()
bf16_peak = torch.cuda.max_memory_allocated() / 1024**2

print(f"FP32: {fp32_peak:.1f} MB | BF16: {bf16_peak:.1f} MB | 节省: {(1-bf16_peak/fp32_peak)*100:.1f}%")

6. 性能分析入门

写出能跑的训练代码只是第一步，跑得快才是 AI Infra 的核心追求。torch.profiler 能告诉你每个操作花了多少时间、GPU 利用率如何、哪里在空等数据。

6.1 torch.profiler 基本用法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch
from torch.profiler import profile, record_function, ProfilerActivity

model = torch.nn.Linear(1024, 1024).cuda()
x = torch.randn(64, 1024).cuda()

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True, profile_memory=True) as prof:
    with record_function("forward_pass"):
        output = model(x)
    with record_function("backward_pass"):
        output.sum().backward()

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

6.2 如何读懂 profiler 输出

列名 含义 关注点
Name 操作名称 找到耗时最多的操作
CPU total CPU 端总耗时 远大于 CUDA total 说明存在 CPU 瓶颈
CUDA total GPU 端总耗时 核心指标
# of Calls 调用次数 高频小操作可能累积成瓶颈

常见性能问题模式：

  1. CPU 时间远大于 CUDA 时间：GPU 在等 CPU（数据预处理、Python 开销）
  2. 大量小 CUDA kernel：launch overhead 累积，考虑 torch.compile 或算子融合
  3. 频繁 CPU-GPU 同步：.item()、print(tensor) 会触发同步，阻塞 GPU 流水线

6.3 常见性能问题

数据加载瓶颈：增加 num_workers，使用 prefetch_factor 预取数据。

CPU-GPU 数据搬运：避免循环中反复 .cuda()，一次性搬运到 GPU。

隐式同步：loss.item() 触发 CPU-GPU 同步，应每 N 步才记录一次。

1
2
3
4
5
6
7
8
9
10
# 不好：每步都同步
for batch in dataloader:
    loss = train_step(batch)
    print(f"loss: {loss.item()}")       # 每步同步！

# 好：每 100 步记录一次
for i, batch in enumerate(dataloader):
    loss = train_step(batch)
    if i % 100 == 0:
        print(f"step {i}, loss: {loss.item()}")

代码示例：用 profiler 分析一个训练 step

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import torch
import torch.nn as nn
from torch.profiler import profile, record_function, ProfilerActivity

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = nn.Sequential(nn.Linear(1024, 4096), nn.ReLU(),
                      nn.Linear(4096, 4096), nn.ReLU(),
                      nn.Linear(4096, 10)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
data = torch.randn(128, 1024, device=device)
labels = torch.randint(0, 10, (128,), device=device)

# warmup（避免首次调用的初始化开销）
for _ in range(3):
    criterion(model(data), labels).backward()
    optimizer.step(); optimizer.zero_grad()

# profiler 分析
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True, profile_memory=True) as prof:
    with record_function("data_to_gpu"):
        d = torch.randn(128, 1024).to(device)
        l = torch.randint(0, 10, (128,)).to(device)
    with record_function("forward"):
        loss = criterion(model(d), l)
    with record_function("backward"):
        loss.backward()
    with record_function("optimizer_step"):
        optimizer.step(); optimizer.zero_grad()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
prof.export_chrome_trace("trace.json")   # 可在 chrome://tracing 中可视化

从输出中可以快速判断：backward 是否比 forward 慢（正常约 2 倍），数据搬运是否占过多时间，优化器步骤是否有异常耗时。

📝 总结

层次 内容 核心能力
Tensor 数据表示与操作 形状操作、设备管理、dtype 选择
autograd 自动微分 计算图、梯度计算、梯度清零
nn.Module 模型组织 定义模型、参数管理、常用层
训练循环 端到端训练 DataLoader、训练流程、checkpoint
GPU 训练 硬件加速 混合精度、显存管理
性能分析 瓶颈定位 profiler 使用、常见瓶颈识别

这些构成了理解后续分布式训练、CUDA 编程和推理优化的必要基础。掌握这些内容后，你在阅读 DeepSpeed、Megatron-LM 等框架的代码时，不会因为 PyTorch 基础不牢而卡住。

🎯 自我检验清单

完成本文学习后，你应该能够：

  • 能用 torch.randn、torch.zeros 等方法创建任意形状的 Tensor，并熟练使用 view、permute、squeeze 等操作变换形状
  • 能解释 requires_grad=True 的作用，手动构建计算图并调用 backward() 获取梯度
  • 能解释为什么每个训练 step 都需要调用 optimizer.zero_grad()，以及梯度累积的原理
  • 能继承 nn.Module 实现自定义模型，并使用 parameters() 和 state_dict() 管理参数
  • 能写出完整的训练循环：DataLoader → forward → loss → backward → optimizer.step → checkpoint
  • 能使用 torch.autocast 实现 bf16/fp16 混合精度训练，并解释 GradScaler 的作用
  • 能使用 torch.profiler 分析训练 step，读懂输出并识别数据加载和 CPU-GPU 传输瓶颈
  • 能解释 fp32、fp16、bf16 三种精度的区别及选择依据

📚 参考资料