add kv_cache to LLM

[vchiley]

This pr include past_key_values (ie kv_cache) in the LLM so that inference can be accelerated.
We also become explicit about how we apply padding_mask for querys and keys.

Shoutout: @dakinggg for working through some of this with me.

cc @dskhudia @alextrott16 @samhavens for after training / inference

[honglu2875]

@vchiley Cached kv values shift the positions as well. Maybe you want to shift the position embeddings in the following?

examples/examples/llm/src/models/mosaic_gpt.py

Line 148 in e98e675

pos = torch.arange(0, S, dtype=torch.long,

Compare with this in HF
https://github.com/huggingface/transformers/blob/60d51ef5123d949fd8c59cd4d3254e711541d278/src/transformers/models/gpt2/modeling_gpt2.py#L801

In our fork of mosaic models, we have the kv cache and the relevant part looks like the following:

        if past_key_values is None:
            past_key_values = [None] * self.cfg.n_layers
            past_position = 0
        else:
            assert len(past_key_values) == self.cfg.n_layers
            # get the key tensor whose spec should be (batch, seq, n_head, head_dim), and
            # collect the `seq`, so that we shift the position embedding later.
            past_position = past_key_values[0][0].size(1)

        tok_emb = self.transformer.wte(input_ids)  # type: ignore
        if self.alibi:
            x = tok_emb
        else:
            if S + past_position > self.cfg.max_seq_len:
                raise ValueError(
                    f'Cannot forward input with past sequence length {past_position} and current sequence length '
                    f'{S + 1}, this model only supports total sequence length <= {self.cfg.max_seq_len}.'
                )
            pos = torch.arange(past_position, S + past_position, dtype=torch.long,
                               device=input_ids.device).unsqueeze(0)
            pos_emb = self.transformer.wpe(pos)  # type: ignore
            x = tok_emb + pos_emb

[dakinggg]

LGTM, can you train a model and make sure nothing is broken?

[dskhudia]

Could you explain the reason for separating out query_padding_mask?

[vchiley]

@dskhudia this formulates a generic attn fn
the queries are potentially not the same as the keys/value and will need their own padding_mask.
This is useful for left padded inputs since past tokens influence future tokens, we NEED to mask them out.

[vchiley]

Note: we should have a conversation about if all raise Error should be changed to assert (see here)