🔴 VLM: compile compatibility

[zucchini-nlp]

What does this PR do?

As per title, adds flags in VLMs when needed, removes test skips and makes sure VLMs are compile compatible. Also for BLIP models adds new cache format in OPT which is one of backbones. Now all official BLIP models can support static cache and thus compile

NOTE:

• Tests with -k compile_forward and -k static_ were run for all models and are passing
• Regarding executorch which I also checked, the model can be exported and run a forward pass. But the generation won't work and probably would need smth similar to what we do when exporting VLMs in ONNX. Still need to dig more into that in later PRs

---

How to run compile and export for VLMs:

import requests
from PIL import Image

import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration
from transformers.generation import GenerationConfig
from transformers.cache_utils import StaticCache
from transformers.integrations.executorch import (
    TorchExportableModuleWithStaticCache,
    convert_and_export_with_cache,
)

model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
model = LlavaForConditionalGeneration.from_pretrained(
    model_id, 
    torch_dtype="float16", 
    device_map="cuda:0",
)
processor = AutoProcessor.from_pretrained(model_id)

conversation = [
    {

      "role": "user",
      "content": [
          {"type": "text", "text": "What are these?"},
          {"type": "image"},
        ],
    },
]
prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"
raw_image = Image.open(requests.get(image_file, stream=True).raw)
inputs = processor(images=raw_image, text=prompt, return_tensors='pt').to(0, torch.float16)

# Run with static cache which compiles the forward in decoding phase for you
output = model.generate(**inputs, max_new_tokens=20, cache_implementation="static")
print(processor.decode(output[0][2:], skip_special_tokens=True))



# Try to export with `torch.export`. NOTE: TorchExportableModuleWithStaticCache is not ready for VLMs
# and as mentioned above, VLMs might need to export 3 different modules as in ONNX. One for text embedding,
# one for vision backbone and one for the LM backbone with simple decoding token-by-token
max_generation_length = 1000
model = LlavaForConditionalGeneration.from_pretrained(
    model_id,
    device_map="cuda:0",
    torch_dtype="float16",
    attn_implementation="sdpa",
    generation_config=GenerationConfig(
        use_cache=True,
        cache_implementation="static",
        max_length=max_generation_length,
        cache_config={
            "batch_size": 1,
            "max_cache_len": max_generation_length,
        },
    ),
)

# Adapted from `TorchExportableModuleWithStaticCache` with minor changes
class TorchExportableModuleWithStaticCacheForVLM(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.static_cache = StaticCache(
            config=self.model.config.get_text_config(),
            batch_size=self.model.generation_config.cache_config.batch_size,
            max_cache_len=self.model.generation_config.cache_config.max_cache_len,
            dtype=self.model.dtype,
            device=model.device,
        )
        self.is_causal = any(("CausalLM" in arch or "ConditionalGeneration" in arch) for arch in self.model.config.architectures)
        if self.is_causal:
            causal_mask = torch.triu(
                torch.full(
                    (
                        self.model.generation_config.cache_config.batch_size,
                        1,
                        self.static_cache.max_cache_len,
                        self.static_cache.max_cache_len
                    ),
                    fill_value=torch.finfo(self.model.dtype).min,
                    dtype=self.model.dtype,
                    device=model.device,
                )
            )
            self.register_buffer("mask", causal_mask, persistent=False)

    def forward(
        self,
        input_ids: torch.Tensor,
        cache_position: torch.Tensor,
        pixel_values: torch.Tensor,
    ):
        _, seqlen = input_ids.shape
        attn_mask = self.mask[:, :, cache_position, :] if self.is_causal else None
        outs = self.model(
            input_ids=input_ids,
            attention_mask=attn_mask,
            position_ids=cache_position.unsqueeze(0),
            pixel_values=pixel_values,
            cache_position=cache_position,
            past_key_values=self.static_cache,
            use_cache=True,
        )
        return outs.logits

cache_position = torch.arange(inputs.input_ids.shape[1], dtype=torch.long, device=model.device)
export_inputs = {"input_ids": inputs.input_ids, "cache_position": cache_position, "pixel_values": inputs.pixel_values}

with torch.no_grad():
    exported_program = torch.export.export(
        TorchExportableModuleWithStaticCacheForVLM(model),
        args=(),
        kwargs=export_inputs,
        strict=True,
    )

torch.export.save(exported_program, "exported_llava.pt2")
exported_program = torch.export.load("exported_llava.pt2")
out = exported_program.module().forward(
    input_ids=inputs.input_ids,
    pixel_values=inputs.pixel_values,
    cache_position=cache_position,
)

---

Benchmark on "llava-hf/llava-onevision-qwen2-7b-ov-hf" using the same script we use for LLMs + dummy image in inputs

Fixes #29891

[zucchini-nlp]

I dont know why but OPT model works as decoder-only but the attention is written as cross-attention (not used anywhere in codebase). So we need to support somehow BC while using the new DynamicCache

As a workaround I simply added a check on cache instance. Another possibility is to accept and return only the correct cache (self or cross attn) but that means all encoder-decoder models will need a change thus breaking BC

[gante]

Very much copy-paste from somewhere else, see this comment

IMO, we can make our maintenance easier and assume no encoder-decoder stuff :) But don't spend more time here, eventually this will be rewritten with modular

[zucchini-nlp]

Ready for review failing test is flaky otherwise everything is passing on my end, including slow test for compile/StaticCache

[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.

[gante]

@zucchini-nlp

In the PR header we can read

all VLMs have dynamic control in prepare_inputs_for_generation and thus skip test_compile_forward which compiles the model for pre-fill phase. But the test for decoding stage compile is green therefore I'm leaving the flag as True

test_compile_forward was the old name for test_generate_compile_model_forward if I'm not mistaken, back when it also did end-to-end compilation tests. We no longer have the end-to-end compilation tests, so this part of the PR header is no longer accurate, correct?

[gante]

The failing test seems related to this PR :D

[ArthurZucker]

Very nice but 🔴 there are a few breaking changes so let's be careful!
And do you have some benches / perf imporvements to share (making sure reduce overhead is working etc_

[ArthurZucker]

yeah, tho calling .continguous works as well

[zucchini-nlp]

Soooo, here is the correct eval with llava-ov-7b. One thing to note is that VLMs will not benefit from torch when the context length is very high, like in case of videos or high-res images. That's the reason vanilla llava got speed up from first try, while llava-ov needed a few runs to notice how many input tokens we had

I will make fixup and merge this, because VLMs with less tokens per image get speedups