Ei>SSKrSSKJr SSKJr SSKJrJr \R4Sjr S\R4Sjr S\R4Sjr \R\R4Sjr \R\R4Sjr\R4S jr\R4S jr\R4S jrSS\R4S jr\R\R4S jr"SS\5r"SS\5r"SS\5r"SS\5r"SS\5r"SS\5r"SS\5r"SS\5r"SS\5r"S S!\5rg)"N)Function)groupReduceOpc.[RXU5$)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)tensorsrcrs \/home/mande/repo/quber/.venv/lib/python3.13/site-packages/torch/distributed/nn/functional.py broadcastr s"   C //c.[RXU5$)a4 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)r dstrs r gatherr s ==V ,,r c0[R"X/UQ76$)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 rs r scatterr/s" >># /w //r c.[RXX05$)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)r roprs r reducerCs$ ==% 00r c2[R"X#U/UQ76$)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_listrrs r reduce_scatterrXs  F @Z @@r c,[RX5$)z 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)r rs r all_gatherr!ks   E **r c.[RXU5$)av 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_tensorrs r _all_gather_baser&zsB    U CCr c0[R"X /UQ76$)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_listrs r all_to_allr+s ??5 I7H IIr c0[RX@X#U5$)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_sizesrs r all_to_all_singler1s4   )e r c.[RXU5$)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)r rrs r all_reducer4s$   Bv ..r c4\rSrSr\S5r\S5rSrg)rcXlX l[R"US9UlUR 5n[R "X1US9 U$Nr)r rdistget_rankrankcloner )ctxr rr s r forward_Broadcast.forwards> ==u- v%0 r c[RUR[RUR U5nURUR :waUR5 SSU4$N)rrr rSUMrr<zero_)r> grad_outputgxs r backward_Broadcast.backwardsJ]]377HLL#))[ I 77chh  HHJdBr N__name__ __module__ __qualname____firstlineno__ staticmethodr?rG__static_attributes__rIr r rrs(  r rc4\rSrSr\S5r\S5rSrg)rcdXlX l[[R"US95Vs/sHn[ R "U5PM nnUR5n[R"US9U:Xa[R"X5XS9 O[R"USXS9 [U5$s snfr8) rrranger:get_world_sizetorch zeros_like contiguousr;rtuple)r>rrr i tensor_lists r r?_Gather.forwards /4D4G4Ge4T.U .UE  V $.U  ""$ ==u % , KKS > KKc 7[!! s B-cbS[R"URUR/UQ764-$NNN)rrrr)r> grad_outputss r rG_Gather.backward s(x~~cggsyyP<PRRRr rINrJrIr r rrs*""$SSr rc4\rSrSr\S5r\S5rSrg)ric(^XlX l[U4SjT55(de[R"TS5n[ R "US9U:Xa![ R"U[T5XS9 U$[ R"USXS9 U$)Nc3j># UH(oR5TSR5:Hv M* g7f)rN)size).0trs r #_Scatter.forward..s%B'Q668wqz00's03rr9) r rallrVrWr:r;rlist)r>r rrrs ` r r?_Scatter.forwards} B'BBBBB!!'!*- ==u % , LLg A  LLs 8 r c^S[RURURU5-$r^)rrr rr>rEs r rG_Scatter.backward s$gmmCGGSYY LLLr rINrJrIr r rrs*  MMr rc4\rSrSr\S5r\S5rSrg)ri&chXlX0lUR5n[R"XAX#S9 U$)Nrr)r rr=r:r)r>r rrr s r r?_Reduce.forward's,  FB4 r c`S[RURURU54-$N)NNN)rrr rrns r rG_Reduce.backward0s)"Z%5%5cggsyy+%V$XXXr rINrJrIr r rr&s*YYr rc4\rSrSr\S5r\S5rSrg)ri6cX lUR5n[SU55n[R"U[ U5XS9 U$)Nc3@# UHoR5v M g7frBrXrfrgs r rh*_Reduce_Scatter.forward..=s!L:KQ,,..:Krr)rrXrYr:rrk)r>rrr r*s r r?_Reduce_Scatter.forward7sF ""$!!L:K!LL FD):$;P r cHS[RURU5-$ru)r rrrns r rG_Reduce_Scatter.backwardAs "J$4$4SYY $LLLr rINrJrIr r rr6s*MMr rc4\rSrSr\S5r\S5rSrg)r iGcUR5nXl[[R"US95Vs/sHn[ R "U5PM nn[R"XBUS9 [U5$s snfr8) rXrrTr:rUrV empty_liker!rY)r>rr _out_tensor_lists r r?_AllGather.forwardHso""$ .3D4G4Ge4T.U .UE  V $.U   u=_%%  s A8ct[R"URS9[RR[RR 4;ak[R "URS9n[R"X5n[R"[RURU/UQ76nSU4$UVs/sHn[R"U5PM nn[R"URU/UQ76n[R"[R"U5SS9nSU4$s snf)Nr9r)dim)r: get_backendrBackendNCCLXCCLr;rVrrrrrCr(sumstack)r>r`r<rFr r[gxss r rG_AllGather.backwardVs   #)) ,1B1BDLLDUDU0V V==syy1D!!,"45B &&x||SYYR\RBbzCOO,5++F3,KO//#))[H<HC5;;s+3BbzPs D5rINrJrIr r r r Gs( & &  r r c4\rSrSr\S5r\S5rSrg)r#iecXX0l[R"XR5US9 U$r8)rr:r&rX)r>r$r%rs r r?_AllGatherBase.forwardfs(  m-D-D-FeTr ct[R"URS9[RR[RR 4;a[R "URS9n[UR55nUSU-S:wa[SUSU35eUS[R "URS9-US'[R"X1RURS9n[R"XA[R UR5 O [S5eSUS4$)Nr9rzTensor with dimensions: z8 does not have first dimension divisible by world_size: devicedtypezBackend not supported!)r:rrrrrrUrkre RuntimeErrorrVemptyrr_reduce_scatter_baserrC)r>rE world_sizeout_sizerFs r rG_AllGatherBase.backwardms   #)) ,1B1BDLLDUDU0V V,,399=JK,,./H{Z'1,".xj9IIS V#1+)<)<399)MMHQK!3!3;;L;LB  % %bx||SYY O78 8b$r rINrJrIr r r#r#es(   r r#c4\rSrSr\S5r\S5rSrg)r(ic:Xl[[R"US95Vs/sHoCUR 5PM snUl[R "US9n[SU55n[R"US9[RRLaO[[R"US95H-nSnXE:Xa [U5n[R"X$XdUS9 M/ O[R"U[U5US9 [U5$s snf)Nr9c3@# UHoR5v M g7frBrzr{s r rh$_AlltoAll.forward..s81 r})rrTr:rUreinput_tensor_size_listr;rYrrGLOOrkrr+)r>rrrrZmy_rankto_sends r r?_AlltoAll.forwards ',T-@-@u-M'N& 'N!AJOO 'N& "--e,888   % (DLL,=,= =4..U;<<"7mG _/5I = OOW   _%%%& sDc URVs/sH2n[R"X!SRUSRS9PM4 nnS[ R "URU/UQ76-$s snf)Nrrr_)rrVrrrr(rr)r>r`rer[s r rG_AlltoAll.backwardsv 22  3 KK!_33<?;P;P 3  ioociiT|TTT  s9A/rINrJrIr r r(r(s*&&,UUr r(c4\rSrSr\S5r\S5rSrg)r-icXlUR5UlX@lX0l[ R "UUUUUS9 U$)N)r/r0r)rre input_sizer/r0r:r1)r>rrr/r0r.s r r?_AlltoAllSingle.forwardsF !2 2   1/   r c [R"URURURS9nS[ R URUURURUR554-$)Nr)NNNN) rVrrrrr-rrr/r0rX)r>rEr s r rG_AlltoAllSingle.backwardss NN;#5#5[=N=N (  ! ! &&%%&&(  +   r rINrJrIr r r-r-s(      r r-c4\rSrSr\S5r\S5rSrg)r3icX lXlUR[RS9n[ R "X1US9 U$)N) memory_formatrr)rrr=rVcontiguous_formatr:r4)r>rrr s r r?_AllReduce.forwards7 E,C,CD U3 r c`S[RURURU54-$r^)r3rrrrns r rG_AllReduce.backwards)z// ;OQQQr rINrJrIr r r3r3s*RRr r3)rVtorch.distributed distributedr:torch.autogradrrrWORLDr rrrCrrr!r&r+r1r4rrrrrr r#r(r-r3rIr r rsF # ."'0( -%++0($< !DH=BKK J& ++ >#,,ekk/*  ,ShS4MxM( Yh Y MhM"< X :!U!UH h D R Rr