o 0 iB @sddlZddlZddlZddlZddlmZmZddlmZddl m Z ddl m Z ddl mZmZerMddlZddlmZd ZejrLddlZd Znd Ze eZd d Zed ZdedefddZddZ d/de ede de efddZ!d/de ede de efddZ"dej#fddZ$ddZ%ed Z&d!d"Z'd#d$Z(de efd%d&Z) ' d0ded(e de fd)d*Z*ed+d,e  d1de ed(e de fd-d.Z+dS)2N)contextmanagerredirect_stdout)StringIO)Optional)logging)is_torch_availablerequires) save_fileFTcCs trtjs dStjdkS)z7Return True if rank=0 or we aren't running distributed.Tr)_torch_distributed_availabletorch distributedZis_initializedZget_rankrrn/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/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_diff7s rcCstr dt|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|r t|}n4|r6|ds|d7}|r!tj||n|}td| i|d|}n t d|d|dt |j t |j|d }|jtjtjtjhvrz|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 ValueErrorrrr Zfloat16Zfloat32Zbfloat16updaterr"r#r$r%)r!rrrZ value_outfilepathoutrrr_serialize_tensor_like_ioFs.    r3cst|ttfrfddt|DSt|tr&fdd|DSt|dr4t|jdSt|t j rBt|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_rrr _serialize_io).0ivrrrrr sz!_serialize_io..c s,i|]\}}|t|d|dqSr4r7)r9kr;r<rr sz!_serialize_io..rr6) isinstancelisttuple enumeratedictitemshasattrr3rr Tensorrr)r!rrrrr<rr8ws     r8r!c Csxtjdddt$}t|t||}Wdn1s"wYWdn1s1wYt|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   r'cCs4|dr|dd|dD]}t|qdSdS)Nchildrenoutputs)getpopprune_outputs_if_childrennodechildrrrrRs    rRz (.*)\.(\d+)$csHt|dd}|r|dsdS|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 s&z!is_layer_block..)LAYER_SUFFIX_REmatchrPgroupany)rTr_rr[ris_layer_blocks   rbcs~|dsdSddt|dD}t|dkr2dd|ddDfd dt|dD|d<|dD]}t|q6dS) 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)rbr9r:rUrrrr=s z-prune_intermediate_layers..rXcSsg|]\}}|qSrr)r9r:r5rrrr=srcsg|] \}}|vr|qSrrrc to_removerrr=s)rPrClenprune_intermediate_layers)rTZ layer_blocksrUrrerrhs    rhc s.|r+ztj|ddtj||jd}Wnty*}z td|d|d}~ww|jd}td|d|d}|d }t |j t |d }t j |j |d d Wdn1s_wYfd dt t |j }|t |d }t j ||d d WddS1swYdS)NTexist_okZ _debug_tree"Unexpected or existing debug_path=rYzWriting model trace at z.jsonz_FULL_TENSORS.jsonz _SUMMARY.jsonwrX)indentcsJfdd|di|di|dgD]}|qdS)NcsTt|tr|dd|D]}|qdSt|tr&|D]}|qdSdS)Nr!)r@rDrQvaluesrA)valr;itemcleanrrrrs      z:log_model_debug_trace..strip_values..cleaninputsrOrNrZrS strip_valuesrqrrus   z+log_model_debug_trace..strip_values)r)makedirsr*r+_debugger_module_dump_name Exceptionr/loggerinforR _call_treeopenjsondumploadsdumps)rmodelbasee full_pathZ summary_pathfZ tree_copyrrtrlog_model_debug_traces,     "rrYdo_prune_layersc sjjddgd_g__r3z tjddWnty2}z tdd|d}~wwfdd} D]\}}|d krHq?||d|q?j 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 associated 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)Ncstr'||dfddDtdddgd}j|t|i|}Wdn1s=wYtrtddDd krWd|d <n t|d d|d <j}|d sr|d jrjd d ||S)Nargskwargscs&i|]}t|dkr||qS)r)rgr9r>Z dict_inputsrrr?6s&zY_attach_debugger_logic..wrap_forward..wrapped_forward.._inputsr6rcss|]}dVqdS)rNr)r9r5rrrr]GszX_attach_debugger_logic..wrap_forward..wrapped_forward..rrO_outputsrNrd) rr8_debugger_model_call_stackappendr Zno_gradsumZnamed_childrenrQ)inpskwsrTr2finished)rrrmodule orig_forwardrrrwrapped_forward2s@     zE_attach_debugger_logic..wrap_forward..wrapped_forward)forward functoolswraps)rrr)rrr)rrrr wrap_forward/s 'z,_attach_debugger_logic..wrap_forwardrWcstrdt||ddddgd}j||i|}trnjrnt|dd|d<j}|djd<|djd<|d jd <fd d tjDrhtjt d |S) Nz (top-level)rrr6rrrOrsrNcs$g|]}j|sj|dqS)N)r{rQr)rrrr=s$zG_attach_debugger_logic..top_wrapped_forward..)rr) rr8rrrQr{rAkeysrhr)rrZtop_noder2r class_namerrrZreal_top_forwardrrrtop_wrapped_forwardes:     z3_attach_debugger_logic..top_wrapped_forward) __class____name__r{rrwr)rvrxr/ named_modulesrrr) rrrrrrname submodulerrrr_attach_debugger_logics(. %r)r )backendsc csldd|D}|j||<t||||z|VW|D]\}}||_qdS|D]\}}||_q-w)a # Model addition debugger - context manager for model adders This context manager is a power user tool intended for model adders. It tracks all forward calls within a model forward and logs a slice of each input and output on a nested JSON file. If `use_repr=True` (the default), the JSON file will record a `repr()`-ized version of the tensors as a list of strings. If `use_repr=False`, the full tensors will be stored in separate SafeTensors files and the JSON file will provide a relative path to that file. To note, this context manager enforces `torch.no_grad()`. ## Usage add the context manager to a model to debug ```python import torch from PIL import Image from transformers import LlavaProcessor, LlavaForConditionalGeneration, model_addition_debugger_context torch.random.manual_seed(673) # load pretrained model and processor model_id = "llava-hf/llava-1.5-7b-hf" processor = LlavaProcessor.from_pretrained(model_id) model = LlavaForConditionalGeneration.from_pretrained(model_id) # create random image input random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy()) # prompt prompt = "Describe this image." # process inputs inputs = processor(text=prompt, images=random_image, return_tensors="pt") # call forward method (not .generate!) with model_addition_debugger_context(model, debug_path="Your_debug_path", do_prune_layers=False): output = model.forward(**inputs) ``` cSsi|]\}}||jqSr)r)r9r5mrrrr?sz3model_addition_debugger_context..N)rrrrE)rrrrZ orig_forwardsZmodule_instanceZforward_methodrrrmodel_addition_debugger_contexts3 r)NTN)rYTT)NTT),rr}r)re contextlibrriortypingrutilsrZutils.import_utilsrr r Zsafetensors.torchr r r Z is_availableZtorch.distributed.tensorZ get_loggerrryrcompilerstrrrboolr3r8rGr'rRr^rbrhrrrrrrrsx        1+ -