centos pytorch如何多卡训练

在 CentOS 上进行 PyTorch 多卡训练，常见做法有 DataParallel（不推荐） 和 DistributedDataParallel（推荐）。下面按实用角度给你一套完整说明。

一、环境与准备（CentOS）

1. 查看 GPU

nvidia-smi

确认多张卡（如 0,1,2,3）

2. 安装 PyTorch（GPU 版）

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证：

import torch
print(torch.cuda.device_count())

二、方式一：DistributedDataParallel（✅ 推荐）

1. 启动方式（torchrun）

CentOS 下最推荐用 torchrun

torchrun \
  --nproc_per_node=4 \
  train.py

• 4 = 你机器上的 GPU 数量
• 单机多卡不用改 IP

2. 训练代码模板（train.py）

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

def main():
    dist.init_process_group(backend="nccl")

    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])

    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    for x, y in loader:
        x, y = x.cuda(local_rank), y.cuda(local_rank)
        loss = model(x, y)
        loss.backward()
        optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

3. 注意事项（CentOS 常见坑）

• ✅ 必须用 nccl
• ✅ batch_size 是每张卡的
• ✅ DistributedSampler 必须加
• ✅ 数据不要重复 shuffle

三、方式二：DataParallel（⚠️ 不推荐）

model = torch.nn.DataParallel(model.cuda())

缺点：

• 主卡显存爆
• 速度慢
• 多进程支持差

只适合快速测试

四、多机多卡（扩展）

如果是多台 CentOS 服务器：

torchrun \
  --nnodes=2 \
  --node_rank=0 \
  --master_addr=192.168.1.1 \
  --master_port=29500 \
  --nproc_per_node=4 \
  train.py

五、常见问题

1. NCCL 报错

export NCCL_DEBUG=INFO

2. 权限问题

ulimit -n 65535

3. 冻结代码 / 死锁

• 检查 barrier()
• 避免在主进程外 print 太多

六、总结

如果你愿意，可以告诉我：

• GPU 数量
• PyTorch 版本
• 是否是容器 / 集群 / Slurm

我可以给你完全贴合你环境的多卡训练脚本。