Output dimension of AutoModelForCausalLM
The output dimension of models for causal LM is(batch_size, sequence_length, config.vocab_size). I don’t understand why this is the case. I would expect the outputs to be(batch_size, config.vocab_size), i.e. the logits for next token prediction.Indeed, thegeneratemethod always usesnext_token_logits = outputs.logits[:, -1, :], i.e. only the last tokens logits for next token prediction. However, the logits for all tokens of the sentence are computed, which very quickly blows the memory up for largesequence_lengthormax_new_tokens.For example, for the Bloom family, thevocab_sizeis 250880 (which is one of the largest vocab size). This means that even using bloom-560M (which is a small model), inference on a batch size of 64 with a prompt of say 500 tokens (or a small prompt with a largemax_new_tokens = 500) will take AT LEAST64 * 500 * 250880 * 2 / 1024**3 = 14.95GiB (the factor 2 is in case we loaded the model in float16, otherwise multiply this number by 2 again). And this is not even taking into account the memory needed for intermediate results in theforwardmethod.So we are basically blowing up the memory very quickly for large sequences or largemax_new_tokens, despite the fact that we only need the logits for the next token predictions and not the tokens we already have!Is there any way to modify theforwardmethod of models for causal LM to only output logits of dimension(batch_size, config.vocab_size)? Or is it a property of the underlying models and that would need retraining for scratch?
2024-07-02T06:52:14Z
Causal LM predicts the next tokengiven each token.  This meansoutputs.logits[:, 0, :]is computed given token 0,outputs.logits[:, 1, :]is computed given tokens 0 and 1, and so on…and finallyoutputs.logits[:, -1, :]is computed given the entire sequence (and that’s why it’s the next predicted token).  We do this by using masked attention - at each time step, the model can only compute outputs given past tokens.As far as I know, there isn’t a way to change this without training a model from scratch.  I’m not even sure if it would result in that much compute savings (aside from the size of the output layer).  Remember we still need to attend to all previous tokens when we calculate the next likely token, so it’s not like we can throw those calculations away.