Windows 11 配 WDDM 驱动,二十四 GB 的 4090 在 torch 里 total_memory 只有二十二 GB 出头,按标称容量算好的 batch 一跑就 OOM;同一张卡在 WSL2 里可用量又多了一点,团队对按哪个数字规划显存有分歧。
一开始怀疑是「torch 取的是 cudaMemGetInfo 返回的当前可用值」,查了一圈,方向不对。
真正麻烦的地方在于,几个原因的现象几乎一样。WDDM 会为桌面合成和驱动预留一部分显存。少数情况下是Windows 的显存共享机制让总量按可用量上报。
少数情况下是TCC 模式只在部分数据中心卡上支持。
第一件事是打印设备上报的总量与当前可用量
python -c "import torch;print(torch.cuda.get_device_properties(0).total_memory/2**30, torch.cuda.mem_get_info())"
然后确认驱动模式与显存总量
nvidia-smi --query-gpu=name,driver_version,memory.total,memory.used --format=csv
再往下按实际可用量重算 batch 上限
python -c "import torch;free,total=torch.cuda.mem_get_info();print(round(free/2**30,1),round(total/2**30,1))"
剩下的就是换到 WSL2 里跑同一份脚本对比
wsl -d Ubuntu-22.04 -- python -c "import torch;print(torch.cuda.get_device_properties(0).total_memory/2**30)"
改完先别撤,确认这几项。脚本按二十二 GB 规划后不再 OOM,batch 与序列长度都在预算内,WSL2 下可用量比原生 Windows 略高。
这里还有个坑。消费级卡在 Windows 上跑不掉图形占用,规划显存要按可用量算而不是标称容量,能省掉一轮 OOM。改动前把原配置备份一份,出问题能马上回滚,这比任何技巧都实用。
