基于Transformer框架实现微调后Qwen/DeepSeek模型的非流式批量推理

下面我将介绍如何使用Hugging Face的Transformer框架对微调后的Qwen或DeepSeek模型进行非流式批量推理。

一、准备工作

首先确保已安装必要的库：

复制代码

pip install transformers torch

二、批量推理实现代码

复制代码

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from typing import List

class BatchInference:
    def __init__(self, model_path: str, device: str = "cuda" if torch.cuda.is_available() else "cpu"):
        """
        初始化模型和tokenizer
        :param model_path: 微调后的模型路径(Hugging Face模型ID或本地路径)
        :param device: 推理设备(cpu/cuda)
        """
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=torch.float16 if "cuda" in device else torch.float32,
            device_map="auto"
        )
        self.model.eval()
        
        # 设置pad_token_id为eos_token_id如果不存在pad_token
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
    
    def batch_predict(self, prompts: List[str], max_new_tokens: int = 512, batch_size: int = 4, **generate_kwargs) -> List[str]:
        """
        批量推理方法
        :param prompts: 输入提示列表
        :param max_new_tokens: 生成的最大token数
        :param batch_size: 批量大小
        :param generate_kwargs: 额外的生成参数
        :return: 生成结果列表
        """
        all_results = []
        
        # 分批次处理
        for i in range(0, len(prompts), batch_size):
            batch_prompts = prompts[i:i + batch_size]
            
            # 编码输入
            inputs = self.tokenizer(
                batch_prompts,
                padding=True,
                truncation=True,
                return_tensors="pt",
                max_length=1024  # 可根据需要调整
            ).to(self.device)
            
            # 生成输出
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    pad_token_id=self.tokenizer.pad_token_id,
                    **generate_kwargs
                )
            
            # 解码输出并移除输入部分
            batch_results = []
            for j in range(len(outputs)):
                output = outputs[j]
                input_length = inputs["input_ids"][j].shape[0]
                generated = output[input_length:]
                batch_results.append(self.tokenizer.decode(generated, skip_special_tokens=True))
            
            all_results.extend(batch_results)
        
        return all_results

# 使用示例
if __name__ == "__main__":
    # 替换为你的模型路径(本地或Hugging Face模型ID)
    MODEL_PATH = "Qwen/Qwen-7B-Chat"  # 或"deepseek-ai/deepseek-llm-7b"等
    
    # 初始化推理器
    inferencer = BatchInference(MODEL_PATH)
    
    # 示例输入
    prompts = [
        "请解释一下量子计算的基本原理",
        "写一首关于春天的诗",
        "如何用Python实现快速排序?",
        "Transformer模型的主要创新点是什么?"
    ]
    
    # 批量推理
    results = inferencer.batch_predict(
        prompts,
        max_new_tokens=256,
        batch_size=2,  # 根据GPU内存调整
        temperature=0.7,
        top_p=0.9
    )
    
    # 打印结果
    for prompt, result in zip(prompts, results):
        print(f"输入: {prompt}\n输出: {result}\n{'-'*50}")

三、关键点说明

设备管理:
- 自动检测并使用可用的GPU
- 支持半精度(fp16)推理以节省显存
批量处理:
- 将输入分成小批次处理，避免内存不足
- 自动填充(padding)使批次内样本长度一致
生成参数:
- max_new_tokens: 控制生成的最大长度
- temperature和top_p: 控制生成的随机性
- 可通过generate_kwargs传递其他生成参数
内存优化:
- 使用torch.no_grad()减少内存消耗
- 根据GPU内存调整batch_size

四、进阶优化

使用Flash Attention (如果模型支持):

model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
torch_dtype=torch.float16,
use_flash_attention_2=True, # 启用Flash Attention
device_map="auto"
)

2.量化推理 (减少显存使用):

复制代码

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    trust_remote_code=True,
    torch_dtype=torch.float16,
    load_in_8bit=True,  # 8位量化
    device_map="auto"
)

3.使用vLLM等优化库 (对于生产环境):

复制代码

from vllm import LLM, SamplingParams

llm = LLM(model=MODEL_PATH)
sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=256)
outputs = llm.generate(prompts, sampling_params)

以上代码提供了基于Transformer框架的Qwen/DeepSeek模型批量推理基础实现，可根据实际需求进行调整和优化。