118_LLM模型量化与压缩：从理论到2025年实践技术详解

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

def quantize_weights(model, bits=8):
    """对模型权重进行量化
    
    Args:
        model: 要量化的PyTorch模型
        bits: 量化位宽
    
    Returns:
        quantized_model: 量化后的模型
    """
    # 保存原始模型的状态字典
    original_state_dict = model.state_dict()
    
    # 创建量化后的状态字典
    quantized_state_dict = {}
    
    # 量化范围
    qmin = -2 ** (bits - 1)
    qmax = 2 ** (bits - 1) - 1
    
    # 对每个权重进行量化
    for name, param in original_state_dict.items():
        if 'weight' in name:
            # 计算量化参数
            rmin = param.min().item()
            rmax = param.max().item()
            scale = (rmax - rmin) / (qmax - qmin)
            
            # 对称量化
            scale = max(abs(rmin), abs(rmax)) / (2 ** (bits - 1) - 1)
            zero_point = 0
            
            # 量化
            q_param = torch.round(param / scale + zero_point)
            q_param = torch.clamp(q_param, qmin, qmax)
            
            # 保存量化后的参数和量化参数
            quantized_state_dict[name + '_quantized'] = q_param.to(torch.int8)
            quantized_state_dict[name + '_scale'] = torch.tensor(scale)
            quantized_state_dict[name + '_zero_point'] = torch.tensor(zero_point)
        else:
            # 其他参数保持不变
            quantized_state_dict[name] = param
    
    # 创建量化后的模型
    quantized_model = AutoModelForCausalLM.from_config(model.config)
    
    # 这里需要注意：实际上需要实现自定义的量化模型来加载量化后的权重
    # 为了简化，这里我们只是演示量化过程
    
    return quantized_model, quantized_state_dict

