NameError: free variable 'state_dict' referenced before assignment in enclosing scope
Author: AllentDan
Pull request: no
Labels: name: Accelerate

### System Info

- `transformers` version: 4.42.1
- Platform: Linux-3.10.0-957.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.10.10
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.2
- Accelerate version: 0.31.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

shard = {tensor: state_dict[tensor] for tensor in tensors}
NameError: free variable 'state_dict' referenced before assignment in enclosing scope
```python
model_path = 'liuhaotian/llava-v1.5-7b'
dst_path = 'llava-v1.5-7b-test'

import torch
from transformers import AutoModelForCausalLM

from llava.model.language_model.llava_llama import LlavaConfig
config = LlavaConfig.from_pretrained(model_path, trust_remote_code=True)

model = AutoModelForCausalLM.from_config(config,
                                            trust_remote_code=True)
from accelerate import load_checkpoint_and_dispatch
load_checkpoint_and_dispatch(
    model=model,
    checkpoint=model_path,
    device_map={'': 'cpu'},
    no_split_module_classes=['CLIPEncoderLayer'],
    dtype=torch.half)

model.save_pretrained(dst_path,
                      max_shard_size='2GB',
                      safe_serialization=True)

```
Seems the latest transformers cannot work with accelerate properly.

### Expected behavior

It worked fine with previous versions of transformers.

Comments:
Hi @AllentDan, could you share the values in `LlavaConfig`?

cc @muellerzr @SunMarc  

https://huggingface.co/liuhaotian/llava-v1.5-7b/blob/main/config.json
```
{
  "_name_or_path": "llava-v1.5-7b",
  "architectures": [
    "LlavaLlamaForCausalLM"
  ],
  "bos_token_id": 1,
  "eos_token_id": 2,
  "freeze_mm_mlp_adapter": false,
  "freeze_mm_vision_resampler": false,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "image_aspect_ratio": "pad",
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_length": 4096,
  "max_position_embeddings": 4096,
  "mm_hidden_size": 1024,
  "mm_projector_type": "mlp2x_gelu",
  "mm_resampler_type": null,
  "mm_use_im_patch_token": false,
  "mm_use_im_start_end": false,
  "mm_vision_select_feature": "patch",
  "mm_vision_select_layer": -2,
  "mm_vision_tower": "openai/clip-vit-large-patch14-336",
  "model_type": "llava",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 32,
  "pad_token_id": 0,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "tie_word_embeddings": false,
  "torch_dtype": "float16",
  "transformers_version": "4.31.0",
  "tune_mm_mlp_adapter": false,
  "tune_mm_vision_resampler": false,
  "unfreeze_mm_vision_tower": false,
  "use_cache": true,
  "use_mm_proj": true,
  "vocab_size": 32000
}
```
And in python, the value is:
```python
LlavaConfig {
  "_name_or_path": "llava-v1.5-7b",
  "architectures": [
    "LlavaLlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "freeze_mm_mlp_adapter": false,
  "freeze_mm_vision_resampler": false,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "image_aspect_ratio": "pad",
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_length": 4096,
  "max_position_embeddings": 4096,
  "mlp_bias": false,
  "mm_hidden_size": 1024,
  "mm_projector_type": "mlp2x_gelu",
  "mm_resampler_type": null,
  "mm_use_im_patch_token": false,
  "mm_use_im_start_end": false,
  "mm_vision_select_feature": "patch",
  "mm_vision_select_layer": -2,
  "mm_vision_tower": "openai/clip-vit-large-patch14-336",
  "model_type": "llava_llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 32,
  "pad_token_id": 0,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": null,
  "rope_theta": 10000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "float16",
  "transformers_version": "4.42.1",
  "tune_mm_mlp_adapter": false,
  "tune_mm_vision_resampler": false,
  "unfreeze_mm_vision_tower": false,
  "use_cache": true,
  "use_mm_proj": true,
  "vocab_size": 32000
}
```

The bug comes with `load_checkpoint_and_dispatch`. If I remove it, `save_pretrained` will work fine.

this issue can be reproduced on latest transformers version 4.42.4

@zhulinJulia24 could you try running with an install from source: `pip install git+https://github.com/huggingface/transformers`? The fixed has been merged into main, but isn't part of stable release on pypi

> pip install git+https://github.com/huggingface/transformers

It works by using main code. Thanks!