pytorch 多机单卡分布式训练配置笔记

记录通过torchrun进行pytorch的分布式训练配置方法，示例代码为基本的分布式训练框架代码，无实际功能

环境

操作系统：Ubuntu 22.04
Python环境：anaconda 23.11.0、Python 3.8
pytorch：2.1.2

编写代码

将代码保存为main.py
模型训练代码写到train函数中，本示例仅做打印输出

import argparse
import os
import sys
import tempfile
from time import sleep
from urllib.parse import urlparse

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import datetime
from torch.nn.parallel import DistributedDataParallel as DDP

def train():
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    while True:
        print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...")
        sleep(1)

def run():
    env_dict = {
        key: os.environ[key]
        for key in ("MASTER_ADDR", "MASTER_PORT", "WORLD_SIZE", "LOCAL_WORLD_SIZE")
    }
    print(f"[{os.getpid()}] Initializing process group with: {env_dict}")
    dist.init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=30))
    #dist.init_process_group(backend="nccl")
    train()
    dist.destroy_process_group()


if __name__ == "__main__":
    run()

运行

以torchrun方式运行,此命令无需单独安装，随pytorch一同安装

1. 运行主节点
   torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="127.0.0.1" --master_port=1234 main.py
2. 运行工作节点
   torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="127.0.0.1" --master_port=1234 main.py

参数解释

• nproc_per_node 每个节点线程数
• nnodes 总节点数
• node_rank 当前节点编号，主节点必须为0
• master_addr 主节点ip地址
• master_port 主节点端口号

上述示例命令先运行了主节点，指定每个节点创建四个线程，一共有2个节点，主节点地址为127.0.0.1端口为1234。主启动完成后再启动工作节点。主节点不会执行训练。只有工作节点进行训练