def main():
    # 加载预训练模型
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    # 量化模型
    quantized_model, quantized_state_dict = quantize_weights(model, bits=8)
    
    print(f"原始模型参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
    print(f"量化后INT8权重大小约为: {sum(p.numel() for p in model.parameters() if 'weight' in p.name) / 4e6:.2f}MB")
    print("注意：这里只是演示量化过程，实际应用中需要实现自定义的量化模型")

if __name__ == "__main__":
    main()

import torch
import numpy as np
from transformers import AutoModelForCausalLM

def gptq_quantize_block(block, bits=4, percdamp=0.01):
    """对单个权重块进行GPTQ量化
    
    Args:
        block: 要量化的权重块 (shape: [out_features, in_features])
        bits: 量化位宽
        percdamp: 阻尼百分比，用于数值稳定性
    
    Returns:
        q_block: 量化后的权重块
        scales: 量化比例
        zeros: 量化零点
    """
    out_features, in_features = block.shape
    
    # 计算量化参数
    n = 2 ** bits - 1
    
    # 计算缩放因子
    damp = percdamp * block.abs().mean()
    
    # 初始化Hessian近似
    H = torch.zeros((in_features, in_features), device=block.device)
    
    # 量化结果
    q_block = torch.zeros_like(block, dtype=torch.int8)
    scales = torch.zeros(out_features, device=block.device)
    zeros = torch.zeros(out_features, device=block.device)
    
    # 逐列量化
    for i in range(out_features):
        # 获取当前行
        row = block[i, :].unsqueeze(0)
        
        # 计算量化范围
        row_min = row.min().item()
        row_max = row.max().item()
        
        # 对称量化
        scale = max(abs(row_min), abs(row_max)) / (2 ** (bits - 1) - 1)
        
        if scale == 0:
            q_row = torch.zeros_like(row, dtype=torch.int8)
        else:
            # 量化
            q_row = torch.round(row / scale)
            q_row = torch.clamp(q_row, -2 ** (bits - 1), 2 ** (bits - 1) - 1).to(torch.int8)
        
        # 保存结果
        q_block[i, :] = q_row.squeeze(0)
        scales[i] = scale
        zeros[i] = 0  # 对称量化，零点为0
    
    return q_block, scales, zeros

def gptq_quantize_model(model, bits=4):
    """对模型进行GPTQ量化
    
    Args:
        model: 要量化的PyTorch模型
        bits: 量化位宽
    
    Returns:
        quantized_weights: 量化后的权重字典
    """
    quantized_weights = {}
    
    # 遍历模型的所有参数
    for name, param in model.named_parameters():
        if 'weight' in name and param.ndim == 2:  # 仅量化2D权重
            print(f"量化: {name}")
            
            # 对权重块进行量化
            q_weight, scales, zeros = gptq_quantize_block(param.data, bits=bits)
            
            # 保存量化结果
            quantized_weights[name + '_quantized'] = q_weight
            quantized_weights[name + '_scales'] = scales
            quantized_weights[name + '_zeros'] = zeros
        else:
            # 其他参数保持不变
            quantized_weights[name] = param.data
    
    return quantized_weights

def main():
    # 加载模型（为了演示，使用小型模型）
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    # 量化模型
    quantized_weights = gptq_quantize_model(model, bits=4)
    
    print("量化完成！")
    print(f"注意：这是简化的GPTQ实现，实际实现需要更复杂的优化和Hessian近似")

if __name__ == "__main__":
    main()

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def estimate_importance(model, tokenizer, calibration_texts, n_samples=128):
    """估计权重的重要性
    
    Args:
        model: 要分析的模型
        tokenizer: 分词器
        calibration_texts: 校准文本
        n_samples: 样本数量
    
    Returns:
        importance: 权重重要性字典
    """
    importance = {}
    
    # 选择校准样本
    samples = calibration_texts[:n_samples]
    
    # 启用梯度计算
    for param in model.parameters():
        param.requires_grad = True
    
    # 遍历校准样本
    for text in samples:
        # 编码文本
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        
        # 前向传播
        outputs = model(**inputs, labels=inputs["input_ids"])
        loss = outputs.loss
        
        # 反向传播计算梯度
        loss.backward()
    
    # 计算权重重要性（使用梯度的绝对值作为重要性度量）
    for name, param in model.named_parameters():
        if param.grad is not None:
            importance[name] = param.grad.abs().detach()
    
    # 禁用梯度计算
    for param in model.parameters():
        param.requires_grad = False
    
    return importance

def awq_quantize_model(model, importance, bits=4, keep_ratio=0.1):
    """使用AWQ方法量化模型
    
    Args:
        model: 要量化的模型
        importance: 权重重要性字典
        bits: 量化位宽
        keep_ratio: 保留高精度的权重比例
    
    Returns:
        quantized_weights: 量化后的权重字典
    """
    quantized_weights = {}
    
    # 遍历模型的所有参数
    for name, param in model.named_parameters():
        if 'weight' in name and param.ndim == 2 and name in importance:
            print(f"量化: {name}")
            
            # 获取权重重要性
            imp = importance[name]
            
            # 确定要保留的权重索引
            num_keep = int(param.numel() * keep_ratio)
            _, keep_idx = torch.topk(imp.view(-1), num_keep)
            
            # 量化参数
            qmin = -2 ** (bits - 1)
            qmax = 2 ** (bits - 1) - 1
            
            # 创建量化后的权重
            q_weight = torch.zeros_like(param, dtype=torch.int8)
            scales = torch.ones(param.shape[0], device=param.device)
            
            # 逐行量化
            for i in range(param.shape[0]):
                row = param[i, :]
                row_imp = imp[i, :]
                
                # 计算量化范围（排除重要权重）
               不重要_idx = torch.ones_like(row_imp, dtype=bool)
               不重要_idx.view(-1)[keep_idx[keep_idx < row.numel()]] = False
                
                if不重要_idx.any():
                    # 使用不重要权重计算缩放因子
                    row_min = row[不重要_idx].min().item()
                    row_max = row[不重要_idx].max().item()
                    scale = (row_max - row_min) / (qmax - qmin)
                    
                    if scale == 0:
                        scale = 1e-8
                    
                    scales[i] = scale
                    
                    # 量化整行
                    q_row = torch.round(row / scale)
                    q_row = torch.clamp(q_row, qmin, qmax).to(torch.int8)
                    
                    # 保留重要权重的原始精度（通过调整量化后的值实现）
                    for idx in keep_idx:
                        if idx < row.numel() and idx // param.shape[1] == i:
                            col_idx = idx % param.shape[1]
                            q_row[col_idx] = torch.round(row[col_idx] / scale)
                
                q_weight[i, :] = q_row
            
            # 保存量化结果
            quantized_weights[name + '_quantized'] = q_weight
            quantized_weights[name + '_scales'] = scales
            quantized_weights[name + '_keep_idx'] = keep_idx
        else:
            # 其他参数保持不变
            quantized_weights[name] = param.data
    
    return quantized_weights

def main():
    # 加载模型和分词器
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 校准文本（实际应用中应使用更多样的文本）
    calibration_texts = ["这是一个校准样本。" for _ in range(32)]
    
    print("估计权重重要性...")
    importance = estimate_importance(model, tokenizer, calibration_texts)
    
    print("使用AWQ量化模型...")
    quantized_weights = awq_quantize_model(model, importance, bits=4)
    
    print("量化完成！")
    print(f"注意：这是简化的AWQ实现，实际实现需要更复杂的重要性评估和量化策略")

if __name__ == "__main__":
    main()

import torch
import numpy as np
from transformers import AutoModelForCausalLM

def pq_plus_quantize_weight(weight, num_subspaces=8, codebook_size=16):
    """使用PQ+量化单个权重矩阵
    
    Args:
        weight: 权重矩阵 (shape: [out_features, in_features])
        num_subspaces: 子空间数量
        codebook_size: 码本大小
    
    Returns:
        codes: 量化后的编码
        codebooks: 码本
        indices: 子空间划分索引
    """
    out_features, in_features = weight.shape
    
    # 确保输入特征维度可以被子空间数量整除
    assert in_features % num_subspaces == 0, "in_features must be divisible by num_subspaces"
    subspace_dim = in_features // num_subspaces
    
    # 初始化码本
    codebooks = torch.zeros((num_subspaces, codebook_size, subspace_dim), device=weight.device)
    codes = torch.zeros((out_features, num_subspaces), dtype=torch.int8, device=weight.device)
    
    # 逐个子空间进行量化
    for i in range(num_subspaces):
        # 提取子空间数据
        start_idx = i * subspace_dim
        end_idx = (i + 1) * subspace_dim
        subspace_data = weight[:, start_idx:end_idx]
        
        # 使用K-means聚类创建码本
        # 简化版实现，使用随机初始化的码本
        # 实际应用中应使用更复杂的聚类算法
        codebook = torch.randn((codebook_size, subspace_dim), device=weight.device)
        
        # 量化：找到每个向量最近的码本向量
        for j in range(out_features):
            vec = subspace_data[j, :]
            distances = torch.norm(codebook - vec, dim=1)
            code_idx = torch.argmin(distances)
            codes[j, i] = code_idx
        
        # 更新码本（简化版）
        for k in range(codebook_size):
            indices = codes[:, i] == k
            if indices.any():
                codebook[k, :] = subspace_data[indices, :].mean(dim=0)
        
        codebooks[i, :, :] = codebook
    
    # 生成子空间划分索引
    indices = torch.arange(0, in_features, device=weight.device)
    indices = indices.reshape(num_subspaces, subspace_dim)
    
    return codes, codebooks, indices

def pq_plus_quantize_model(model, num_subspaces=8, codebook_size=16):
    """使用PQ+量化模型
    
    Args:
        model: 要量化的模型
        num_subspaces: 子空间数量
        codebook_size: 码本大小
    
    Returns:
        quantized_weights: 量化后的权重字典
    """
    quantized_weights = {}
    
    # 遍历模型的所有参数
    for name, param in model.named_parameters():
        if 'weight' in name and param.ndim == 2:
            print(f"量化: {name}")
            
            # 对权重进行PQ+量化
            codes, codebooks, indices = pq_plus_quantize_weight(
                param.data, 
                num_subspaces=num_subspaces, 
                codebook_size=codebook_size
            )
            
            # 保存量化结果
            quantized_weights[name + '_codes'] = codes
            quantized_weights[name + '_codebooks'] = codebooks
            quantized_weights[name + '_indices'] = indices
        else:
            # 其他参数保持不变
            quantized_weights[name] = param.data
    
    return quantized_weights

def main():
    # 加载模型
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    # 量化模型
    quantized_weights = pq_plus_quantize_model(model, num_subspaces=8, codebook_size=16)
    
    print("量化完成！")
    print(f"注意：这是简化的PQ+实现，实际实现需要更复杂的聚类算法和优化策略")

if __name__ == "__main__":
    main()

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

class DynamicQuantizationWrapper(nn.Module):
    """动态量化包装器
    
    此包装器为模型添加动态量化功能，在推理过程中根据输入动态调整量化参数。
    """
    def __init__(self, model, quantize_bits=8, dynamic_range=True):
        super().__init__()
        self.model = model
        self.quantize_bits = quantize_bits
        self.dynamic_range = dynamic_range
        self.qmin = -2 ** (quantize_bits - 1)
        self.qmax = 2 ** (quantize_bits - 1) - 1
        
        # 存储激活量化参数
        self.activation_scales = {}
        self.activation_zeros = {}
        
        # 对模型权重进行量化
        self._quantize_weights()
    
    def _quantize_weights(self):
        """量化模型权重"""
        # 保存量化后的权重
        self.quantized_weights = {}
        self.weight_scales = {}
        self.weight_zeros = {}
        
        # 遍历模型的所有参数
        for name, param in self.model.named_parameters():
            if 'weight' in name:
                # 计算量化参数
                rmin = param.min().item()
                rmax = param.max().item()
                
                # 对称量化
                scale = max(abs(rmin), abs(rmax)) / (2 ** (self.quantize_bits - 1) - 1)
                zero_point = 0
                
                # 量化
                q_param = torch.round(param / scale + zero_point)
                q_param = torch.clamp(q_param, self.qmin, self.qmax)
                
                # 保存量化后的参数和量化参数
                self.quantized_weights[name] = q_param.to(torch.int8)
                self.weight_scales[name] = scale
                self.weight_zeros[name] = zero_point
    
    def _dynamic_quantize_activation(self, activation, name):
        """动态量化激活值"""
        # 计算量化参数
        rmin = activation.min().item()
        rmax = activation.max().item()
        
        # 对称量化
        scale = max(abs(rmin), abs(rmax)) / (2 ** (self.quantize_bits - 1) - 1)
        if scale == 0:
            scale = 1e-8
        zero_point = 0
        
        # 保存量化参数
        self.activation_scales[name] = scale
        self.activation_zeros[name] = zero_point
        
        # 量化
        q_activation = torch.round(activation / scale + zero_point)
        q_activation = torch.clamp(q_activation, self.qmin, self.qmax)
        
        # 反量化（在实际应用中，这里应该直接使用量化值进行计算）
        dequant_activation = (q_activation - zero_point) * scale
        
        return dequant_activation
    
    def forward(self, **inputs):
        """前向传播，动态量化激活值"""
        # 这里需要修改模型的前向传播以支持动态量化
        # 为了简化，我们仅返回原始模型的输出
        # 实际实现中，需要修改模型的中间层以插入动态量化操作
        
        # 保存原始的前向传播函数
        original_forward = self.model.forward
        
        # 在推理过程中，我们可以通过钩子（hooks）来量化激活值
        # 这里省略具体实现
        
        # 调用原始前向传播
        outputs = original_forward(**inputs)
        
        return outputs

def main():
    # 加载模型和分词器
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 创建动态量化模型
    quantized_model = DynamicQuantizationWrapper(model, quantize_bits=8)
    
    # 准备输入
    text = "这是一个测试文本。"
    inputs = tokenizer(text, return_tensors="pt")
    
    # 推理
    with torch.no_grad():
        outputs = quantized_model(**inputs)
    
    print("动态量化模型推理完成！")
    print(f"注意：这是简化的动态量化实现，实际实现需要修改模型的前向传播以支持中间激活值的量化")

if __name__ == "__main__":
    main()

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

def evaluate_attention_heads_importance(model, tokenizer, calibration_texts):
    """评估注意力头的重要性
    
    Args:
        model: 要评估的模型
        tokenizer: 分词器
        calibration_texts: 校准文本
    
    Returns:
        head_importance: 注意力头重要性字典
    """
    head_importance = {}
    
    # 启用梯度计算
    for param in model.parameters():
        param.requires_grad = True
    
    # 遍历校准文本
    for text in calibration_texts:
        # 编码文本
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        
        # 前向传播
        outputs = model(**inputs, labels=inputs["input_ids"])
        loss = outputs.loss
        
        # 反向传播计算梯度
        loss.backward()
    
    # 计算注意力头重要性
    # 这里使用注意力头权重的梯度范数作为重要性度量
    for name, param in model.named_parameters():
        if 'attn' in name and 'weight' in name and param.grad is not None:
            # 假设参数形状为 [hidden_size, ...]
            # 计算每个头的重要性
            hidden_size = param.shape[0]
            num_heads = model.config.num_attention_heads
            head_size = hidden_size // num_heads
            
            # 重塑参数以分离不同的头
            reshaped_param = param.reshape(num_heads, head_size, -1)
            reshaped_grad = param.grad.reshape(num_heads, head_size, -1)
            
            # 计算每个头的梯度范数
            head_grads = torch.norm(reshaped_grad, dim=(1, 2))
            
            # 保存重要性分数
            layer_idx = name.split('.')[2]  # 假设name格式为 'transformer.h.{layer_idx}.attn...'
            if layer_idx not in head_importance:
                head_importance[layer_idx] = torch.zeros(num_heads, device=param.device)
            head_importance[layer_idx] += head_grads
    
    # 禁用梯度计算
    for param in model.parameters():
        param.requires_grad = False
    
    return head_importance

def prune_attention_heads(model, head_importance, prune_ratio=0.2):
    """剪枝注意力头
    
    Args:
        model: 要剪枝的模型
        head_importance: 注意力头重要性字典
        prune_ratio: 剪枝比例
    
    Returns:
        pruned_model: 剪枝后的模型
        pruned_heads: 被剪枝的头索引
    """
    pruned_heads = {}
    
    # 遍历每一层
    for layer_idx, importance in head_importance.items():
        # 确定要剪枝的头数量
        num_heads = len(importance)
        num_prune = int(num_heads * prune_ratio)
        
        if num_prune > 0:
            # 选择重要性最低的头
            _, indices = torch.topk(importance, k=num_heads - num_prune, largest=True)
            prune_indices = [i for i in range(num_heads) if i not in indices.tolist()]
            
            pruned_heads[int(layer_idx)] = prune_indices
    
    # 应用剪枝
    # 注意：实际应用中，这里需要使用模型库提供的剪枝API
    # 以下是示意代码
    pruned_model = model
    
    # 记录剪枝信息
    pruned_model.config.pruned_heads = pruned_heads
    
    return pruned_model, pruned_heads

def main():
    # 加载模型和分词器
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # 校准文本
    calibration_texts = ["这是一个校准样本。" for _ in range(32)]
    
    print("评估注意力头重要性...")
    importance = evaluate_attention_heads_importance(model, tokenizer, calibration_texts)
    
    print("剪枝注意力头...")
    pruned_model, pruned_heads = prune_attention_heads(model, importance, prune_ratio=0.2)
    
    print(f"剪枝完成！剪枝的注意力头: {pruned_heads}")
    print(f"注意：这是简化的注意力头剪枝实现，实际应用中需要使用模型库提供的剪枝API")

if __name__ == "__main__":
    main()

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM, AutoTokenizer, get_linear_schedule_with_warmup

def distill_model(teacher_model, student_model, tokenizer, train_texts, epochs=3, batch_size=4, lr=5e-5, temperature=2.0, alpha=0.5):
    """知识蒸馏函数
    
    Args:
        teacher_model: 教师模型
        student_model: 学生模型
        tokenizer: 分词器
        train_texts: 训练文本
        epochs: 训练轮数
        batch_size: 批次大小
        lr: 学习率
        temperature: 蒸馏温度
        alpha: 软目标权重
    
    Returns:
        distilled_model: 蒸馏后的学生模型
    """
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    teacher_model.to(device)
    student_model.to(device)
    
    # 设置模型模式
    teacher_model.eval()  # 教师模型保持评估模式
    student_model.train()  # 学生模型设置为训练模式
    
    # 准备数据加载器
    # 简化版，实际应用中应使用DataLoader
    def get_batch(texts, batch_size):
        for i in range(0, len(texts), batch_size):
            yield texts[i:i+batch_size]
    
    # 定义优化器和学习率调度器
    optimizer = optim.AdamW(student_model.parameters(), lr=lr)
    total_steps = (len(train_texts) // batch_size) * epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=0, num_training_steps=total_steps
    )
    
    # 定义损失函数
    soft_loss_fn = nn.KLDivLoss(reduction="batchmean")
    hard_loss_fn = nn.CrossEntropyLoss()
    
    # 训练循环
    for epoch in range(epochs):
        epoch_loss = 0
        
        for batch_texts in get_batch(train_texts, batch_size):
            # 准备输入
            inputs = tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True)
            inputs = {k: v.to(device) for k, v in inputs.items()}
            
            # 教师模型前向传播（不计算梯度）
            with torch.no_grad():
                teacher_logits = teacher_model(**inputs).logits
                teacher_probs = nn.functional.softmax(teacher_logits / temperature, dim=-1)
            
            # 学生模型前向传播
            student_outputs = student_model(**inputs)
            student_logits = student_outputs.logits
            student_probs = nn.functional.log_softmax(student_logits / temperature, dim=-1)
            
            # 计算软目标损失
            soft_loss = soft_loss_fn(student_probs, teacher_probs) * (temperature ** 2)
            
            # 计算硬目标损失
            labels = inputs["input_ids"]
            shift_logits = student_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            hard_loss = hard_loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            
            # 组合损失
            loss = alpha * soft_loss + (1 - alpha) * hard_loss
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()
            
            epoch_loss += loss.item()
        
        print(f"Epoch {epoch+1}/{epochs}, Loss: {epoch_loss / (len(train_texts) // batch_size)}")
    
    return student_model

