[ImageGPT] Cross-attention crashes, ignores encoder_attention_mask, and doubles its cache

[blipbyte]

System Info

- `transformers` version: 5.18.0.dev0
- Platform: Linux-7.0.0-31-generic-x86_64-with-glibc2.39
- Python version: 3.10.20
- Huggingface_hub version: 1.15.0
- Safetensors version: 0.8.0
- Accelerate version: 1.13.0
- Accelerate config: not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.12.0+cu130 (CUDA)
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: no, the script below is CPU-only
- GPU type: NVIDIA GeForce RTX 3060

Who can help?

@zucchini-nlp @vasqu

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

ImageGPTModel accepts encoder_hidden_states and encoder_attention_mask and builds the
encoder mask at modeling_imagegpt.py:526, but that path is broken four ways. GPT-2, which it was
hand-copied from, passes all four. CPU-only, no checkpoint, on main at 606e6e8bb5:

import torch
from transformers import (
    DynamicCache,
    EncoderDecoderCache,
    GPT2Config,
    GPT2Model,
    ImageGPTConfig,
    ImageGPTModel,
)

torch.manual_seed(0)
SEQ, ENC, DIM = 7, 11, 32


def build(kind):
    config_cls, model_cls = (ImageGPTConfig, ImageGPTModel) if kind == "imagegpt" else (GPT2Config, GPT2Model)
    config = config_cls(vocab_size=99, n_embd=DIM, n_layer=2, n_head=4, n_positions=64,
                        n_inner=37, add_cross_attention=True, bos_token_id=98, eos_token_id=98)
    return model_cls(config).eval()


