如何处理训练过程中出现OOM（显存不足）

最新推荐文章于 2025-04-22 10:21:11 发布

昇思MindSpore

最新推荐文章于 2025-04-22 10:21:11 发布

阅读量3.8k

点赞数 3

文章标签：深度学习人工智能

原文链接：https://bbs.huaweicloud.com/forum/forum.php?mod=viewthread&tid=169771

版权

本文探讨了GPU网络训练中常见的OOM问题，分析了过大batchsize、动态shape输入、tensor缓存和显存泄漏等可能原因，提供了排查和修复方法，强调了输出tensor管理和算子优化的重要性。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

【现象描述】

GPU上网络运行过程中出现OOM（显存不足）

【原因分析】

出现该现象，可能原因如下：

1.训练batchsize过大。

2.输入数据的shape是变化的（输入数据动态shape）。

3.输出结果的tensor保存起来了。

4.网络中的算子里出现显存泄漏（算子里每次launch都申请显存，并且不释放）。

【解决方法】

步骤1：排查训练的batchsize是否过大，可以逐步缩小batch，如果觉得batch不够大，还是出现OOM，则排除这个原因。

步骤2：确认输入数据是否是动态shape，当前动态shape还不支持输入shape可变，因为针对这种场景，每次step训练都会编译新图，导致显存不断申请最后OOM，可以通过图的个数是不是不断増长来判断。

步骤3：输出tensor里会挂接device上的地址，方便print时同步device数据输出，在tensor析构的时候会释放device地址，因此如果每次step训练都把输出tensor保存起来的话，导致tensor不会析构，随着训练step的增加则会OOM。

步骤4：前面3个步骤确认都没有问题的话，则有可能是算子实现bug，可以排查下是否有新增算子，算子里是否有申请显存的操作导致显存泄漏。

案例代码：

```

bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &outputs, void *stream_ptr) override {
T *input = GetDeviceAddress<T>(inputs, 0);
S *indices = GetDeviceAddress<S>(inputs, 1);
T *updates = GetDeviceAddress<T>(inputs, 2);
T *output = GetDeviceAddress<T>(outputs, 0);

const size_t indices_len = sizeof(S) * out_strides_.size();
void *indices_stride_work = device::gpu::GPUMemoryAllocator::GetInstance().AllocTensorMem(indices_len);
if (indices_stride_work == nullptr) {
MS_LOG(EXCEPTION) << "Failed to alloc indices_stride_work, size: " << indices_len;
}

```

---》每次launch都调用内存池接口AllocTensorMem申请内存，导致随着训练step增加出现OOM

【建议与总结】

1.输出tensor一般不建议保存起来，如果需要保存，可以保存asnumpy后的对象，tensor经过asnumpy后拷贝device到host后会释放掉device的地址。

2.算子实现逻辑不应该出现申请显存的逻辑，一般都可以通过workspace来实现临时device地址存储的目的。

【相关参考文档】

无