1. 多GPU环境下的设备管理挑战

第一次在实验室接触8卡GPU服务器时，我像发现新大陆一样兴奋。但很快就被现实打脸——明明指定了GPU 1训练模型，日志却显示占用率最高的是GPU 0。这种设备分配混乱的问题，相信很多刚接触多GPU开发的同学都遇到过。

PyTorch默认会使用GPU 0作为主设备，当我们不显式指定时，所有.cuda()操作都会跑到这张卡上。这就好比你去餐厅吃饭，服务员不问你要坐哪个位置，直接把你带到第一张空桌。在单GPU环境下没问题，但在多GPU场景下就会引发三个典型问题：

1. 资源争抢：多个训练任务挤在同一张GPU，显存很快爆掉
2. 负载不均：有些GPU跑满100%利用率，有些却在摸鱼
3. 调试困难：日志中的cuda:0让人分不清实际用的是哪张卡

# 典型的问题代码示例
model = MyModel()
model.cuda()  # 默认使用GPU 0
data = data.cuda()  # 也跑到GPU 0上

2. torch.cuda.set_device的局限性与替代方案

2.1 为什么官方不推荐set_device

官方文档中明确标注"不鼓励使用torch.cuda.set_device"，这句话背后其实有血泪教训。去年我们团队就踩过一个坑：在分布式训练脚本里混用set_device和DataParallel，导致GPU内存泄漏。根本原因在于：

• 作用域不明确：set_device会影响整个进程的设备状态
• 线程不安全：多线程环境下可能引发设备切换竞争
• 与并行API冲突：和nn.DataParallel/DistributedDataParallel配合使用时容易出错

# 危险的使用示例
torch.cuda.set_device(0)  # 全局设置
model = MyModel().cuda()  # 使用GPU 0

# 另一个线程可能同时执行
torch.cuda.set_device(1)  # 导致设备状态混乱

2.2 更优雅的替代方案

经过多次实践验证，我总结出三个更可靠的设备管理方案：

方案一：显式设备指定（推荐）

device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
model = MyModel().to(device)
data = data.to(device)

方案二：环境变量控制

# 启动时指定可见GPU
CUDA_VISIBLE_DEVICES=2,3 python train.py

方案三：上下文管理器

with torch.cuda.device(1):
    model = MyModel().cuda()  # 仅在上下文内使用GPU 1

这三种方式的对比：

方法	作用域	线程安全	兼容并行训练	易用性
set_device	全局	否	差	★★☆
显式to(device)	局部对象	是	优	★★★
CUDA_VISIBLE_DEVICES	进程级	是	良	★★☆
上下文管理器	代码块	是	良	★★★

3. 分布式训练中的最佳实践

3.1 DataParallel的隐藏陷阱

很多教程教我们用nn.DataParallel实现单机多卡训练，但很少有人提到它的三个致命缺陷：

1. 主卡瓶颈：GPU 0要汇总所有梯度，容易成为性能瓶颈
2. 显存不均：主卡需要额外存储完整模型参数，占用更多显存
3. 设备绑定：必须依赖set_device或默认设备

# 有问题的DataParallel用法
model = MyModel()
model = nn.DataParallel(model).cuda()  # 默认占用所有可见GPU

改进方案是配合环境变量使用：

# 只使用GPU 1,2做DataParallel训练
CUDA_VISIBLE_DEVICES=1,2 python train.py

3.2 DistributedDataParallel的正确姿势

在帮某电商客户优化推荐系统训练时，我们发现切换到DistributedDataParallel后训练速度提升了3倍。关键配置要点：

1. 每个进程独占GPU：避免设备竞争
2. 初始化后端：通常用NCCL效率最高
3. 数据分片：保证各GPU处理不同数据

# 分布式训练示例
import torch.distributed as dist

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)  # 这里set_device是安全的

def cleanup():
    dist.destroy_process_group()

4. 工程化部署的实用技巧

4.1 环境隔离方案

在容器化部署时，我习惯用docker的--gpus参数实现设备隔离：

# 为容器分配特定GPU
docker run --gpus '"device=2,3"' my_training_image

4.2 自动化设备选择

分享一个实用工具函数，可以自动选择空闲GPU：

def auto_select_gpu(min_memory=5000):
    """选择显存大于阈值的GPU"""
    for i in range(torch.cuda.device_count()):
        mem = torch.cuda.mem_get_info(i)[0] / 1024**2
        if mem > min_memory:
            return torch.device(f'cuda:{i}')
    raise RuntimeError('No available GPU')

4.3 监控与调试

这些命令可以帮助排查设备问题：

# 查看GPU使用情况
nvidia-smi -l 1  

# PyTorch内查看设备信息
python -c "import torch; print(f'可用GPU: {torch.cuda.device_count()}张')"

在模型开发中，我习惯在训练开始时打印完整的设备配置：

def print_device_info():
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA可用: {torch.cuda.is_available()}")
    print(f"当前设备: {torch.cuda.current_device()}")
    print(f"设备名称: {torch.cuda.get_device_name()}")
    print(f"设备内存: {torch.cuda.get_device_properties(0).total_memory/1024**3:.2f}GB")

记住一个原则：在多GPU环境下，显式指定设备永远比依赖默认行为更可靠。就像老司机常说的——不要让你的代码猜该用哪张卡，明确告诉它。