Failed to load Llama2 with customized device_map
Author: polarispw
Pull request: no

### System Info

- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.24.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): 2.13.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

@ArthurZucker;  @younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Load Llama2 on a PC with 4060(8GB) and 32GB RAM
2. Passing device_map to AutoModel.from_pretrained()
3. Using "auto" and GPU memory + RAM is not enough, or just manually locating some layers on the disk
4. Receive the following info:
```
 "E:\anaconda\envs\transformers\lib\site-packages\transformers\modeling_utils.py", line 3633, in _load_pretrained_model
    offload_index = {
  File "E:\anaconda\envs\transformers\lib\site-packages\transformers\modeling_utils.py", line 3636, in <dictcomp>
    if param_device_map[p] == "disk"
  KeyError: 'lm_head.weight'
```

I have checked it and found there were no items in para_device_map, which is created by `expand_device_map()` in line 3621.
```
def expand_device_map(device_map, param_names):
    """
    Expand a device map to return the correspondance parameter name to device.
    """
    new_device_map = {}
    for module, device in device_map.items():
        new_device_map.update({p: device for p in param_names if p == module or p.startswith(f"{module}.")})
    return new_device_map
```
**The key is that the update func failed to extract p and device:** All the param_names have "model" as prefix and the lm_head was not included.
```
device_map:
{'embed_tokens': 'cpu', 'layers.0': 0, 'layers.1': 0, 'layers.2': 0, 'layers.3': 0, 'layers.4': 'cpu', 'layers.5': 'cpu', 'layers.6': 'cpu', 'layers.7': 'cpu', 'layers.8': 'cpu', 'layers.9': 'cpu', 'layers.10': 'cpu', 'layers.11': 'cpu', 'layers.12': 'cpu', 'layers.13': 'cpu', 'layers.14': 'cpu', 'layers.15': 'cpu', 'layers.16': 'cpu', 'layers.17': 'cpu', 'layers.18': 'cpu', 'layers.19': 'cpu', 'layers.20': 'cpu', 'layers.21': 'cpu', 'layers.22': 'cpu', 'layers.23': 'cpu', 'layers.24': 'cpu', 'layers.25': 'disk', 'layers.26': 'disk', 'layers.27': 'disk', 'layers.28': 'disk', 'layers.29': 'disk', 'layers.30': 'disk', 'layers.31': 'disk', 'norm': 'cpu'}

para_names:
['lm_head.weight', 'model.embed_tokens.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.rotary_emb.inv_freq', 'model.layers.0.self_attn.v_proj.weight', ...}
```

### Expected behavior

I wonder if there are some methods to avoid it

Comments:
Hi @polarispw, thanks for raising this issue! 

I think this bug - specifically the logic in `expand_device_map` - should be resolved with the merging of #27204

If you modify `expand_device_map` to:

```
def expand_device_map(device_map, param_names):
    """
    Expand a device map to return the correspondance parameter name to device.
    """
    new_device_map = {}
    for module, device in device_map.items():
        new_device_map.update(
            {p: device for p in param_names if p == module or p.startswith(f"{module}.") or module == ""}
        )
    return new_device_map
```

are you able to successfully run your code? 

Thanks for your answers, but I don't think it works because those module names in `device_map` are still mismatched with those in `param_names`.
By the way, I created my own `device_map` by using the following func:
```
def generate_device_map(model, model_config, start_layer=0, no_split_module_classes=["LlamaDecoderLayer"]):
    device_map = infer_auto_device_map(model, no_split_module_classes=no_split_module_classes)
    my_device_map = device_map
    my_device_map["embed_tokens"] = "cpu"
    my_device_map["norm"] = "cpu"
    start = start_layer  # from 0 to 7
    gpu_layers = range(start * 4, start * 4 + 4)
    for i in range(model_config.num_hidden_layers):
        if i in gpu_layers:
            my_device_map[f"layers.{i}"] = 0
        elif my_device_map[f"layers.{i}"] == 0:
            my_device_map[f"layers.{i}"] = "cpu"
    check_device_map(model, my_device_map)
    return my_device_map
```
It means the names in `device_map` comes from the `accelerate.utils.infer_auto_device_map`. I can avoid this issue by:
```
def expand_device_map(device_map, param_names):
    """
    Expand a device map to return the correspondance parameter name to device.
    """
    new_device_map = {}
    for module, device in device_map.items():
        new_device_map.update({p: device for p in param_names if p == module or p.startswith(f"{module}.") or \
                                                  p.startswith(f"model.{module}.")}) # adding here to adopt prefix
    return new_device_map


def get_disk_only_shard_files(device_map, sharded_metadata):
    """
    Returns the list of shard files containing only weights offloaded to disk.
    """
    files_content = collections.defaultdict(list)
    for weight_name, filename in sharded_metadata["weight_map"].items():
        while len(weight_name) > 0 and weight_name not in device_map:
            weight_name = weight_name.split("model.")[-1]    # adding here to remove prefix
            weight_name = ".".join(weight_name.split(".")[:-1])
        files_content[filename].append(device_map[weight_name])

    return [fname for fname, devices in files_content.items() if set(devices) == {"disk"}]
```
But it will definitely affect other models' loading. So maybe you can connect with the Accelerate team and ask for aligned names  after `infer_auto_device_map`?

@polarispw Thanks for such a detailed explanation - I was hoping there might be an easy fix, but alas it's never that easy! 

As it relates to accelerate, passing this to @muellerzr and @pacman100 

cc @SunMarc 

Hi @polarispw, this happens because you are using `AutoModel` and not `AutoModelForCausalLM`. The weights that are saved correspond to `AutoModelForCausalLM` model, hence the `lm_head` and the prefix `model.` for the parameters. I don't know if this is what you had in mind but you won't be able to fine-tune your model when offloaded weights to disk. If your use case is just for inference, I suggest you to switch to `AutoModelForCausalLM`. 

Right now, there is indeed an issue with offloading params to disk using `safetensors` + loading weights from a derived model into the base model. Does it makes sense to fix this issue considering that we can only do inference when offloading to disk ? @amyeroberts If yes, I know how to fix this and I can submit a PR. 

@SunMarc that is the point and I am just loading the model to check some param matrices inside. I focus on the base model.

I'd be happy if the auto conversion from derived to base model could be supported, as it might be more in line with the user's intuitive thoughts. It seems that when mapping only includes GPU and CPU, the logic does not go this way and causes an error

@SunMarc Yes - I think it makes sense to fix this, in particular for loading weights into the base model.