ids = torch.randint(0, 90, (2, SEQ))
encoder_states = torch.randn(2, ENC, DIM)
encoder_mask = torch.ones(2, ENC, dtype=torch.long)
encoder_mask[:, ENC // 2:] = 0                  # second half of the encoder output is padding
garbage = encoder_states.clone()
garbage[:, ENC // 2:] = 1e6                     # what the mask hides must not reach the output

print("1. default call (use_cache defaults to True, as generate() does)")
for kind in ("gpt2", "imagegpt"):
    try:
        with torch.no_grad():
            build(kind)(input_ids=ids, encoder_hidden_states=encoder_states, encoder_attention_mask=encoder_mask)
        print(f"   {kind:9} ok")
    except Exception as e:
        print(f"   {kind:9} {type(e).__name__}: {e}")

print("\n2. is the encoder padding mask honoured? (use_cache=False, to get past 1)")
for kind in ("gpt2", "imagegpt"):
    model = build(kind)
    with torch.no_grad():
        a = model(input_ids=ids, encoder_hidden_states=encoder_states,
                  encoder_attention_mask=encoder_mask, use_cache=False).last_hidden_state
        b = model(input_ids=ids, encoder_hidden_states=garbage,
                  encoder_attention_mask=encoder_mask, use_cache=False).last_hidden_state
    print(f"   {kind:9} padding ignored: {torch.equal(a, b)}   max|delta| = {(a - b).abs().max():.4g}")

print("\n3. same, plus a padded decoder attention_mask (the ordinary batched call)")
decoder_mask = torch.tensor([[1] * SEQ, [1] * (SEQ - 2) + [0, 0]])
for kind in ("gpt2", "imagegpt"):
    model = build(kind)
    try:
        with torch.no_grad():
            model(input_ids=ids, attention_mask=decoder_mask, encoder_hidden_states=encoder_states,
                  encoder_attention_mask=encoder_mask, use_cache=False)
        print(f"   {kind:9} ok")
    except Exception as e:
        print(f"   {kind:9} {type(e).__name__}: {e}")

print("\n4. one cached decoding step (pass an EncoderDecoderCache, to get past 1)")
for kind in ("gpt2", "imagegpt"):
    model = build(kind)
    cache = EncoderDecoderCache(DynamicCache(), DynamicCache())
    with torch.no_grad():
        out = model(input_ids=ids, past_key_values=cache, encoder_hidden_states=encoder_states,
                    encoder_attention_mask=encoder_mask, use_cache=True)
        model(input_ids=ids[:, -1:], past_key_values=out.past_key_values,
              encoder_hidden_states=encoder_states, encoder_attention_mask=encoder_mask, use_cache=True)
    kept = out.past_key_values.cross_attention_cache.layers[0].keys.shape[-2]
    print(f"   {kind:9} cross-attention cache holds {kept} keys for {ENC} encoder positions")

1. default call (use_cache defaults to True, as generate() does)
   gpt2      ok
   imagegpt  UnboundLocalError: local variable 'is_updated' referenced before assignment

2. is the encoder padding mask honoured? (use_cache=False, to get past 1)
   gpt2      padding ignored: True   max|delta| = 0
   imagegpt  padding ignored: False   max|delta| = 4.652

3. same, plus a padded decoder attention_mask (the ordinary batched call)
   gpt2      ok
   imagegpt  RuntimeError: The size of tensor a (11) must match the size of tensor b (7) at non-singleton dimension 3

4. one cached decoding step (pass an EncoderDecoderCache, to get past 1)
   gpt2      cross-attention cache holds 11 keys for 11 encoder positions
   imagegpt  cross-attention cache holds 22 keys for 11 encoder positions

Two commits did this; a third fixed it everywhere but here.

#38635 (c8524aeb07) rewrote ImageGPTAttention.forward by hand. One hunk deleted
attention_mask = encoder_attention_mask — cases 2 and 3 — and added a cross-attention cache
write with no reuse guard, so from the second decoding step the cache appends to itself — case 4.

#39754 (ccb2e0e03b) added that guard to GPT-2 thirteen days later, and to
decision_transformer through its # Copied from. ImageGPT carries neither, so it was skipped.

#40811 (7a1aeec36e) then swapped ImageGPT's EncoderDecoderCache for a plain DynamicCache,
leaving is_updated unbound at line 232 — case 1, which now fires first and hides the rest.

Expected behavior

All four should behave as GPT-2 does: no crash, masked encoder positions not moving the output,
and a cache the length of the encoder.

Not only reachable by hand: VisionEncoderDecoderConfig sets add_cross_attention=True on the
decoder for you, and its docstring shows the same.

decoder_config = ImageGPTConfig(..., is_decoder=True, add_cross_attention=True)
model = VisionEncoderDecoderModel(
    encoder=ViTModel(encoder_config), decoder=ImageGPTForCausalImageModeling(decoder_config)
)
model(pixel_values=pixel_values, decoder_input_ids=ids)      # UnboundLocalError
model.generate(pixel_values=pixel_values, max_new_tokens=3)  # RuntimeError: 10 vs 2

A GPT-2 decoder in the same harness works. No released openai/imagegpt-* checkpoint sets the
flag and nobody has reported this in fourteen months; from v4.54.1 to v5.0.0 case 4 raised
nothing and silently doubled the cache.

The contract is on record: on #35430 @Cyrilvallez wrote that in cross-attention "the mask in use
is encoder_attention_mask". #47946 then fixed case 2 for GPT-2. ImageGPT has no cross-attention
test, so nothing mechanical was catching it.

#45773 fixes the same crash for whisper, moonshine, pix2struct and audioflamingo3, not ImageGPT.

Is cross-attention on ImageGPT in scope? If not, a line saying so closes this.

I have all three fixes and three regression tests, red on main and green with them. Glad to
open the PR if a maintainer wants it.

[vasqu]

Argh I really want to avoid the encoder-decoder models that are created by the EncoderDecoder classes

But yes it seems right for now to fix. I would just warn: its not a high priority and that this feature might get deleted in the future as it's not used too much and creates a lot of maintenance on us

[YeonwooSung]

@vasqu one related model — not asking to expand #48694.

GPTBigCode has the same cross-attention family (is_updated is only bound for EncoderDecoderCache, while the model builds a DynamicCache; encoder_attention_mask never reaches attention_interface; the 2D unsqueeze is not assigned). There is no modular file.

Unlike ImageGPT, the block never forwards layer_past to cross-attn, so the default path never hits the UnboundLocalError. It dies earlier: c_attn is split as (head_dim, head_dim), so n_head != 1 is the old #24516 RuntimeError (closed, unfixed), and default multi_query=True (public StarCoder-style configs) raises NotImplementedError for cross-attn entirely.

Not opening a GPTBigCode issue or PR. Public checkpoints are decoder-only + MQA, and you already said this EncoderDecoder path is low priority / may be deleted. Is GPTBigCode in scope after #48694, or should it be left as a dead path?

[zucchini-nlp]

We need to do old models deprecation at some point so quite sure ImageGPT has low downloads and will be deprecated. Anyway, agreed with @vasqu that we want to deprecate add_cross_attention config attr in the future, it is a real maintenance burden