PyTorch单机多卡训练(无废话)

目前大家基本都在使用DistributedDataParallel(简称DDP)用来训练,该方法主要用于分布式训练,但也可以用在单机多卡。

第一步:初始化分布式环境,主要用来帮助进程间通信

复制代码
torch.distributed.init_process_group(backend='nccl')

第二步:负责创建 args.local_rank 变量,并接受 torch.distributed.launch 注入的值

归根到底是创建一个变量,来接收torch.distributed.launch 注入。

目前代码中常见的两种方式:

复制代码
local_rank = int(os.environ["LOCAL_RANK"])   # 这种是从自定义config文件中获取LOCAL_RANK

另外一种是parser:

复制代码
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int, default=-1)
args = parser.parse_args()

第三步:每个进程根据自己的local_rank设置应该使用的GPU

复制代码
torch.cuda.set_device(args.local_rank)
torch.manual_seed(hps.train.seed)   # 非必要,尽可能固定种子

第四步:分布式数据和模型

复制代码
# 分布式数据
train_sampler = DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size) 
# 分布式模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

第五步:运行

nproc_per_node表示每个node有多少个进程,每个GPU对应一个进程。

nnodes表示使用几个节点,一个节点对应一台电脑。

复制代码
python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py 

解释

  1. torch.distributed.launch参数解析(终端运行命令的参数)

    --> python -m torch.distributed.launch --help

    usage: launch.py [-h] [--nnodes NNODES] [--node_rank NODE_RANK]
    [--nproc_per_node NPROC_PER_NODE] [--master_addr MASTER_ADDR] [--master_port MASTER_PORT]
    [--use_env] [-m] [--no_python] [--logdir LOGDIR]
    training_script ...

  • nnodes:节点的数量,通常一个节点对应一个主机,方便记忆,直接表述为主机
  • node_rank:节点的序号,从0开始
  • nproc_per_node:一个节点中显卡的数量
  • master_addr:master节点的ip地址,也就是0号主机的IP地址,该参数是为了让 其他节点 知道0号节点的位,来将自己训练的参数传送过去处理
  • master_port:master节点的port号,在不同的节点上master_addr和master_port的设置是一样的,用来进行通信
  1. torch.ditributed.launch相关环境变量解析(代码中os.environ中的参数)
  • WORLD_SIZE:os.environ["WORLD_SIZE"]所有进程的数量
  • LOCAL_RANK:os.environ["LOCAL_RANK"]每张显卡在自己主机中的序号,从0开始
  • RANK:os.environ["RANK"]进程的序号,一般是1个gpu对应一个进程
  1. 一般WORLD_SIZE和RANK和LOCAL_RANK在一起声明

    rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])
    n_gpus = dist.get_world_size()

  2. 多机和单机的一些概念问题

  • rank
    多机多卡:代表某一台机器
    单机多卡:代表某一块GPU
  • world_size
    多机多卡:代表有几台机器
    单机多卡:代表有几块GPU
  • local_rank
    多机多卡:代表某一块GPU的编号
    单机多卡:代表某一块GPU的编号

参考文献

1.https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel

  1. https://github.com/jia-zhuang/pytorch-multi-gpu-training

  2. https://blog.csdn.net/weixin_44966641/article/details/121872773

  3. https://blog.csdn.net/magic_ll/article/details/122359490

相关推荐
esmap1 分钟前
ESMAP 智慧消防解决方案:以数字孪生技术构建全域感知消防体系,赋能消防安全管理智能化升级
人工智能·物联网·3d·编辑器·智慧城市
LaughingZhu6 分钟前
Product Hunt 每日热榜 | 2026-02-08
大数据·人工智能·经验分享·搜索引擎·产品运营
sensen_kiss6 分钟前
INT303 Coursework1 爬取影视网站数据(如何爬虫网站数据)
爬虫·python·学习
芷栀夏15 分钟前
CANN ops-math:筑牢 AI 神经网络底层的高性能数学运算算子库核心实现
人工智能·深度学习·神经网络
用户51914958484516 分钟前
CVE-2025-47812:Wing FTP Server 高危RCE漏洞分析与利用
人工智能·aigc
阿里云大数据AI技术21 分钟前
【AAAI2026】阿里云人工智能平台PAI视频编辑算法论文入选
人工智能
玄同76523 分钟前
我的 Trae Skill 实践|使用 UV 工具一键搭建 Python 项目开发环境
开发语言·人工智能·python·langchain·uv·trae·vibe coding
Yorlen_Zhang33 分钟前
Python Tkinter Text 控件完全指南:从基础编辑器到富文本应用
开发语言·python·c#
苍何44 分钟前
腾讯重磅开源!混元图像 3.0 图生图真香!
人工智能
千里马也想飞1 小时前
人工智能在医疗领域的应用与研究论文写作实操:AI辅助快速完成框架+正文创作
人工智能