ubuntu pytorch怎么多卡训练

在 Ubuntu 上用 PyTorch 做多卡训练，常见做法有以下几种，从简单到进阶：

一、最常用：torch.nn.DataParallel（简单但不推荐大模型）

适合快速上手、卡数不多（≤2~4）

import torch
import torch.nn as nn

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

⚠️ 缺点：

• 主卡负载高
• 多进程通信效率低
• 不适合大模型 / 多机

二、推荐：torch.nn.parallel.DistributedDataParallel（DDP）

官方推荐，性能最好

1️⃣ 单机多卡（最常见）

启动方式（推荐 torchrun）

torchrun \
  --nproc_per_node=4 \
  train.py

（4 = GPU 数量）

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

# 初始化
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# 数据
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=32)

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

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

✅ 优点：

• 每张卡负载均衡
• 支持多机
• 训练速度最快

三、多机多卡（简单说明）

# 机器1
torchrun --nnodes=2 --node_rank=0 --master_addr=192.168.1.1 --master_port=29500 --nproc_per_node=8 train.py

# 机器2
torchrun --nnodes=2 --node_rank=1 --master_addr=192.168.1.1 --master_port=29500 --nproc_per_node=8 train.py

四、查看 GPU 是否可用

nvidia-smi
python -c "import torch; print(torch.cuda.device_count())"

五、常见坑

❗ 必须用 DistributedSampler
❗ batch_size 是每张卡的
❗ 保存模型只用一个进程：

if local_rank == 0:
    torch.save(model.state_dict(), "model.pth")

六、如果你用 HuggingFace / Lightning

• Transformers

accelerate launch train.py

• PyTorch Lightning

Trainer(accelerator="gpu", devices=4, strategy="ddp")

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

• 你是 单机几卡
• 用的是 原生 PyTorch / HF / Lightning
• 模型大小（小模型 / LLM）

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