PyTorch Distributed 与 Collective Communication
PyTorch Distributed 与 Collective Communication torch distributed torch distributed以进程组(process group, PG) 作为整个通信系统，每个PG有自己的world_size与成员process，每个process拥有自己的rank，以此为基础进行通信。通信后端支持nccl(n卡GPU), gloo(CPU), MPI(各种Message Passing Interface是西安)以及自定义PG后端, torch 2.x还支持ucc(Nvidia推出的统一collective通信库, xccl(XPU Collective Communication Library, 围绕Intel GPU/XPU). torch distributed通常需要定义一些进程内环境变量： MASTER_ADDR: 表示RANK0所在机器的地址(IP/hostname) WORLD_SIZE: PG成员数量 RANK: 当前process的全局ID(0~WORLD_SIZE-1) LOCAL_RANK: process在本机内的编号 LOCAL_RANK 通常用来设置初始时torch cuda设置device，比如: local_rank=init(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) torch.distributed里和PG相关的操作主要有: init_process_group : 初始化默认PG destroy_process_group: 摧毁释放指定PG,未指定则释放默认PG new_group : 在默认PG的基础上，按照指定ranks创建子PG is_initialized: 默认PG是否被初始化 get_backend: 获取指定PG的backend get_world_size:获取指定PG的world_size MASTER_PORT: RANK0上的一个空闲窗口 与communication相关的操作主要有: torch.distributed.broadcast(tensor, src, group=None) torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=None) torch.distributed.reduce(tensor, dst, op=..., group=None) torch.distributed.all_gather(tensor_list, tensor, group=None) torch.distributed.gather(tensor, dst, group=None) torch.distributed.scatter(tensor, scatter_list, src, group=None) torch.distributed.reduce_scatter(output, input_list, op=..., group=None) torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None) torch.distributed.all_gather_object(object_list, obj, group=None) torch.distributed.broadcast_object_list(obj_list, src, group=None) 除去上述这8中collective(还有一些和python object相关)，还支持P2P通信: ...