最近在研究 PyTorch 训练显存 offload 相关的技术, 发现了一些有意思的现象和问题, 也踩了一些有意思的坑
model.to 的实现
在 PyTorch 中, 如果我们想把模型挪到某个设备中, 我们会调用 model.to
这个方法, 例如如果想把 model 从 GPU 挪到 CPU,
我们会使用 model.to('cuda')
实现模型的原地 device offload, 也就是说, 调用这个方法并不需要你写成
model = model.to('cuda')
。但是我们知道如果是一个 tensor 的话是需要写成 tensor = tensor.to('cpu')
的。
那么为什么模型就无需这样操作呢? torch 是如何实现一个 tensor 原地替换呢?
注意到, pytorch 有一个有意思的方法叫 torch.utils.swap_tensors
, 正如文档里描述, 我们可以这样的使用它来实现不改tensor reference的情况下改里面的storage
1 | import torch |
那么 model 就可以使用类似的方法来实现原地替换, 参考 文档。
我们也可以直接从 torch 源码中找到对应的实现
需要注意的是, 如果一个 tensor 被别人存了一个 slice, 那么 swap_tensors 会报错, 因此当我们想使用 swap_tensors 的时候, 不要做一些 slice 相关操作
奇妙的 64MiB 的显存占用
我们书写一个非常简单的程序如下
1 | import torch |
运行代码得到了下面的输出
1 | [start] Current CUDA allocated 0.00 MiB, reserved 0.00 MiB. |
看起来非常奇怪, 因为我们只分配了很小的三个 tensor, 分别是 a
, b
和 loss
, 为什么经过一次 matmul 之后,
显存涨了 32MiB, 然后经过一次 backward 之后, 显存又涨了 32MiB? 而且这部分增长的 64MiB 显存即便是
使用 torch.cuda.empty_cache()
也无法清理, 到底是哪里占了这些显存?
在 torch 的 code 中我们发现 torch 在 matmul 和 bwd 的时候需要给 cublas 一个 32M 的 workspace。
这个问题还是很隐蔽的, 其中最大的问题是会导致 torch 显存释放的时候出现碎片。
如果第一次 fwd 或者 bwd 是在代码比较靠后的位置做的, 那么前面申请的那些显存如果释放的话极有可能会导致 torch 出现较大的显存碎片,
这两个 32MiB 极端情况可能会占着一个非常大的 torch Segment 而无法释放, 因此我们需要额外关注这个问题。
一个简单的解决办法是把上面这段 demo 代码在你的代码最开头运行一次, 相当于做一次 fwd 和 bwd 的 warmup, 这样就能极大的避免显存碎片
即可以直接在代码最开头执行下面的代码:
1 | torch.matmul( |