centos pytorch如何多卡

在 CentOS 上使用 PyTorch 进行多卡训练，常见方式有两种：DataParallel（DP） 和 DistributedDataParallel（DDP）。生产环境强烈推荐 DDP。

下面按“环境准备 → 单机多卡 → 多机多卡”说明。

一、环境准备（CentOS）

1. 确认 GPU 和驱动

nvidia-smi

能看到多张卡即正常。

2. 安装 CUDA / cuDNN

确保 CUDA 版本与 PyTorch 匹配：

nvcc -V

3. 安装 PyTorch（示例）

pip install torch torchvision torchaudio

或指定 CUDA 版本：

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

二、单机多卡（最常见）

方式 1：DataParallel（不推荐，简单）

import torch
import torch.nn as nn

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

⚠️ 缺点：

• 主卡负载高
• 不支持多机
• 性能不如 DDP

方式 2：DistributedDataParallel（✅ 推荐）

1. 启动方式（PyTorch ≥ 1.9）

torchrun \
  --nproc_per_node=4 \
  train.py

4 = 4 张 GPU

2. 示例代码（train.py）

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
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("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    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()

if __name__ == "__main__":
    main()

三、多机多卡（CentOS 集群）

1. 机器互通

• 主机名 / IP 可互通
• 防火墙放行端口（如 29500）

2. 启动命令

主节点（node0）

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

从节点（node1）

torchrun \
  --nnodes=2 \
  --node_rank=1 \
  --nproc_per_node=4 \
  --master_addr=192.168.1.10 \
  --master_port=29500 \
  train.py

四、常见问题（CentOS）

1. NCCL 报错

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

2. 显存不均衡

• 使用 DDP
• 减小 batch size
• 使用 gradient_checkpointing

3. 多用户共享机器

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun ...

五、推荐总结

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

• GPU 数量
• CentOS 版本
• PyTorch 版本
• 是否是 SLURM / 集群

我可以直接给你可运行的启动脚本。