跳到主要内容

WSL2 CUDA禁用共享内存

· 阅读需 4 分钟
RibomBalt
RibomBalt
CTF enthusiastist, GeoPhysics PhD, Amateur coder

TL;DR;​

  • 对于双显卡设备，当N卡显存不够用时，会使用集显提供的共享显存，会显著降低性能。
  • 可以通过torch.cuda.set_per_process_memory_fraction(1.0, 0)在torch模块中禁用共享显存。
  • 禁用共享显存会导致专用显存报错时直接退出而不是保持运行，出现共享显存问题首先应该调参减小模型显存占用。

背景​

这个问题是在玩fish-speech TTS模型时遇到的。我的NVIDIA RTX 4060Ti只有8GB显存，对模型微调时显存爆了，任务管理器可以看到使用了共享显存，并且运行速度非常慢。当我减小batch_size，减少显存占用后，所有显存均使用的N卡的显存，这时运行速度是正常的。

诊断​

首先，通过问QwQ-32B，得到了如下测试代码，可以用nvcc -o test test.cu编译：

#include <stdio.h>
#include <unistd.h>
int main() {
    size_t size = 1024 * 1024 * 1024;  // 超出显存
    float *d_ptr;

    for (size_t i=0; i<10; i++) {
        cudaMallocManaged(&d_ptr, size);    // 抛出错误而非回退
        printf("alloc GB: %u\n", i+1);
        usleep(2000000);
    }

    return 0;
}

上述代码在WSL2下运行时会只使用共享显存，即使仍然有大量的专用显存，这一点非常奇怪，暂时不明原因。AI还告诉我似乎有两个环境变量CUDA_VISIBLE_DEVICES和CUDA_MANAGED_FORCE_DEVICE_ALLOC 可以限制只使用设备内存，我找到了对应的CUDA文档说明，但我测试下来WSL2下不生效。

值得注意的是，CUDA设备是32bit的（这一点存疑，但是编译期确实出现了overflow），size_t是unsigned long int，是32位的，所以单次分配不能超过4GB，这里是每次分配1GB实现的。cudaMallocManaged分配的内存将由CUDA自主管理，不需要我们释放。

不过后来我意识到其实更应该直接在torch中测试，而不是手动编译CUDA kernel。可以用更简单的代码：

import torch
x = torch.ones((12, 1024, 1024, 1024), dtype=torch.uint8, device='cuda')
print(x.device)

当分配显存大于剩余专用显存时，如果禁用fallback，这里会报错；否则会使用共享显存正常运行退出。

解决方案​

Windows下运行CUDA​

首先，可以通过修改NVIDIA控制面板，在全局/指定程序设置的CUDA - Sysmem Fallback Policy （系统内存回退政策），改为偏好无回退（这个选项需要中译中，英文应该是no fallback，也就是设备显存爆了不会妥协用共享显存的意思）。在GPT-SoVITS的文档也提到了这个事情。

这个方案适用于Windows下运行的CUDA程序（包括调用了CUDA的torch），适用于Windows端的代码/整合包。但我更倾向于在WSL2中跑模型，因此这个方案不适用。事实上，N卡控制面板无法选中WSL中的程序。而在我之前的测试中，WSL2上运行的CUDA并不遵守控制面板的设置，这个问题其他人也遇到过，参加微软WSL仓库的 issue，目前仍是open状态。

WSL2中解决方案​

上述issue中提到，torch中可以通过设置解决：

torch.cuda.set_per_process_memory_fraction(1.0, 0)

推测纯CUDA C代码也可以通过某些设置解决（看起来就是调用setMemoryFraction函数），对Python torch模块来说，直接在第一次分配显存前加这句话就可以了。

Comment​

灵魂之问：为什么要关闭共享显存？

关闭共享显存意味着显存不够时会直接报错退出而不是用共享显存低速运行。事实上直接报错和共享显存下运行对于模型训练来说几乎是同等严重程度的不可用状态，某种程度上前者可能更糟糕一些（因为程序稳定性差了，想象一下最后1%突然爆显存了，是比较慢的跑完好，还是直接报错前面白跑比较好）。

出现这个问题本质上还是显存不够用了，需要做的是减小在运行程序的显存占用，包括关闭无关程序，调整训练超参。比如我是在微调时遇到这个问题的，解决方案是把batch_size减小，当独显显存够用之后这个选项其实是没有影响的。目前没有测试过这个选项能否让模型更优先地使用设备专用显存（我推测不会）。