大语言模型推理优化踩坑记录

作者: 醉月思发布日期: 2024-07-12 分类: 技术实践

这次做大语言模型推理优化，从量化到 LoRA，。

推理性能评估

基准测试

import time
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def benchmark_inference(prompt, max_length=100):
    """基准测试推理性能"""
    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # 记录时间
    start_time = time.time()
    
    # 推理
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=max_length,
            do_sample=True
        )
    
    end_time = time.time()
    
    # 计算指标
    inference_time = end_time - start_time
    tokens_generated = outputs.shape[1] - inputs.input_ids.shape[1]
    tokens_per_second = tokens_generated / inference_time
    
    return {
        "inference_time": inference_time,
        "tokens_generated": tokens_generated,
        "tokens_per_second": tokens_per_second
    }

# 运行基准测试
results = benchmark_inference("The future of AI is")
print(f"Inference time: {results['inference_time']:.2f}s")
print(f"Tokens per second: {results['tokens_per_second']:.2f}")

性能分析

import torch
from transformers import pipeline

# 使用 pipeline 进行性能分析
generator = pipeline(
    "text-generation",
    model=model_name,
    device=0  # 使用 GPU
)

# 分析内存使用
def analyze_memory():
    if torch.cuda.is_available():
        print(f"GPU Memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
        print(f"GPU Memory Cached: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

analyze_memory()

量化优化

动态量化

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载模型
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 只量化 Linear 层
    dtype=torch.qint8
)

# 测试量化后的性能
def test_quantized_model(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    start_time = time.time()
    with torch.no_grad():
        outputs = quantized_model.generate(
            inputs.input_ids,
            max_length=100
        )
    end_time = time.time()
    
    print(f"Quantized inference time: {end_time - start_time:.2f}s")
    print(f"Generated text: {tokenizer.decode(outputs[0], skip_special_tokens=True)}")

test_quantized_model("The future of AI is")

静态量化

import torch.quantization as quant

# 准备校准数据
def prepare_calibration_data():
    calibration_data = []
    for text in sample_texts:  # 使用代表性的文本样本
        inputs = tokenizer(text, return_tensors="pt")
        calibration_data.append(inputs)
    return calibration_data

# 校准模型
def calibrate_model(model, calibration_data):
    model.eval()
    
    # 配置量化
    model.qconfig = quant.get_default_qconfig('fbgemm')
    quant.prepare(model, inplace=True)
    
    # 校准
    with torch.no_grad():
        for data in calibration_data:
            model(**data)
    
    # 转换为量化模型
    quant.convert(model, inplace=True)
    return model

# 使用静态量化
calibration_data = prepare_calibration_data()
quantized_model = calibrate_model(model, calibration_data)

INT8 量化

from optimum.bettertransformer import BetterTransformer

# BetterTransformer 优化
model = BetterTransformer.transform(model, attn_implementation="sdpa")

# INT8 量化
from optimum import INCModel

model = INCModel.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto"
)

# 测试 INT8 量化
def test_int8_model(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    start_time = time.time()
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=100
        )
    end_time = time.time()
    
    print(f"INT8 inference time: {end_time - start_time:.2f}s")
    print(f"Generated text: {tokenizer.decode(outputs[0], skip_special_tokens=True)}")

test_int8_model("The future of AI is")

LoRA 微调

LoRA 基础

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, TrainingArguments, Trainer

# 加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto"
)

# 配置 LoRA
lora_config = LoraConfig(
    r=8,  # LoRA 秩
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],  # 目标模块
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# 添加 LoRA 适配器
model = get_peft_model(base_model, lora_config)

# 查看可训练参数
model.print_trainable_parameters()

LoRA 微调

from datasets import Dataset
from transformers import Trainer

# 准备训练数据
train_data = [
    {"input": "What is AI?", "output": "Artificial Intelligence is..."},
    {"input": "How does ML work?", "output": "Machine Learning works by..."},
    # 更多训练样本
]

# 转换为 Dataset 格式
def prepare_dataset(data):
    inputs = [f"Input: {item['input']}\nOutput: " for item in data]
    targets = [item['output'] for item in data]
    
    return Dataset.from_dict({
        "input_ids": [tokenizer.encode(inp, truncation=True) for inp in inputs],
        "labels": [tokenizer.encode(tar, truncation=True) for tar in targets]
    })

train_dataset = prepare_dataset(train_data)

# 配置训练参数
training_args = TrainingArguments(
    output_dir="./lora_model",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    learning_rate=2e-4,
    fp16=True,
    logging_steps=10,
    save_steps=100,
)

# 创建 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

# 微调模型
trainer.train()

