译者:univeryinli
torch.distributed
支持三个后端,每个后端具有不同的功能。下表显示哪些功能可用于CPU/CUDA张量。仅当用于构建PyTorch的实现支持时,MPI才支持CUDA。
后端 | gloo | mpi | nccl |
---|---|---|---|
设备 | CPU | GPU | CPU |
— | — | — | — |
发送 | ✓ | ✘ | ✓ |
接收 | ✓ | ✘ | ✓ |
广播 | ✓ | ✓ | ✓ |
all_reduce | ✓ | ✓ | ✓ |
reduce | ✓ | ✘ | ✓ |
all_gather | ✓ | ✘ | ✓ |
收集 | ✓ | ✘ | ✓ |
分散 | ✓ | ✘ | ✓ |
屏障 | ✓ | ✘ | ✓ |
目前PyTorch分发版仅支持Linux。默认情况下,Gloo和NCCL后端构建并包含在PyTorch的分布之中(仅在使用CUDA构建时为NCCL)。MPI是一个可选的后端,只有从源代码构建PyTorch时才能包含它。(例如,在安装了MPI的主机上构建PyTorch)
在过去,我们经常被问到:“我应该使用哪个后端?”。