从零到一：DeepSeek如何在竞争激烈的AI领域脱颖而出

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

# 知识蒸馏实现示例
class DeepSeekDistillation:
    def __init__(self, teacher_model="deepseek-ai/deepseek-llm-67b-base", 
                student_model="deepseek-ai/deepseek-llm-7b-base",
                alpha=0.5, temperature=2.0):
        # 加载教师模型（大模型）
        self.teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model)
        self.teacher_model = AutoModelForCausalLM.from_pretrained(
            teacher_model, torch_dtype=torch.bfloat16, device_map="auto"
        )
        self.teacher_model.eval()  # 设为评估模式
        
        # 加载学生模型（小模型）
        self.student_tokenizer = AutoTokenizer.from_pretrained(student_model)
        self.student_model = AutoModelForCausalLM.from_pretrained(
            student_model, torch_dtype=torch.bfloat16, device_map="auto"
        )
        
        # 蒸馏超参数
        self.alpha = alpha  # 硬标签和软标签的权重平衡
        self.temperature = temperature  # 软化logits的温度参数
        
    def distillation_loss(self, student_logits, teacher_logits, labels):
        """
        计算蒸馏损失：结合硬标签交叉熵和软标签KL散度
        """
        # 硬标签损失 - 学生模型预测与真实标签的交叉熵
        hard_loss = F.cross_entropy(student_logits, labels)
        
        # 软标签损失 - 学生模型与教师模型输出分布的KL散度
        soft_student = F.log_softmax(student_logits / self.temperature, dim=-1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=-1)
        soft_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean')
        
        # 综合损失
        loss = (1 - self.alpha) * hard_loss + self.alpha * soft_loss * (self.temperature ** 2)
        return loss
    
    def train_step(self, batch_inputs, batch_labels):
        """执行一步蒸馏训练"""
        # 获取教师模型输出（无梯度计算）
        with torch.no_grad():
            teacher_outputs = self.teacher_model(batch_inputs, labels=batch_labels)
            teacher_logits = teacher_outputs.logits
        
        # 获取学生模型输出
        student_outputs = self.student_model(batch_inputs, labels=batch_labels)
        student_logits = student_outputs.logits
        
        # 计算蒸馏损失
        loss = self.distillation_loss(student_logits, teacher_logits, batch_labels)
        
        return loss