NameError: free variable 'state_dict' referenced before assignment in enclosing scope

[AllentDan]

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?:
• Using GPU in script?:
• GPU type: NVIDIA A100-SXM4-80GB

Who can help?

No response

Information

• The official example scripts
• 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

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.

[amyeroberts]

Hi @AllentDan, could you share the values in LlavaConfig?

cc @muellerzr @SunMarc

[AllentDan]

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:

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
}

[AllentDan]

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

[zhulinJulia24]

this issue can be reproduced on latest transformers version 4.42.4

[amyeroberts]

@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

[zhulinJulia24]

pip install git+https://github.com/huggingface/transformers

It works by using main code. Thanks!