def main():
    # 加载教师模型和分词器
    teacher_model_name = "gpt2-large"
    teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model_name)
    
    # 创建学生模型（较小的模型）
    student_model_name = "gpt2"
    student_model = AutoModelForCausalLM.from_pretrained(student_model_name)
    
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
    
    # 准备训练数据（简化版，实际应用中应使用更多样的数据集）
    train_texts = ["这是一个训练样本。" for _ in range(128)]
    
    print("开始知识蒸馏...")
    distilled_model = distill_model(
        teacher_model, student_model, tokenizer, train_texts,
        epochs=3, batch_size=4, lr=5e-5, temperature=2.0, alpha=0.5
    )
    
    print("知识蒸馏完成！")
    print(f"注意：这是简化的知识蒸馏实现，实际应用中需要更复杂的训练策略和更大的数据集")

if __name__ == "__main__":
    main()

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM

def low_rank_decomposition(weight, rank_ratio=0.5):
    """低秩分解函数
    
    Args:
        weight: 要分解的权重矩阵
        rank_ratio: 目标秩与原始矩阵最小维度的比例
    
    Returns:
        u: 左低秩矩阵
        v: 右低秩矩阵
    """
    # 执行奇异值分解
    u, s, v = torch.svd_lowrank(weight, q=min(weight.shape))
    
    # 计算目标秩
    target_rank = int(min(weight.shape) * rank_ratio)
    
    # 截断奇异值和奇异向量
    u_truncated = u[:, :target_rank]
    s_truncated = s[:target_rank]
    v_truncated = v[:, :target_rank]
    
    # 重构低秩矩阵
    u_scaled = u_truncated @ torch.diag(torch.sqrt(s_truncated))
    v_scaled = torch.diag(torch.sqrt(s_truncated)) @ v_truncated.t()
    
    return u_scaled, v_scaled

