如何指定PyTorch脚本使用特定GPU单元？解决CUDA显存不足问题

问题描述 #

• 使用Kohya Trainer训练模型时触发显存不足错误：

显存不足错误：尝试分配2.00 MiB显存失败（GPU 1；总容量23.65 GiB；已分配144.75 MiB；剩余2.81 MiB；PyTorch总共预留146.00 MiB）。若预留显存远大于已分配显存，可尝试设置max_split_size_mb避免内存碎片，详见PyTorch内存管理文档及PYTORCH_CUDA_ALLOC_CONF配置。

• 排查发现脚本默认使用GPU1，但GPU1显存紧张，GPU0资源充足，切换GPU操作均失败：
  1. 修改text_encoder.to("cuda")为text_encoder.to("cuda:0")后，脚本仍占用GPU1
  2. 设置环境变量CUDA_VISIBLE_DEVICES=0时，出现错误：

  运行时错误：CUDA错误：无效的设备序号

系统GPU配置 #

执行以下Python代码检测GPU：

import torch
for i in range(torch.cuda.device_count()):
    print(torch.cuda.get_device_properties(i))

输出结果：

_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', 主版本号=8, 次版本号=9, 总显存=24217MB, 多处理器数量=128)
_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', 主版本号=8, 次版本号=9, 总显存=24217MB, 多处理器数量=128)

解决方案 #

全局强制指定默认GPU #

在脚本开头加入以下代码，让PyTorch所有操作默认使用GPU0：

import torch
torch.cuda.set_device(0)

统一修改脚本中设备指定 #

搜索脚本内所有.to("cuda")或.to(device)的代码行，替换为.to("cuda:0")，或定义全局设备变量统一管理：

device = torch.device("cuda:0")
text_encoder.to(device)
unet.to(device)
vae.to(device)
# 其他需指定设备的组件同理调整

正确配置环境变量 #

按系统类型设置环境变量后启动脚本：

• Linux/macOS终端：

CUDA_VISIBLE_DEVICES=0 python train_network.py

• Windows命令提示符：

set CUDA_VISIBLE_DEVICES=0 && python train_network.py

• Windows PowerShell：

$env:CUDA_VISIBLE_DEVICES=0; python train_network.py

若仍报设备序号错误，重启终端重试，排查是否有其他程序篡改环境变量。

清理GPU1占用进程 #

执行nvidia-smi命令查看GPU1的占用进程，关闭无关进程释放显存资源。

内容的提问来源于stack exchange，提问作者Raptor