问题背景

最近在搭建多智能体强化学习的训练环境时，遇到了各种PyTorch多GPU的问题。经过一番折腾，总结了一些常见问题和解决方案。

😅 我踩过的坑

刚开始天真地以为只要有多张GPU，PyTorch就会自动利用，结果发现事情远没有那么简单…

🔧 常见问题及解决方案

1. CUDA版本不匹配

问题描述：

1
RuntimeError: CUDA error: no kernel image is available for execution on the device

解决方案：

1
2
3
4
5
6
7
8
9
10
11
12
13
# 查看CUDA版本
nvidia-smi
nvcc --version

# 卸载现有PyTorch
pip uninstall torch torchvision torchaudio

# 安装对应CUDA版本的PyTorch
# 对于CUDA 11.8
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 对于CUDA 12.1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

验证安装：

1
2
3
4
import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"CUDA version: {torch.version.cuda}")
print(f"GPU count: {torch.cuda.device_count()}")

2. 内存不足问题

问题描述：

1
RuntimeError: CUDA out of memory

解决方案：

方法1：减小batch size

1
2
3
4
5
# 之前
batch_size = 64

# 修改后
batch_size = 32  # 或者更小

方法2：梯度累积

1
2
3
4
5
6
7
8
9
10
11
12
accumulation_steps = 4
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 归一化损失
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

方法3：清理GPU缓存

1
2
3
4
5
6
7
import torch
import gc

# 在训练循环中定期清理
if batch_idx % 100 == 0:
    torch.cuda.empty_cache()
    gc.collect()

3. DataParallel vs DistributedDataParallel

**问题：**选择哪种并行方式？

DataParallel（简单但效率低）：

1
2
3
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)
model.to(device)

DistributedDataParallel（推荐）：

1
2
3
4
5
6
7
8
9
10
11
12
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化进程组
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

# 包装模型
model = model.to(rank)
model = DDP(model, device_ids=[rank])

4. 混合精度训练

**问题：**训练速度慢，GPU利用率不高

解决方案：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from torch.cuda.amp import autocast, GradScaler

model = model.to(device)
scaler = GradScaler()

for inputs, labels in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

5. 环境变量配置

常用的CUDA环境变量：

1
2
3
4
5
6
7
8
# 控制可见GPU
export CUDA_VISIBLE_DEVICES=0,1,2,3

# 内存增长策略
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

# 调试信息
export CUDA_LAUNCH_BLOCKING=1

🛠️ 调试技巧

1. 监控GPU使用情况

1
2
3
4
5
6
# 实时监控
watch -n 1 nvidia-smi

# 或者使用Python
import GPUtil
GPUtil.showUtilization()

2. 检查内存泄漏

1
2
3
4
5
6
7
8
9
import torch

def check_memory():
    if torch.cuda.is_available():
        print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
        print(f"Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

# 在训练循环中定期调用
check_memory()

3. 性能分析

1
2
3
4
5
6
7
8
9
10
11
12
import torch.profiler as profiler

with profiler.profile(
    activities=[profiler.ProfilerActivity.CPU, profiler.ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    # 你的训练代码
    pass

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

📝 最佳实践总结

  1. 环境配置：

    • 确保CUDA、PyTorch版本匹配
    • 使用conda管理环境更稳定

  2. 内存管理：

    • 适当的batch size
    • 定期清理GPU缓存
    • 使用混合精度训练

  3. 并行策略：

    • 小规模实验用DataParallel
    • 大规模训练用DistributedDataParallel

  4. 调试习惯：

    • 先在单GPU上验证代码
    • 使用小数据集测试并行效果
    • 监控GPU利用率和内存使用

希望这些经验能帮大家少踩一些坑！ 😄

🔗 参考资源