机器学习基础系列——LoRA与参数高效微调

import numpy as np

print("全量微调的问题:")
print("=" * 50)
print("• 存储开销: 每个任务需要完整模型副本")
print("• 计算成本: 更新所有参数需要大量显存")
print("• 灾难遗忘: 可能破坏预训练知识")
print()

# 模型参数量示例
models = {
    'BERT-base': 110e6,
    'BERT-large': 340e6,
    'GPT-2': 1.5e9,
    'LLaMA-7B': 7e9,
    'LLaMA-65B': 65e9,
}

print("模型参数量:")
for name, params in models.items():
    memory = params * 4 / 1e9  # float32
    print(f"  {name}: {params/1e6:.0f}M 参数, 约{memory:.1f}GB 显存")

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    
    class LoRALayer(nn.Module):
        """LoRA层实现"""
        
        def __init__(self, in_features, out_features, rank=4, alpha=1.0):
            super().__init__()
            
            self.rank = rank
            self.alpha = alpha
            self.scaling = alpha / rank
            
            # 原始权重（冻结）
            self.weight = nn.Parameter(
                torch.randn(out_features, in_features) * 0.02,
                requires_grad=False
            )
            
            # LoRA分解矩阵
            self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.02)
            self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
            
            # 只有LoRA参数需要梯度
        
        def forward(self, x):
            # 原始输出
            out = F.linear(x, self.weight)
            
            # LoRA增量
            lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B)
            
            return out + lora_out * self.scaling
        
        def merge_weights(self):
            """合并权重用于推理"""
            merged = self.weight + self.lora_B @ self.lora_A * self.scaling
            return merged
        
        def num_trainable_params(self):
            """可训练参数数量"""
            return self.lora_A.numel() + self.lora_B.numel()
        
        def num_total_params(self):
            """总参数数量"""
            return self.weight.numel()
    
    # 测试
    in_dim, out_dim = 768, 768
    rank = 8
    
    lora = LoRALayer(in_dim, out_dim, rank=rank)
    
    total = lora.num_total_params()
    trainable = lora.num_trainable_params()
    
    print(f"原始参数: {total:,}")
    print(f"LoRA参数: {trainable:,}")
    print(f"参数比例: {trainable/total*100:.2f}%")
    
except ImportError:
    print("PyTorch未安装")

class LoRALayerNumpy:
    """NumPy实现的LoRA层"""
    
    def __init__(self, in_features, out_features, rank=4, alpha=1.0):
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        
        # 原始权重
        self.weight = np.random.randn(out_features, in_features) * 0.02
        
        # LoRA矩阵
        self.lora_A = np.random.randn(rank, in_features) * 0.02
        self.lora_B = np.zeros((out_features, rank))
    
    def forward(self, x):
        # x: [batch, in_features]
        # 原始输出
        out = x @ self.weight.T
        
        # LoRA增量: x -> A -> B
        lora_out = (x @ self.lora_A.T) @ self.lora_B.T
        
        return out + lora_out * self.scaling
    
    def merge_weights(self):
        """合并权重"""
        delta_W = self.lora_B @ self.lora_A * self.scaling
        return self.weight + delta_W

# 测试
lora_np = LoRALayerNumpy(768, 768, rank=8)
x = np.random.randn(4, 768)
out = lora_np.forward(x)
print(f"输出形状: {out.shape}")

# 参数效率
total_params = 768 * 768
lora_params = 8 * 768 * 2
print(f"参数节省: {(1 - lora_params/total_params)*100:.1f}%")

try:
    class Adapter(nn.Module):
        """Adapter模块"""
        
        def __init__(self, hidden_size, bottleneck_size=64):
            super().__init__()
            
            self.down_proj = nn.Linear(hidden_size, bottleneck_size)
            self.up_proj = nn.Linear(bottleneck_size, hidden_size)
            self.act = nn.GELU()
        
        def forward(self, x):
            # 下投影 -> 激活 -> 上投影 -> 残差
            residual = x
            x = self.down_proj(x)
            x = self.act(x)
            x = self.up_proj(x)
            return x + residual
    
    
    class TransformerWithAdapter(nn.Module):
        """带Adapter的Transformer层"""
        
        def __init__(self, hidden_size, num_heads, bottleneck_size=64):
            super().__init__()
            
            # 原始组件（冻结）
            self.self_attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
            self.ffn = nn.Sequential(
                nn.Linear(hidden_size, hidden_size * 4),
                nn.GELU(),
                nn.Linear(hidden_size * 4, hidden_size)
            )
            self.norm1 = nn.LayerNorm(hidden_size)
            self.norm2 = nn.LayerNorm(hidden_size)
            
            # Adapter（可训练）
            self.adapter_attn = Adapter(hidden_size, bottleneck_size)
            self.adapter_ffn = Adapter(hidden_size, bottleneck_size)
        
        def forward(self, x):
            # Self-attention + Adapter
            attn_out, _ = self.self_attn(x, x, x)
            attn_out = self.adapter_attn(attn_out)  # Adapter
            x = self.norm1(x + attn_out)
            
            # FFN + Adapter
            ffn_out = self.ffn(x)
            ffn_out = self.adapter_ffn(ffn_out)  # Adapter
            x = self.norm2(x + ffn_out)
            
            return x
    
    print("Adapter特点:")
    print("  • 插入到Transformer层中")
    print("  • 瓶颈结构压缩参数")
    print("  • 残差连接保持稳定性")
    
except NameError:
    print("需要先导入PyTorch")

