Add model visual debugger

[molbap]

What does this PR do?

This PR introduces a tool to trace ALL model inputs and outputs to a json nested format. Main values, shapes, dtypes are all outputted.

Usage

It's a simple context manager to add before calling a forward on your inputs. To note, inference_mode is enforced here.

import torch
from PIL import Image
from transformers import LlavaProcessor, LlavaForConditionalGeneration
from transformers.model_debugging_utils import  model_addition_debugger_context

# 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, low_cpu_mem_usage=True)

# create random image input
torch.random.manual_seed(673)
random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy())

# prompt
prompt = "<image>Describe this image."

# process inputs
inputs = processor(text=prompt, images=random_image, return_tensors="pt")

with model_addition_debugger_context(model=model, debug_path="model_debug"):
    output = model.forward(**inputs)

Why??

Because when porting models to transformers, even from python to python, model adders often have to do a lot of manual operations, involving saving and loading tensors, comparing dtypes, etc. This small tool can hopefully shave off some time.

Example output

Here, you can see two jsons of the same model, where the sole difference is the epsilon of a layer normalization. I added the decorator, and we can see with string matching that the outputs start to differ. I'll include a snippet with difflib too I think to make it simpler.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

super nice! Just missing something like register_model_for_debug(LlamaModel) !

[molbap]

Can remove this/move it to another PR, with #36827 done we should be able to get rid of this check in make fixup

[ArthurZucker]

would put this in the model_debugging_utils

[molbap]

added in the doc :) not sure of the output cause the builder is broken currently