深度学习模型量化完全指南：精通 INT8、INT4、GPTQ、AWQ、GGUF

import torch
import numpy as np

# 查看各数据类型每个元素占用的内存大小
x_fp32 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.float32)
x_fp16 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.float16)
x_bf16 = torch.tensor([1.5, -2.3, 0.7], dtype=torch.bfloat16)

print(f"FP32: {x_fp32.element_size()} bytes per element")  # 4 bytes
print(f"FP16: {x_fp16.element_size()} bytes per element")  # 2 bytes
print(f"BF16: {x_bf16.element_size()} bytes per element")  # 2 bytes

# 模型内存计算示例(7B 参数模型)
params = 7e9
fp32_memory_gb = params * 4 / 1e9
fp16_memory_gb = params * 2 / 1e9
int8_memory_gb = params * 1 / 1e9
int4_memory_gb = params * 0.5 / 1e9

print(f"\n7B 模型内存需求:")
print(f"FP32: {fp32_memory_gb:.1f} GB")   # 28.0 GB
print(f"FP16: {fp16_memory_gb:.1f} GB")   # 14.0 GB
print(f"INT8: {int8_memory_gb:.1f} GB")   # 7.0 GB
print(f"INT4: {int4_memory_gb:.1f} GB")   # 3.5 GB

q = clamp(round(x / scale) + zero_point, q_min, q_max)

x_approx = scale * (q - zero_point)

import torch
import numpy as np

def symmetric_quantize(x: torch.Tensor, num_bits: int = 8):
    """对称量化实现"""
    q_max = 2 ** (num_bits - 1) - 1  # INT8 时为 127
    q_min = -q_max  # -127

    # 计算 scale
    max_abs = x.abs().max()
    scale = max_abs / q_max

    # 量化
    q = torch.clamp(torch.round(x / scale), q_min, q_max).to(torch.int8)

    return q, scale

def asymmetric_quantize(x: torch.Tensor, num_bits: int = 8):
    """非对称量化实现"""
    q_max = 2 ** num_bits - 1  # UINT8 时为 255
    q_min = 0

    # 计算 scale 和 zero_point
    min_val = x.min()
    max_val = x.max()
    scale = (max_val - min_val) / (q_max - q_min)
    zero_point = q_min - torch.round(min_val / scale)
    zero_point = torch.clamp(zero_point, q_min, q_max).to(torch.int32)

    # 量化
    q = torch.clamp(torch.round(x / scale) + zero_point, q_min, q_max).to(torch.uint8)

    return q, scale, zero_point

def dequantize(q: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor = None):
    """反量化"""
    if zero_point is None:
        return scale * q.float()
    return scale * (q.float() - zero_point.float())

# 测试
x = torch.randn(100)
print(f"原始数据范围: [{x.min():.4f}, {x.max():.4f}]")

# 对称量化
q_sym, scale_sym = symmetric_quantize(x)
x_reconstructed_sym = dequantize(q_sym, scale_sym)
error_sym = (x - x_reconstructed_sym).abs().mean()
print(f"对称量化平均误差: {error_sym:.6f}")

# 非对称量化
q_asym, scale_asym, zp_asym = asymmetric_quantize(x)
x_reconstructed_asym = dequantize(q_asym, scale_asym, zp_asym)
error_asym = (x - x_reconstructed_asym).abs().mean()
print(f"非对称量化平均误差: {error_asym:.6f}")

import torch

def per_group_quantize(weight: torch.Tensor, group_size: int = 128, num_bits: int = 4):
    """Per-Group 量化实现"""
    rows, cols = weight.shape

    # 按分组切分
    weight_grouped = weight.reshape(-1, group_size)

    # 每个分组的最大/最小值
    max_vals = weight_grouped.max(dim=1, keepdim=True)[0]
    min_vals = weight_grouped.min(dim=1, keepdim=True)[0]

    q_max = 2 ** num_bits - 1  # INT4 时为 15

    # 计算 scale
    scales = (max_vals - min_vals) / q_max
    zero_points = torch.round(-min_vals / scales)

    # 量化
    q = torch.clamp(torch.round(weight_grouped / scales) + zero_points, 0, q_max)

    # 反量化
    weight_dequant = scales * (q - zero_points)
    weight_dequant = weight_dequant.reshape(rows, cols)

    return q, scales, zero_points, weight_dequant

# 示例：Transformer 权重量化
weight = torch.randn(4096, 4096)  # Llama 风格权重
q, scales, zp, weight_dequant = per_group_quantize(weight, group_size=128, num_bits=4)

error = (weight - weight_dequant).abs().mean()
print(f"Per-Group INT4 量化平均误差: {error:.6f}")
print(f"压缩率: {weight.element_size() * weight.numel() / (q.numel() / 2 + scales.numel() * 4):.2f}x")

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from datasets import load_dataset