def apply_low_rank_decomposition(model, rank_ratio=0.5):
    """对模型应用低秩分解
    
    Args:
        model: 要分解的模型
        rank_ratio: 目标秩与原始矩阵最小维度的比例
    
    Returns:
        decomposed_model: 分解后的模型
    """
    decomposed_model = model
    
    # 遍历模型的所有层
    for name, module in list(decomposed_model.named_modules()):
        if hasattr(module, 'mlp') and hasattr(module.mlp, 'c_fc') and hasattr(module.mlp, 'c_proj'):
            # 对前馈网络的第一层应用低秩分解
            fc_weight = module.mlp.c_fc.weight.data
            
            # 执行低秩分解
            u, v = low_rank_decomposition(fc_weight, rank_ratio)
            
            # 创建新的低秩层
            in_features = fc_weight.shape[1]
            out_features = fc_weight.shape[0]
            rank = u.shape[1]
            
            # 替换原始层
            module.mlp.c_fc = nn.Sequential(
                nn.Linear(in_features, rank, bias=False),
                nn.Linear(rank, out_features, bias=module.mlp.c_fc.bias is not None)
            )
            
            # 初始化权重
            module.mlp.c_fc[0].weight.data = u
            module.mlp.c_fc[1].weight.data = v
            
            if module.mlp.c_fc[1].bias is not None:
                module.mlp.c_fc[1].bias.data = module.mlp.c_fc.bias.data
            
            print(f"分解层: {name}.mlp.c_fc")
            print(f"原始维度: {fc_weight.shape}, 分解维度: ({u.shape}, {v.shape})")
    
    return decomposed_model

