deepSpeed多机多卡训练服务器之间,和服务器内两个GPU是怎么通信
DeepSpeed 在多机多卡训练时,主要依赖 NCCL 和 PyTorch Distributed 进行通信。具体来说,分为服务器之间和服务器内两种情况:
1. 服务器之间的通信(跨节点通信)
DeepSpeed 采用 NCCL(NVIDIA Collective Communications Library)作为主要的通信后端,结合 PyTorch Distributed 实现数据同步。其核心原理如下:
- 进程组通信:DeepSpeed 使用
torch.distributed
进行 进程组(Process Group) 管理,每个 GPU 都是一个独立进程,每个服务器上的所有进程形成一个集群。 - NCCL 后端通信:
- 主要使用 AllReduce、AllGather 和 ReduceScatter 进行梯度同步和参数聚合。
- NCCL 通过 InfiniBand 或 RDMA 提供高带宽、低延迟的通信能力。
- 默认使用 TCP 或 RDMA 进行数据传输,具体取决于网络配置。
- Rendezvous 机制:DeepSpeed 通过
torch.distributed.launch
或deepspeed.init_distributed()
自动建立通信,确保所有进程正确连接。
具体的通信流程
- 参数广播(Broadcast):所有进程从 rank=0 的主进程获取初始参数,并进行同步。
- 梯度同步(Gradient AllReduce):
- 每个 GPU 计算梯度后,使用 NCCL AllReduce 聚合所有梯度。
- 计算完成后,将梯度更新到本地参数中,确保所有节点的模型权重一致。
- 参数更新(Optimizer Step):
- DeepSpeed 可以选择 ZeRO 优化器,将优化计算分布到多个 GPU 上,以减少显存占用。
- 如果使用 ZeRO-3,DeepSpeed 还会分片存储模型参数,在参数更新时需要通信。
2. 服务器内的通信(单节点多 GPU 通信)
单机多卡情况下,DeepSpeed 仍然依赖 NCCL 进行 GPU 间的通信。核心机制如下:
- NCCL Ring-AllReduce:所有 GPU 通过 PCIe 或 NVLink 互连,形成环状拓扑,每个 GPU 负责传输数据的一部分,从而实现梯度聚合。
- Shared Memory:如果多张 GPU 之间有 NVLink,NCCL 会利用 CUDA IPC(Inter-Process Communication) 共享显存,提高数据传输效率。
- ZeRO Offload:如果显存不足,可以选择将部分计算转移到 CPU 进行。
具体通信方式
- GPU 之间数据同步
- 默认采用 NCCL AllReduce 实现所有 GPU 之间的梯度聚合。
- 如果使用 ZeRO-2 或 ZeRO-3,DeepSpeed 还会对优化器状态和参数进行分片存储和交换。
- 数据并行(Data Parallelism)
- 每个 GPU 负责不同的数据批次,训练过程中保持模型参数一致。
- DeepSpeed 通过 NCCL 进行梯度同步,避免参数漂移。
3. 总结
通信层级 | 通信方式 | 传输介质 |
---|---|---|
服务器之间 | NCCL + PyTorch Distributed | TCP / InfiniBand / RDMA |
服务器内部 | NCCL AllReduce / CUDA IPC | PCIe / NVLink |
DeepSpeed 主要依靠 NCCL 进行 GPU 之间的梯度同步,并通过 ZeRO 进行参数优化和分片,以减少显存占用,提高训练效率。在多机多卡环境下,它依赖 torch.distributed 来初始化分布式进程,并使用 InfiniBand/RDMA 进行高效的数据传输。