分布式训练下的多进程环境
核心就是创建一个包含多个进程的进程组，进程之间可以相互通信
使用torch.distributed 在单机多进程环境中启动多个进程，实现分布式训练。
首先，在单个机器节点上生成两个进程，每个进程通过相同的主机IP地址和端口号进行通信。
然后，设置分布式环境，初始化进程组，旨在允许进程通过共享位置来相互通信。
- 导入需要的模块
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
os：设置环境变量。torch.distributed：用于分布式通信。torch.multiprocessing：用于多进程启动，与 Python 原生的multiprocessing类似，但兼容 CUDA。
- 启动多个进程
if __name__ == "__main__":size = 2 # 总进程数（即 world size）processes = []mp.set_start_method("spawn") # 设置启动方法，推荐使用 "spawn" 以兼容性更好for rank in range(size):p = mp.Process(target=init_process, args=(rank, size, run)) # 为每个 rank 创建一个进程p.start()processes.append(p)for p in processes:p.join() # 等待所有子进程完成
- 启动两个子进程（
size = 2），每个进程都会执行init_process。 - 每个进程分配一个唯一的
rank，从0到size-1。 - 设置
spawn模式更安全，尤其是多平台/多 GPU 情况下。
- 初始化进程组
def init_process(rank, size, fn, backend="gloo"):os.environ['MASTER_ADDR'] = "127.0.0.1" # 主节点 IPos.environ['MASTER_PORT'] = '29500' # 主节点端口dist.init_process_group(backend, rank=rank, world_size=size) # 初始化分布式进程组fn(rank, size) # 执行实际的任务函数
- 设置通信的主节点地址和端口（因为我们在单机上跑，IP 就用本地地址）。
- 使用
dist.init_process_group创建一个通信上下文，这样不同进程之间可以互相通信。 - 初始化后执行传入的
fn函数（即run函数）。
- 定义任务逻辑
def run(rank, size):"""具体要执行的分布式函数"""pass
- 每个进程在初始化好分布式环境后会执行这里的逻辑。
- 你可以在这里写
all_reduce、broadcast、send/recv等具体的分布式通信逻辑。 rank是当前进程的编号，size是总进程数。
简单总结一下：
几点：
创建进程，设置主机地址和端口号，然后初始化进程组
来源于PyTorch文档：https://docs.pytorch.org/tutorials/intermediate/dist_tuto.html#collective-communication