PyTorch加载Mistral-7B：如何仅加载一次避免每次推理重复加载？

方案1：交互式循环复用加载好的模型 #

修改脚本，让模型只加载一次，之后进入循环接收输入快速推理，和Notebook的使用逻辑一致：

修改if __name__ == "__main__"部分的代码：

if __name__ == "__main__":
    # 原有模型加载逻辑不变
    torch.cuda.empty_cache()
    model_name = "mistralai/Mistral-7B-Instruct-v0.2"
    compute_dtype = getattr(torch, "float16")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16
    )
    base_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    
    # 启动循环，支持多次推理
    print("模型加载完成，输入prompt即可推理，输入exit退出")
    while True:
        eval_prompt = input(">>> ")
        if eval_prompt.lower() == "exit":
            break
        result = evaluate_model(base_model, tokenizer, eval_prompt)
        print("推理结果:\n", result, "\n")

运行python run_mycode.py后，模型仅加载一次，后续每次输入prompt都能直接执行推理。

方案2：封装成API服务（适合批量/远程调用） #

用FastAPI把模型封装成HTTP接口，服务启动时加载一次模型，后续所有请求复用已加载的模型：

1. 先安装依赖：pip install fastapi uvicorn
2. 修改脚本如下：

from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')
import pandas as pd
import torch
import csv
from datasets import Dataset, load_dataset
from huggingface_hub import notebook_login
from peft import LoraConfig, PeftModel
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
    logging,
)
from trl import SFTTrainer
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

# 全局加载模型和tokenizer，服务启动时仅执行一次
torch.cuda.empty_cache()
model_name = "mistralai/Mistral-7B-Instruct-v0.2"
compute_dtype = getattr(torch, "float16")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
base_model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# 定义请求格式
class PromptRequest(BaseModel):
    prompt: str

def evaluate_model(model, tokenizer, eval_prompt):
    model_input = tokenizer(eval_prompt, return_tensors="pt").to("cpu")
    model.eval()
    with torch.no_grad():
        return tokenizer.decode(model.generate(**model_input, max_new_tokens=256, pad_token_id=2)[0], skip_special_tokens=True)

# 推理接口
@app.post("/infer")
async def infer(request: PromptRequest):
    result = evaluate_model(base_model, tokenizer, request.prompt)
    return {"result": result}

3. 启动服务：uvicorn run_mycode:app --host 0.0.0.0 --port 8000
4. 调用接口：用curl或Postman发送POST请求到http://localhost:8000/infer，请求体为{"prompt": "你的推理内容"}，每次请求无需重新加载模型。

方案3：利用本地缓存减少下载耗时 #

如果加载慢是因为重复下载模型分片，确保Hugging Face缓存生效：

• 模型默认缓存到~/.cache/huggingface/hub目录，首次下载后后续加载直接读取本地文件
• 也可以手动指定缓存路径，在from_pretrained中添加cache_dir="./mistral_local_cache"，将模型缓存到项目本地目录，避免跨环境重复下载

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