o h~4@sXddlZddlZddlmZddlmZddlmZmZm Z m Z m Z ddl Z ddl mZddl mZddlmZgZdeded eeed feed fffd d Zd e e jdeded eeeffddZdeed fdeed fd eeed feeefffddZe deeeZe de jeZe dede jde d eefddZ!e dede jde d eefddZ!ddZ!dMdeded e d dfd!d"Z"d#e jd$e j#d dfd%d&Z$d#e jfd'd(Z%e d)Z&e d*eeee'eeeZ(e d+ee jge&fd,e jd e&fd-d.Z)e d+ee jgefd,e(d e(fd/d.Z)d0d.Z)deed fde eeefde jde d eeed feeeefd fff d1d2Z* dNd3ej+d4ee jd5e d6fd7d8Z, dMd9ej-d3ej+d:e.d;e.dd?Z/d3ej+d@ee jd:e.d;e.d df dAdBZ0dCeeefdDedEed dfdFdGZ1d#e jd e fdHdIZ2dJeej-d eej-fdKdLZ3dS)ON) OrderedDict) Container)AnyCallableOptionaloverloadTypeVar)nn)PackedSequenceargskwargsreturn.cOsBg}t|}|D]\}}||||q t|t|fS)a Turn argument list into separate key list and value list (unpack_kwargs does the opposite). Inspiration: https://github.com/facebookresearch/fairscale/blob/eeb6684/fairscale/internal/containers.py#L70 Usage:: kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4) assert kwarg_keys == ("a", "b") assert flat_args == (1, 2, 3, 4) args, kwargs = unpack_kwargs(kwarg_keys, flat_args) assert args == (1, 2) assert kwargs == {"a": 3, "b": 4} Returns: Tuple[Tuple[Any, ...], Tuple[str, ...]]: The first tuple element gives gives both positional args and kwarg values, where the positional args proceed kwarg values and kwarg values are ordered consistently with the kwarg keys. The second tuple element gives the kwarg keys. The second tuple element's length is at most the first tuple element's length. )listitemsappendtuple)r r kwarg_keys flat_argskvrk/var/www/html/construction_image-detection-poc/venv/lib/python3.10/site-packages/torch/distributed/utils.py _pack_kwargss   rdtypecs>dur||fSdtjdtjffdd }t||t||fS)z Cast floating point tensors in ``args`` and ``kwargs`` to ``input_dtype``. This respects the existing ``requires_grad`` on the tensors. Nxr cs"t|r |jkr |S|SN)torchis_floating_pointrto)rrrrcast_fn;s z%_cast_forward_inputs..cast_fn)rTensor_apply_to_tensors)rr r r rrr_cast_forward_inputs.s r#rrcCsvt|t|ksJdt|dt|t|dkr |ifS|dt| }tt||t| d}||fS)zSee _pack_kwargs.ztoo many keys z vs. rN)lendictzip)rrr r rrr_unpack_kwargsCs r'STinputs target_device!use_side_stream_for_tensor_copiescCdSrrr*r+r,rrr _recursive_toUr/cCr-rrr.rrrr/[r0cs*fddz |}Wd|Sdw)z-Recursively moves input to the target_device.cs ttjtfrttrjjnj}|krfSs#fStt|jd}|jdks3|dur9fSddl m }|}| | }Wdn1sXwY|j 5| }||ttry|j|nt|tjsJ||Wd|fSWd|fS1swY|fSddlm}|rfddttDSttrtdkrtttSttrtdkrddttDSttrtdkrfddttDSgS) Ncpur) _get_stream_is_namedtuplecsg|]}t|qSrtype).0r objrr z1_recursive_to..to_map..cSsg|]}t|qSr)rr7irrrr:scsg|]}t|qSrr5r<r8rrr:r;) isinstancerr!r datadevicergetattrr6torch.nn.parallel._functionsr2streamindexcurrent_stream wait_stream record_stream torch.nn.parallel.scatter_gatherr4r&maprr$rr%r)r9r@ device_modr2rCoutputrEr4r+to_mapr,r8rrMdsL              z_recursive_to..to_mapNr)r*r+r,resrrLrr/as / Tcondsraise_assertion_errorcCs(|st|t|rt|dSdS)zwAlternate to ``assert`` when in the backward context to print the error message ``s`` since otherwise, it is swallowed.N)print traceback print_stackAssertionError)rOrPrQrrr _p_assertsrVtensorsizecCstAtjjs5||k}|s=|}t|dkd| |WddSWddSWddS1sHwYdS)z Allocate storage for ``tensor`` with the given size. Returns: bool: ``True`` if this method allocated storage and ``False`` if the storage was already allocated. rzCTensor storage should have been resized to be 0 but got PLACEHOLDErN) rno_grad distributed_functional_collectivesis_torchdynamo_compiling_typed_storage_sizenumelrV_resize_)rWrXalready_allocatedtensor_storage_sizerrr_alloc_storages   "rcc CstItjjs=|dk}|sEt|dkd|d|d|j | dWddSWddSWddS1sPwYdS)z Frees the underlying storage of ``tensor``. Returns: bool: ``True`` if the method freed the storage and ``False`` if the storage was already freed. rzVFreeing a tensor's storage is unsafe when it is not the sole occupant storage offset: z storage size: z tensor shape: N) rrYrZr[r\r]r^rVstorage_offsetshaper`)rW already_freedrrr _free_storages(    "rgQRfn containercCr-rrrjrkrrrr"r0r"cCr-rrrlrrrr"scsfdd|S)zFRecursively apply to all tensor in different kinds of container types.cs ddlm}t|tjr|St|dr0t|fddtD}tjfi|St|t rJ| }| D] \}}|||<q=|St|t rV|j |St|trffdd| DS||ryfdd|D}t||St|tttfrt|fd d|DS|S) Nrr3__dataclass_fields__cs i|] }|jt|jqSr)namerA)r7f)applydcrr sz4_apply_to_tensors..apply..csi|] \}}||qSrr)r7keyvaluerprrrrsc3|]}|VqdSrrr7elrurr z3_apply_to_tensors..apply..c3rvrrrwrurrryrz)rHr4r>rr!hasattr dataclassesreplacefieldsr __class__rr r?r%r6rrset)rr4changesodrsrtrNrprj)rqrrps2          z _apply_to_tensors..applyrrlrrrr"scCs|rt|||ng}|rt|||ng}t|t|kr/|ddtt|t|Dnt|t|krI|ddtt|t|Dt|t|fS)NcSsg|]}dqS)rrr7_rrrr:z_to_kwargs..cSsg|]}iqSrrrrrrr:r)r/r$extendranger)r*r r+r, moved_inputs moved_kwargsrrr _to_kwargss&$r process_grouptensorsloggerz dist.LoggercCst|||Sr)dist_verify_params_across_processes)rrrrrr$_verify_param_shape_across_processessrmodulebroadcast_bucket_sizesrcparams_and_buffers_to_ignorebroadcast_buffersc Csjg}|D]\}}||vr||q|r,|D]\}} ||vr+|| qt||||dS)ag Sync ``module``'s parameters and buffers state. Syncs ``module``'s parameters and buffers state so that all ranks contain the same module state across all ranks. Note that this API assumes that all parameter shapes are consistent before running the synchronization. This can be checked with ``_verify_param_shape_across_processes``. N)named_parametersrdetach named_buffers_sync_params_and_buffers) rrrrrr module_statesrnparambufferrrr_sync_module_statessrrcCs$t|dkrt||||dSdS)zfSynchronize ``module_states`` (list of tensors) across all processes by broadcasting them from rank 0.rN)r$r_broadcast_coalesced)rrrrrrrr:s r state_dict old_prefix new_prefixcCsX||krtdt|D]}||sq||t|d}||||<||=qdS)a Replace all keys that match a given old_prefix with a new_prefix (in-place). Usage:: state_dict = {"layer.xyz": torch.tensor(1)} replace_by_prefix_(state_dict, "layer.", "module.layer.") assert state_dict == {"module.layer.xyz": torch.tensor(1)} z*old_prefix and new_prefix must be distinctN) ValueErrorrkeys startswithr$)rrrrsnew_keyrrr_replace_by_prefixGs  rcCs|dkS)Nr)untyped_storagedata_ptr)rWrrr_data_ptr_allocated_srmodulescCs`g}dd|D}|D]"}d}|D]\}}||uo||v}|r%d}nq|r-||q |S)z Returns the modules in ``modules`` that are root modules (i.e. parent-less) with respect to the set ``modules``. In other words, these are the modules in ``modules`` that are the not child of any other module in ``modules``. cSsi|] }|t|qSr)rr)r7rrrrrrksz%_get_root_modules..TF)rr)r root_modulesmodule_to_modulescandidate_moduleis_root_moduler_modulesis_child_modulerrr_get_root_modulescs" r)Tr)4r|rS collectionsrcollections.abcrtypingrrrrrrtorch.distributedrZrr torch.nn.utils.rnnr __all__rstrrrr#r%r'rr(r!r)r@boolr/rVSizercrgrhrrir"r ProcessGrouprModuleintrrrrrrrrrs     .       9  $# "      "