把模型搬到 GPU 上跑，很多人第一反应是给张量加一句.cuda()，然后盯着显存占用和利用率来回看。真跑起来才发现，显存爆了、view报错了、GPU 利用率在低位和满载之间反复横跳。这些问题大多不是模型结构写错了，而是 PyTorch 数据结构没有按 GPU 的脾气组织。GPU 加速的本质，是让数据以连续、对齐、类型匹配的方式进入显存，再让计算核按可合并访存的模式去读。张量的device、dtype、stride、layout这几个属性，决定了数据在显存里怎么摆，也决定了计算核是跑满还是空转。
这篇内容面向刚接触 GPU 训练、或者训练脚本已经能跑但速度不理想的读者。我会从 CPU 和 GPU 的内存模型差异讲起，拆开 PyTorch 张量的几个关键属性，再给出一套可以直接照着改的数据管道和排查流程。不堆概念，重点放在“为什么这样摆数据”和“改哪几行代码”上，尽量让每个步骤都能在你的项目里复现。
1. 先想清楚：GPU 加速到底加速了什么
1.1 CPU 张量与 GPU 张量的内存模型差异
CPU 上的张量，数据放在主机内存里，张量对象本身也在主机内存里，对象里存着指向存储区的指针、形状、步长这些元数据。调用.cuda()或者.to('cuda')之后，PyTorch 会在显存里新开一块存储区，把数据从主机内存拷过去，然后返回一个新的张量对象。原来的 CPU 张量还在，两个张量共享同一份逻辑数据，但物理位置完全不同。
这个拷贝动作是有成本的。主机内存和显存之间隔着一条总线，数据要经过它搬过去，带宽比显存内部低一个数量级。所以我在改训练脚本时，会先看数据搬运动作发生在哪里：是在每个 batch 里反复搬小张量，还是一次性把大张量搬过去。前者会把大量时间耗在等待拷贝上，后者才是 GPU 该干的活。
训练循环里还有一个容易被忽略的点：拷贝默认是同步的。也就是说，CPU 发出拷贝指令后会等它完成，才继续往下走。PyTorch 提供了non_blocking=True参数，配合页锁定内存，可以让拷贝和计算重叠起来。页锁定内存的作用是把这块主机内存锁定，不让操作系统把它换到磁盘上，这样直接内存访问才能高效工作。没有页锁定内存，异步拷贝基本等于没开。
1.2 连续性与 stride：计算核眼里的数据长什么样
张量在显存里是一段线性地址，形状是逻辑视图，真正决定元素位置的是stride。一个形状为(4, 6)的二维张量，默认步长是(6, 1)，意思是沿第一维走一步跳过 6 个元素，沿第二维走一步跳过 1 个元素。计算核读数据时按步长算地址，如果步长和内存布局匹配，连续的元素能被合并成一次宽访存；如果步长乱了，计算核就得跳着读，访存效率会明显下降。
转置操作最能说明问题。对(4, 6)的张量做.t()，形状变成(6, 4)，步长变成(1, 6)，数据本身没有动，只是换了个看法。这时候它不再是连续张量，is_contiguous()返回False。很多算子要求输入连续，view更是直接要求连续，否则报错。reshape看起来更宽容，但它内部可能会偷偷做一次拷贝，把非连续张量变成连续的，这个拷贝发生在显存里，代价不小。
我习惯在关键节点打印一下x.stride()和x.is_contiguous()，尤其是做完转置、切片、permute之后。卷积网络里permute用得多，稍不注意就把连续张量变成了非连续，后面接一个view就炸。与其报错后到处找，不如提前看一眼。
1.3 dtype 与带宽：省一半显存不等于快一倍
数据类型直接决定每个元素占几个字节。float32占 4 字节，float16和bfloat16占 2 字节，int64占 8 字节。显存容量和带宽都是按字节算的，把float32换成float16，理论上显存占用和访存量都能减半。但实际训练里，混合精度不是简单地把所有张量转成半精度就完事。
模型参数、梯度、优化器状态通常要保持float32，前向和反向计算用半精度，这就是混合精度的基本盘。PyTorch 的autocast会自动给合适的算子选半精度，同时保留一部分对数值敏感的算子用float32。自己手动把所有张量.half()，很容易在损失函数或者归一化层出问题，梯度下溢之后损失变成NaN，排查起来很费时间。
还有一个常被忽略的点是索引张量。PyTorch 的嵌入层、gather、index_select默认要求索引是int64，一个大规模词表的索引张量可能比想象中占得多。如果显存紧张，先看索引张量是不是不必要地用了int64，再看中间激活值。不是所有张量都值得转半精度，索引和掩码这类张量转错了类型反而会报错。
2. PyTorch 数据结构的关键属性与 GPU 适配
2.1 device、dtype、layout 三件套
拿到一个张量，我一般先看三个属性：在哪个设备上、什么数据类型、什么内存布局。device决定数据在主机内存还是显存，dtype决定每个元素占多大，layout决定是稠密布局还是稀疏布局。三者不一致，轻则触发隐式拷贝，重则直接报错。
| 属性 | 常见取值 | 对 GPU 的影响 | 典型问题 |
|---|---|---|---|
device | cpu、cuda:0 | 决定数据是否在显存 | 模型在 GPU、标签在 CPU |
dtype | float32、float16、int64 | 影响显存占用与访存带宽 | 索引类型不对导致算子报错 |
layout | torch.strided、torch.sparse_coo | 影响计算核的读取方式 | 稀疏结构用错场景反而更慢 |
这三个属性不是孤立的。一个float16的连续张量在 GPU 上很舒服，一个float32的非连续张量在 GPU 上就要多花拷贝和访存的开销。我见过不少脚本，模型参数是float32，输入数据也是float32，但中间某个预处理步骤把张量转成了float64，结果整个计算图被迫用双精度跑，速度直接掉一截。排查这种问题，打印dtype比看代码快。
2.2 非连续张量：view 为什么报错
view报错是 PyTorch 新手最常见的错误之一，错误信息通常说形状不兼容。根因往往不是形状算错了，而是张量不连续。view要求张量在内存里是连续的，这样才能按新的形状重新解释这段内存。非连续张量没法直接重新解释，只能先拷贝成连续张量，或者改用reshape。
import torch x = torch.randn(4, 6, device='cuda') y = x.t() print(x.stride()) # (6, 1) print(y.stride()) # (1, 6) print(y.is_contiguous()) # False # y.view(24) 会报错 z = y.reshape(24) # 可能触发一次隐式拷贝 print(z.is_contiguous()) # True这段代码在 GPU 上跑，reshape触发的拷贝发生在显存里，数据量小的时候感觉不到，数据量大的时候就是实打实的开销。我的做法是：在需要view的地方，先确认前一步操作没有破坏连续性；如果确实需要转置后再展平，用contiguous()显式拷贝一次，心里有数，比让reshape偷偷拷贝强。
还有一种情况是切片。对一个连续张量做x[:, ::2]，得到的张量步长变了，不再连续。后面如果接一个需要连续输入的算子，就会触发拷贝。切片在数据预处理里很常见，尤其是做数据增强或者通道选择的时候，切完记得看一眼连续性。
2.3 稀疏、嵌套、量化结构的适用边界
PyTorch 不只有稠密张量。稀疏张量适合大部分元素为零的场景，比如大规模推荐系统里的特征矩阵；嵌套张量适合变长序列，比如一批长度不一的文本；量化张量把权重压成低精度整数，适合推理阶段省显存。这些结构在 GPU 上都有对应的计算核，但不是所有场景都划算。
稀疏张量在元素极度稀疏时能省显存，但稀疏算子的实现成熟度不如稠密算子，小规模稀疏矩阵用稀疏布局反而更慢。嵌套张量对变长序列友好，但某些算子不支持，调试时容易卡在类型转换上。量化张量主要用于推理，训练阶段用起来限制多。我的建议是：除非数据本身有明确的稀疏或变长特征，否则先用稠密张量把流程跑通，再考虑换结构。换结构之前，先用小规模数据对比一下耗时，别凭感觉选。
2.4 显存分配器与缓存命中
PyTorch 在 GPU 上有一套缓存分配器。申请显存时，它先看缓存里有没有合适大小的块，有就直接复用，没有才向驱动申请。释放张量时，显存不一定立刻还给驱动，而是回到缓存里等着下次复用。这个机制让反复申请释放同尺寸张量变得很快，但也带来一个现象：torch.cuda.memory_allocated()显示占用量不高，torch.cuda.memory_reserved()显示缓存占了很多。
很多人一看到显存不够就调torch.cuda.empty_cache()，这个操作会把缓存里的空闲块还给驱动，但不会解决碎片问题，频繁调用还会拖慢速度。真正该做的是看显存分配的快照，确认是大块申请失败还是碎片导致。如果模型里有动态形状，每步申请的显存大小不一样，缓存命中率低，碎片就容易积累。固定输入形状、避免在循环里创建尺寸变化的临时张量，比反复清缓存有效得多。
3. 实操：把训练管道改造成 GPU 友好结构
3.1 建立可复现的基线
改任何性能问题之前，先有一个能重复跑的基线。我一般会固定随机种子、固定输入形状、固定 batch size，然后用 GPU 事件计时，而不是用time.time()。time.time()测的是主机时间，GPU 是异步执行的，不同步的话测出来的数字没有意义。
import torch def benchmark(fn, warmup=10, iters=50): for _ in range(warmup): fn() torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(iters): fn() end.record() torch.cuda.synchronize() return start.elapsed_time(end) / iters x = torch.randn(64, 3, 224, 224, device='cuda') model = torch.nn.Conv2d(3, 32, 3).cuda() print(benchmark(lambda: model(x)))预热很重要。第一次跑某个算子时，PyTorch 可能要做内存分配、算子选择、编译准备，这些开销不该算进稳定耗时里。预热十次左右，再正式计时，数字才稳定。基线建立之后，每次只改一个变量，改完再测，才能知道是哪一步起了作用。
3.2 DataLoader 的页锁定内存与异步拷贝
数据加载是最容易拖慢 GPU 的环节。GPU 算得再快，数据供不上就只能等。DataLoader的pin_memory=True会把 batch 放进页锁定内存，这样拷贝到 GPU 时可以用异步方式。配合non_blocking=True，拷贝能和计算重叠。
from torch.utils.data import DataLoader loader = DataLoader( dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True, drop_last=True, ) for x, y in loader: x = x.to('cuda', non_blocking=True) y = y.to('cuda', non_blocking=True) # 前向、反向num_workers不是越大越好。工作进程太多会争抢 CPU 和内存带宽，反而拖慢数据准备。我一般从 4 开始试，观察 GPU 利用率；如果 GPU 利用率上不去、数据加载时间占比高，再加到 8。persistent_workers=True可以避免每个 epoch 重建工作进程，epoch 数多的时候能省下不少启动开销。
还有一个细节：non_blocking=True只在页锁定内存上才真正异步。如果pin_memory=False，这个参数基本没效果，拷贝还是同步的。所以这两个参数要成对出现，别只开一个。
3.3 模型参数布局与 channels_last
卷积网络对内存布局敏感。默认的NCHW布局在 GPU 上做卷积时，通道维度的访存模式不一定最优。channels_last布局把通道放到最后，让相邻通道的数据在内存里挨着，卷积核读取时更容易合并访存。PyTorch 提供了memory_format参数，可以整体转换模型和输入。
model = model.to(memory_format=torch.channels_last) x = x.to(memory_format=torch.channels_last) # 训练循环里保持输入布局一致 out = model(x)这个改动对卷积层多、通道数大的模型效果明显，对全连接层为主或者 Transformer 结构不一定有收益。我通常先在一两个 batch 上对比耗时，确认有提升再全局改。改完之后要检查中间张量有没有被某些算子转回默认布局，必要的时候在关键位置重新转一次。
模型参数本身也有一些结构上的讲究。比如把多个小矩阵拼成一个大矩阵做批量矩阵乘法，比循环调用小算子效率高；把频繁访问的常量张量提前放到 GPU 上，避免在循环里反复创建。这些改动不大，但累积起来对吞吐有影响。
3.4 混合精度与梯度缩放
混合精度是省显存、提吞吐的常用手段，但用法要对。PyTorch 的autocast负责给算子选精度，GradScaler负责放大梯度，防止半精度下溢。
scaler = torch.cuda.amp.GradScaler() for x, y in loader: x = x.to('cuda', non_blocking=True) y = y.to('cuda', non_blocking=True) optimizer.zero_grad(set_to_none=True) with torch.cuda.amp.autocast(): out = model(x) loss = criterion(out, y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()set_to_none=True比把梯度置零更省显存，因为它直接把梯度张量释放掉，而不是填零。这个细节在大模型上能省下可观的显存。GradScaler会自动调整缩放系数，不需要手动调。如果训练中出现NaN，先看是不是某个算子在半精度下不稳定，可以把那部分计算强制转成float32，而不是直接放弃混合精度。
3.5 用 profiler 找出真正的拷贝点
性能问题不能靠猜。PyTorch 自带的 profiler 能同时记录 CPU 和 GPU 活动，看清楚时间花在计算、拷贝还是同步上。
from torch.profiler import profile, ProfilerActivity with profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3), on_trace_ready=torch.profiler.tensorboard_trace_handler('./prof_log'), record_shapes=True, ) as prof: for step, (x, y) in enumerate(loader): if step >= 6: break x = x.to('cuda', non_blocking=True) y = y.to('cuda', non_blocking=True) out = model(x) loss = criterion(out, y) loss.backward() optimizer.step() optimizer.zero_grad(set_to_none=True) prof.step()看结果时重点看三类：显存拷贝占了多大比例、有没有频繁的同步调用、计算核的耗时是否被小算子切碎。如果拷贝占比高，回到数据管道检查页锁定内存和异步拷贝；如果同步调用多，检查代码里有没有频繁.item()、.cpu()或者打印张量；如果小算子太多，考虑合并操作或者换更高效的实现。
4. 显存估算与 batch size 选择
4.1 显存占用的四个部分
训练时的显存占用大致分四块：模型参数、梯度、优化器状态、中间激活值。前 three 块可以按参数量估算，第四块跟 batch size、输入尺寸、模型结构都有关，最难精确算。
| 占用项 | 估算方式 | 备注 |
|---|---|---|
| 模型参数 | 参数量 × 每参数字节数 | float32为 4 字节 |
| 梯度 | 与参数同量级 | 通常和参数一样的精度 |
| 优化器状态 | 与优化器类型有关 | 自适应优化器通常存两份状态 |
| 中间激活 | 与 batch、序列长度、通道数相关 | 建议实测，公式只能估数量级 |
以float32训练为例，一个参数量为 1 亿的模型，参数占 0.4 GB 左右，梯度再占 0.4 GB。如果用自适应优化器，每个参数通常还要存两份状态，又是 0.8 GB 左右。三块加起来 1.6 GB，这还没算中间激活和临时缓冲区。实际跑起来，显存占用往往是这个数字的两三倍。
4.2 估算实例与参数推导
假设输入是(32, 3, 224, 224)，模型第一层卷积把 3 通道变成 64 通道，卷积核 3×3。这一层的参数是64 × 3 × 3 × 3 + 64，约 1792 个参数，可以忽略不计。真正占显存的是激活值：输出特征图是(32, 64, 224, 224)，float32下约 32 × 64 × 224 × 224 × 4 字节，接近 411 MB。这只是一层的输出，深层网络累加起来，激活值很容易超过参数本身。
所以调 batch size 时，不能只按参数量算。我一般先跑一个最小 batch，记录torch.cuda.max_memory_allocated()，再按比例往上推。推到显存占用的八成左右就停，留出余量给临时缓冲和碎片。预留太少，训练跑着跑着就 OOM；预留太多，batch size 上不去，吞吐吃亏。
4.3 batch size 与吞吐的拐点
batch size 增大，GPU 利用率通常先升后平。太小的 batch 喂不饱计算单元，太大的 batch 可能让寄存器压力上升、缓存命中下降，吞吐不再增长。找到拐点的办法是固定其他条件，只改 batch size，测每个 epoch 的耗时和显存占用。
如果显存不够大 batch，可以用梯度累积。把多个小 batch 的梯度累加在一起，再更新一次参数，效果接近大 batch，显存占用却按小 batch 算。累积步数不要太多，否则训练动态会变，学习率也要相应调整。
5. 常见问题与排查实录
5.1 OOM 与碎片化
OOM 是最常见的问题，但原因不一定是 batch 太大。可能是某个中间张量没释放，可能是动态形状导致缓存命中率低，也可能是碎片让大块申请失败。排查时先看torch.cuda.memory_summary()，确认是活跃显存高还是缓存占得多。如果是缓存占得多但活跃显存不高，说明碎片或缓存没复用，考虑固定输入形状、减少动态创建张量。如果是活跃显存高，用 profiler 看哪一步显存涨得最快。
注意：不要在训练循环里频繁调用
torch.cuda.empty_cache()。这个操作会打断缓存复用，反而让后续申请更慢。只在确认需要释放缓存、且不在性能敏感路径上时用。
5.2 device 不一致
模型在 GPU、输入在 CPU，或者模型的一部分在 GPU、另一部分在 CPU，都会报 device 不一致。多卡训练时更容易出现，比如某张卡上的参数没有正确同步。我的习惯是定义一个device变量，所有.to()都引用它，不在代码里散落写'cuda'或'cuda:0'。加载 checkpoint 时也指定map_location，避免把 CPU 上的权重直接塞进 GPU 模型。
5.3 隐式同步与隐式拷贝
GPU 是异步执行的，但有些操作会强制同步：.item()、.cpu()、print(tensor)、torch.cuda.synchronize()。在训练循环里频繁调用这些，GPU 就会不停等 CPU，利用率上不去。隐式拷贝则常发生在dtype不一致的运算里，比如float32和float64混算，PyTorch 会把低精度转成高精度，多一次拷贝和计算。
| 现象 | 可能原因 | 排查方式 | 处理建议 |
|---|---|---|---|
| 利用率周期性掉到低位 | 数据加载慢或同步点多 | profiler 看拷贝与同步占比 | 开页锁定内存、减少.item() |
| 显存缓慢增长 | 中间张量被引用未释放 | 查全局列表、日志缓存 | 及时del，避免持有计算图 |
| 报错说形状不兼容 | 非连续张量用了view | 打印stride和连续性 | 用contiguous()或reshape |
损失变NaN | 半精度下溢或学习率过大 | 关闭混合精度对比 | 用梯度缩放，检查数据 |
5.4 数值异常
半精度训练时出现NaN或inf，先确认GradScaler有没有正确使用。如果用了缩放还是异常，逐层检查，把可疑的归一化层或损失计算强制转成float32。数据本身也可能有问题，比如某个 batch 里有极端值，前向传播后梯度爆炸。加一个梯度裁剪通常能缓解，但根源还是要查数据。
5.5 多卡训练的数据结构问题
多卡训练时，每个进程拿到的是不同的数据子集，模型参数需要同步。如果自定义了采样器或者数据划分逻辑，要确保每个进程的 batch 结构一致，否则同步时会卡住。drop_last=True能避免最后一个不完整的 batch 导致各卡步数不一致。另外，多卡下每个进程都会占一份显存，估算单卡显存时要按单进程算，不要按总量算。
6. 我踩过的坑和几条硬经验
6.1 看起来对但很慢的写法
第一种是循环里创建张量。比如每个 step 都torch.zeros(...)创建一个同样大小的缓冲区，缓存分配器虽然能复用，但频繁申请释放仍有开销。更好的做法是提前创建好，用.zero_()原地清零。第二种是频繁在 CPU 和 GPU 之间搬小张量，比如把损失值搬到 CPU 记录日志，每个 step 都搬一次，同步开销比计算还大。可以累积几个 step 再一起搬。第三种是滥用permute之后不处理连续性，后面接一个隐式拷贝的算子，白花一次显存带宽。
6.2 调试顺序与工具
我现在遇到 GPU 性能问题，顺序基本固定：先看显存占用和 batch size 是否匹配，再看数据加载是不是瓶颈，然后用 profiler 看拷贝和同步占比，最后才怀疑算子实现。这个顺序能排除掉大部分低级问题。工具上，torch.cuda.memory_summary()看显存分布，profiler 看时间分布，nvidia-smi这类系统工具看整体利用率。三者结合，基本能定位到具体环节。
还有一条经验：改性能之前先保证正确性。非连续张量、混合精度、多卡同步这些改动，都可能引入不易察觉的数值问题。每改一步，跑一个小规模的全量训练，确认损失曲线正常，再继续优化。否则优化到最后，速度上去了，模型效果却不对，回头查的成本更高。
最后分享一个我常用的检查动作：在每个 epoch 开始时打印一次关键张量的device、dtype、shape和is_contiguous()，只在调试阶段开。这四个信息能覆盖大部分数据结构相关的问题，比出事后翻代码快得多。等流程稳定了，再把这些打印去掉，别让它们留在正式训练循环里拖慢速度。
