NCCL集合通信算子DEMO及性能测试

本文主要是介绍NCCL集合通信算子DEMO及性能测试，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

以下代码用于测试NCCL算子的性能及正确性

一.复现代码

tee ccl_benchmark.py <<-'EOF'
import os
import torch
import argparse
import torch.distributed as dist
from torch.distributed import ReduceOp
from datetime import datetime
import time
import argparse
import numpy as np
dev_type="cuda"class Timer:def __init__(self,duration):        self.duration=durationdef __enter__(self):dist.barrier()self.beg= datetime.now().timestamp() * 1e6def __exit__(self, exc_type, exc_val, exc_tb):dist.barrier()self.end=datetime.now().timestamp() * 1e6self.duration.append(self.end-self.beg)op_mapping={}
class ccl_benchmark:def __init__(self,func):global op_mapping  op_mapping[func.__name__]=funcself.func=funcdef __call__(self,*args,**kwargs):return self.func(*args,**kwargs)@ccl_benchmark
def all_gather(shape,device,rank,world_size,iters=5):'''将每个rank input_tensor的数据在dim 0维度拼接在一起'''duration=[]input_tensor=(torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100+rank)).to(device)gather_list=[torch.zeros((shape[0]//world_size,shape[1]),dtype=torch.int64).to(device) for _ in range(world_size)]for _ in range(iters):with Timer(duration):dist.all_gather(gather_list,input_tensor)   output=torch.cat(gather_list,dim=0)gt=[torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100+i) for i in range(world_size)]gt=torch.cat(gt,dim=0)return duration,(output.cpu()==gt).all()@ccl_benchmark
def scatter(shape,device,rank,world_size,iters=5):'''将每个rank从scatter_list[rank]取数据到output_tensor'''duration=[]output_tensor=torch.zeros((shape[0]//world_size,shape[1]),dtype=torch.int64).to(device)scatter_list=[(torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*i).to(device) for i in range(world_size)]for _ in range(iters):with Timer(duration):if rank == 0:dist.scatter(output_tensor,scatter_list=scatter_list,src =0)else:dist.scatter(output_tensor,src  = 0)gt=torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*rankreturn duration,(output_tensor.cpu()==gt).all()@ccl_benchmark
def gather(shape,device,rank,world_size,iters=5):'''将每个rank input_tensor的数据在dim 0维度拼接在一起 只在批定的rank做'''duration=[]input_tensor=(torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100+rank)).to(device)gather_list=[torch.zeros((shape[0]//world_size,shape[1]),dtype=torch.int64).to(device) for _ in range(world_size)]for _ in range(iters):with Timer(duration):if rank == 0:dist.gather(input_tensor,gather_list=gather_list,dst=0)else:dist.gather(input_tensor,dst=0)ret=Trueif rank==0:output=torch.cat(gather_list,dim=0)gt=[torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100+i) for i in range(world_size)]gt=torch.cat(gt,dim=0)ret=(output.cpu()==gt).all()return duration,ret@ccl_benchmark
def reduce(shape,device,rank,world_size,iters=5):'''将每个rank input_tensor的数据在dim 0维度拼接在一起 只在批定的rank做'''duration=[]   for _ in range(iters):input_tensor=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+rank)).to(device)# input_tensor的内容会被修改,所以放在循环里with Timer(duration):dist.reduce(input_tensor,dst=0,op=dist.ReduceOp.SUM)ret=Trueif rank==0:gt=[torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+i) for i in range(world_size)]gt_=gt[0]       for i in range(1,world_size):gt_=gt_+gt[i]ret=(input_tensor.cpu()==gt_).all()return duration,ret@ccl_benchmark
def broadcast(shape,device,rank,world_size,iters=5):'''将src的rank的数据广播到其它rank'''duration=[]   for _ in range(iters):input_tensor=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+rank)).to(device)with Timer(duration):dist.broadcast(input_tensor,src=0)gt=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+0)).to('cpu')ret=(input_tensor.cpu()==gt).all()return duration,ret@ccl_benchmark
def p2p(shape,device,rank,world_size,iters=5):'''将src的rank的数据广播到其它rank'''duration=[]   for _ in range(iters):input_tensor=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+rank)).to(device)with Timer(duration):if rank!=0:dist.recv(input_tensor,rank-1)               if rank!=world_size-1:               dist.send(input_tensor,dst=rank+1)   gt=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+0)).to('cpu')ret=(input_tensor.cpu()==gt).all()return duration,ret@ccl_benchmark
def all_reduce(shape,device,rank,world_size,iters=5):'''将每个rank input_tensor的数据在dim 0维度拼接在一起'''duration=[]   for _ in range(iters):input_tensor=(torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+rank)).to(device)# input_tensor的内容会被修改,所以放在循环里with Timer(duration):dist.all_reduce(input_tensor,op=dist.ReduceOp.SUM)gt=[torch.ones((shape[0],shape[1]),dtype=torch.int64)*(100+i) for i in range(world_size)]gt_=gt[0]       for i in range(1,world_size):gt_=gt_+gt[i]ret=(input_tensor.cpu()==gt_).all()return duration,ret@ccl_benchmark
def reduce_scatter(shape,device,rank,world_size,iters=5):''''''duration=[]output_tensor=torch.zeros((shape[0]//world_size,shape[1]),dtype=torch.int64).to(device)input_list=[(torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100*rank)+chunk_id).to(device) for chunk_id in range(world_size)]for _ in range(iters):with Timer(duration):dist.reduce_scatter(output_tensor,input_list=input_list,op=dist.ReduceOp.SUM)gt_list=[(torch.ones((shape[0]//world_size,shape[1]),dtype=torch.int64)*(100*rk)+rank).to('cpu') for rk in range(world_size)]gt_=gt_list[0]       for i in range(1,world_size):gt_=gt_+gt_list[i]    return duration,(output_tensor.cpu()==gt_).all()def main():dist.init_process_group(backend='nccl')if not torch.distributed.is_initialized():returnparser = argparse.ArgumentParser(description='test')parser.add_argument('--shape', type=str, default="(1024,8192)", help='Number of epochs to train.')parser.add_argument('--iters', type=int, default=5, help='Number of epochs to train.')parser.add_argument('--op', type=str, default="", help='Number of epochs to train.')args = parser.parse_args()global op_mappingif args.op in op_mapping:torch.manual_seed(1)world_size = torch.distributed.get_world_size()rank = torch.distributed.get_rank()local_rank=int(os.environ['LOCAL_RANK'])torch.cuda.set_device(local_rank)device = torch.device(dev_type,local_rank)shape=eval(args.shape)duration,passed=op_mapping[args.op](shape,device,rank,world_size,args.iters)time.sleep(0.1*rank)print("rank:{} op:{} shape:{} iters:{} mean(us):{:.3f} passed:{}".format(rank,args.op,shape,args.iters,np.mean(duration[len(duration)//2:]),passed))dist.destroy_process_group()if __name__=='__main__':main()EOFexport NCCL_DEBUG=error
export NCCL_SOCKET_IFNAME=ens8
export NCCL_IB_DISABLE=1  
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=all_gather --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=scatter --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=gather --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=reduce --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=broadcast --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=p2p --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=all_reduce --shape="(1024,4096)" --iters=5
torchrun -m --nnodes=1 --nproc_per_node=4 ccl_benchmark --op=reduce_scatter --shape="(1024,4096)" --iters=5

这篇关于NCCL集合通信算子DEMO及性能测试的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---