一、单卡训练基础: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 的层次结构:

GPU 架构层次图 GPU(整体) SM 0(流式多处理器) CUDA Core CUDA Core CUDA Core CUDA Core CUDA Core CUDA Core Shared Memory(共享内存) SM 1(流式多处理器) CUDA Core CUDA Core CUDA Core CUDA Core CUDA Core CUDA Core Shared Memory(共享内存) Global Memory(全局内存 / 显存)

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 单卡训练,说白了就是五步走:

  1. 数据加载:用 DataLoader 把数据喂给模型。
  2. 模型定义:继承 nn.Module,定义网络结构。
  3. 损失函数和优化器:比如 CrossEntropyLoss 和 SGD。
  4. 训练循环:前向传播 → 计算损失 → 反向传播 → 更新参数。
  5. 评估和保存:验证集上跑一下,保存模型权重。

一个标准的训练循环:

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 训练流程、性能分析工具——这四个点吃透了,单卡训练你就基本入门了。下一章我们会聊多卡训练,但前提是你能把单卡跑得又快又稳。

可可精品资料库