PyTorch DDP快速上手附代码

如标题所示,博主有使命 想要快速上手DDP以测试RDMA参与下的耽误带宽.
环境准备:

验证PyTorch:

1. python -c "import torch; print(torch.__version__, torch.cuda.device_count())"

复制代码

核心概念:

进程组：每个GPU对应一个独立进程
​AllReduce：跨卡聚合数据（如梯度）
​torchrun：官方保举启动工具

测试两卡all-reduce

看看博主的卡:

1. import torch
2. import torch.distributed as dist
4. dist.init_process_group(backend='nccl')
5. rank = dist.get_rank()
6. world_size = dist.get_world_size()
8. print(f'I am rank: {rank}  of world {world_size}')
10. input_tensor = torch.rand([1024, 1024, 10], dtype=torch.float).to('cuda:%d' % rank)
11. input_tensor.fill_(1.0)
13. dist.all_reduce(input_tensor)
15. print(input_tensor[0][0])

复制代码

效果:

还想进一步测试差别buffer size下的耽误和带宽环境:

1. import os
2. import torch
3. import torch.distributed as dist
4. import time
6. os.environ['MASTER_ADDR'] = '127.0.0.1'
7. os.environ['MASTER_PORT'] = '29500'
9. dist.init_process_group(backend='nccl')
10. rank = dist.get_rank()
11. torch.cuda.set_device(rank)
12. world_size = dist.get_world_size()
14. buffer_sizes = [1, 1024, 2048, 4096, 8192,
15.                 256 * 1024, 512 * 1024,
16.                 1024 * 1024, 4 * 1024 * 1024,
17.                 256 * 1024 * 1024]
19. def benchmark(size):
20.     input_tensor = torch.ones(size, dtype=torch.float32).to(f'cuda:{rank}')
21.    
22.     torch.cuda.synchronize()
23.     start = time.time()
24.    
25.     dist.all_reduce(input_tensor)
26.    
27.     torch.cuda.synchronize()
28.     elapsed = time.time() - start
29.    
30.     if rank == 0:
31.         data_size = size * 4  # float32占4字节
32.         bandwidth = (data_size / elapsed) / (1024**2)  # MB/s
33.         print(f"Size: {size/1024:.1f} KB\tLatency: {elapsed*1000:.3f} ms\tBandwidth: {bandwidth:.2f} MB/s")
35. if rank == 0:
36.     print("Buffer Size (KB)\tLatency (ms)\tBandwidth (MB/s)")
37.     print("------------------------------------------------")
39. for size in buffer_sizes:
40.     benchmark(size)
42. dist.destroy_process_group()

复制代码

效果:

至此 单机两卡h20测试乐成 带宽耽误符合预期