linux pytorch怎么多卡训练

在 Linux 下使用 PyTorch 进行多卡训练，常见方式有以下几种，按使用频率和推荐程度排序：

一、torch.nn.DataParallel（最简单，但不推荐用于多机/大项目）

适合单机多卡、快速上手。

import torch
import torch.nn as nn

model = MyModel()
model = nn.DataParallel(model)   # 默认使用所有 GPU
model = model.cuda()

训练代码不变，只是 loss 会在多卡上自动平均。

⚠️ 缺点：

• 主卡瓶颈明显
• 不支持多机
• 效率不如 DistributedDataParallel

二、torch.nn.parallel.DistributedDataParallel（✅ 官方推荐）

适合单机多卡 / 多机多卡，性能最好。

1️⃣ 启动方式（Linux 常用）

单机多卡（推荐）

torchrun --nproc_per_node=4 train.py

多机（示例）

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

2️⃣ 训练代码示例（关键部分）

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    torch.cuda.set_device(rank)

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

    dataset = MyDataset()
    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    loader = torch.utils.data.DataLoader(
        dataset, batch_size=32, sampler=sampler
    )

    for epoch in range(10):
        sampler.set_epoch(epoch)
        for x, y in loader:
            x, y = x.cuda(), y.cuda()
            loss = model(x, y)
            loss.backward()
            optimizer.step()

三、使用 Hugging Face Trainer（最简单高阶方案）

如果你用 transformers：

training_args = TrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    ddp_find_unused_parameters=False,
    local_rank=-1  # torchrun 会自动设置
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

启动：

torchrun --nproc_per_node=4 train.py

四、常见问题排查

✅ 查看 GPU

nvidia-smi

✅ 设置可见 GPU

export CUDA_VISIBLE_DEVICES=0,1,2,3

✅ 后端选择

• Linux + NVIDIA → nccl
• CPU → gloo

五、总结建议

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

• 是 单机几卡 还是 多机
• 用的 PyTorch 版本
• 是否是 NLP / CV / 自定义模型

我可以直接给你一份可运行模板。