def collect_calibration_data(model_name: str, num_samples: int = 128):
    """收集校准数据"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 通常使用 WikiText-2 或 C4 数据集
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

    texts = []
    for item in dataset:
        if len(item['text'].strip()) > 100:
            texts.append(item['text'].strip())
        if len(texts) >= num_samples:
            break

    # 分词
    encoded = [
        tokenizer(text, return_tensors="pt", max_length=2048, truncation=True)
        for text in texts
    ]

    return encoded

# 用校准数据收集激活统计信息
def collect_activation_stats(model, calibration_data, layer_name: str):
    """收集特定层的激活统计信息"""
    stats = {"min": float("inf"), "max": float("-inf"), "histogram": []}

    def hook_fn(module, input, output):
        with torch.no_grad():
            act = output.detach().float()
            stats["min"] = min(stats["min"], act.min().item())
            stats["max"] = max(stats["max"], act.max().item())

    # 注册钩子
    target_layer = dict(model.named_modules())[layer_name]
    handle = target_layer.register_forward_hook(hook_fn)

    # 运行校准数据
    model.eval()
    with torch.no_grad():
        for batch in calibration_data[:32]:
            model(**batch)

    handle.remove()
    return stats

class MinMaxCalibrator:
    """最小-最大校准器"""

    def __init__(self):
        self.min_val = float("inf")
        self.max_val = float("-inf")

    def update(self, tensor: torch.Tensor):
        self.min_val = min(self.min_val, tensor.min().item())
        self.max_val = max(self.max_val, tensor.max().item())

    def compute_scale_zp(self, num_bits: int = 8, symmetric: bool = True):
        q_max = 2 ** (num_bits - 1) - 1 if symmetric else 2 ** num_bits - 1

        if symmetric:
            max_abs = max(abs(self.min_val), abs(self.max_val))
            scale = max_abs / q_max
            zero_point = 0
        else:
            scale = (self.max_val - self.min_val) / q_max
            zero_point = -round(self.min_val / scale)

        return scale, zero_point

import numpy as np
from scipy import stats

class HistogramCalibrator:
    """基于直方图的校准器(最小化 KL 散度)"""

    def __init__(self, num_bins: int = 2048):
        self.num_bins = num_bins
        self.histogram = None
        self.bin_edges = None

    def update(self, tensor: torch.Tensor):
        data = tensor.detach().float().numpy().flatten()

        if self.histogram is None:
            self.histogram, self.bin_edges = np.histogram(data, bins=self.num_bins)
        else:
            new_hist, _ = np.histogram(data, bins=self.bin_edges)
            self.histogram += new_hist

    def compute_optimal_range(self, num_bits: int = 8):
        """寻找最小化 KL 散度的最优范围"""
        num_quantized_bins = 2 ** num_bits - 1

        best_kl = float("inf")
        best_threshold = None

        # 遍历不同的 threshold
        for i in range(num_quantized_bins, len(self.histogram)):
            # 把直方图压缩为 num_quantized_bins 个区间
            reference = self.histogram[:i].copy().astype(float)
            reference /= reference.sum()

            # 计算 KL 散度(近似)
            quantized = np.zeros(i)
            bin_size = i / num_quantized_bins

            for j in range(num_quantized_bins):
                start = int(j * bin_size)
                end = int((j + 1) * bin_size)
                quantized[start:end] = reference[start:end].sum() / (end - start)

            # 处理值为 0 的区间
            quantized = np.where(quantized == 0, 1e-10, quantized)
            reference_clipped = np.where(reference == 0, 1e-10, reference)

            kl = stats.entropy(reference_clipped, quantized)

            if kl < best_kl:
                best_kl = kl
                best_threshold = self.bin_edges[i]

        return -best_threshold, best_threshold

import torch
import math
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_perplexity(model, tokenizer, text: str, device: str = "cuda"):
    """计算困惑度"""
    encodings = tokenizer(text, return_tensors="pt")
    input_ids = encodings.input_ids.to(device)

    max_length = 1024
    stride = 512

    nlls = []
    prev_end_loc = 0

    for begin_loc in range(0, input_ids.size(1), stride):
        end_loc = min(begin_loc + max_length, input_ids.size(1))
        trg_len = end_loc - prev_end_loc

        input_ids_chunk = input_ids[:, begin_loc:end_loc]
        target_ids = input_ids_chunk.clone()
        target_ids[:, :-trg_len] = -100

        with torch.no_grad():
            outputs = model(input_ids_chunk, labels=target_ids)
            neg_log_likelihood = outputs.loss

        nlls.append(neg_log_likelihood)
        prev_end_loc = end_loc

        if end_loc == input_ids.size(1):
            break

    ppl = torch.exp(torch.stack(nlls).mean())
    return ppl.item()

# 各模型 PPL 对比示例
# FP16: PPL ≈ 5.68
# INT8: PPL ≈ 5.71(约增加 0.5%)
# INT4 (GPTQ): PPL ≈ 5.89(约增加 3.7%)
# INT4 (naive): PPL ≈ 6.52(约增加 14.8%)

import torch
import torch.nn as nn
import torch.nn.functional as F

class FakeQuantize(nn.Module):
    """伪量化模块"""

    def __init__(self, num_bits: int = 8, symmetric: bool = True):
        super().__init__()
        self.num_bits = num_bits
        self.symmetric = symmetric

        self.register_buffer('scale', torch.tensor(1.0))
        self.register_buffer('zero_point', torch.tensor(0))
        self.register_buffer('fake_quant_enabled', torch.tensor(1))

        if symmetric:
            self.q_min = -(2 ** (num_bits - 1))
            self.q_max = 2 ** (num_bits - 1) - 1
        else:
            self.q_min = 0
            self.q_max = 2 ** num_bits - 1

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.fake_quant_enabled[0] == 0:
            return x

        # 更新 scale(滑动平均)
        if self.training:
            with torch.no_grad():
                if self.symmetric:
                    max_abs = x.abs().max()
                    new_scale = max_abs / self.q_max
                else:
                    new_scale = (x.max() - x.min()) / (self.q_max - self.q_min)

                # 用指数滑动平均更新 scale
                self.scale.copy_(0.9 * self.scale + 0.1 * new_scale)

        # 伪量化：先量化再反量化
        x_scaled = x / self.scale
        x_clipped = torch.clamp(x_scaled, self.q_min, self.q_max)
        x_rounded = torch.round(x_clipped)
        x_dequant = x_rounded * self.scale

        return x_dequant

class STERound(torch.autograd.Function):
    """round() 的 Straight-Through Estimator"""

    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, grad_output):
        # 反向传播时直接让梯度穿过 round()(用恒等函数近似)
        return grad_output

class STEClamp(torch.autograd.Function):
    """clamp() 的 Straight-Through Estimator"""

    @staticmethod
    def forward(ctx, x, min_val, max_val):
        ctx.save_for_backward(x)
        ctx.min_val = min_val
        ctx.max_val = max_val
        return torch.clamp(x, min_val, max_val)

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        # 只在 clamp 范围内传递梯度
        grad = grad_output * ((x >= ctx.min_val) & (x <= ctx.max_val)).float()
        return grad, None, None

class QATLinear(nn.Module):
    """应用了 QAT 的 Linear 层"""

    def __init__(self, in_features, out_features, num_bits=8):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.weight_fake_quant = FakeQuantize(num_bits=num_bits)
        self.act_fake_quant = FakeQuantize(num_bits=num_bits, symmetric=False)

    def forward(self, x):
        # 激活量化
        x_q = self.act_fake_quant(x)
        # 权重量化
        w_q = self.weight_fake_quant(self.linear.weight)
        # FP32 运算(实际部署中为 INT8)
        return F.linear(x_q, w_q, self.linear.bias)

# QAT 训练工作流
import torch.optim as optim
from torch.quantization import prepare_qat, convert

def train_qat_model(model, train_loader, num_epochs=10):
    """QAT 训练示例"""

    # 准备 QAT
    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
    model_prepared = prepare_qat(model.train())

    optimizer = optim.Adam(model_prepared.parameters(), lr=1e-5)

    for epoch in range(num_epochs):
        for batch in train_loader:
            inputs, labels = batch
            outputs = model_prepared(inputs)
            loss = F.cross_entropy(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    # 转换为 INT8 模型
    model_prepared.eval()
    model_quantized = convert(model_prepared)

    return model_quantized

import torch
from torch.ao.quantization import (
    get_default_qconfig,
    get_default_qat_qconfig,
    prepare,
    prepare_qat,
    convert
)

# 静态量化(PTQ)
def static_quantization_example():
    """静态量化示例"""
    model = MyModel()
    model.eval()

    # 后端设置(fbgemm: x86，qnnpack: ARM)
    model.qconfig = get_default_qconfig('fbgemm')

    # 准备校准
    model_prepared = prepare(model)

    # 用校准数据收集统计信息
    with torch.no_grad():
        for data in calibration_loader:
            model_prepared(data)

    # 转换为 INT8 模型
    model_quantized = convert(model_prepared)

    return model_quantized

# 动态量化(对 LSTM、Linear 效果好)
def dynamic_quantization_example():
    """动态量化示例"""
    model = MyModel()

    model_quantized = torch.quantization.quantize_dynamic(
        model,
        {nn.Linear, nn.LSTM},  # 要量化的层类型
        dtype=torch.qint8
    )

    return model_quantized

from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
from torch.ao.quantization import QConfigMapping

def fx_quantization_example(model, calibration_data):
    """FX 图模式量化"""
    model.eval()

    # 配置 QConfig
    qconfig_mapping = QConfigMapping().set_global(
        get_default_qconfig('fbgemm')
    )

    # 示例输入
    example_inputs = (torch.randn(1, 3, 224, 224),)

    # 基于 FX 图准备
    model_prepared = prepare_fx(
        model,
        qconfig_mapping,
        example_inputs
    )

    # 校准
    with torch.no_grad():
        for batch in calibration_data:
            model_prepared(batch)

    # 转换
    model_quantized = convert_fx(model_prepared)

    return model_quantized

argmin_Q ||WX - QX||_F^2

# GPTQ 核心算法实现(简化版)
import torch
import math

def gptq_quantize_weight(weight: torch.Tensor,
                          hessian: torch.Tensor,
                          num_bits: int = 4,
                          group_size: int = 128,
                          damp_percent: float = 0.01):
    """
    用 GPTQ 算法量化权重

    Args:
        weight: [out_features, in_features] 权重矩阵
        hessian: [in_features, in_features] 海森矩阵 (H = 2 * X @ X.T)
        num_bits: 量化位数
        group_size: 分组大小
        damp_percent: 用于稳定海森矩阵的阻尼比例
    """
    W = weight.clone().float()
    n_rows, n_cols = W.shape

    # 海森矩阵阻尼(数值稳定性)
    H = hessian.clone().float()
    dead_cols = torch.diag(H) == 0
    H[dead_cols, dead_cols] = 1
    W[:, dead_cols] = 0

    damp = damp_percent * H.diag().mean()
    H.diagonal().add_(damp)

    # 用 Cholesky 分解求海森矩阵的逆
    H_inv = torch.linalg.cholesky(H)
    H_inv = torch.cholesky_inverse(H_inv)
    H_inv = torch.linalg.cholesky(H_inv, upper=True)

    Q = torch.zeros_like(W)
    Losses = torch.zeros_like(W)

    q_max = 2 ** (num_bits - 1) - 1

    for col_idx in range(n_cols):
        w_col = W[:, col_idx]  # 当前列的权重
        h_inv_diag = H_inv[col_idx, col_idx]  # 海森逆矩阵的对角元素

        # 按分组计算 scale
        if group_size != -1 and col_idx % group_size == 0:
            group_end = min(col_idx + group_size, n_cols)
            w_group = W[:, col_idx:group_end]
            max_abs = w_group.abs().max(dim=1)[0].unsqueeze(1)
            scale = max_abs / q_max
            scale = torch.clamp(scale, min=1e-8)

        # 量化
        q_col = torch.clamp(torch.round(w_col / scale.squeeze()), -q_max, q_max)
        q_col = q_col * scale.squeeze()
        Q[:, col_idx] = q_col

        # 量化误差
        err = (w_col - q_col) / h_inv_diag
        Losses[:, col_idx] = err ** 2 / 2

        # 把误差传播到剩余权重上(核心步骤！)
        W[:, col_idx + 1:] -= err.unsqueeze(1) * H_inv[col_idx, col_idx + 1:].unsqueeze(0)

    return Q, Losses


def collect_hessian(model_layer, calibration_data, device='cuda'):
    """用校准数据收集海森矩阵"""
    hessians = {}

    def make_hook(name):
        def hook(module, input, output):
            inp = input[0].detach().float()
            if inp.dim() == 3:
                inp = inp.reshape(-1, inp.size(-1))

            if name not in hessians:
                hessians[name] = torch.zeros(inp.size(1), inp.size(1), device=device)

            hessians[name] += 2 * inp.T @ inp
        return hook

    handles = []
    for name, module in model_layer.named_modules():
        if isinstance(module, torch.nn.Linear):
            handles.append(module.register_forward_hook(make_hook(name)))

    with torch.no_grad():
        for batch in calibration_data:
            model_layer(batch.to(device))

    for h in handles:
        h.remove()

    return hessians

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch

def quantize_with_gptq(
    model_name: str,
    output_dir: str,
    bits: int = 4,
    group_size: int = 128
):
    """用 AutoGPTQ 对模型进行量化"""

    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)

    # 量化配置
    quantize_config = BaseQuantizeConfig(
        bits=bits,              # 4 或 8
        group_size=group_size,  # 推荐 128
        damp_percent=0.01,      # 海森矩阵阻尼
        desc_act=False,         # 激活重排序(提升质量，降低速度)
        sym=True,               # 对称量化
        true_sequential=True    # 逐层顺序量化
    )

    # 加载模型
    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config=quantize_config
    )

    # 准备校准数据
    from datasets import load_dataset
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

    calibration_data = []
    for text in dataset["text"][:128]:
        if len(text.strip()) > 50:
            encoded = tokenizer(
                text.strip(),
                return_tensors="pt",
                max_length=2048,
                truncation=True
            )
            calibration_data.append(encoded["input_ids"].squeeze())

    # 执行 GPTQ 量化
    print(f"开始 GPTQ {bits}bit 量化...")
    model.quantize(calibration_data)

    # 保存
    model.save_quantized(output_dir, use_safetensors=True)
    tokenizer.save_pretrained(output_dir)

    print(f"量化完成: {output_dir}")
    return model, tokenizer


def load_gptq_model(model_dir: str, device: str = "cuda"):
    """加载 GPTQ 量化模型"""

    model = AutoGPTQForCausalLM.from_quantized(
        model_dir,
        device=device,
        use_triton=False,       # 是否使用 Triton 内核
        disable_exllama=False,  # 使用 ExLlama 内核(提升速度)
        inject_fused_attention=True,
        inject_fused_mlp=True
    )

    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    return model, tokenizer

# 使用示例
# model, tokenizer = quantize_with_gptq("meta-llama/Llama-2-7b-hf", "./llama2-7b-gptq-4bit")
# model, tokenizer = load_gptq_model("./llama2-7b-gptq-4bit")

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

def quantize_with_awq(
    model_name: str,
    output_dir: str,
    bits: int = 4,
    group_size: int = 128
):
    """用 AutoAWQ 对模型进行量化"""

    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True
    )

    model = AutoAWQForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
        use_cache=False
    )

    # AWQ 量化配置
    quant_config = {
        "zero_point": True,   # 非对称量化
        "q_group_size": group_size,
        "w_bit": bits,
        "version": "GEMM"     # GEMM 或 GEMV(适合小批量)
    }

    # 执行量化
    print(f"开始 AWQ {bits}bit 量化...")
    model.quantize(tokenizer, quant_config=quant_config)

    # 保存
    model.save_quantized(output_dir)
    tokenizer.save_pretrained(output_dir)

    print(f"AWQ 量化完成: {output_dir}")
    return model

def load_awq_model(model_dir: str, device: str = "cuda"):
    """加载 AWQ 量化模型"""

    model = AutoAWQForCausalLM.from_quantized(
        model_dir,
        fuse_layers=True,       # 层融合以提升速度
        trust_remote_code=True,
        safetensors=True
    )

    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    return model, tokenizer

# 与 Hugging Face transformers 集成
from transformers import AutoModelForCausalLM

def load_awq_with_transformers(model_dir: str):
    """用 transformers 加载 AWQ 模型"""
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    return model, tokenizer

GGUF 文件结构：
┌─────────────────────────────┐
│ 魔数 (GGUF)                  │
│ 版本                         │
│ 张量数量                     │
│ 元数据 KV 对                 │
│  - 模型架构                  │
│  - 上下文长度                │
│  - 注意力头数                │
│  - 嵌入维度                  │
├─────────────────────────────┤
│ 张量信息 (名称、类型、形状)   │
├─────────────────────────────┤
│ 张量数据                     │
└─────────────────────────────┘

# 克隆并构建 llama.cpp
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp

# 支持 CUDA 的构建
cmake -B build -DGGML_CUDA=ON
cmake --build build --config Release -j $(nproc)

# 仅 CPU 构建
cmake -B build
cmake --build build --config Release -j $(nproc)

# 把 HuggingFace 模型转换为 GGUF
python convert_hf_to_gguf.py \
    --model meta-llama/Llama-2-7b-hf \
    --outfile llama2-7b-f16.gguf \
    --outtype f16

# GGUF 量化 (Q4_K_M)
./build/bin/llama-quantize \
    llama2-7b-f16.gguf \
    llama2-7b-q4_k_m.gguf \
    Q4_K_M

# 执行推理
./build/bin/llama-cli \
    -m llama2-7b-q4_k_m.gguf \
    -p "The future of AI is" \
    -n 100 \
    --ctx-size 4096 \
    --threads 8 \
    --n-gpu-layers 35

from llama_cpp import Llama

# 加载模型
llm = Llama(
    model_path="./llama2-7b-q4_k_m.gguf",
    n_ctx=4096,          # 上下文长度
    n_gpu_layers=35,     # 卸载到 GPU 的层数(-1 表示全部)
    n_threads=8,         # CPU 线程数
    verbose=False
)

# 文本生成
output = llm(
    "Once upon a time",
    max_tokens=200,
    temperature=0.7,
    top_p=0.9,
    stop=["</s>", "\n\n"]
)

print(output["choices"][0]["text"])

# 对话补全格式
response = llm.create_chat_completion(
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is machine learning?"}
    ],
    max_tokens=500,
    temperature=0.7
)

print(response["choices"][0]["message"]["content"])

# 流式输出
for chunk in llm.create_chat_completion(
    messages=[{"role": "user", "content": "Tell me a joke"}],
    stream=True
):
    delta = chunk["choices"][0].get("delta", {})
    if "content" in delta:
        print(delta["content"], end="", flush=True)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载 INT8 模型
model_8bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    load_in_8bit=True,
    device_map="auto"
)

# 查看内存占用
def print_model_size(model, label):
    """打印模型内存占用"""
    total_params = sum(p.numel() for p in model.parameters())
    total_bytes = sum(
        p.numel() * p.element_size() for p in model.parameters()
    )
    print(f"{label}: {total_params/1e9:.2f}B params, {total_bytes/1e9:.2f} GB")

print_model_size(model_8bit, "INT8 模型")
# INT8 模型: 6.74B params, ~7.0 GB

import bitsandbytes as bnb
from transformers import BitsAndBytesConfig

# NF4 量化配置(QLoRA)
bnb_config_nf4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # NF4 或 FP4
    bnb_4bit_compute_dtype=torch.bfloat16, # 运算时使用的数据类型
    bnb_4bit_use_double_quant=True,       # 二次量化(连量化常数也量化)
)

# FP4 量化配置
bnb_config_fp4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="fp4",
    bnb_4bit_compute_dtype=torch.float16,
)

# 加载模型
model_4bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config_nf4,
    device_map="auto"
)

print_model_size(model_4bit, "NF4 模型")
# NF4 模型: 6.74B params, ~4.0 GB(含二次量化)

# QLoRA 微调设置
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model_4bit = prepare_model_for_kbit_training(model_4bit)

lora_config = LoraConfig(
    r=64,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model_lora = get_peft_model(model_4bit, lora_config)
model_lora.print_trainable_parameters()
# trainable params: 4,194,304 || all params: 3,504,607,232 || trainable%: 0.1197

import numpy as np
import matplotlib.pyplot as plt

# NF4 量化点可视化
def get_nf4_quantization_points():
    """NF4 的 16 个量化点"""
    # 正态分布的 1/16 分位数
    nf4_points = []
    for i in range(16):
        quantile = (i + 0.5) / 16
        nf4_points.append(scipy.stats.norm.ppf(quantile))

    # 归一化
    max_val = max(abs(p) for p in nf4_points)
    nf4_points = [p / max_val for p in nf4_points]

    return nf4_points

# NF4: [-1.0, -0.6961, -0.5250, -0.3949, -0.2844, -0.1848, -0.0911, 0.0000,
#        0.0796, 0.1609, 0.2461, 0.3379, 0.4407, 0.5626, 0.7230, 1.0]

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def analyze_activation_outliers(model, tokenizer, text: str, threshold: float = 100.0):
    """分析激活异常值"""

    activations = {}

    def make_hook(name):
        def hook(module, input, output):
            act = output.detach().float()
            max_val = act.abs().max().item()
            outlier_ratio = (act.abs() > threshold).float().mean().item()
            activations[name] = {
                "max": max_val,
                "outlier_ratio": outlier_ratio,
                "std": act.std().item()
            }
        return hook

    handles = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            handles.append(module.register_forward_hook(make_hook(name)))

    input_ids = tokenizer(text, return_tensors="pt").input_ids.cuda()

    with torch.no_grad():
        model(input_ids)

    for h in handles:
        h.remove()

    # 按异常值多少对层排序
    sorted_acts = sorted(
        activations.items(),
        key=lambda x: x[1]["max"],
        reverse=True
    )

    print("异常值最大的 Top 10 层:")
    for name, stats in sorted_acts[:10]:
        print(f"  {name}: max={stats['max']:.1f}, outlier_ratio={stats['outlier_ratio']:.3%}")

    return activations

Y = (X * diag(s)^(-1)) * (diag(s) * W)
  = X_smooth * W_smooth

def smooth_quantize(
    model,
    calibration_samples,
    alpha: float = 0.5
):
    """
    应用 SmoothQuant

    Args:
        alpha: 迁移强度(0=仅权重，1=仅激活)
               推荐值：0.5(均等分配)
    """

    # 收集激活统计信息
    act_scales = {}

    def collect_scales(name):
        def hook(module, input, output):
            inp = input[0].detach()
            if inp.dim() == 3:
                inp = inp.reshape(-1, inp.size(-1))

            channel_max = inp.abs().max(dim=0)[0]

            if name not in act_scales:
                act_scales[name] = channel_max
            else:
                act_scales[name] = torch.maximum(act_scales[name], channel_max)
        return hook

    handles = []
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear):
            handles.append(module.register_forward_hook(collect_scales(name)))

    with torch.no_grad():
        for sample in calibration_samples:
            model(**sample)

    for h in handles:
        h.remove()

    # 计算并应用 scale
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) and name in act_scales:
            act_scale = act_scales[name]
            weight_scale = module.weight.abs().max(dim=0)[0]

            # 计算迁移 scale
            smooth_scale = (act_scale ** alpha) / (weight_scale ** (1 - alpha))
            smooth_scale = torch.clamp(smooth_scale, min=1e-5)

            # 把 scale 应用到权重上
            module.weight.data = module.weight.data / smooth_scale.unsqueeze(0)

            # 对前一层(LayerNorm 等)的输出 scale 做逆向调整
            # (实际实现中需要找到前一层并修改)

    return model, act_scales

import torch

def spqr_quantize(weight: torch.Tensor,
                   num_bits: int = 3,
                   outlier_threshold_percentile: float = 1.0):
    """
    SpQR 量化(简化版)

    核心思路：把 top p% 的异常值以 FP16 保存，其余部分做低位量化
    """

    # 计算异常值阈值
    threshold = torch.quantile(weight.abs(), 1 - outlier_threshold_percentile / 100)

    # 异常值掩码
    outlier_mask = weight.abs() > threshold

    # 保存异常值(FP16)
    outlier_values = weight.clone()
    outlier_values[~outlier_mask] = 0

    # 量化剩余部分
    regular_weight = weight.clone()
    regular_weight[outlier_mask] = 0

    # 应用 Per-group 量化
    q_max = 2 ** (num_bits - 1) - 1
    group_size = 16

    rows, cols = regular_weight.shape
    regular_grouped = regular_weight.reshape(-1, group_size)

    max_abs = regular_grouped.abs().max(dim=1, keepdim=True)[0]
    scales = max_abs / q_max
    scales = torch.clamp(scales, min=1e-8)

    q = torch.clamp(torch.round(regular_grouped / scales), -q_max, q_max).to(torch.int8)
    regular_dequant = (scales * q.float()).reshape(rows, cols)

    # 最终重建
    reconstructed = regular_dequant + outlier_values

    error = (weight - reconstructed).abs().mean().item()

    # 计算内存占用
    outlier_memory = outlier_mask.sum().item() * 2  # FP16 = 2 bytes
    regular_memory = (~outlier_mask).sum().item() * (num_bits / 8)
    total_memory = outlier_memory + regular_memory
    original_memory = weight.numel() * weight.element_size()
    compression_ratio = original_memory / total_memory

    print(f"异常值比例: {outlier_mask.float().mean():.2%}")
    print(f"平均重建误差: {error:.6f}")
    print(f"压缩率: {compression_ratio:.2f}x")

    return q, scales, outlier_values, outlier_mask

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import psutil
import GPUtil

def benchmark_quantization(model, tokenizer, device="cuda", num_runs=50):
    """量化模型基准测试"""

    prompt = "The history of artificial intelligence began"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    # 内存占用
    if device == "cuda":
        torch.cuda.synchronize()
        gpu = GPUtil.getGPUs()[0]
        memory_used_gb = gpu.memoryUsed / 1024
    else:
        memory_used_gb = psutil.virtual_memory().used / 1e9

    # 预热
    with torch.no_grad():
        for _ in range(5):
            outputs = model.generate(
                **inputs,
                max_new_tokens=50,
                do_sample=False
            )

    # 测量速度
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.time()

    with torch.no_grad():
        for _ in range(num_runs):
            outputs = model.generate(
                **inputs,
                max_new_tokens=50,
                do_sample=False
            )

    if device == "cuda":
        torch.cuda.synchronize()
    elapsed = time.time() - start

    avg_time = elapsed / num_runs
    tokens_per_second = 50 / avg_time

    return {
        "memory_gb": memory_used_gb,
        "avg_time_ms": avg_time * 1000,
        "tokens_per_second": tokens_per_second
    }

# 结果示例(基于 A100 80GB，Llama-2-7B)
benchmark_results = {
    "FP16": {"memory_gb": 13.5, "tokens_per_second": 52.3, "ppl": 5.68},
    "INT8 (bitsandbytes)": {"memory_gb": 7.8, "tokens_per_second": 38.1, "ppl": 5.71},
    "INT4 GPTQ": {"memory_gb": 4.5, "tokens_per_second": 65.2, "ppl": 5.89},
    "INT4 AWQ": {"memory_gb": 4.3, "tokens_per_second": 68.7, "ppl": 5.86},
    "Q4_K_M (GGUF)": {"memory_gb": 4.1, "tokens_per_second": 45.2, "ppl": 5.91},  # CPU
    "INT4 NF4": {"memory_gb": 4.0, "tokens_per_second": 31.5, "ppl": 5.94},
}

print("=" * 80)
print(f"{'方法':<25} {'内存(GB)':<12} {'Tok/s':<12} {'PPL':<8}")
print("=" * 80)
for method, stats in benchmark_results.items():
    print(f"{method:<25} {stats['memory_gb']:<12.1f} {stats['tokens_per_second']:<12.1f} {stats['ppl']:<8.2f}")

def recommend_quantization(
    task: str,
    model_size_b: float,
    gpu_memory_gb: float,
    cpu_only: bool = False,
    fine_tuning_needed: bool = False
):
    """根据任务和环境推荐量化方案"""

    recommendations = []

    if cpu_only:
        recommendations.append({
            "method": "GGUF Q4_K_M",
            "reason": "针对 CPU 推理优化，基于 llama.cpp",
            "library": "llama-cpp-python"
        })
        return recommendations

    if fine_tuning_needed:
        recommendations.append({
            "method": "bitsandbytes NF4 + QLoRA",
            "reason": "可微调，额外约 4GB 内存即可训练 LoRA 适配器",
            "library": "bitsandbytes + peft"
        })
        return recommendations

    # 计算内存需求
    fp16_memory = model_size_b * 2  # FP16 = 每参数 2 bytes
    int8_memory = model_size_b * 1  # INT8 = 每参数 1 byte
    int4_memory = model_size_b * 0.5  # INT4 = 每参数 0.5 bytes

    if fp16_memory <= gpu_memory_gb * 0.8:
        recommendations.append({
            "method": "FP16(基础)",
            "reason": "内存充裕，质量最佳",
            "memory_gb": fp16_memory
        })

    if int8_memory <= gpu_memory_gb * 0.8:
        if task in ["chat", "completion", "summarization"]:
            recommendations.append({
                "method": "AWQ INT8",
                "reason": "质量与速度的最佳平衡",
                "library": "autoawq",
                "memory_gb": int8_memory
            })

    if int4_memory <= gpu_memory_gb * 0.8:
        recommendations.append({
            "method": "AWQ INT4",
            "reason": "高速推理，质量出色",
            "library": "autoawq",
            "memory_gb": int4_memory
        })
        recommendations.append({
            "method": "GPTQ INT4",
            "reason": "最佳 INT4 质量，量化速度较慢",
            "library": "auto-gptq",
            "memory_gb": int4_memory
        })

    return recommendations

# 使用示例
recommendations = recommend_quantization(
    task="chat",
    model_size_b=7.0,
    gpu_memory_gb=16.0,
    fine_tuning_needed=False
)

for rec in recommendations:
    print(f"\n方法: {rec['method']}")
    print(f"理由: {rec['reason']}")
    if 'library' in rec:
        print(f"库: {rec['library']}")
    if 'memory_gb' in rec:
        print(f"预计内存: {rec['memory_gb']:.1f} GB")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from awq import AutoAWQForCausalLM
import json
import os

class QuantizationPipeline:
    """统一的量化流水线"""

    def __init__(self, model_name: str, output_base_dir: str):
        self.model_name = model_name
        self.output_base_dir = output_base_dir
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        os.makedirs(output_base_dir, exist_ok=True)

    def quantize_gptq(self, bits: int = 4, group_size: int = 128):
        """GPTQ 量化"""
        output_dir = os.path.join(self.output_base_dir, f"gptq-{bits}bit")

        config = BaseQuantizeConfig(
            bits=bits,
            group_size=group_size,
            sym=True,
            desc_act=False
        )

        model = AutoGPTQForCausalLM.from_pretrained(
            self.model_name,
            quantize_config=config
        )

        # 校准数据
        calibration_data = self._prepare_calibration_data()

        model.quantize(calibration_data)
        model.save_quantized(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        print(f"GPTQ {bits}bit 已保存: {output_dir}")
        return output_dir

    def quantize_awq(self, bits: int = 4, group_size: int = 128):
        """AWQ 量化"""
        output_dir = os.path.join(self.output_base_dir, f"awq-{bits}bit")

        model = AutoAWQForCausalLM.from_pretrained(
            self.model_name,
            low_cpu_mem_usage=True
        )

        quant_config = {
            "zero_point": True,
            "q_group_size": group_size,
            "w_bit": bits,
            "version": "GEMM"
        }

        model.quantize(self.tokenizer, quant_config=quant_config)
        model.save_quantized(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        print(f"AWQ {bits}bit 已保存: {output_dir}")
        return output_dir

    def _prepare_calibration_data(self, num_samples: int = 128):
        """准备校准数据"""
        from datasets import load_dataset

        dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

        data = []
        for text in dataset["text"]:
            if len(text.strip()) > 50:
                encoded = self.tokenizer(
                    text.strip(),
                    return_tensors="pt",
                    max_length=2048,
                    truncation=True
                )
                data.append(encoded["input_ids"].squeeze())
                if len(data) >= num_samples:
                    break

        return data

    def evaluate_all(self, test_text: str = None):
        """评估所有量化模型"""
        if test_text is None:
            from datasets import load_dataset
            dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
            test_text = " ".join(dataset["text"][:10])

        results = {}

        # FP16 基准线
        model_fp16 = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16,
            device_map="auto"
        )

        # 计算 PPL
        from transformers import pipeline

        # 输出各模型的评估结果
        print("\n=== 量化评估结果 ===")
        print(f"模型: {self.model_name}")
        print(f"{'方法':<20} {'PPL':<10} {'内存(GB)':<12}")
        print("-" * 42)

        return results


# 执行完整流水线
pipeline = QuantizationPipeline(
    model_name="meta-llama/Llama-2-7b-hf",
    output_base_dir="./quantized_models"
)

# GPTQ 4bit 量化
gptq_dir = pipeline.quantize_gptq(bits=4)

# AWQ 4bit 量化
awq_dir = pipeline.quantize_awq(bits=4)