IЦi>$hSSKrSSKJrJrJrJrJrJrJrJ r J r SSK r SSK J r SSKJr SSKJr SSKJrJr Sr"SS \5rS \S \S \4S jr"SS5r\"\S9SS\ R0S\\\ RR0S \ R04Sjj5rS\ R0S \4Sjrg)N) AnycastDictIterableListNoReturnOptionalSetTuple)_State)DistributedDataParallel) _get_registrycontractc ^\rSrSrSU4Sjjr\4S\RS\\RS\\RS\ SS4 S jjr SS jr S\RS\\RSS4S jr SS jrSS jrS\RS\\S4S\\ \4S\4SjrS\RS\\R*S\R*S\R*4SjrSrU=r$)_ReplicateStatereturnNc>[TU]5 [R"5UlSUl[R"5UlURUl/UlSUl SUl 0Ul /Ul g)NF) super__init__nn ParameterListmodulehas_initialized _param_list _orig_module _param_names_no_sync _init_args _init_kwargs_comm_hook_args)self __class__s f/var/www/html/ai-image-ml/venv/lib/python3.13/site-packages/torch/distributed/_composable/replicate.pyr_ReplicateState.__init__sk !#!1!1!3 %*-/-=-=-?!KK')# 59,.*,rignored_modulesignored_paramsprefixcl[U5(agX;agU[:waUS3O[nURSS9HFupgXs;dM URR U5 UR R UU35 MH UR 5HupURU UUUU3S9 M g)N.F)recurse)r+)_is_fully_sharded_ROOT_MODULE_PREFIXnamed_parametersrappendrnamed_children_collect_params) r$rr)r*r+recurse_prefixnpname child_modules r&r4_ReplicateState._collect_params s V $ $   $ #&99vhaL?R ++E+:DA&  ''*!!((N+;A3)?@; #)"7"7"9 D  ()$0 ! #:r(cN^[R"SS9U4Sj5nU"5 g)NT) recursivec>TRceTR"TR0TRD6 TR5 STl0Tlg)N)r!initr"register_comm_hook)r$sr& _lazy_init-_ReplicateState.lazy_init.._lazy_initAsI??. .. IIt <$*;*; <  # # % DO "D r()torch_disable_dynamo)r$rAs` r& lazy_init_ReplicateState.lazy_init@s'    . # / #  r(c UR(agSUlXlUVVs1sHoDR5HoUiM M nnnUR5H4n[ U5(dMUR UR55 M6 SSKJn U"XS9 URXU5 SU;a]USbAUSn [U [R5(aU RS:XaSUS'O U /US'OSUS'URS5 [UR40UD6Ul["R$"UR 5[&R)UR5lgs snnf)NTr)_localize_dtensor)r* device_idcpu device_ids)rr parametersmodulesr/update%torch.distributed.tensor.parallel.ddprHr4 isinstancerCdevicetypepopr r_ddpweakrefref replicatestate _ddp_weakref) r$rr)kwargsmr7r* submodulerHrIs r&r?_ReplicateState.initKs6    # %4M_lln!n!_M)I ++%%i&:&:&<=* L&@ VnE & k".";/ i669>>U;R+/F<(-6;F<('+|$ JJ{ #+D,<,<GG 4;KK 4J  $1;Ns"E1cURH!upURR"U0UD6 M# URR5 gN)r#rTr@clear)r$ comm_args comm_kwargss r&r@"_ReplicateState.register_comm_hookus>&*&:&: "I II ( () C{ C'; ""$r(cXlX lgr_)r!r"r$argsrZs r&record_init_args _ReplicateState.record_init_argszs "r(rf.rZcUR(dUR(aUR5 UR(+URlURR "U0UD6$r_)r!r"rEr rTrequire_backward_grad_sync _pre_forward)r$rrfrZs r&forward_pre_hook _ReplicateState.forward_pre_hook~sJ ??d// NN 37==/@ ,yy%%t6v66r(inputoutputc8URRU5$r_)rT _post_forward)r$rrnros r&forward_post_hook!_ReplicateState.forward_post_hooks yy&&v..r() r#rTr!r"r rrrrrrN)__name__ __module__ __qualname____firstlineno__rr0rModuler Parameterstrr4rEr?r@rgr rrrlrCTensorrr__static_attributes__ __classcell__)r%s@r&rrs -&*  RYYBLL)     @ (K (KRYY(K  (KT% #7ii7',S#X7@DS#X7 7/ /U\\"/ /  //r(rrfrZrc[S5e)NzGDDP does not support deepcopy. Please use state dict for serialization.)AssertionError)rfrZs r&unimplemented_deepcopyrs Q r(c6\rSrSrSrS\SS4SjrS SjrSrg) DDPcLURSnUR"U/UQ70UD6$)z Override ``__new__`` to remove the DDP class and directly construct the original class for cases like indexing into a container module. )__mro____new__)clsrfrZorig_clss r&r DDP.__new__s,;;q>:4:6::r(requires_gradient_syncrNcBU(+[RU5lg)z Sets if the module should sync gradients. This can be used to implement gradient accumulation without communication. Args: requires_gradient_sync (bool): Whether to reduce gradients for the module's parameters. N)rWrXr )r$rs r&set_requires_gradient_syncDDP.set_requires_gradient_syncs.D)C &r(cb[RU5RRX45 gr_)rWrXr#r2res r&r@DDP.register_comm_hooks!--44d^Dr(r>rt) rurvrwrxrboolrr@r}r>r(r&rrs#; D D$ DEr(r) state_clsrr)c ,[RRS5 SU;aB[US[[R 45(d[ S[US535e[U5(a [ S5eUc0nO [U5n[[[RU55nURURSS9 UR!SS5nUbFS S KJn UR'U5nXd:wa*S S KJnJn URU5 UR/U5 UR/UR05 UR2"X40UD6 UR4n S [60n [S U R83[:U 4U 5n XlU$)zReplicates a module Args: module (torch.nn.Module): module to replicate Example:: >>> # xdoctest: +REQUIRES(module:torch._C._distributed_c10d) >>> module = nn.Linear(3, 3) >>> replicate(module) ztorch.distributed.replicaterIz6Expected device_id to be int or torch.device, but got zGCannot apply `replicate()` on a Module already managed by `fully_shard`NT) with_kwargs device_meshr)_mesh_resources)rH_reconstruct_dtensor __deepcopy__r)rC_C_log_api_usage_oncerPintrQ RuntimeErrorrRr/setrrrWrXregister_forward_pre_hookrlgettorch.distributed.device_meshr get_root_meshrOrHrregister_forward_hookrrrgr%rrur) rr)rZrXrr root_meshrHrrdctnew_clss r&rWrWs  HH  !>?f&-U\\/BCC{ 3457     U  o. )//&"9 :E $$U%;%;$N**]D1KA#11+>   #   , ,-A B  ( (): ;   !8!89 6=f=   C 1 2CS'#sS9G Mr(c*[U5nUcgSU;$)z+Check if module is marked with fully_shard.F fully_shard)r)rregistrys r&r/r/s V$H H $$r(r_)rUtypingrrrrrrr r r rCtorch.nnr#torch.distributed._composable_stater torch.nn.parallelr rrr0rrrryrWrr/r>r(r&rsRRR 65-{/f{/|#EE4 O$<@E IIEhuxx78EYY E%EP%bii%D%r(