# 保存模型
model.save_pretrained("./lora_model")

使用 LoRA 模型

# 加载基础模型
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto"
)

# 加载 LoRA 适配器
from peft import PeftModel

model = PeftModel.from_pretrained(base_model, "./lora_model")

# 推理
def inference_with_lora(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=100,
            do_sample=True,
            temperature=0.7
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

result = inference_with_lora("What is AI?")
print(result)

其他优化技术

KV Cache

# KV Cache 是 Transformer 模型中的缓存机制
# 它缓存每个 token 的 Key 和 Value，避免重复计算

# 使用缓存加速推理
def generate_with_cache(prompt, max_tokens=100):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            inputs.input_ids,
            max_length=inputs.input_ids.shape[1] + max_tokens,
            use_cache=True,  # 启用 KV Cache
            do_sample=True
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

Flash Attention

from transformers import AutoModelForCausalLM

# 使用 Flash Attention
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    use_flash_attention_2=True,
    device_map="auto"
)

# Flash Attention 可以显著加速注意力计算

Batch Inference

# 批量推理
def batch_inference(prompts, batch_size=4):
    all_results = []
    
    for i in range(0, len(prompts), batch_size):
        batch = prompts[i:i+batch_size]
        
        # Tokenize
        inputs = tokenizer(
            batch,
            padding=True,
            return_tensors="pt",
            truncation=True
        )
        
        # 推理
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_length=inputs.input_ids.shape[1] + 50,
                use_cache=True
            )
        
        # 解码
        for output in outputs:
            all_results.append(
                tokenizer.decode(output, skip_special_tokens=True)
            )
    
    return all_results

踩过的坑

坑一：量化后质量下降

量化后模型输出质量明显下降。

解决：使用量化感知训练或选择合适的量化策略。

# 量化感知训练
quant_config = {
    "quantization_config": {
        "load_in_8bit": True,
        "llm_int8_threshold": 6.0,
        "llm_int8_has_fp16_weight": True
    }
}

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    **quant_config
)

坑二：LoRA 微调效果不明显

LoRA 微调后效果不明显。

解决：调整 LoRA 参数和训练策略。

lora_config = LoraConfig(
    r=16,  # 增加秩
    lora_alpha=64,  # 增加 alpha
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # 增加目标模块
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)

坑三：内存溢出

模型太大，内存溢出。

解决：使用模型并行或梯度检查点。

# 使用模型并行
from accelerate import Accelerator

accelerator = Accelerator()

# 梯度检查点
model.gradient_checkpointing_enable()

# 8-bit 量化
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,
    device_map="auto"
)

写在最后

LLM 推理优化这东西，不只是技术，是成本和体验。

优化了：

推理速度
内存占用
推理成本

带来了：

模型质量可能下降
复杂度增加
学习成本

优化之前先评估：

质量要求
成本预算
性能要求
团队能力

量化适合：

对精度要求不高
需要快速推理
资源有限

LoRA 适合：

需要微调
保持模型能力
资源有限

不是所有场景都需要优化，有时候原始模型就够用。

这次 LLM 推理优化花了一个月，从量化到 LoRA。优化完成后，推理速度提升了 4 倍，内存占用减少了 50%，成本降低了 60%。

版权声明: 本文首发于指尖魔法屋-大语言模型推理优化踩坑记录（https://blog.thinkmoon.cn/post/90-llm-inference-optimization-quantization-lora-practice/）转载或引用必须申明原指尖魔法屋来源及源地址！

大模型微调折腾手记

结果光显存这块就踩了一堆坑：先是用全参数微调直接 OOM，换 LoRA 还是顶不住，最后上 Q-LoRA 才勉强在 24G 显存上跑起来。

📁 技术实践

📅 2024-02-07

AI LoRA：这次怎么落地的

边界条件才会告诉你方案能不能留。从原理到实践，搞清楚 LoRA 的工作机制，把那些踩过的坑和…

📁 技术实践

📅 2024-04-17

把量化换到剪枝时踩过的坑

项目里那个 7B 参数的模型训练好了，部署到推理服务器上时，CPU 占用直接拉满，延迟干到了 3 秒多。

📁 技术实践

📅 2022-01-10

大语言模型推理优化踩坑记录

推理性能评估

基准测试

性能分析

量化优化

动态量化

静态量化

INT8 量化

LoRA 微调

LoRA 基础

LoRA 微调

使用 LoRA 模型

其他优化技术

KV Cache

Flash Attention

Batch Inference

踩过的坑

坑一：量化后质量下降

坑二：LoRA 微调效果不明显

坑三：内存溢出

写在最后

评论

相关阅读