o ho;@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 s/var/www/html/construction_image-detection-poc/venv/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)output 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)output_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 rSUMrr0zero_)r2 grad_outputgxr r r backwards  z_Broadcast.backwardN__name__ __module__ __qualname__ staticmethodr3r9r 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)r2rrr tensor_listr rEr r3s z_Gather.forwardcGsdtj|j|jg|RfSNNN)rrrr)r2 grad_outputsr r r r9 sz_Gather.backwardNr:r r r r r s  r c@r+)rcst||_||_tfddDsJtd}tj|d|kr/tj|t||d|Stj|d||d|S)Nc3s$|] }|dkVqdS)rNsizerCtrr r s"z#_Scatter.forward..rr-) r rallr@rAr.r/rlist)r2r rrrr rTr r3sz_Scatter.forwardcCsdt|j|j|SrM)r rr rr2r7r r r r9sz_Scatter.backwardNr:r r r r rs   rc@r+)rcCs*||_||_|}tj||||d|S)Nrr)r rr1r.r)r2r rrrr r r r3"s z_Reduce.forwardcCdt|j|j|fSN)NNN)rrr rrXr r r r9*z_Reduce.backwardNr:r r r r r!  rc@r+)rcGs:||_|}tdd|D}tj|t|||d|S)Ncs|]}|VqdSr4rJrRr r r rU5z*_Reduce_Scatter.forward..rY)rrJrKr.rrW)r2rrrr"r r r r30s z_Reduce_Scatter.forwardcCsdt|j|Sr[)rrrrXr r r r99sz_Reduce_Scatter.backwardNr:r r r r r/s  rc@r+)rcsD||_fddttj|dD}tj||dt|S)Ncr?r r@ empty_like)rC_rEr r rFErGz&_AllGather.forward..r-)rJrrHr.rIrrK)r2rrout_tensor_listr rEr r3?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 ra)rCrr r r rFUsz'_AllGather.backward..r)dim)r. get_backendrBackendNCCLr/r@rbrrrr5r sumstack)r2rOr0r8rLgxsr r r r9Lsz_AllGather.backwardNr:r r r r r>s   rc@r+)rcCs||_tj|||d|Sr,)rr.rrJ)r2rrrr r r r3\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.rfrrgrhrIrWrQ RuntimeErrorr@emptyrmrn_reduce_scatter_baserr5)r2r7 world_sizeout_sizer8r r r r9bs    z_AllGatherBase.backwardNr:r r r r r[s  rc@r+)r cs||_fddttj|dD|_tj|d}tddDtj|dtjj urPttj|dD]}d}||kr@t }tj |||||dq4t|Stj |t |dt|S)Ncsg|]}|qSr rPrBrTr r rFzs z%_AlltoAll.forward..r-csr^r4r_rRr r r rU~r`z$_AlltoAll.forward..) rrHr.rIinput_tensor_size_listr/rKrfrgGLOOrWrr#)r2rrdrmy_rankrDto_sendr rTr r3ws&  z_AlltoAll.forwardcs.fdd|jD}dtj|j|gRS)Ncs(g|]}tj|djdjdqS)rrl)r@rprmrn)rCrQrOr r rFs z&_AlltoAll.backward..rN)rtr rr)r2rOrLr rxr r9s z_AlltoAll.backwardNr:r r r r r vs  r c@r+)r$cCs4||_||_||_||_tj|||||d|S)N)r&r'r)rrQ input_sizer&r'r.r()r2rrr&r'r%r r r r3s z_AlltoAllSingle.forwardc Cs8tj|j|j|jd}dt|j||j|j | fS)Nrl)NNNN) r@rpryrmrnr$rrr&r'rJ)r2r7rr r r r9s z_AlltoAllSingle.backwardNr:r r r r r$s  r$c@r+)r)cCs.||_||_|jtjd}tj|||d|S)N) memory_formatrY)rrr1r@contiguous_formatr.r*)r2rrrr r r r3s z_AllReduce.forwardcCrZrM)r)rrrrXr r r r9r\z_AllReduce.backwardNr:r r r r r)r]r))r@torch.distributed distributedr.torch.autogradrrrWORLDr rrr5rrrrr#r(r*rr rrrrrr r$r)r r r r s6  $ #