a hB@sddlZddlZddlZddlZddlmZmZddlmZddl m Z ddl m Z m Z e rddlZddlmZejrddlZdZnd Zdd lmZeeZd d Zed ZeedddZddZd,e ee e edddZ!d-e ee e edddZ"ej#dddZ$ddZ%edZ&ddZ'd d!Z(d"d#Z)d.e ee e e d%d&d'Z*e d(d)ed/e ee e e e d%d*d+Z+dS)0N)contextmanagerredirect_stdout)StringIO)Optional)is_torch_availablerequires) save_fileTF)loggingcCs trtjsdStjdkS)z7Return True if rank=0 or we aren't running distributed.Tr)_torch_distributed_availabletorch distributedZis_initializedZget_rankrr^/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/model_debugging_utils.py _is_rank_zero,srzobject at 0x[0-9A-Fa-f]+)x_strreturncCs td|S)z Replace memory addresses in an object's repr with a stable placeholder so that beautiful JSON diffs won't be ruined by ephemeral addresses. zobject at 0xXXXXXXXX)MEMORY_ADDRESS_REGEXsub)rrrr_sanitize_repr_for_diff6srcCstrdt|jSdS)z@Return a stable string representation for a DTensor-like object.zDTensor (rank0) -> zDTensor(non-rank0))rrepr _local_tensor)xrrr _dtensor_repr>sr debug_pathuse_repr path_to_valuec Cstjdd|rt|}nh|rl|ds0|d7}|rBtj||n|}td| i|d|}nt d|d|dt |j t |j|d }|jtjtjtjhvr|tt |tt |tt |tt |d |S) a Converts Tensors and DTensors to a JSON-serializable dictionary representation. Args: value: Any Python object, often including torch Tensors, lists, dicts, etc. debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files. use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensor as the `value` property in the asscoiated FULL_TENSORS.json file, or to store the full tensors in separate SafeTensors file and store the relative path to that file in the `value` property in the dictionary. path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full tensor value if `use_repr=False`. Returns: A nested Python structure (list, dict, or sanitized string) that is safe to json.dump. T)sci_modez .safetensorsdataz./z use_repr=z and path_to_value=z cannot both be falsy.)shapedtypevalue)meanstdminmax)r set_printoptions _repr_to_listendswithospathjoinr contiguousdetachcpu ValueErrorrr r!Zfloat16Zfloat32Zbfloat16updaterr#r$r%r&)r"rrrZ value_outfilepathoutrrr_serialize_tensor_like_ioEs.    r4cst|ttfr(fddt|DSt|trLfdd|DSt|drht|jdSt|t j rt|dSt t |S)a Recursively build a JSON-serializable Python structure from `value`. Tensors and DTensors become either sanitized repr strings, or are saved to disk as SafeTensors files and their relative paths are recorded in the returned Python structure. Lists/tuples/dicts are recursed into. All memory addresses are replaced with a stable placeholder. Args: value: Any Python object, often including torch Tensors, lists, dicts, etc. debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files. use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the `value` property in the asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors files and store the relative path to that file in the `value` property. path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full tensor value if `use_repr=False`. Returns: A nested Python structure (list, dict, or sanitized string) that is safe to json.dump. c s*g|]"\}}t|d|dqS_r _serialize_io).0ivrrrrr sz!_serialize_io..c s,i|]$\}}|t|d|dqSr5r7)r9kr;r<rr sz!_serialize_io..rr) isinstancelisttuple enumeratedictitemshasattrr4rr Tensorrr)r"rrrrr<rr8vs    r8)r"c CsxtjdddtH}t| t||}Wdn1sD0YWdn1sb0Yt|S)z Converts a tensor into a sanitized multi-line string representation. Args: value (`torch.Tensor`): The tensor to represent. Returns: `list[str]`: List of string lines representing the tensor. Tx)rZ linewidthN)r r'rrprintgetvaluer splitlines)r"bufrawrrrr(s Dr(cCs0|dr,|dd|dD] }t|qdS)Nchildrenoutputs)getpopprune_outputs_if_childrennodechildrrrrRs   rRz (.*)\.(\d+)$csHt|dd}|r |ds$dS|dtfdd|dDS)z Checks whether a node represents a layer block with submodules. Args: node (`dict`): A node from the call tree. Returns: `bool`: Whether the node is a layer block. module_pathrNFc3s&|]}dd|ddvVqdS).rVrWNrP)r9rUnumberrr z!is_layer_block..)LAYER_SUFFIX_REmatchrPgroupany)rTr`rr[ris_layer_blocks  rccs~|dsdSddt|dD}t|dkrddd|ddDfd dt|dD|d<|dD] }t|qldS) z Recursively removes intermediate layers from the tree to improve readability. Keeps at least the first and last layers if many consecutive layers are present. Args: node (`dict`): The root or subnode to prune recursively. rNNcSs g|]\}}t|r||fqSr)rcr9r:rUrrrr=r^z-prune_intermediate_layers..rXcSsg|] \}}|qSrr)r9r:r6rrrr=r^rcsg|]\}}|vr|qSrrrdZ to_removerrr=r^)rPrClenprune_intermediate_layers)rTZ layer_blocksrUrrfrrhs   rhc s<|rdz&tj|ddtj||jd}Wqnty`}ztd|d|WYd}~qnd}~00n |jd}td|d|d}|d }t |j t |d "}t j |j |d d Wdn1s0Yfd dt t |j }|t |d  }t j ||d d Wdn1s.0YdS)NTexist_okZ _debug_tree"Unexpected or existing debug_path=rYzWriting model trace at z.jsonz_FULL_TENSORS.jsonz _SUMMARY.jsonwrX)indentcsJfdd|di|di|dgD] }|q8dS)NcsNt|tr.|dd|D] }|qnt|trJ|D] }|q.strip_values..cleaninputsrOrNrZrS strip_valuesrqrrus  z+log_model_debug_trace..strip_values)r*makedirsr+r,_debugger_module_dump_name Exceptionr0loggerinforR _call_treeopenjsondumploadsdumps)rmodelbasee full_pathZ summary_pathfZ tree_copyrrtrlog_model_debug_traces$*   0  rrY)rdo_prune_layersrc sjjddgd_g__rrztjddWn6typ}ztdd|WYd}~n d}~00fdd} D]&\}}|d krq||d|qj t fd d }|_ dS) a Attaches a debugging wrapper to every module in the model. This records structured inputs and outputs during the forward pass into a call tree. Args: model (`PreTrainedModel`, `nn.Module`): Model to wrap. debug_path (`str`): Optional directory to dump debug JSON files. do_prune_layers (`bool`, *optional*, defaults to `True`): Whether to prune intermediate layers. use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the `value` property in the asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors files and store the relative path to that file in the `value` property. NrVrsrOrNTrirkrYcs0jtfdd}|_dS)NcstrN||dfddDtdddgd}j|t|i|}Wdn1sz0YtrtddDd krd|d <nt|d d|d <j}|d s|d jrjd d ||S)Nargskwargscs&i|]}t|dkr||qS)r)rgr9r>Z dict_inputsrrr?5r^zY_attach_debugger_logic..wrap_forward..wrapped_forward.._inputsrrcss|] }dVqdS)rNr)r9r6rrrr]Fr^zX_attach_debugger_logic..wrap_forward..wrapped_forward..rrO_outputsrNre) rr8_debugger_model_call_stackappendr Zno_gradsumZnamed_childrenrQ)inpskwsrTr3finished)rrrmodule orig_forwardrrrwrapped_forward1s>   ,    zE_attach_debugger_logic..wrap_forward..wrapped_forward)forward functoolswraps)rrr)rrr)rrrr wrap_forward.s'z,_attach_debugger_logic..wrap_forwardrWcstr>dt||ddddgd}j||i|}tr܈jrt|dd|d<j}|djd<|djd<|d jd <fd d tjDrtjt d |S) Nz (top-level)rrrrrrOrsrNcs$g|]}j|sj|dqS)N)r{rQr)rrrr=r^zG_attach_debugger_logic..top_wrapped_forward..)rr) rr8rrrQr{rAkeysrhr)rrZtop_noder3r class_namerrrZreal_top_forwardrrrtop_wrapped_forwardds:      z3_attach_debugger_logic..top_wrapped_forward) __class____name__r{rrwr*rvrxr0 named_modulesrrr) rrrrrrname submodulerrrr_attach_debugger_logics$(.%r)r )backendsc csldd|D}|j||<t||||z"|VW|D]\}}||_q.N)rrrrE)rrrrZ orig_forwardsZmodule_instanceZforward_methodrrrmodel_addition_debugger_contexts3  r)NTN)NTN)rYTT)NTT),rr}r*re contextlibrriortypingrZutils.import_utilsrrr Zsafetensors.torchr r Z is_availableZtorch.distributed.tensorr utilsr Z get_loggerrryrcompilerstrrrboolr4r8rGr(rRr_rcrhrrrrrrrsb        1+ -