def main():
    # 加载模型
    model_name = "gpt2"
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    # 计算原始模型参数量
    original_params = sum(p.numel() for p in model.parameters())
    
    # 应用低秩分解
    decomposed_model = apply_low_rank_decomposition(model, rank_ratio=0.5)
    
    # 计算分解后模型参数量
    decomposed_params = sum(p.numel() for p in decomposed_model.parameters())
    
    # 计算压缩率
    compression_ratio = original_params / decomposed_params
    
    print(f"原始模型参数量: {original_params / 1e6:.2f}M")
    print(f"分解后模型参数量: {decomposed_params / 1e6:.2f}M")
    print(f"压缩率: {compression_ratio:.2f}x")
    print(f"注意：这是简化的低秩分解实现，实际应用中需要更复杂的分解策略和微调过程")

if __name__ == "__main__":
    main()

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoConfig

def create_shared_parameter_model(base_model_name, share_layers=True, share_attention=True):
    """创建共享参数模型
    
    Args:
        base_model_name: 基础模型名称
        share_layers: 是否共享层参数
        share_attention: 是否共享注意力头参数
    
    Returns:
        shared_model: 共享参数模型
    """
    # 加载基础配置
    config = AutoConfig.from_pretrained(base_model_name)
    
    # 创建模型
    model = AutoModelForCausalLM.from_pretrained(base_model_name)
    
    if share_layers:
        print("共享层参数...")
        # 获取第一层作为共享层模板
        if hasattr(model, 'transformer') and hasattr(model.transformer, 'h'):
            layers = model.transformer.h
            
            # 对于偶数层，共享奇数层的参数
            for i in range(1, len(layers), 2):
                layers[i].load_state_dict(layers[i-1].state_dict(), strict=False)
                print(f"层 {i} 共享层 {i-1} 的参数")
    
    if share_attention:
        print("共享注意力头参数...")
        # 遍历每一层，共享注意力头参数
        if hasattr(model, 'transformer') and hasattr(model.transformer, 'h'):
            for layer_idx, layer in enumerate(model.transformer.h):
                if hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
                    # 获取注意力投影层
                    attn_proj = layer.attn.c_attn
                    
                    # 假设注意力头数量为 config.num_attention_heads
                    num_heads = config.num_attention_heads
                    hidden_size = config.hidden_size
                    head_size = hidden_size // num_heads
                    
                    # 确保维度正确
                    if attn_proj.weight.shape[0] == 3 * hidden_size:  # Q, K, V 投影
                        # 共享Q、K、V投影的参数
                        # 获取Q投影权重
                        q_weight = attn_proj.weight[:hidden_size, :].clone()
                        
                        # 将K、V投影权重设置为Q投影权重
                        attn_proj.weight.data[hidden_size:2*hidden_size, :] = q_weight
                        attn_proj.weight.data[2*hidden_size:, :] = q_weight
                        
                        # 如果有偏置，也共享偏置
                        if attn_proj.bias is not None:
                            q_bias = attn_proj.bias[:hidden_size].clone()
                            attn_proj.bias.data[hidden_size:2*hidden_size] = q_bias
                            attn_proj.bias.data[2*hidden_size:] = q_bias
                        
                        print(f"层 {layer_idx} 的注意力头参数已共享")
    
    return model

