1. 多GPU环境配置基础

当你第一次面对服务器上琳琅满目的GPU设备时,可能会感到既兴奋又困惑。我清楚地记得第一次使用8卡A100服务器时的场景——就像走进了一个糖果店的孩子,却不知道从哪颗糖开始品尝。PyTorch的多GPU管理其实就像给这些"糖果"分配不同的品尝顺序,而CUDA_VISIBLE_DEVICES就是你的分配工具。

理解设备编号规则是第一步。假设服务器装有4块GPU,系统默认编号为0到3。但这里有个关键细节容易被忽视:当你设置CUDA_VISIBLE_DEVICES=2,3时,在程序中它们会变成设备0和1。这种重新编号机制经常让新手踩坑,我就曾因此浪费了半天调试时间。

临时设置环境变量是最快上手的方式。在Linux终端中,只需执行:

export CUDA_VISIBLE_DEVICES=0,1  # 使用前两块GPU

Windows用户则应该使用:

set CUDA_VISIBLE_DEVICES=0,1

但在实际项目中,我建议采用更稳妥的Python代码设置方式:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,2'  # 使用第1和第3块GPU

这种方式特别适合在Jupyter Notebook中进行快速实验,避免了每次重启kernel都要重新设置环境变量的麻烦。

2. 深入理解设备可见性机制

很多开发者以为CUDA_VISIBLE_DEVICES只是简单的过滤器,其实它的工作机制要精妙得多。这个环境变量实际上创建了一个虚拟的设备映射层,重新定义了GPU的物理编号。这种设计带来了灵活性,但也埋下了一些陷阱。

让我分享一个实际案例:在一次多任务调度系统中,我们同时运行三个训练任务,分别设置CUDA_VISIBLE_DEVICES为"0,1"、"2,3"和"1,3"。结果第三个任务意外抢占了第一个任务的资源,导致OOM错误。这是因为我们没有考虑到设备编号的重新映射规则。

正确的设备隔离应该这样做:

# 任务1
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
device = torch.device('cuda:0')  # 实际使用物理GPU 0

# 任务2
os.environ['CUDA_VISIBLE_DEVICES'] = '2,3'
device = torch.device('cuda:0')  # 实际使用物理GPU 2

对于长期运行的训练任务,我强烈推荐在.bashrc中设置永久变量。但要注意这会影响所有在该用户下运行的程序:

echo 'export CUDA_VISIBLE_DEVICES=0,1' >> ~/.bashrc
source ~/.bashrc

3. PyTorch原生设备控制API

除了环境变量,PyTorch提供了一组更灵活的设备控制API。torch.cuda模块就像是一个精细的GPU操作面板,让你可以实时查看和调整设备状态。

我最常用的几个工具方法包括:

# 获取当前活跃设备
current = torch.cuda.current_device()
print(f"当前使用设备:{current}")

# 查询设备数量
count = torch.cuda.device_count()
print(f"可用设备数:{count}")

# 获取设备名称
name = torch.cuda.get_device_name(0)
print(f"设备0名称:{name}")

手动设置设备虽然灵活,但在大型项目中容易出错。我建议采用这种模式:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

这种写法既保证了代码在没有GPU的环境下也能运行,又明确了设备选择逻辑。

4. DataParallel实战技巧

torch.nn.DataParallel是PyTorch中最简单的多GPU并行方案,但也是最容易用错的。经过多次踩坑,我总结出了几个关键要点:

首先,设备ID的连续性要求是个大坑。DataParallel要求device_ids必须是连续的,比如[0,1,2]。如果设置[0,2]就会报错。解决方案是:

# 正确做法
model = nn.DataParallel(model, device_ids=[0,1,2])

# 错误示范(会报错)
model = nn.DataParallel(model, device_ids=[0,2])

其次,主设备的选择需要特别注意。DataParallel总是使用device_ids[0]作为主设备,所有梯度计算都会汇总到这里。如果主设备负载过高,可以这样平衡:

# 将较大batch的数据放在次要设备上
inputs = inputs.to(f'cuda:{model.device_ids[1]}')

最后,我强烈建议在使用DataParallel时添加这个异常处理:

try:
    output = model(input)
except RuntimeError as e:
    if 'out of memory' in str(e):
        print('遇到OOM错误,尝试减小batch size')
    raise

5. 异构GPU环境下的优化策略

现实中的服务器往往配置不同型号的GPU,这种异构环境需要特殊处理。我曾管理过混合了Titan RTX和2080Ti的集群,总结出以下经验:

首先,通过性能分析确定各设备的实际算力:

from torch.utils.benchmark import Timer

def benchmark_device(device_id):
    device = f'cuda:{device_id}'
    x = torch.randn(10000, 10000, device=device)
    timer = Timer(
        stmt='x @ x',
        globals={'x': x},
        label=f'Device {device_id}'
    )
    print(timer.timeit(100))

然后,根据性能差异动态分配batch size:

class HeterogeneousDataParallel(nn.Module):
    def __init__(self, model, device_ratios):
        super().__init__()
        self.devices = list(device_ratios.keys())
        self.ratios = list(device_ratios.values())
        self.models = [model.to(f'cuda:{d}') for d in self.devices]
        
    def forward(self, x):
        chunks = torch.split(x, [int(r*x.size(0)) for r in self.ratios])
        results = [m(c.to(f'cuda:{d}')) 
                  for m, c, d in zip(self.models, chunks, self.devices)]
        return torch.cat(results, dim=0)

6. 常见问题与调试技巧

多GPU编程中90%的问题都集中在设备通信和内存管理上。这里分享几个实用的调试命令:

查看各GPU内存使用情况:

print(torch.cuda.memory_summary())

手动清空缓存(解决内存碎片问题):

torch.cuda.empty_cache()

检查设备间的P2P通信是否启用:

for i in range(torch.cuda.device_count()):
    for j in range(torch.cuda.device_count()):
        print(f'{i}->{j}:', torch.cuda.can_device_access_peer(i,j))

当遇到神秘的CUDA错误时,这个调试流程屡试不爽:

  1. 首先检查torch.cuda.is_available()
  2. 确认CUDA_VISIBLE_DEVICES设置正确
  3. 使用nvidia-smi查看实际GPU状态
  4. 逐步缩小问题范围,从最小示例开始测试

7. 进阶:自定义设备分配策略

对于需要精细控制的大型项目,可以构建自己的设备管理器。这是我常用的一个模板:

class GPUManager:
    def __init__(self):
        self.lock = threading.Lock()
        self.devices = list(range(torch.cuda.device_count()))
        self.allocated = set()
    
    def acquire(self, n=1):
        with self.lock:
            available = [d for d in self.devices if d not in self.allocated]
            if len(available) < n:
                raise RuntimeError(f'需要{n}个GPU,但只有{len(available)}个可用')
            selected = available[:n]
            self.allocated.update(selected)
            return selected
    
    def release(self, devices):
        with self.lock:
            self.allocated.difference_update(devices)

# 使用示例
manager = GPUManager()
try:
    my_devices = manager.acquire(2)
    os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, my_devices))
    # 执行训练代码
finally:
    manager.release(my_devices)

在多团队共享的GPU集群中,这种资源管理机制可以避免冲突。我还经常配合使用Slurm等作业调度系统,实现更复杂的分配策略。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