try:
    class PrefixTuning(nn.Module):
        """Prefix Tuning实现"""
        
        def __init__(self, num_layers, hidden_size, prefix_length=10):
            super().__init__()
            
            self.prefix_length = prefix_length
            
            # 可训练的前缀
            # [num_layers, 2, prefix_length, hidden_size]
            # 2表示key和value
            self.prefix_tokens = nn.Parameter(
                torch.randn(num_layers, 2, prefix_length, hidden_size) * 0.02
            )
        
        def get_prefix(self, batch_size, layer_idx):
            """获取特定层的前缀"""
            # [2, prefix_length, hidden_size]
            prefix = self.prefix_tokens[layer_idx]
            
            # 扩展到batch
            # [batch, 2, prefix_length, hidden_size]
            prefix = prefix.unsqueeze(0).expand(batch_size, -1, -1, -1)
            
            prefix_key = prefix[:, 0]    # [batch, prefix_length, hidden_size]
            prefix_value = prefix[:, 1]  # [batch, prefix_length, hidden_size]
            
            return prefix_key, prefix_value
    
    print("Prefix Tuning:")
    print("  • 在attention的key/value前添加可训练前缀")
    print("  • 不修改模型结构")
    print("  • 参数量与序列长度无关")
    
except NameError:
    print("需要先导入PyTorch")

try:
    class PromptTuning(nn.Module):
        """Prompt Tuning实现"""
        
        def __init__(self, num_tokens, hidden_size, init_from_vocab=False):
            super().__init__()
            
            self.num_tokens = num_tokens
            
            # 可训练的软提示
            self.soft_prompt = nn.Parameter(
                torch.randn(num_tokens, hidden_size) * 0.02
            )
        
        def forward(self, input_embeds):
            """
            input_embeds: [batch, seq_len, hidden_size]
            """
            batch_size = input_embeds.shape[0]
            
            # 扩展软提示到batch
            prompt = self.soft_prompt.unsqueeze(0).expand(batch_size, -1, -1)
            
            # 拼接: [prompt | input]
            return torch.cat([prompt, input_embeds], dim=1)
    
    print("Prompt Tuning vs Prefix Tuning:")
    print("  • Prompt Tuning: 在输入嵌入层添加")
    print("  • Prefix Tuning: 在每层attention添加")
    
except NameError:
    print("需要先导入PyTorch")

print("QLoRA关键技术:")
print("=" * 50)
print("• 4-bit量化: 将模型权重量化到4位")
print("• Double Quantization: 量化常数也量化")
print("• Paged Optimizers: 防止显存溢出")
print("• NormalFloat4: 对正态分布优化的数据类型")
print()

# 显存估算
def estimate_memory(num_params, bits=16, lora_rank=8, lora_ratio=0.1):
    """估算显存使用"""
    
    # 基础模型
    base_memory = num_params * bits / 8 / 1e9
    
    # LoRA参数（全精度）
    lora_params = num_params * lora_ratio * lora_rank * 2 / num_params
    lora_memory = lora_params * 16 / 8 / 1e9
    
    return base_memory, lora_memory

# 7B模型对比
params_7b = 7e9

mem_fp16, _ = estimate_memory(params_7b, bits=16)
mem_int8, _ = estimate_memory(params_7b, bits=8)
mem_int4, lora_mem = estimate_memory(params_7b, bits=4)

print("7B模型显存估算:")
print(f"  FP16全量微调: ~{mem_fp16*2:.0f}GB")  # 梯度和优化器状态
print(f"  INT8 + LoRA: ~{mem_int8 + lora_mem:.0f}GB")
print(f"  INT4 (QLoRA): ~{mem_int4 + lora_mem:.0f}GB")

try:
    class LoRAConfig:
        """LoRA配置"""
        
        def __init__(self):
            # 常用配置
            self.rank = 8              # 秩，越大容量越大
            self.alpha = 16            # 缩放因子
            self.dropout = 0.1         # Dropout
            self.target_modules = [    # 应用LoRA的模块
                'q_proj', 'k_proj', 'v_proj', 'o_proj',  # Attention
                'gate_proj', 'up_proj', 'down_proj'      # FFN
            ]
    
    def apply_lora(model, config):
        """应用LoRA到模型"""
        
        for name, module in model.named_modules():
            if any(target in name for target in config.target_modules):
                if isinstance(module, nn.Linear):
                    # 替换为LoRA版本
                    lora_module = LoRALayer(
                        module.in_features,
                        module.out_features,
                        rank=config.rank,
                        alpha=config.alpha
                    )
                    # 复制原始权重
                    lora_module.weight.data = module.weight.data.clone()
                    
                    # 替换模块（伪代码）
                    # setattr(parent, child_name, lora_module)
        
        # 冻结非LoRA参数
        for name, param in model.named_parameters():
            if 'lora' not in name:
                param.requires_grad = False
        
        return model
    
    print("LoRA最佳实践:")
    print("  • rank通常4-16即可")
    print("  • alpha通常设为rank的2倍")
    print("  • 主要应用于Attention层")
    print("  • 推理时合并权重消除开销")
    
except NameError:
    print("需要先导入PyTorch")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model_id = "meta-llama/Meta-Llama-3-8B"

# 1. 4-bit 量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype="float16",
    bnb_4bit_use_double_quant=True,
)

# 2. 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

# 3. 准备量化训练
model = prepare_model_for_kbit_training(model)

# 4. LoRA 配置
config = LoraConfig(
    r=16, 
    lora_alpha=32, 
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], 
    lora_dropout=0.05, 
    bias="none", 
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, config)
model.print_trainable_parameters()