def main():
    # 模型名称
    model_name = "gpt2"
    
    # 创建共享参数模型
    shared_model = create_shared_parameter_model(model_name, share_layers=True, share_attention=True)
    
    # 计算原始模型参数量（理论值）
    original_params = 124000000  # GPT-2 参数量约为1.24亿
    
    # 计算共享参数后的参数量（近似值）
    # 假设共享了50%的层参数和注意力头参数
    shared_params = original_params * 0.75  # 简化计算
    
    print(f"估计原始模型参数量: {original_params / 1e6:.2f}M")
    print(f"估计共享参数后模型参数量: {shared_params / 1e6:.2f}M")
    print(f"参数减少比例: {(1 - shared_params / original_params) * 100:.2f}%")
    print(f"注意：这是简化的参数共享实现，实际应用中需要根据模型架构调整共享策略")

if __name__ == "__main__":
    main()

# 安装GPTQ工具
pip install git+https://github.com/IST-DASLab/gptq

# 安装依赖
pip install torch transformers accelerate datasets

from transformers import AutoModelForCausalLM, AutoTokenizer
from gptq import GPTQQuantizer

# 加载模型和分词器
model_name = "facebook/opt-1.3b"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 创建量化器实例
quantizer = GPTQQuantizer(
    bits=4,           # 量化位宽
    group_size=128,   # 分组大小
    damp_percent=0.01, # 阻尼百分比
    desc_act=True     # 是否使用描述性激活量化
)

