如何处理GPU训练显存不足[memory isn&apos；t enough][alloc failed][out of memory]

【现象描述】

GPU上网络运行过程中出现OOM（显存不足），报错日志中有如下信息：

cudaMalloc failed, ret[2], out of memory

或者：

memory isn't enough and alloc failed

【原因分析】

网络训练过错中出现显存不足的原因很多，后续遇到我会持续补充，当前遇到的一些原因如下：

1.卡被占用，导致可用显存变小了。

2.网络训练batchsize过大。

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

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

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

【排查步骤和解决方法】

步骤1：执行前使用命令nvidia-smi查看卡的使用状态，是否有被占用以及被占用的大小，剩下可用显存有多大。

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

步骤3：确认输入数据是否是动态shape，当前动态shape还不支持输入shape可变，因为针对这种场景，每次step训练都会编译新图，导致显存不断申请最后OOM，可以通过图的个数是不是不断増长来判断。（context.set_context(save_graphs=True)，然后观察*_validate*.ir这一类ir是不是一直在增多）

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

案例代码：

losses=[]
for i in range(steps):
    loss = train_network()
    losses.append(loss)

该案例代码将图的执行结果loss放到全局losses中存储，因此每次循环都会保存loss的device地址，随着训练次数的增加可能会导致OOM。如果有存储图输出结果的诉求，可以转化成asnumpy后保存，改成loss = train_network().asnumpy()，loss经过asnumpy后拷贝device到host后会释放掉device的地址。

步骤5：前面步骤确认都没有问题的话，则有可能是算子实现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，可以搜索代码查看算子的Launch中是否有调用AllocTensorMem函数。

步骤6：前面的几种可能原因都排查没有问题的话，需要MindSpore开发人员详细分析了，context.set_context(save_graphs=True)同时export GLOG_v=1保存执行日志，将保存的图和日志打包发给开发人员排查确认。

【建议与总结】

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

2.算子实现逻辑不应该出现申请显存的逻辑，一般都可以通过workspace来实现临时device地址存储的目的。上面的案例修改方案参考：https://gitee.com/mindspore/mindspore/pulls/25633