一、单卡训练基础:GPU架构简介、CUDA编程模型、PyTorch单卡训练流程、性能分析工具(nvidia-smi)
好,咱们正式开始。这一章是地基,我把它叫做「从零到一」的关键一步。
很多人一上来就搞分布式,结果单卡都没跑明白,排查问题的时候一脸懵。我个人习惯是:先把单卡吃透,再谈集群。你想想看,连一辆车都开不稳,怎么能去开列车队?
1.1 GPU架构简介:别把它当黑盒
GPU 不是一块简单的「加速卡」。它本质上是一个超大规模的并行计算单元。
我当年刚接触 GPU 时,以为它就是 CPU 的升级版,结果写出来的代码跑得比 CPU 还慢。后来才明白,GPU 和 CPU 的设计哲学完全不同。
- CPU:擅长复杂逻辑控制,核心少但单核能力强,适合串行任务。
- GPU:擅长简单并行计算,核心多但单核能力弱,适合大规模数据并行。
以 NVIDIA 的架构为例,几个关键概念你得记住:
| 概念 | 说明 | 我的理解 |
|---|---|---|
| SM(Streaming Multiprocessor) | GPU 的核心计算单元,包含多个 CUDA Core | 相当于 GPU 里的「车间」 |
| CUDA Core | 执行浮点运算的基本单元 | 车间里的「工人」 |
| Global Memory | 显存,所有线程都能访问 | 公共仓库,但速度慢 |
| Shared Memory | 同一个 Block 内的线程共享 | 车间内部的小仓库,速度快 |
| Register | 每个线程私有的寄存器 | 工人手里的工具,最快 |
核心要点:GPU 的性能瓶颈往往不在算力,而在显存带宽和内存访问模式。说白了,数据搬来搬去的时间,比计算本身还长。
这里我画了一张图,帮你理清 GPU 的层次结构:
1.2 CUDA编程模型:理解「线程层次」
CUDA 编程模型的核心,就是线程的组织方式。说白了,你得告诉 GPU:谁干什么活,谁和谁是一伙的。
三个层次:
- Grid(网格):一个 Kernel 启动时,对应一个 Grid。Grid 包含多个 Block。
- Block(线程块):一组线程,可以协作。同一个 Block 内的线程共享 Shared Memory。
- Thread(线程):最小的执行单元。每个线程有自己的寄存器和局部内存。
我的经验:Block 的大小不是随便设的。我刚开始时设了 1024 个线程一个 Block,结果发现 Shared Memory 不够用。后来改成 256,性能反而提升了。为什么?因为每个 Block 能用的资源是有限的,Block 太大反而导致并行度下降。
一个简单的 CUDA 核函数示例:
// 向量加法 Kernel
__global__ void vecAdd(float* A, float* B, float* C, int N) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < N) {
C[i] = A[i] + B[i];
}
}
// 调用方式
int threadsPerBlock = 256;
int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock;
vecAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, N);
这里 blockIdx.x 是 Block 的索引,threadIdx.x 是线程在 Block 内的索引,blockDim.x 是 Block 的大小。这三个变量帮你定位到具体的线程。
注意:Kernel 是异步执行的。调用 cudaDeviceSynchronize() 才能确保 GPU 上的任务完成。我曾经踩过这个坑——CPU 侧已经打印了结果,但 GPU 还没算完,打印出来的全是垃圾值。
1.3 PyTorch单卡训练流程:从数据到模型
好了,理论说完了,咱们来点实际的。PyTorch 单卡训练,说白了就是五步走:
- 数据加载:用
DataLoader把数据喂给模型。 - 模型定义:继承
nn.Module,定义网络结构。 - 损失函数和优化器:比如
CrossEntropyLoss和SGD。 - 训练循环:前向传播 → 计算损失 → 反向传播 → 更新参数。
- 评估和保存:验证集上跑一下,保存模型权重。
一个标准的训练循环:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
model = MyModel().cuda() # 把模型搬到 GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
for epoch in range(num_epochs):
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.cuda(), target.cuda() # 数据也要搬到 GPU
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')
关键点:.cuda() 是 PyTorch 把数据或模型搬到 GPU 的方法。但注意,不是所有操作都能自动搬到 GPU——比如 NumPy 数组就不行。我见过有人把 NumPy 数组直接传给模型,结果报错半天没反应过来。
我个人习惯在训练前加一句 torch.cuda.empty_cache(),清一下显存缓存,避免之前残留的数据影响。
1.4 性能分析工具:nvidia-smi 和更多
训练跑起来了,但怎么知道它跑得好不好?这时候就需要工具了。
nvidia-smi 是最基础的 GPU 监控工具。直接在终端敲:
nvidia-smi
你会看到:
- GPU 利用率(%):是不是跑满了?
- 显存使用(MiB):有没有爆显存?
- 温度(°C):是不是快烧了?
- 功耗(W):有没有达到 TDP 限制?
避坑指南:我曾经遇到一个情况,GPU 利用率只有 30%,但显存快满了。后来发现是 DataLoader 的 num_workers 设得太小,CPU 来不及喂数据,GPU 一直在空等。把 num_workers 从 0 改成 4,利用率直接飙到 95%。
除了 nvidia-smi,还有几个工具我经常用:
| 工具 | 用途 | 我的评价 |
|---|---|---|
| nvidia-smi | 实时监控 GPU 状态 | 必备,简单粗暴 |
| nvtop | 类似 htop 的 GPU 监控 | 可视化更好,推荐 |
| PyTorch Profiler | 分析每个操作的耗时 | 定位瓶颈的神器 |
| Nsight Systems | 全系统性能分析 | 专业级,但学习成本高 |
用 PyTorch Profiler 的简单示例:
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CUDA],
record_shapes=True
) as prof:
output = model(data)
loss = criterion(output, target)
loss.backward()
print(prof.key_averages().table(sort_by="cuda_time_total"))
它会告诉你哪个操作最耗时。嗯,这里要注意:有时候最耗时的不是 forward,而是 to(device) 的数据搬运操作。所以尽量在加载数据时就放到 GPU 上,别在循环里反复搬。
警告:nvidia-smi 显示的显存使用量,不一定等于你模型实际占用的显存。PyTorch 有缓存分配器,会预分配一些显存。用 torch.cuda.memory_summary() 可以看更详细的信息。
好了,这一章的内容就是这些。GPU 架构、CUDA 编程模型、PyTorch 训练流程、性能分析工具——这四个点吃透了,单卡训练你就基本入门了。下一章我们会聊多卡训练,但前提是你能把单卡跑得又快又稳。