# 准备校准数据
calibration_texts = ["这是校准样本1。", "这是校准样本2。", ...]
calibration_tokens = tokenizer(calibration_texts, return_tensors="pt", padding=True, truncation=True)

# 执行量化
quantized_model = quantizer.quantize(model, calibration_tokens["input_ids"])

# 保存量化后的模型
quantized_model.save_pretrained("opt-1.3b-gptq-4bit")
tokenizer.save_pretrained("opt-1.3b-gptq-4bit")

# 安装AWQ工具
pip install git+https://github.com/mit-han-lab/llm-awq

# 安装依赖
pip install torch transformers accelerate

from transformers import AutoModelForCausalLM, AutoTokenizer
from awq import AutoAWQForCausalLM

# 加载模型和分词器
model_name = "facebook/opt-1.3b"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 加载并量化模型
model = AutoAWQForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=False,  # 不需要先加载8位模型
    torch_dtype=torch.float16
)

# 量化模型
model.quantize(
    tokenizer,
    bits=4,          # 量化位宽
    group_size=128,  # 分组大小
    zero_point=True  # 是否使用零点量化
)

# 保存量化后的模型
model.save_quantized("opt-1.3b-awq-4bit")
tokenizer.save_pretrained("opt-1.3b-awq-4bit")

