LlavaForConditionalGeneration._merge_input_ids_with_image_features is incorrect

[fpgaminer]

System Info

commit: 38f9f10

Who can help?

@amyeroberts @qubvel

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

The function _merge_input_ids_with_image_features of class LlavaForConditionalGeneration gives incorrect results if it has to pad the result (which arises when the batch has different amounts of image tokens in each row).

transformers/src/transformers/models/llava/modeling_llava.py

Line 297 in 38f9f10

def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels):

I created the following reproduction example to show the failure mode simply. Note that it uses a copy of the function to add print statements for debugging, but the code is otherwise an identical copy at HEAD.

def _merge_input_ids_with_image_features(self, image_features, inputs_embeds, input_ids, attention_mask, labels):
	num_images, num_image_patches, embed_dim = image_features.shape
	batch_size, sequence_length = input_ids.shape
	left_padding = not torch.sum(input_ids[:, -1] == torch.tensor(self.pad_token_id))
	# 1. Create a mask to know where special image tokens are
	special_image_token_mask = input_ids == self.config.image_token_index
	num_special_image_tokens = torch.sum(special_image_token_mask, dim=-1)
	# Compute the maximum embed dimension
	max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length
	batch_indices, non_image_indices = torch.where(input_ids != self.config.image_token_index)

	# 2. Compute the positions where text should be written
	# Calculate new positions for text tokens in merged image-text sequence.
	# `special_image_token_mask` identifies image tokens. Each image token will be replaced by `nb_text_tokens_per_images - 1` text tokens.
	# `torch.cumsum` computes how each image token shifts subsequent text token positions.
	# - 1 to adjust for zero-based indexing, as `cumsum` inherently increases indices by one.
	new_token_positions = torch.cumsum((special_image_token_mask * (num_image_patches - 1) + 1), -1) - 1
	nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1]
	print(f"nb_image_pad: {nb_image_pad}")
	if left_padding:
		new_token_positions += nb_image_pad[:, None]  # offset for left padding
	text_to_overwrite = new_token_positions[batch_indices, non_image_indices]

	# 3. Create the full embedding, already padded to the maximum position
	final_embedding = torch.zeros(
		batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device
	)
	final_attention_mask = torch.zeros(
		batch_size, max_embed_dim, dtype=attention_mask.dtype, device=inputs_embeds.device
	)
	if labels is not None:
		final_labels = torch.full(
			(batch_size, max_embed_dim), self.config.ignore_index, dtype=input_ids.dtype, device=input_ids.device
		)
	# In case the Vision model or the Language model has been offloaded to CPU, we need to manually
	# set the corresponding tensors into their correct target device.
	target_device = inputs_embeds.device
	batch_indices, non_image_indices, text_to_overwrite = (
		batch_indices.to(target_device),
		non_image_indices.to(target_device),
		text_to_overwrite.to(target_device),
	)
	attention_mask = attention_mask.to(target_device)

	print(f"max_embed_dim: {max_embed_dim}")
	print(f"batch_indices: {batch_indices}")
	print(f"non_image_indices: {non_image_indices}")
	print(f"text_to_overwrite: {text_to_overwrite}")

	# 4. Fill the embeddings based on the mask. If we have ["hey" "<image>", "how", "are"]
	# we need to index copy on [0, 577, 578, 579] for the text and [1:576] for the image features
	final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices]
	final_attention_mask[batch_indices, text_to_overwrite] = attention_mask[batch_indices, non_image_indices]
	if labels is not None:
		final_labels[batch_indices, text_to_overwrite] = labels[batch_indices, non_image_indices]

	# 5. Fill the embeddings corresponding to the images. Anything that is not `text_positions` needs filling (#29835)
	image_to_overwrite = torch.full(
		(batch_size, max_embed_dim), True, dtype=torch.bool, device=inputs_embeds.device
	)
	image_to_overwrite[batch_indices, text_to_overwrite] = False
	print(f"image_to_overwrite cumsum: {image_to_overwrite.cumsum(-1) - 1}")
	image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device)
	print(f"image_to_overwrite: {image_to_overwrite}")

	if image_to_overwrite.sum() != image_features.shape[:-1].numel():
		raise ValueError(
			f"The input provided to the model are wrong. The number of image tokens is {torch.sum(special_image_token_mask)} while"
			f" the number of image given to the model is {num_images}. This prevents correct indexing and breaks batch generation."
		)

	final_embedding[image_to_overwrite] = image_features.contiguous().reshape(-1, embed_dim).to(target_device)
	final_attention_mask |= image_to_overwrite
	position_ids = (final_attention_mask.cumsum(-1) - 1).masked_fill_((final_attention_mask == 0), 1)

	# 6. Mask out the embedding at padding positions, as we later use the past_key_value value to determine the non-attended tokens.
	batch_indices, pad_indices = torch.where(input_ids == self.pad_token_id)
	indices_to_mask = new_token_positions[batch_indices, pad_indices]

	final_embedding[batch_indices, indices_to_mask] = 0

	if labels is None:
		final_labels = None

	return final_embedding, final_attention_mask, final_labels, position_ids


