1. 问题复现：当PyTorch对多GPU视而不见时

上周调试模型时遇到了诡异现象：实验室那台8卡服务器用nvidia-smi明明显示所有GPU都空闲，但我的PyTorch代码死活只能用到0号卡。更糟的是，当尝试用torch.cuda.set_device(4)指定其他卡时，直接报错"invalid device ordinal"。这种单卡囚徒困境相信不少人都遇到过——服务器硬件资源充足，但框架却像戴了眼罩，只能摸到最近的那张显卡。

通过torch.cuda.device_count()验证，返回值果然是1。这意味着PyTorch的CUDA后端虽然能正常工作，但其可见设备范围被某种机制限制了。有趣的是，这种情况通常与CUDA驱动无关——既然nvidia-smi能识别全部显卡，说明驱动层是正常的，问题出在框架与驱动之间的通信层。

2. 环境变量陷阱：CUDA_VISIBLE_DEVICES的生效逻辑

2.1 变量设置的三种途径

排查多GPU识别问题首先要理解CUDA_VISIBLE_DEVICES这个关键环境变量。它能控制CUDA程序可见的GPU范围，但设置方式不同会导致生效时机差异：

1. 终端临时设置

   CUDA_VISIBLE_DEVICES=0,1 python train.py
   

   这种方式仅对当前终端会话有效，适合快速测试不同显卡组合

2. 代码内硬编码

   import os
   os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
   

   注意！这段代码必须放在所有CUDA相关操作之前，包括import torch。我曾在import torch之后设置，结果完全无效

3. Shell配置文件（推荐）
   在~/.bashrc或~/.zshrc末尾添加：

   export CUDA_VISIBLE_DEVICES=0,1,2,3
   

   记得用source ~/.bashrc使其生效

2.2 变量冲突的经典案例

最隐蔽的坑莫过于多层级设置的冲突。比如：

• .bashrc里设置了export CUDA_VISIBLE_DEVICES=0
• 代码里又写os.environ["CUDA_VISIBLE_DEVICES"]="0,1"
• 实际运行时框架可能只会看到0号卡

建议用以下命令检查最终生效值：

import os
print("Effective CUDA_VISIBLE_DEVICES:", os.environ.get("CUDA_VISIBLE_DEVICES"))

3. 组合拳解决方案：环境变量与代码的协同配置

经过多次测试，我发现最可靠的配置方案是双重确认法：

1. 基础配置
   在.bashrc设置默认可见设备（比如所有卡）：

   export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
   

2. 代码层验证
   在Python脚本开头添加设备检查：

   import torch
   print("可见GPU数量:", torch.cuda.device_count())
   print("当前设备:", torch.cuda.current_device())
   

3. 运行时灵活覆盖
   若需要临时调整可见设备，可以用：

   CUDA_VISIBLE_DEVICES=2,3 python train.py
   

这种组合方式既保证了默认情况下的全卡可用，又保留了临时调整的灵活性。实测在PyTorch 1.8+和CUDA 11.x环境下稳定有效。

4. 深度排查：当常规方法失效时

4.1 检查CUDA与PyTorch版本匹配

版本不兼容是隐形杀手，用以下命令检查：

nvcc --version  # CUDA编译器版本
python -c "import torch; print(torch.__version__)"  # PyTorch版本
python -c "import torch; print(torch.version.cuda)"  # PyTorch编译时的CUDA版本

这三个版本应该相互兼容。常见问题包括：

• 用conda安装的PyTorch可能自带CUDA运行时，与系统CUDA冲突
• pip安装时误选错误版本（如cpu版本）

4.2 设备锁定检测

有些情况下GPU会被进程隐式锁定：

fuser -v /dev/nvidia*  # 查看占用GPU的进程

如果发现未知进程占用，可以用kill -9 [PID]终止（需管理员权限）

4.3 容器环境特殊处理

在Docker中使用GPU时需要特别注意：

1. 必须添加--gpus all参数
2. 确保容器内安装了匹配的CUDA驱动
3. 检查NVIDIA Container Toolkit是否安装：

   docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
   

5. 多GPU编程的最佳实践

成功识别全卡后，建议采用这些模式避免常见问题：

5.1 设备选择的安全写法

import torch
device_id = 4  # 目标设备号
assert device_id < torch.cuda.device_count(), f"设备{device_id}不可用"
device = torch.device(f"cuda:{device_id}")

5.2 自动负载均衡

用循环自动分配任务到不同GPU：

devices = [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
for batch_idx, data in enumerate(dataloader):
    target_device = devices[batch_idx % len(devices)]
    data = data.to(target_device)
    # 后续计算...

5.3 内存优化技巧

即使多卡可用，单卡内存不足仍可能报错。可以：

• 启用torch.backends.cudnn.benchmark = True加速卷积运算
• 使用梯度检查点技术：

  from torch.utils.checkpoint import checkpoint
  def custom_forward(x):
      # 定义前向计算
      return x
  output = checkpoint(custom_forward, input)
  

记得在代码部署前用torch.cuda.empty_cache()清理缓存，这对长期运行的训练任务尤其重要。当所有方法都试过仍不奏效时，终极方案是重建CUDA环境：彻底卸载驱动和CUDA后，按照PyTorch官网推荐的版本组合重新安装。