https://ccalwayscc.com/tags/pytorch/Recent content in PyTorch on ccHugozh-CNSat, 01 Aug 2026 00:00:00 +0800https://ccalwayscc.com/posts/torch-distributed/Sat, 01 Aug 2026 00:00:00 +0800https://ccalwayscc.com/posts/torch-distributed/<h1 id="pytorch-distributed-与-collective-communication">PyTorch Distributed 与 Collective Communication</h1> <h2 id="torch-distributed">torch distributed</h2> <p>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).</p> <p>torch distributed通常需要定义一些进程内环境变量：</p> <ul> <li><code>MASTER_ADDR</code>: 表示RANK0所在机器的地址(IP/hostname)</li> <li><code>WORLD_SIZE</code>: PG成员数量</li> <li><code>RANK</code>: 当前process的全局ID(0~WORLD_SIZE-1)</li> <li><code>LOCAL_RANK</code>: process在本机内的编号</li> </ul> <p><code>LOCAL_RANK</code> 通常用来设置初始时torch cuda设置device，比如:</p> <div class="highlight"><pre tabindex="0" style="color:#f8f8f2;background-color:#272822;-moz-tab-size:4;-o-tab-size:4;tab-size:4;-webkit-text-size-adjust:none;"><code class="language-python" data-lang="python"><span style="display:flex;"><span>local_rank<span style="color:#f92672">=</span>init(os<span style="color:#f92672">.</span>environ[<span style="color:#e6db74">&#34;LOCAL_RANK&#34;</span>]) </span></span><span style="display:flex;"><span>torch<span style="color:#f92672">.</span>cuda<span style="color:#f92672">.</span>set_device(local_rank) </span></span></code></pre></div><p>torch.distributed里和PG相关的操作主要有:</p> <ul> <li><code>init_process_group</code> : 初始化默认PG</li> <li><code>destroy_process_group</code>: 摧毁释放指定PG,未指定则释放默认PG</li> <li><code>new_group</code> : 在默认PG的基础上，按照指定ranks创建子PG</li> <li><code>is_initialized</code>: 默认PG是否被初始化</li> <li><code>get_backend</code>: 获取指定PG的backend</li> <li><code>get_world_size</code>:获取指定PG的world_size</li> <li><code>MASTER_PORT</code>: RANK0上的一个空闲窗口</li> </ul> <p>与communication相关的操作主要有:</p> <ul> <li><code>torch.distributed.broadcast(tensor, src, group=None)</code></li> <li><code>torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=None)</code></li> <li><code>torch.distributed.reduce(tensor, dst, op=..., group=None)</code></li> <li><code>torch.distributed.all_gather(tensor_list, tensor, group=None)</code></li> <li><code>torch.distributed.gather(tensor, dst, group=None)</code></li> <li><code>torch.distributed.scatter(tensor, scatter_list, src, group=None)</code></li> <li><code>torch.distributed.reduce_scatter(output, input_list, op=..., group=None)</code></li> <li><code>torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None)</code></li> <li><code>torch.distributed.all_gather_object(object_list, obj, group=None)</code></li> <li><code>torch.distributed.broadcast_object_list(obj_list, src, group=None)</code></li> </ul> <p>除去上述这8中collective(还有一些和python object相关)，还支持P2P通信:</p>