o Zho;@stddlZddlmZddlmZddlmZmZejfddZ dejfddZ dejfdd Z ej ejfd d Z ej ejfd d ZejfddZejfddZejfddZddejfddZej ejfddZGdddeZGdddeZGdddeZGdddeZGd d!d!eZGd"d#d#eZGd$d%d%eZGd&d'd'eZGd(d)d)eZGd*d+d+eZdS),N)Function)groupReduceOpcCt|||S)a Broadcasts the tensor to the whole group. ``tensor`` must have the same number of elements in all processes participating in the collective. Arguments: tensor (Tensor): Data to be sent if ``src`` is the rank of current process. src (int): Source rank. group (ProcessGroup, optional): The process group to work on. Returns: Tensor: Received tensor from the broadcast op. ) _Broadcastapply)tensorsrcrr N/var/www/auris/lib/python3.10/site-packages/torch/distributed/nn/functional.py broadcast sr cCr)aT Gathers a list of tensors in a single process. Arguments: tensor (Tensor): Input tensor. dst (int, optional): Destination rank (default is 0). group (ProcessGroup, optional): The process group to work on. Returns: tuple[Tensor]: List of appropriately-sized tensors with the gathered data. )_Gatherr)rdstrr r r gather s rcCstj||g|RS)a Scatters a list of tensors to all processes in a group. Each process will receive exactly one tensor and store its data in the ``tensor`` argument. Arguments: tensors (list[Tensor]): List of tensors to scatter on the source rank. Receivers must pass ``None`. src (int, optional): Source rank (default is 0). group (ProcessGroup, optional): The process group to work on. Returns: Tensor: Output tensor from the scatter operation. )_Scatterr)tensorsr rr r r scatter/srcCst||||S)a Reduces the tensor data across all machines. Only the process with rank ``dst`` is going to receive the final result. Arguments: tensor (Tensor): Input of the collective. dst (int): Destination rank. op (optional): One of the values from ``torch.distributed.ReduceOp`` enum. Specifies an operation used for element-wise reductions. group (ProcessGroup, optional): The process group to work on. Returns: Tensor: Output of the collective. )_Reducer)rroprr r r reduceCsrcCstj|||g|RS)a Reduces, then scatters a list of tensors to all processes in a group. Arguments: output (Tensor): Output tensor. input_list (list[Tensor]): List of tensors to reduce and scatter. op (optional): One of the values from ``torch.distributed.ReduceOp`` enum. Specifies an operation used for element-wise reductions. group (ProcessGroup, optional): The process group to work on. Returns: Tensor: Output of the collective. )_Reduce_Scatterr)outputZ input_listrrr r r reduce_scatterXsrcCs t||S)a Gathers tensors from the whole group in a list. Arguments: tensor (Tensor): Tensor to be broadcast from current process. group (ProcessGroup, optional): The process group to work on. Returns: tuple([Tensor]): Output of the collective. ) _AllGatherr)rrr r r all_gatherks rcCst|||S)a Single tensor all gather. Gathers a single tensor from all ranks, and puts them in a single output tensor. Args: output_tensor (Tensor): Output tensor. It should contain correctly-sized tensors to be used for output of the collective. input_tensor (Tensor): Tensor to be broadcast from current process. group (ProcessGroup, optional): The process group to work on. If None, the default process group will be used. Examples: >>> # All tensors below are of torch.int64 dtype. >>> # We have 2 process groups, 2 ranks. >>> # xdoctest: +SKIP("incorrect want text") >>> output_tensor = torch.zeros(2, dtype=torch.int64) >>> output_tensor [tensor([0, 0])] # Rank 0 and 1 >>> tensor = torch.arange(1, dtype=torch.int64) + 1 + rank >>> tensor tensor([1]) # Rank 0 tensor([2]) # Rank 1 >>> dist.all_gather_base(output_tensor, tensor) >>> output_tensor tensor([1,2]) # Rank 0 tensor([1,2]) # Rank 1 .. warning:: `_all_gather_base` is experimental and subject to change. It is the caller's responsibility to ensure the output_tensor is correctly sized. )_AllGatherBaser) output_tensor input_tensorrr r r _all_gather_basezs!rcCstj||g|RS)a Each process scatters list of input tensors to all processes in a group and return gathered list of tensors in output list. Arguments: output_tensor_list (list[Tensor]): list of tensors to gather one per rank. input_tensor_list (list[Tensor]): List of tensors to scatter one per rank. group (ProcessGroup, optional): The process group to work on. Returns: tuple([Tensor]): Output of the collective. ) _AlltoAllr)Zoutput_tensor_listinput_tensor_listrr r r all_to_alls r!cCst|||||S)a Each process splits input tensor and then scatters the split list to all processes in a group. Then concatenate the received tensors from all the processes in the group and return single output tensor. Arguments: output (Tensor): Gathered concatenated output tensor. input (Tensor): Input tensor to scatter. output_split_sizes: (list[Int], optional): Output split sizes for dim 0 if specified None or empty, dim 0 of ``output`` tensor must divide equally by ``world_size``. input_split_sizes: (list[Int], optional): Input split sizes for dim 0 if specified None or empty, dim 0 of ``input`` tensor must divide equally by ``world_size``. Returns: Tensor: Output of the collective. )_AlltoAllSingler)rinputoutput_split_sizesinput_split_sizesrr r r all_to_all_singles r&cCr)a& Reduces the tensor data across all machines in such a way that all get the final result. After the call the returned tensor is going to be bitwise identical in all processes. Arguments: tensor (Tensor): Input of the collective. op (optional): One of the values from ``torch.distributed.ReduceOp`` enum. Specifies an operation used for element-wise reductions. group (ProcessGroup, optional): The process group to work on. Returns: Tensor: Output of the collective ) _AllReducer)rrrr r r all_reducesr(c@$eZdZeddZeddZdS)rcCs6||_||_tj|d|_|}tj|||d|SNr)r rdistget_rankrankcloner )ctxr rrr r r forwards z_Broadcast.forwardcCs4t|jtj|j|}|j|jkr|dd|fSN)rrr rSUMrr.Zzero_)r0 grad_outputgxr r r backwards  z_Broadcast.backwardN__name__ __module__ __qualname__ staticmethodr1r6r r r r rs   rc@r))r csv||_||_fddttj|dD}tj|d|kr.tj|||dt|Stjd||dt|S)Ncg|]}tqSr )torch zeros_like.0irr r  z#_Gather.forward..r+) rrranger,get_world_size contiguousr-rtuple)r0rrr tensor_listr rBr r1s z_Gather.forwardcGsdtj|j|jg|RfSNNN)rrrr)r0 grad_outputsr r r r6 sz_Gather.backwardNr7r r r r r s  r c@r))rcst||_||_tfddDsJtd}tj|d|kr/tj|t||d|Stj|d||d|S)Nc3s$|] }|dkVqdS)rNsizer@trr r s"z#_Scatter.forward..rr+) r rallr=r>r,r-rlist)r0r rrrr rQr r1sz_Scatter.forwardcCsdt|j|j|SrJ)r rr rr0r4r r r r6sz_Scatter.backwardNr7r r r r rs   rc@r))rcCs*||_||_|}tj||||d|S)Nrr)r rr/r,r)r0r rrrr r r r1"s z_Reduce.forwardcCdt|j|j|fSN)NNN)rrr rrUr r r r6*z_Reduce.backwardNr7r r r r r!  rc@r))rcGs:||_|}tdd|D}tj|t|||d|S)Ncs|]}|VqdSr2rGrOr r r rR5z*_Reduce_Scatter.forward..rV)rrGrHr,rrT)r0rrrr r r r r10s z_Reduce_Scatter.forwardcCsdt|j|SrX)rrrrUr r r r69sz_Reduce_Scatter.backwardNr7r r r r r/s  rc@r))rcsD||_fddttj|dD}tj||dt|S)Ncr<r r= empty_like)r@_rBr r rCErDz&_AllGather.forward..r+)rGrrEr,rFrrH)r0rrout_tensor_listr rBr r1?s z_AllGather.forwardcGstj|jdtjjur*tj|jd}t||}tj t j |j|g|R}d|fSdd|D}t j |j|g|R}tj t|dd}d|fS)Nr+cSsg|]}t|qSr r^)r@rr r r rCUsz'_AllGather.backward..r)dim)r, get_backendrBackendNCCLr-r=r_rrrr3rsumstack)r0rLr.r5rIZgxsr r r r6Lsz_AllGather.backwardNr7r r r r r>s   rc@r))rcCs||_tj|||d|Sr*)rr,rrG)r0rrrr r r r1\sz_AllGatherBase.forwardcCstj|jdtjjurLtj|jd}t|}|d|dkr*td|d||dtj|jd|d<t j ||j |j d}t ||tj|jntdd|dfS)Nr+rzTensor with dimensions: z8 does not have first dimension divisible by world_size: devicedtypezBackend not supported!)r,rcrrdrerFrTrN RuntimeErrorr=emptyrirjZ_reduce_scatter_baserr3)r0r4Z world_sizeZout_sizer5r r r r6bs    z_AllGatherBase.backwardNr7r r r r r[s  rc@r))rcs||_fddttj|dD|_tj|d}tddDtj|dtjj urPttj|dD]}d}||kr@t }tj |||||dq4t|Stj |t |dt|S)Ncsg|]}|qSr rMr?rQr r rCzs z%_AlltoAll.forward..r+csr[r2r\rOr r r rR~r]z$_AlltoAll.forward..) rrEr,rFinput_tensor_size_listr-rHrcrdZGLOOrTrr!)r0rrarZmy_rankrAto_sendr rQr r1ws&  z_AlltoAll.forwardcs.fdd|jD}dtj|j|gRS)Ncs(g|]}tj|djdjdqS)rrh)r=rlrirj)r@rNrLr r rCs z&_AlltoAll.backward..rK)rmrrr)r0rLrIr ror r6s z_AlltoAll.backwardNr7r r r r rvs  rc@r))r"cCs4||_||_||_||_tj|||||d|S)N)r$r%r)rrN input_sizer$r%r,r&)r0rrr$r%r#r r r r1s z_AlltoAllSingle.forwardc Cs8tj|j|j|jd}dt|j||j|j | fS)Nrh)NNNN) r=rlrprirjr"rrr$r%rG)r0r4rr r r r6s z_AlltoAllSingle.backwardNr7r r r r r"s  r"c@r))r'cCs.||_||_|jtjd}tj|||d|S)N)Z memory_formatrV)rrr/r=Zcontiguous_formatr,r()r0rrrr r r r1s z_AllReduce.forwardcCrWrJ)r'rrrrUr r r r6rYz_AllReduce.backwardNr7r r r r r'rZr')r=Ztorch.distributed distributedr,Ztorch.autogradrrrZWORLDr rrr3rrrrr!r&r(rr rrrrrrr"r'r r r r s6  $ #