Attention中存在的问题-错位对齐

import os import torch import sys import types import importlib.util from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig # --------------------- 配置项（根据硬件调整） --------------------- MODEL_PATH = r"E:\LocalAI\AIModel\DeepSeek\DeepSeekV2LiteChat" # 模型路径 USE_4BIT = True # 启用4-bit量化（显存需求降至约10GB） USE_8BIT = False # 8-bit量化（显存需求约20GB） MAX_NEW_TOKENS = 200 # 生成文本最大长度（降低内存压力） DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # --------------------- 量化配置（4-bit/8-bit） --------------------- quantization_config = None if USE_4BIT or USE_8BIT: quantization_config = BitsAndBytesConfig( load_in_4bit=USE_4BIT, load_in_8bit=USE_8BIT, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", # 4-bit量化类型 bnb_8bit_quant_type="nf8" if USE_8BIT else None ) # --------------------- 加载模型（跳过Flash Attention） --------------------- if sys.platform == "win32": if "flash_attn" not in sys.modules: # 创建一个 dummy 的 module spec dummy_spec = importlib.util.spec_from_loader("flash_attn", loader=None) dummy_module = types.ModuleType("flash_attn") dummy_module.__spec__ = dummy_spec # 根据模型代码中可能调用的函数名称添加 dummy 实现（仅返回 None） dummy_module.flash_attn_unpadded = lambda *args, **kwargs: None dummy_module.flash_attn_varlen_qkvpacked_func = lambda *args, **kwargs: None sys.modules["flash_attn"] = dummy_module # --------------------- 加载分词器 --------------------- tokenizer = AutoTokenizer.from_pretrained( MODEL_PATH, trust_remote_code=True, padding_side="left", # 左对齐避免生成错位 use_fast=True # 强制使用快速分词器 ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 确保有pad token # --------------------- 加载模型 --------------------- model = AutoModelForCausalLM.from_pretrained( MODEL_PATH, quantization_config=quantization_config, trust_remote_code=True ) # --------------------- 显存监控函数 --------------------- def print_gpu_memory(): if DEVICE == "cuda": allocated = torch.cuda.memory_allocated() / 1024**3 reserved = torch.cuda.memory_reserved() / 1024**3 print(f"[显存] 已分配: {allocated:.2f} GB | 保留: {reserved:.2f} GB") else: print("当前使用CPU运行") # --------------------- 轻量生成函数 --------------------- def light_generate(prompt: str): print("\n" + "="*30 + " 生成开始 " + "="*30) # 精简输入处理 inputs = tokenizer( f"<|startoftext|>User: {prompt}\nAssistant:", return_tensors="pt", max_length=512, # 限制输入长度 truncation=True ).to(DEVICE) # 低资源生成参数 outputs = model.generate( inputs.input_ids, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=0.3, # 降低随机性 top_p=0.9, repetition_penalty=1.1, # 减少重复 pad_token_id=tokenizer.eos_token_id ) # 清理显存 torch.cuda.empty_cache() # 解码结果 generated = tokenizer.decode( outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True ) return generated.strip() # --------------------- 测试运行 --------------------- if __name__ == "__main__": print_gpu_memory() # 初始显存状态 # 测试样例（短文本验证基础功能） test_prompts = [ "请输出我问过你的问题" ] for prompt in test_prompts: print(f"\n[输入] {prompt}") try: reply = light_generate(prompt) print(f"[输出] {reply}") except RuntimeError as e: if "CUDA out of memory" in str(e): print("[错误] 显存不足！请尝试减小 MAX_NEW_TOKENS 或启用4-bit量化") else: print(f"[错误] {str(e)}") print_gpu_memory() # 每次生成后显存状态 这是我的代码