File size: 2,142 Bytes
0695d0b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import torch
from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
from PIL import Image
# ===================== 1. 基础配置 =====================
# 你的本地路径
local_model_path = "/home/nashen/deepseek-ocr/DeepSeek-OCR-master/DeepSeek-OCR-vllm/model/"
device = "cuda" if torch.cuda.is_available() else "cpu"
# ===================== 2. 正确的分层量化配置 (BitsAndBytesConfig) =====================
# 核心修改：不再使用字典，而是创建一个 Config 对象
# 逻辑：默认全部 4bit 量化，然后通过 skip_modules 保护视觉层
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4", # 推荐 nf4
bnb_4bit_use_double_quant=True, # 双重量化省显存
bnb_4bit_compute_dtype=torch.bfloat16,
# 【重点】这里实现你的“分层”逻辑
# 强制跳过视觉层(vision_model)和输出层(lm_head)，保持 bf16 原生精度
llm_int8_skip_modules=[
"vision_model", # 视觉编码器
"lm_head", # 输出层
"embed_tokens", # 词嵌入
"output" # 部分架构的输出层
]
)
# ===================== 3. 加载模型 =====================
print(f"🚀 正在从本地加载模型: {local_model_path}")
try:
tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(
local_model_path,
trust_remote_code=True,
# 核心修改：这里使用 quantization_config 参数，而不是 **kwargs 解包
quantization_config=quantization_config,
device_map="auto"
)
# 开启 KV Cache 优化
model.config.use_cache = True
model.eval()
print("\n✅ DeepSeek-OCR 本地模型加载成功！")
print("✅ 混合精度策略已应用：")
print(" - 视觉层/输出层：BF16 (原画质，不量化)")
print(" - 语言层：NF4 (显存压缩)")
# 测试一下显存占用
print(f"✅ 模型加载完毕。你可以尝试运行推理了。")
except Exception as e:
print(f"\n❌ 发生错误: {e}")
|