linux pytorch怎样多卡

在 Linux 下使用 PyTorch 进行多卡（多 GPU）训练，常见方式有以下几种，从简单到进阶：

一、最常用：torch.nn.DataParallel（简单但不推荐新项目）

适合单机多卡、快速上手，但效率不如 DistributedDataParallel。

import torch
import torch.nn as nn

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

注意：

• 主卡（GPU 0）负载更高
• 不支持多机
• PyTorch 官方更推荐下面这种方式

二、推荐方式：DistributedDataParallel（DDP，单机/多机都支持）

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

启动方式（推荐 torchrun）

torchrun \
  --nproc_per_node=4 \
  train.py

train.py 示例

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(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 = torch.utils.data.distributed.DistributedSampler(dataset)
    loader = torch.utils.data.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()

✅ 优点：

• 速度快
• 负载均衡
• 支持多机

三、多机多卡（DDP + 网络）

假设：

• 2 台机器
• 每台 4 张卡

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

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

四、查看 GPU 使用情况

nvidia-smi

五、常见问题

❓ 只用了一张卡？

• 没用 DistributedSampler
• batch size 没乘 GPU 数
• 使用了 DataParallel 但代码写错

❓ CUDA out of memory

• 减小 batch_size
• 使用 gradient_accumulation
• 使用 torch.cuda.empty_cache()

六、总结建议

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

• 你是 单机几卡
• 用的是 CNN / Transformer / LLM
• 是否用 HuggingFace

我可以直接给你一份可运行的多卡训练代码。