torch-xla在多GPU下NCCL ncclUniqueId的生成过程

money0534

已于 2024-09-04 11:17:10 修改

阅读量482

点赞数 3

分类专栏： deep learning 文章标签： XLA torch-xla NCCL GRPC distributed

于 2024-03-07 16:28:05 首次发布

本文链接：https://blog.csdn.net/liuzonrze/article/details/136536462

版权

deep learning 专栏收录该内容

12 篇文章 0 订阅

订阅专栏

本文基于torch-xla2.0，结论如下：

ncclUniqueId就是当前主机的IP和端口号；
在一个通信组内，每个rank使用相同的ncclUniqueId；
多个通信组时，各个通信组内ncclUniqueId相同，每个通信组的ncclUniqueId不同；

生成过程：在rank 0上创建MeshService，也就是GRPC的服务端，每个rank创建一个MeshClient，即GRPC的客户端，客户端通过GRPC请求向服务端获取ncclUniqueId（因为多进程环境，所以请求天然是多线程的，响应也同理，NcclUidManager::GetNcclUniqueUid线程安全）；由于一个通信组内的请求参数replicas（即rank列表，例如2个rank时：0,1）相同，所以总是得到相同的ncclUniqueId。

这是一个典型的ncclUniqueId生成和应用场景。对于不能提供CPU侧通信机制的场景，可以使用环境变量NCCL_COMM_ID实现，且每个进程的需要相同的环境变量值（实测不同时，无法建立通信）。
参考NCCL issues 730 @sjeaugey commented on Oct 10, 2022：

NCCL_COMM_ID is the current way for NCCL to initialize without a CPU-side communication mechanism. Having a CPU-side communication mechanism is arguably expensive but it usually needed anyway on most parallel computing applications (MPI can be used for that purpose).