tokenizer = AutoTokenizer.from_pretrained("llava-hf/llava-1.5-7b-hf", use_fast=True)

model = LlavaForConditionalGeneration.from_pretrained("llava-hf/llava-1.5-7b-hf", device_map=0, torch_dtype=torch.float16)
model.eval()
assert isinstance(model, LlavaForConditionalGeneration)


model._merge_input_ids_with_image_features = types.MethodType(_merge_input_ids_with_image_features, model)


prompt1 = "System prompt<image>Caption 1<pad><pad><pad><pad><pad>"
prompt2 = "System prompt<image>Caption 2<image>Caption 3"           # 12 tokens

print("`System Prompt` tokens: ", tokenizer.encode("System prompt", add_special_tokens=False, truncation=False))
print("`Caption 1` tokens: ", tokenizer.encode("Caption 1", add_special_tokens=False, truncation=False))

# Tokenize the prompt
input_ids = [
	tokenizer.encode(prompt1, return_tensors="pt", add_special_tokens=False, truncation=False).squeeze(0),
	tokenizer.encode(prompt2, return_tensors="pt", add_special_tokens=False, truncation=False).squeeze(0),
]
input_ids = torch.stack(input_ids)
assert isinstance(input_ids, torch.Tensor)

print(f"Input IDs: {input_ids.shape}")
print(f"Input IDs: {input_ids}")


gc.collect()
torch.cuda.empty_cache()

with torch.no_grad():
	input_embeddings = model.get_input_embeddings()(input_ids.to('cuda'))

	embedded_images = torch.randn(3, 5, 4096, device='cuda', dtype=torch.float16)
	attention_mask = torch.ones(2, 12, device='cuda', dtype=torch.bool)
	result = model._merge_input_ids_with_image_features(embedded_images, input_embeddings, input_ids, attention_mask, None)[0]
	print(result.shape)

	print("`System Prompt` diff: ", (result[0, :2] - input_embeddings[0, :2]).abs().max())
	print("`image` diff: ", (result[0, 2:7] - embedded_images[0]).abs().max())
	print("`Caption 1` diff: ", (result[0, 7:11] - input_embeddings[0, 3:7]).abs().max())

	print("`System Prompt` diff: ", (result[1, :2] - input_embeddings[1, :2]).abs().max())
	print("`image` diff: ", (result[1, 2:7] - embedded_images[1]).abs().max())
	print("`Caption 2` diff: ", (result[1, 7:11] - input_embeddings[1, 3:7]).abs().max())
	print("`image` diff: ", (result[1, 11:16] - embedded_images[2]).abs().max())
	print("`Caption 3` diff: ", (result[1, 16:20] - input_embeddings[1, 8:12]).abs().max())

Running this gives:

`System Prompt` tokens:  [2184, 9508]
`Caption 1` tokens:  [9243, 683, 29871, 29896]
Input IDs: torch.Size([2, 12])
Input IDs: tensor([[ 2184,  9508, 32000,  9243,   683, 29871, 29896, 32001, 32001, 32001,
         32001, 32001],
        [ 2184,  9508, 32000,  9243,   683, 29871, 29906, 32000,  9243,   683,
         29871, 29941]])
nb_image_pad: tensor([4, 0])
max_embed_dim: 20
batch_indices: tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       device='cuda:0')
non_image_indices: tensor([ 0,  1,  3,  4,  5,  6,  7,  8,  9, 10, 11,  0,  1,  3,  4,  5,  6,  8,
         9, 10, 11], device='cuda:0')
text_to_overwrite: tensor([ 0,  1,  7,  8,  9, 10, 11, 12, 13, 14, 15,  0,  1,  7,  8,  9, 10, 16,
        17, 18, 19], device='cuda:0')
image_to_overwrite cumsum: tensor([[-1, -1,  0,  1,  2,  3,  4,  4,  4,  4,  4,  4,  4,  4,  4,  4,  5,  6,
          7,  8],
        [-1, -1,  0,  1,  2,  3,  4,  4,  4,  4,  4,  5,  6,  7,  8,  9,  9,  9,
          9,  9]], device='cuda:0')
image_to_overwrite: tensor([[False, False, False, False, False, False,  True, False, False, False,
         False, False, False, False, False, False,  True,  True,  True,  True],
        [False, False,  True,  True,  True,  True,  True, False, False, False,
         False,  True,  True,  True,  True,  True, False, False, False, False]],
       device='cuda:0')
torch.Size([2, 20, 4096])
`System Prompt` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`image` diff:  tensor(5.0898, device='cuda:0', dtype=torch.float16)
`Caption 1` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`System Prompt` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`image` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`Caption 2` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`image` diff:  tensor(0., device='cuda:0', dtype=torch.float16)
`Caption 3` diff:  tensor(0., device='cuda:0', dtype=torch.float16)

