Aluode/PerceptionLabPortable

1# Copyright 2025 The HuggingFace Inc. team.2# All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16import functools17import json18import os19import re20from contextlib import contextmanager, redirect_stdout21from io import StringIO22from typing import Optional23 24from .utils import logging25from .utils.import_utils import is_torch_available, requires26 27 28if is_torch_available():29    import torch30    from safetensors.torch import save_file31 32    _torch_distributed_available = False33    # Note to code inspectors: this toolbox is intended for people who add models to `transformers`.34    if torch.distributed.is_available():35        import torch.distributed.tensor36 37        _torch_distributed_available = True38else:39    _torch_distributed_available = False40 41 42logger = logging.get_logger(__name__)43 44 45def _is_rank_zero():46    """Return True if rank=0 or we aren't running distributed."""47    if not (_torch_distributed_available and torch.distributed.is_initialized()):48        return True49    return torch.distributed.get_rank() == 050 51 52MEMORY_ADDRESS_REGEX = re.compile(r"object at 0x[0-9A-Fa-f]+")53 54 55def _sanitize_repr_for_diff(x_str: str) -> str:56    """57    Replace memory addresses in an object's repr with a stable placeholder58    so that beautiful JSON diffs won't be ruined by ephemeral addresses.59    """60    return MEMORY_ADDRESS_REGEX.sub("object at 0xXXXXXXXX", x_str)61 62 63def _dtensor_repr(x):64    """Return a stable string representation for a DTensor-like object."""65    if _is_rank_zero():66        return f"DTensor (rank0) -> {repr(x._local_tensor)}"67    return "DTensor(non-rank0)"68 69 70def _serialize_tensor_like_io(71    value, debug_path: Optional[str] = None, use_repr: bool = True, path_to_value: Optional[str] = None72):73    """74    Converts Tensors and DTensors to a JSON-serializable dictionary representation.75 76    Args:77        value: Any Python object, often including torch Tensors, lists, dicts, etc.78        debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files.79        use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensor as the80            `value` property in the asscoiated FULL_TENSORS.json file, or to store the full tensors in separate81            SafeTensors file and store the relative path to that file in the `value` property in the dictionary.82        path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full83            tensor value if `use_repr=False`.84 85    Returns:86        A nested Python structure (list, dict, or sanitized string) that is safe to json.dump.87    """88    torch.set_printoptions(sci_mode=True)89 90    if use_repr:91        value_out = _repr_to_list(value)92    elif path_to_value:93        if not path_to_value.endswith(".safetensors"):94            path_to_value += ".safetensors"95 96        filepath = os.path.join(debug_path, path_to_value) if debug_path else path_to_value97        save_file({"data": value.contiguous().detach().cpu()}, filepath)98        value_out = f"./{path_to_value}"99    else:100        raise ValueError(f"{use_repr=} and {path_to_value=} cannot both be falsy.")101 102    out = {103        "shape": repr(value.shape),104        "dtype": repr(value.dtype),105        "value": value_out,106    }107    if value.dtype in {torch.float16, torch.float32, torch.bfloat16}:108        out.update(109            {110                "mean": _sanitize_repr_for_diff(repr(value.mean())),111                "std": _sanitize_repr_for_diff(repr(value.std())),112                "min": _sanitize_repr_for_diff(repr(value.min())),113                "max": _sanitize_repr_for_diff(repr(value.max())),114            }115        )116    return out117 118 119def _serialize_io(value, debug_path: Optional[str] = None, use_repr: bool = True, path_to_value: Optional[str] = None):120    """121    Recursively build a JSON-serializable Python structure from `value`.122    Tensors and DTensors become either sanitized repr strings, or are saved to disk as SafeTensors files and their123    relative paths are recorded in the returned Python structure.124    Lists/tuples/dicts are recursed into.125    All memory addresses are replaced with a stable placeholder.126 127    Args:128        value: Any Python object, often including torch Tensors, lists, dicts, etc.129        debug_path (`str`, *optional*, defaults to `None`): Directory to dump debug JSON and SafeTensors files.130        use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the131            `value` property in the asscoiated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors132            files and store the relative path to that file in the `value` property.133        path_to_value (`str`, *optional*, defaults to `None`): The file name for the SafeTensors file holding the full134            tensor value if `use_repr=False`.135 136    Returns:137        A nested Python structure (list, dict, or sanitized string) that is safe to json.dump.138    """139    if isinstance(value, (list, tuple)):140        return [141            _serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f"{path_to_value}_{i}")142            for i, v in enumerate(value)143        ]144 145    if isinstance(value, dict):146        return {147            k: _serialize_io(v, debug_path=debug_path, use_repr=use_repr, path_to_value=f"{path_to_value}_{k}")148            for k, v in value.items()149        }150 151    if hasattr(value, "_local_tensor"):152        return _serialize_tensor_like_io(153            value._local_tensor, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value154        )155 156    if isinstance(value, torch.Tensor):157        return _serialize_tensor_like_io(value, debug_path=debug_path, use_repr=use_repr, path_to_value=path_to_value)158 159    return _sanitize_repr_for_diff(repr(value))160 161 162def _repr_to_list(value: torch.Tensor):163    """164    Converts a tensor into a sanitized multi-line string representation.165 166    Args:167        value (`torch.Tensor`): The tensor to represent.168 169    Returns:170        `list[str]`: List of string lines representing the tensor.171    """172    torch.set_printoptions(sci_mode=True, linewidth=120)173    with StringIO() as buf, redirect_stdout(buf):174        print(value)  # to redirected stdout to avoid line splits175        raw = buf.getvalue()176    return _sanitize_repr_for_diff(raw).splitlines()177 178 179def prune_outputs_if_children(node):180    # if there are children, remove this node's "outputs"181    # so we only see outputs at the leaf level182    if node.get("children"):183        node.pop("outputs", None)184        for child in node["children"]:185            prune_outputs_if_children(child)186 187 188LAYER_SUFFIX_RE = re.compile(r"(.*)\.(\d+)$")  # should be generic enough, ends with a number189 190 191def is_layer_block(node):192    """193    Checks whether a node represents a layer block with submodules.194 195    Args:196        node (`dict`): A node from the call tree.197 198    Returns:199        `bool`: Whether the node is a layer block.200    """201    match = LAYER_SUFFIX_RE.match(node.get("module_path", ""))202    if not match or not node.get("children"):203        return False204    number = match.group(2)205    return any(f".{number}." in child.get("module_path", "") for child in node["children"])206 207 208def prune_intermediate_layers(node):209    """210    Recursively removes intermediate layers from the tree to improve readability.211    Keeps at least the first and last layers if many consecutive layers are present.212 213    Args:214        node (`dict`): The root or subnode to prune recursively.215    """216    if not node.get("children"):217        return218    layer_blocks = [(i, child) for i, child in enumerate(node["children"]) if is_layer_block(child)]219 220    if len(layer_blocks) > 2:221        to_remove = [i for i, _ in layer_blocks[1:-1]]222        node["children"] = [child for i, child in enumerate(node["children"]) if i not in to_remove]223 224    for child in node["children"]:225        prune_intermediate_layers(child)226 227 228def log_model_debug_trace(debug_path: Optional[str], model):229    if debug_path:230        try:231            os.makedirs(debug_path, exist_ok=True)232            base = os.path.join(debug_path, model._debugger_module_dump_name + "_debug_tree")233        except Exception as e:234            raise ValueError(f"Unexpected or existing debug_path={debug_path}.") from e235    else:236        base = model._debugger_module_dump_name + "_debug_tree"237 238    logger.info(f"Writing model trace at {base}.json")239    full_path = base + "_FULL_TENSORS.json"240    summary_path = base + "_SUMMARY.json"241 242    prune_outputs_if_children(model._call_tree)243 244    with open(full_path, "w") as f:245        json.dump(model._call_tree, f, indent=2)246 247    # summary-only version for readability - traversing the tree again #TODO optimize?248    def strip_values(node):249        def clean(val):250            if isinstance(val, dict):251                val.pop("value", None)252                for v in val.values():253                    clean(v)254            elif isinstance(val, list):255                for item in val:256                    clean(item)257 258        clean(node.get("inputs", {}))259        clean(node.get("outputs", {}))260 261        for child in node.get("children", []):262            strip_values(child)263 264    tree_copy = json.loads(json.dumps(model._call_tree))  # deep copy265    strip_values(tree_copy)266 267    with open(summary_path, "w") as f:268        json.dump(tree_copy, f, indent=2)269 270 271def _attach_debugger_logic(272    model,273    debug_path: str = ".",274    do_prune_layers: bool = True,275    use_repr: bool = True,276):277    """278    Attaches a debugging wrapper to every module in the model.279 280    This records structured inputs and outputs during the forward pass into a call tree.281 282    Args:283        model (`PreTrainedModel`, `nn.Module`): Model to wrap.284        debug_path (`str`): Optional directory to dump debug JSON files.285        do_prune_layers (`bool`, *optional*, defaults to `True`): Whether to prune intermediate layers.286        use_repr (bool, *optional*, defaults to `True`): Whether to save a `repr()`-ized version of the tensors as the287            `value` property in the associated FULL_TENSORS.json file, or to store full tensors in separate SafeTensors288            files and store the relative path to that file in the `value` property.289    """290    class_name = model.__class__.__name__291 292    # Prepare data structures on the model object293    model._call_tree = {"module_path": class_name, "inputs": None, "outputs": None, "children": []}294    model._debugger_model_call_stack = []295    model._debugger_module_dump_name = class_name  # used for final JSON filename296 297    if debug_path:298        try:299            os.makedirs(debug_path, exist_ok=True)300        except Exception as e:301            raise ValueError(f"Unexpected or existing debug_path={debug_path}.") from e302 303    def wrap_forward(module, full_path):304        orig_forward = module.forward305 306        @functools.wraps(orig_forward)307        def wrapped_forward(*inps, **kws):308            if _is_rank_zero():309                dict_inputs = {"args": inps, "kwargs": kws}310                dict_inputs = {k: dict_inputs[k] for k in dict_inputs if len(dict_inputs[k]) > 0}311                node = {312                    "module_path": full_path,313                    "inputs": _serialize_io(314                        dict_inputs,315                        debug_path=debug_path,316                        use_repr=use_repr,317                        path_to_value=f"{full_path}_inputs",318                    ),319                    "outputs": None,320                    "children": [],321                }322                model._debugger_model_call_stack.append(node)323            with torch.no_grad():324                out = orig_forward(*inps, **kws)325 326            if _is_rank_zero():327                if sum(1 for _ in module.named_children()) > 0:328                    node["outputs"] = None329                else:330                    node["outputs"] = _serialize_io(331                        out,332                        debug_path=debug_path,333                        use_repr=use_repr,334                        path_to_value=f"{full_path}_outputs",335                    )336 337                finished = model._debugger_model_call_stack.pop()338                # prune empty vertices here as well (mostly empty children nodes)339                if not finished["children"]:340                    finished.pop("children")341 342                if model._debugger_model_call_stack:343                    model._debugger_model_call_stack[-1]["children"].append(finished)344            return out345 346        module.forward = wrapped_forward347 348    # wrap all submodules349    for name, submodule in model.named_modules():350        if name == "":351            continue352        wrap_forward(submodule, f"{class_name}.{name}")353 354    # wrap top-level forward355    real_top_forward = model.forward356 357    @functools.wraps(real_top_forward)358    def top_wrapped_forward(*inps, **kws):359        if _is_rank_zero():360            top_node = {361                "module_path": f"{class_name} (top-level)",362                "inputs": _serialize_io(363                    {"args": inps, "kwargs": kws},364                    debug_path=debug_path,365                    use_repr=use_repr,366                    path_to_value=f"{class_name}_inputs",367                ),368                "outputs": None,369                "children": [],370            }371            model._debugger_model_call_stack.append(top_node)372 373        out = real_top_forward(*inps, **kws)374        if _is_rank_zero() and model._debugger_model_call_stack:375            top_node["outputs"] = _serialize_io(376                out,377                debug_path=debug_path,378                use_repr=use_repr,379                path_to_value=f"{class_name}_outputs",380            )381            finished = model._debugger_model_call_stack.pop()382            model._call_tree["inputs"] = finished["inputs"]383            model._call_tree["outputs"] = finished["outputs"]384            model._call_tree["children"] = finished["children"]385            # prune empty stuff for visibility386            [model._call_tree.pop(k, None) for k in list(model._call_tree.keys()) if not model._call_tree[k]]387 388            # prune layers that are not 0 or last389            if do_prune_layers:390                prune_intermediate_layers(model._call_tree)391            # Write final JSON trace here392            log_model_debug_trace(debug_path=debug_path, model=model)393        return out394 395    model.forward = top_wrapped_forward396 397 398@requires(backends=("torch",))399@contextmanager400def model_addition_debugger_context(401    model,402    debug_path: Optional[str] = None,403    do_prune_layers: bool = True,404    use_repr: bool = True,405):406    """407    # Model addition debugger - context manager for model adders408    This context manager is a power user tool intended for model adders.409 410    It tracks all forward calls within a model forward and logs a slice of each input and output on a nested JSON file.411    If `use_repr=True` (the default), the JSON file will record a `repr()`-ized version of the tensors as a list of412    strings. If `use_repr=False`, the full tensors will be stored in separate SafeTensors files and the JSON file will413    provide a relative path to that file.414 415    To note, this context manager enforces `torch.no_grad()`.416 417    ## Usage418 419    add the context manager to a model to debug420 421    ```python422    import torch423 424    from PIL import Image425    from transformers import LlavaProcessor, LlavaForConditionalGeneration, model_addition_debugger_context426 427    torch.random.manual_seed(673)428 429    # load pretrained model and processor430    model_id = "llava-hf/llava-1.5-7b-hf"431    processor = LlavaProcessor.from_pretrained(model_id)432    model = LlavaForConditionalGeneration.from_pretrained(model_id)433 434    # create random image input435    random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy())436 437    # prompt438    prompt = "<image>Describe this image."439 440    # process inputs441    inputs = processor(text=prompt, images=random_image, return_tensors="pt")442 443    # call forward method (not .generate!)444    with model_addition_debugger_context(model, debug_path="Your_debug_path", do_prune_layers=False):445        output = model.forward(**inputs)446    ```447 448    """449    orig_forwards = {m: m.forward for _, m in model.named_modules()}450    orig_forwards[model] = model.forward451    _attach_debugger_logic(model, debug_path, do_prune_layers, use_repr)452    try:453        yield model454    finally:455        for module_instance, forward_method in orig_forwards.items():456            module_instance.forward = forward_method457