大模型微调新范式：QLoRA+MoE混合训练

# 导入必要库
import torch
from torch import nn
from transformers import AutoModelForCausalLM, AutoTokenizer
import bitsandbytes as bnb
from accelerate import Accelerator

# 定义QLoRA配置类
class QLoRAConfig:
    def __init__(self, 
                 r: int = 8,  # LoRA秩大小
                 quantization_bits: int = 8,  # 量化位数
                 alpha: float = 16  # LoRA alpha参数
                ):
        self.r = r
        self.quantization_bits = quantization_bits
        self.alpha = alpha

# 定义MoE配置类
class MoEConfig:
    def __init__(self, 
                 num_experts: int = 8,  # 专家网络数量
                 top_k: int = 2,  # Top-k门控机制中的k值
                 expert_capacity: int = 32  # 每个专家的最大处理容量
                ):
        self.num_experts = num_experts
        self.top_k = top_k
        self.expert_capacity = expert_capacity

# 定义QLoRA+MoE混合模型类
class QLoRAMoEModel(nn.Module):
    def __init__(self, 
                 base_model_name: str,
                 qLora_config: QLoRAConfig,
                 moe_config: MoEConfig
                ):
        super().__init__()
        
        # 加载基础模型并量化
        self.base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
        self.quantize_model(qLora_config.quantization_bits)
        
        # 初始化MoE相关组件
        self.num_experts = moe_config.num_experts
        self.top_k = moe_config.top_k
        
        # 创建专家网络列表
        self.experts = nn.ModuleList([
            self._create_qlora_expert(qLora_config) for _ in range(moe_config.num_experts)
        ])
        
        # 创建门控网络并应用QLoRA
        self.gating_network = self._create_gating_network(qLora_config)
        
    def quantize_model(self, bits: int):
        """对模型进行量化处理"""
        for module in self.base_model.modules():
            if isinstance(module, nn.Linear):
                # 应用bitsandbytes量化
                module.weight = bnb.nn.Int8Params(
                    module.weight.data.contiguous(),
                    has_fp16_weights=False,
                    requires_grad=False
                )
                # 替换为量化线性层
                quantized_linear = bnb.nn.Linear4bit(
                    module.in_features,
                    module.out_features,
                    bias=module.bias is not None,
                    compute_dtype=torch.float16
                )
                quantized_linear.weight = module.weight
                if module.bias is not None:
                    quantized_linear.bias = module.bias
                # 替换原始层
                parent = module.parent
                name = module.name
                setattr(parent, name, quantized_linear)
    
    def _create_qlora_expert(self, qLora_config: QLoRAConfig):
        """创建单个专家网络并应用QLoRA"""
        lora_modules = nn.ModuleDict()
        for name, module in self.base_model.named_modules():
            if isinstance(module, bnb.nn.Linear4bit):
                # 创建LoRA分解矩阵
                lora_A = nn.Parameter(torch.randn(qLora_config.r, module.in_features))
                lora_B = nn.Parameter(torch.randn(module.out_features, qLora_config.r))
                
                # 计算LoRA缩放因子
                scaling = qLora_config.alpha / qLora_config.r
                
                # 将LoRA参数添加到专家网络
                lora_modules[name] = nn.ModuleDict({
                    'lora_A': lora_A,
                    'lora_B': lora_B,
                    'scaling': scaling
                })
        
        return nn.ModuleDict({
            'lora_params': lora_modules,
            'expert_id': len(self.experts)  # 专家唯一标识
        })
    
    def _create_gating_network(self, qLora_config: QLoRAConfig):
        """创建门控网络并应用QLoRA"""
        # 门控网络输入维度为模型隐藏层大小
        hidden_size = self.base_model.config.hidden_size
        
        # 创建门控网络主体（多层感知机）
        gating_network = nn.Sequential(
            bnb.nn.Linear4bit(hidden_size, hidden_size // 2, compute_dtype=torch.float16),
            nn.GELU(),
            bnb.nn.Linear4bit(hidden_size // 2, self.num_experts, compute_dtype=torch.float16)
        )
        
        # 对门控网络应用QLoRA分解
        for name, module in gating_network.named_modules():
            if isinstance(module, bnb.nn.Linear4bit):
                # 创建LoRA分解矩阵
                lora_A = nn.Parameter(torch.randn(qLora_config.r, module.in_features))
                lora_B = nn.Parameter(torch.randn(module.out_features, qLora_config.r))
                
                # 添加到门控网络
                setattr(module, 'lora_A', lora_A)
                setattr(module, 'lora_B', lora_B)
        
        return gating_network
    
    def forward(self, input_ids, attention_mask):
        # 基础模型前向传播
        base_outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True
        )
        
        # 获取最后一层隐藏状态作为门控输入
        gating_input = base_outputs.hidden_states[-1]
        
        # 门控网络计算专家选择概率
        raw_gating_scores = self.gating_network(gating_input)
        gating_scores = torch.softmax(raw_gating_scores, dim=-1)
        
        # Top-k选择专家
        top_k_scores, top_k_indices = torch.topk(gating_scores, k=self.top_k, dim=-1)
        
        # 初始化专家输出集合
        expert_outputs = []
        for expert_id in range(self.num_experts):
            # 获取选择当前专家的样本索引
            selected_samples = torch.where(top_k_indices == expert_id)
            
            # 如果没有样本选择当前专家，则跳过
            if len(selected_samples[0]) == 0:
                continue
            
            # 获取专家网络的QLoRA参数
            expert = self.experts[expert_id]
            lora_params = expert['lora_params']
            
            # 计算QLoRA更新后的权重
            updated_weights = {}
            for layer_name, lora_param in lora_params.items():
                base_weight = getattr(self.base_model, layer_name).weight
                updated_weight = base_weight + lora_param['lora_B'] @ lora_param['lora_A'] * lora_param['scaling']
                updated_weights[layer_name] = updated_weight
            
            # 使用更新后的权重进行前向传播
            expert_output = self._forward_expert(
                input_ids=input_ids[selected_samples],
                attention_mask=attention_mask[selected_samples],
                updated_weights=updated_weights
            )
            
            # 将专家输出添加到集合
            expert_outputs.append((expert_output, selected_samples, top_k_scores[selected_samples]))
        
        # 合并专家输出
        final_output = self._merge_expert_outputs(expert_outputs)
        
        return final_output
    
    def _forward_expert(self, input_ids, attention_mask, updated_weights):
        """使用专家网络特定的QLoRA权重进行前向传播"""
        # 创建临时模型副本并替换权重
        with torch.no_grad():
            expert_model = copy.deepcopy(self.base_model)
            for layer_name, weight in updated_weights.items():
                setattr(expert_model, layer_name + '.weight', weight)
        
        # 使用专家模型进行前向传播
        return expert_model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
    
    def _merge_expert_outputs(self, expert_outputs):
        """合并多个专家网络的输出"""
        # 初始化最终输出张量
        batch_size = expert_outputs[0][1][0].shape[0]
        seq_length = expert_outputs[0][0].logits.shape[1]
        vocab_size = expert_outputs[0][0].logits.shape[2]
        
        final_logits = torch.zeros(batch_size, seq_length, vocab_size, device=expert_outputs[0][0].logits.device)
        
        # 按样本位置填充专家输出
        for expert_output, (batch_indices, seq_indices), scores in expert_outputs:
            final_logits[batch_indices, seq_indices] += expert_output.logits[batch_indices, seq_indices] * scores.unsqueeze(-1)
        
        return final_logits

# 定义训练流程类
class QLoRAMoETrainer:
    def __init__(self, 
                 model: QLoRAMoEModel,
                 train_dataset,
                 eval_dataset,
                 learning_rate: float = 2e-5,
                 num_epochs: int = 3,
                 batch_size: int = 8
                ):
        self.model = model
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.learning_rate = learning_rate
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        
        # 初始化加速器
        self.accelerator = Accelerator()
        
        # 创建优化器，仅优化QLoRA和门控网络参数
        self.optimizer = torch.optim.AdamW(
            [
                {'params': self.model.gating_network.parameters(), 'lr': learning_rate},
                {'params': self._get_qlora_parameters(), 'lr': learning_rate * 0.1}  # LoRA参数学习率稍低
            ],
            weight_decay=0.01
        )
        
        # 准备数据加载器
        self.train_dataloader = self._prepare_dataloader(train_dataset, shuffle=True)
        self.eval_dataloader = self._prepare_dataloader(eval_dataset, shuffle=False)
    
    def _get_qlora_parameters(self):
        """获取所有QLoRA参数"""
        qlora_params = []
        for expert in self.model.experts:
            for layer_params in expert['lora_params'].values():
                qlora_params.append(layer_params['lora_A'])
                qlora_params.append(layer_params['lora_B'])
        return qlora_params
    
    def _prepare_dataloader(self, dataset, shuffle: bool):
        """准备数据加载器"""
        return torch.utils.data.DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            collate_fn=self._collate_fn
        )
    
    def _collate_fn(self, batch):
        """自定义数据整理函数"""
        input_ids = torch.stack([torch.tensor(item['input_ids']) for item in batch])
        attention_mask = torch.stack([torch.tensor(item['attention_mask']) for item in batch])
        labels = torch.stack([torch.tensor(item['labels']) for item in batch])
        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'labels': labels
        }
    
    def train(self):
        """开始训练"""
        # 使用accelerate准备模型、优化器和数据加载器
        self.model, self.optimizer, self.train_dataloader, self.eval_dataloader = self.accelerator.prepare(
            self.model, self.optimizer, self.train_dataloader, self.eval_dataloader
        )
        
        # 训练循环
        for epoch in range(self.num_epochs):
            self.model.train()
            total_loss = 0
            
            for batch in self.train_dataloader:
                outputs = self.model(**batch)
                loss = self._compute_loss(outputs, batch['labels'])
                
                self.accelerator.backward(loss)
                self.optimizer.step()
                self.optimizer.zero_grad()
                
                total_loss += loss.item()
            
            avg_loss = total_loss / len(self.train_dataloader)
            self.accelerator.print(f"Epoch {epoch+1}, Training Loss: {avg_loss:.4f}")
            
            # 评估模型
            self.evaluate()
    
    def _compute_loss(self, outputs, labels):
        """计算损失"""
        return nn.CrossEntropyLoss()(outputs.logits.view(-1, outputs.logits.size(-1)), labels.view(-1))
    
    def evaluate(self):
        """评估模型"""
        self.model.eval()
        total_loss = 0
        
        with torch.no_grad():
            for batch in self.eval_dataloader:
                outputs = self.model(**batch)
                loss = self._compute_loss(outputs, batch['labels'])
                total_loss += loss.item()
        
        avg_loss = total_loss / len(self.eval_dataloader)
        self.accelerator.print(f"Evaluation Loss: {avg_loss:.4f}")
        
        return avg_loss

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def preprocess_squad(example):
    # 分割上下文为块
    context_chunks = []
    for i in range(0, len(tokenizer.tokenize(example['context'])), 512):
        context_chunks.append(tokenizer.decode(
            tokenizer.encode(example['context'])[i:i+512],
            skip_special_tokens=True
        ))
    
    # 生成问题-上下文配对
    question = example['question']
    answer = example['answers'][0] if example['answers'] else None
    
    examples = []
    for context in context_chunks:
        encoding = tokenizer(
            question,
            context,
            max_length=512,
            truncation="only_second",
            padding="max_length"
        )
        
        # 标记答案位置（如果有答案）
        if answer:
            answer_start = context.find(answer['text'])
            if answer_start != -1:
                answer_end = answer_start + len(answer['text'])
                start_positions = encoding.char_to_token(answer_start)
                end_positions = encoding.char_to_token(answer_end)
                examples.append({
                    'input_ids': encoding['input_ids'],
                    'attention_mask': encoding['attention_mask'],
                    'start_positions': start_positions,
                    'end_positions': end_positions
                })
            else:
                # 答案不在当前上下文块
                examples.append({
                    'input_ids': encoding['input_ids'],
                    'attention_mask': encoding['attention_mask'],
                    'start_positions': 0,
                    'end_positions': 0
                })
        else:
            # 无答案情况
            examples.append({
                'input_ids': encoding['input_ids'],
                'attention_mask': encoding['attention_mask'],
                'start_positions': 0,
                'end_positions': 0
            })
    
    return examples

# 应用预处理
from datasets import load_dataset

squad_dataset = load_dataset("squad_v2")
processed_dataset = squad_dataset.map(preprocess_squad, batched=True, remove_columns=squad_dataset['train'].column_names)