The last set of prints shows the difference between what the function outputs and what is expected, by comparing the different pieces of the result to the original inputs. As can be seen here, the result is almost correct but has misplaced the first image.

The reason is this line:

transformers/src/transformers/models/llava/modeling_llava.py

Line 352 in 38f9f10

image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device)

It looks like this line is trying to account for padding in the outputs, but assumes the outputs are going to be left padded, whereas the rest of the function dynamically switches between left and right padding depending on the left_padding variable. If the outputs are being right padded, the comparison on this line is incorrect. That can be seen in the debug output above where image_to_overwrite shows the image being scattered around the tensor.

Importantly, this incorrect behavior can occur for any padding, because left_padding is "automatically" detected from the input. Even if the model is setup for left padding, if a particular batch doesn't have any padding going into this function, the function will assume right padding and thus mangle the output.

One solution is to add an if statement on that line based on left_padding and use a different condition to handle right padding. I think this code (for the right padding case) would work, but I haven't tested it extensively:

	idxs = torch.arange(max_embed_dim, device=image_to_overwrite.device).expand(batch_size, -1)
	image_to_overwrite &= idxs < (max_embed_dim - nb_image_pad[:, None]).to(target_device)

Side note:

I think the current implementation for handling the left padding case could also be replaced with arange based logic. To me it would be more clear, as it indicates the intention of that line explicitly (set all padded indexes to False). Whereas I found the use of image_to_overwrite.cumsum(-1) opaque and difficult to decipher what it was trying to calculate.

Expected behavior

See above

[ArthurZucker]

Hey! Thanks for the report. I think this piece of code has been changed quite a lot since initial porting so @zucchini-nlp will have a better idea of what is going on!

Note that this is the legacy branch, which should not be triggered anymore and we overall deprecate!

[fpgaminer]

Note that this is the legacy branch, which should not be triggered anymore and we overall deprecate!

Ah, I hadn't noticed! Good point, and it looks like this code should disappear soon based on the comment TODO: @raushan retain only the new behavior after v4.47. So maybe this issue isn't worth fixing?

[zucchini-nlp]

Hey! Yes, it is worth fixing as I still didn't update our files on the hub and we fall to the legacy logic. I am planning to gradually update files and in the meanwhile we need to be sure the legacy branch is working correctly

Interesting that the padding="right" is not correct, as it should have been covered when adding the model. I'll get back to this on Monday/Tuesday, thanks

[ArthurZucker]

Yeah I remember that we had some tests for this!

[zucchini-nlp]

Okey, verified that indeed the merging is not correct when we do right padding, and it is really weird noone noticed that earlier. We added handling for padding side only in llava-next because the number of image tokens can be different and needs unpadding. And the other models it was assumed that the following line handles it correctly

if left_padding:
	new_token_positions += nb_image_pad[:, None]  # offset for left padding

Also related to #33662 which I didn't had time to look due to absence of clear reproducer. @fpgaminer would you like to open a PR and fix this in all llava models? We might also need a test with dummy inputs and tiny model, to make sure it doesn't happen in the future :)

[MonolithFoundation]

Got some questions about this new branch:

elif image_features is not None:
            n_image_tokens = (input_ids == self.config.image_token_index).sum().item()
            n_image_features = image_features.shape[0] * image_features.shape[1]

            if n_image_tokens != n_image_features:
                raise ValueError(
                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
                )
            special_image_mask = (
                (input_ids == self.config.image_token_index)
                .unsqueeze(-1)
                .expand_as(inputs_embeds)
                .to(inputs_embeds.device)
            )
            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)

that means image token had to be expanded to image_features_tokens before feeding into forward?

For some dynamic iamge token models, this could probably only knows when it output from vision tower.

If such, wouldn't it's not convenient to expand it first?

[ArthurZucker]

For some dynamic iamge token models, this could probably only knows when it output from vision tower.
I am not sure we have seen such models yet no?

[MonolithFoundation]

Currently, most models with fixed resolution, such as Siglip with a 448 input, can have their tokens pre-calculated by dividing into 14 patches. However, for some models that utilize WindowAtten or dynamic position embedding, users can employ dynamic input. As a result, when using a dynamic input resolution, the output length of tokens can be different.
For such models can reference AIMv2-native resolution version. I think original way is more flexible, you don't need pre pad your <image> tokens in preprocess or somewhere.

[zucchini-nlp]

@MonolithFoundation We already have models with dynamic image lengths like Qwen2-VL and it works by checking the image size after preprocessing is done, on a per sample basis. I guess for models like AIMv2 same would work

The basic intuition behind this modifications was to make VLMs more flexible and reduce bugs coming from incorrect expansion of the attention mask or kv cache during generation, to account for image+text concatenation