import torch
import torch.quantization
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 将模型设置为评估模式
model.eval()

# 应用动态量化
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 要量化的层类型
    dtype=torch.qint8   # 量化数据类型
)

# 保存量化后的模型
torch.save(quantized_model.state_dict(), "gpt2-dynamic-quantized.pth")

import torch
import torch.quantization
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 准备模型进行静态量化
model.eval()
model.qconfig = torch.quantization.get_default_qconfig("fbgemm")
model_prepared = torch.quantization.prepare(model)

# 校准模型
calibration_texts = ["这是校准样本1。", "这是校准样本2。", ...]
for text in calibration_texts:
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        model_prepared(**inputs)

# 转换为量化模型
quantized_model = torch.quantization.convert(model_prepared)

# 保存量化后的模型
torch.save(quantized_model.state_dict(), "gpt2-static-quantized.pth")

# 安装TensorRT
# 请根据您的CUDA版本和平台下载相应的TensorRT安装包
# 参考NVIDIA官方文档：https://developer.nvidia.com/tensorrt

# 安装PyTorch TensorRT集成
pip install torch-tensorrt

import torch
import tensorrt as trt
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 将模型转换为ONNX格式
dummy_input = tokenizer("这是一个测试输入。", return_tensors="pt")
input_names = ["input_ids", "attention_mask"]
output_names = ["logits"]
torch.onnx.export(
    model,
    tuple(dummy_input.values()),
    "gpt2.onnx",
    opset_version=15,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size", 1: "seq_len"}
    }
)

# 创建TensorRT引擎
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析ONNX文件
with open("gpt2.onnx", "rb") as model:
    parser.parse(model.read())

# 配置生成器
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB

# 启用INT8量化
config.set_flag(trt.BuilderFlag.INT8)

# 设置校准器
# 注意：这里需要实现一个自定义的校准器类
# 请参考TensorRT文档实现校准器

# 构建引擎
serialized_engine = builder.build_serialized_network(network, config)

# 保存引擎
with open("gpt2.trt", "wb") as f:
    f.write(serialized_engine)