Streamlit中用bitsandbytes加载4bit HuggingFace模型时quant_state为None报错

问题背景 #

使用Streamlit的@st.cache_resource缓存加载HuggingFace的starchat-beta 4bit量化模型，加载过程正常，但调用model.generate()时触发断言错误quant_state is not None，错误源自bitsandbytes库的matmul_4bit函数。相同代码在Jupyter Notebook中可正常运行。

加载模型的代码如下：

@st.cache_resource
def load_model():
    """Initialize the tokenizer and the AI model."""

    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/starchat-beta")
    print("loaded tokenizer")
    model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/starchat-beta",load_in_4bit=True, device_map="cuda")
    print("loaded models")
    return tokenizer,model

可能的解决方案 #

1. 改用Streamlit Session State缓存模型 #

@st.cache_resource会对缓存对象进行序列化操作，而bitsandbytes的4bit量化模型的底层量化状态无法被正确序列化/反序列化，导致调用时丢失状态。改用Session State直接存储模型对象，避免序列化步骤：

def load_model():
    """Initialize the tokenizer and the AI model."""
    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/starchat-beta")
    print("loaded tokenizer")
    model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/starchat-beta", load_in_4bit=True, device_map="cuda")
    print("loaded models")
    return tokenizer, model

# 初始化时将模型存入Session State
if "tokenizer" not in st.session_state or "model" not in st.session_state:
    st.session_state.tokenizer, st.session_state.model = load_model()

# 调用时从Session State取出
tokenizer = st.session_state.tokenizer
model = st.session_state.model

2. 升级相关依赖版本 #

bitsandbytes、transformers或accelerate的旧版本可能存在与Streamlit缓存兼容的问题，升级到最新版：

pip install --upgrade bitsandbytes transformers accelerate streamlit

3. 在缓存函数内预热模型 #

在加载模型后，执行一次小的generate操作，确保量化状态被正确初始化并保留：

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/starchat-beta")
    print("loaded tokenizer")
    model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/starchat-beta", load_in_4bit=True, device_map="cuda")
    print("loaded models")
    
    # 预热模型，初始化量化状态
    dummy_input = tokenizer("dummy prompt", return_tensors="pt").to("cuda")
    with torch.no_grad():
        model.generate(**dummy_input, max_new_tokens=1)
    
    return tokenizer, model

原因分析 #

Streamlit的@st.cache_resource为了实现缓存，会对对象进行序列化和反序列化。而bitsandbytes的4bit量化模型依赖底层C++实现的quant_state状态，该状态无法被Python序列化机制正确保存，反序列化后状态丢失，导致model.generate()调用时触发断言错误。Jupyter环境中没有经过缓存的序列化步骤，模型状态完整，因此可以正常运行。

内容的提问来源于stack exchange，提问作者Abhilash Pal