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

【现象描述】

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 &inputs, const std::vector &workspace,
              const std::vector &outputs, void *stream_ptr) override {
    T *input = GetDeviceAddress(inputs, 0);
    S *indices = GetDeviceAddress(inputs, 1);
    T *updates = GetDeviceAddress(inputs, 2);
    T *output = GetDeviceAddress(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地址存储的目的。

【相关参考文档】

你可能感兴趣的:(深度学习,人工智能)