PyTorch分布式训练怎么选卡

在 PyTorch 分布式训练里，“选卡”通常分为物理选哪些 GPU和逻辑上怎么分配 GPU 给进程两层意思。下面按常见场景系统讲一下。

一、先明确：你用的是哪种分布式方式？

不同方式“选卡”方法不一样：

1. 单机多卡（最常见）

• DataParallel（不推荐）
• DistributedDataParallel (DDP) ✅ 推荐

1. 多机多卡

• 一定是 DDP + 多进程

1. 模型并行 / 流水并行

• 手动指定 layer 放哪张卡

下面主要讲 DDP 的选卡方式。

二、单机多卡：如何“选卡”

✅ 方法 1：用 CUDA_VISIBLE_DEVICES（最常用）

这是物理选卡的标准做法。

CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py

含义：

• 程序里只能看到 4 张卡
• 逻辑编号：
• cuda:0 → 物理 GPU 0
• cuda:1 → 物理 GPU 1
• …

✅ 适合：

• 只跑部分卡
• 避免占用全部 GPU

✅ 方法 2：DDP 中按 local_rank 选卡（推荐）

import os
import torch
import torch.distributed as dist
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

启动方式：

torchrun --nproc_per_node=4 train.py

✅ 优点：

• 每张卡一个进程
• 自动绑定 GPU
• 不会选错卡

✅ 方法 3：手动指定（不推荐但可用）

gpu_id = [0, 2, 3]
torch.cuda.set_device(gpu_id[local_rank])

⚠️ 容易混乱，除非你非常清楚拓扑结构。

三、多机多卡：怎么选卡

每台机器都：

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

• node_rank 区分机器
• local_rank 区分本机 GPU
• 物理卡通过 CUDA_VISIBLE_DEVICES 控制

四、高级：根据拓扑选卡（NVLink / PCIe）

查看 GPU 拓扑

nvidia-smi topo -m

优先选择：

• ✅ NVLink 连接的卡（通信快）
• ❌ 跨 CPU socket 的卡（慢）

例如：

GPU0 GPU1 GPU2 GPU3
GPU0  X  NV4  SYS SYS

✅ 选 0,1 而不是 0,2

五、常见错误

❌ 错误 1：

device = "cuda:0"

所有进程都用 0 号卡 → OOM
❌ 错误 2：

CUDA_VISIBLE_DEVICES=0,1
torchrun --nproc_per_node=4

只有 2 张卡，却开 4 个进程
✅ 正确：

CUDA_VISIBLE_DEVICES=0,1
torchrun --nproc_per_node=2

六、一句话总结

PyTorch 分布式训练选卡 = CUDA_VISIBLE_DEVICES 控制物理卡 + local_rank 控制逻辑卡
如果你愿意，可以告诉我：

• 单机还是多机？
• 几张卡？
• 使用 torchrun 还是 mp.spawn？

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