LlamaForCausalLM
Author: Shikamaru5
Pull request: no

I'm trying to implement the bitsandbytes library into my script to run a Llama2 model and in this instance it requires that I use a custom device map. I have about 458 layers and most of them need to be sent to the GPU. However, I realized upon the loading the model it seems there's this huge bottleneck of leading every single one of those layers individually to the GPU.
My load time might increase if I could load the layers if not all at the same time, maybe in multiple batches to the GPU. I'm however, uncertain of how to do this as multiple different things that I've tried have not succeeded. I would greatly appreciate if anyone knows how or if I can solve this issue, and thank you for taking the time to read this.

Comments:
Maybe cc @SunMarc :) 

Hi @Shikamaru5, if i understood that correctly, you want to be able to load the parameters contained in the state_dict in batches to the GPU. If this is indeed the case, i don't think we have an easy fix. Here is the [line](see code [here](https://github.com/huggingface/transformers/blob/946bac798caefada3f5f1c9fecdcfd587ed24ac7/src/transformers/modeling_utils.py#L680)) that put the parameters to the GPU in case you want to try something. bitsandbytes takes a lot of time to load because the quantization happens when we move the parameters from the CPU to the GPU. 

Ok I understand, I was just thinking it might be something of value, although I'm uncertain as to how I could do it. I was watching a video on I think it was bettertransformer or Flash attention or something and they were talking about how they do batches of the inputs and perhaps parameters to send to the gpu at once in order to speed up inference and training. It had occurred to me that these sorts of layers:

     device_map = {
         "transformer.word_embeddings": 0,
         "transformer.word_embeddings_layernorm": 0,
         "model.embed_tokens.weight": 0,
         "model.layers.0.self_attn.q_proj.weight": 0,
         "model.layers.0.self_attn.k_proj.weight": 0,
         "model.layers.0.self_attn.v_proj.weight": 0,
         "model.layers.0.self_attn.o_proj.weight": 0,
         "model.layers.0.mlp.gate_proj.weight": 0,
         "model.layers.0.mlp.up_proj.weight": 0,
         "model.layers.0.mlp.down_proj.weight": 0,
         "model.layers.0.input_layernorm.weight": 0,
         "model.layers.0.post_attention_layernorm.weight": 0,
        #rest of layers,
     }

might be able to be grouped together more efficiently, instead of one by one calling : 0, perhaps it could be something like:

     device_map = {
              ("transformer.word_embeddings",
              "transformer.word_embeddings_layernorm",
              "model.embed_tokens.weight",
              "model.layers.0.self_attn.q_proj.weight",
              "model.layers.0.self_attn.k_proj.weight",
              "model.layers.0.self_attn.v_proj.weight",
              "model.layers.0.self_attn.o_proj.weight",
              "model.layers.0.mlp.gate_proj.weight",
              "model.layers.0.mlp.up_proj.weight",
              "model.layers.0.mlp.down_proj.weight",
              "model.layers.0.input_layernorm.weight",
              "model.layers.0.post_attention_layernorm.weight").to("cuda"),
             #rest of the layers,
          }

Although I did try that and it told me that it wasn't set up to handle this. Anyway thanks for the suggestion.

Hi @Shikamaru5 , flash attention should work with device_map out the box. In flash attention, the layers are already on the GPU and we just handle better the communication of the parameters inside the gpu, so that we spend less time storing, reading and writing the keys, queries and values tensors and more time on the computation. Check out this [doc](https://huggingface.co/docs/text-generation-inference/conceptual/flash_attention) that explain better how it works 

Right that would've been nice, but because I'm using 4bit quantization with bitsandbytes and my model needs to offload between gpu and cpu it needs a custom device_map.