LMHead is processing redundant tokens in prefill

[null-pointer-access]

While using GPT2LMHeadModel.generate() and compare its performance with vLLM, I noticed a significant inefficiency in the forward() implementation of many huggingface models. For example, in the GPT2LMHeadModel.forward, self.lm_head is applied to all token hidden states, even when called from the generate() method, where only the logits of the last token are needed for next-token prediction. This computes logits over the entire sequence and can introduce significant overhead.

# src/transformers/models/gpt2/modeling_gpt2.py, line 1233
lm_logits = self.lm_head(hidden_states)

Suggested Fix: add a conditional branch in forward() to slice the hidden states before computing logits if it’s a generation step.

[zucchini-nlp]

Hey @null-pointer-access !

We have a similar way to drop unnecessary tokens in generation in most recent models (e.g. Llama) though the change wasn't propagated to all language models in the library. Would you like to open a PR to add smth similar to below code snippet in GPT models?

transformers/src/transformers/models/llama/modeling_llama.py

Lines 565 to 567 in 2166b6b

	# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
	slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
	logits = self.lm_head(hidden_states[:, slice_indices, :])

[null-pointer-access]

Hey @null-pointer-access !

We have a similar way to drop unnecessary tokens in generation in most recent models (e.g. Llama) though the change wasn't propagated to all language models in the library. Would you like to open a PR to add smth similar to below code snippet in GPT models?

transformers/src/transformers/models/llama/modeling_llama.py

Lines 565 to 567 in 2166b6b

Only compute necessary logits, and do not upcast them to float if we are not computing the loss

slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])

Sure! I'll submit a PR soon

[null-pointer-access]

Fixed in #39016