知识蒸馏
├── 白盒蒸馏（White-box）
│   ├── Logit 蒸馏（软标签）
│   ├── 中间层蒸馏（特征/注意力）
│   └── 注意力蒸馏
└── 黑盒蒸馏（Black-box）
    ├── 响应蒸馏（Response-based）
    │   ├── 指令跟随蒸馏
    │   ├── CoT 蒸馏
    │   └── 多轮对话蒸馏
    └── 行为蒸馏（Behavior-based）
        ├── 排序蒸馏
        └── 反馈蒸馏（RLAIF）

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

class KDLoss(nn.Module):
    """
    经典知识蒸馏损失
    L = α * KL(soft_student || soft_teacher) + (1-α) * CE(hard_label)
    """
    def __init__(self, temperature: float = 4.0, alpha: float = 0.7):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha

    def forward(
        self,
        student_logits: torch.Tensor,   # (B, V)
        teacher_logits: torch.Tensor,    # (B, V)
        labels: torch.Tensor,            # (B,)
    ) -> torch.Tensor:
        # 软标签损失
        soft_student = F.log_softmax(student_logits / self.temperature, dim=-1)
        soft_teacher = F.softmax(teacher_logits / self.temperature, dim=-1)
        kd_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (self.temperature ** 2)

        # 硬标签损失
        ce_loss = F.cross_entropy(student_logits, labels)

        # 加权组合
        total = self.alpha * kd_loss + (1 - self.alpha) * ce_loss
        return total, kd_loss, ce_loss

def llm_logit_distillation(
    student_model,
    teacher_model,
    input_ids: torch.Tensor,       # (B, seq_len)
    attention_mask: torch.Tensor,
    temperature: float = 2.0,
    alpha: float = 0.5,
):
    """
    LLM 逐 token Logit 蒸馏
    """
    # Teacher 前向（不计算梯度）
    with torch.no_grad():
        teacher_outputs = teacher_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        teacher_logits = teacher_outputs.logits  # (B, seq_len, vocab_size)

    # Student 前向
    student_outputs = student_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
    )
    student_logits = student_outputs.logits  # (B, seq_len, vocab_size)

    # 逐 token 计算 KD loss
    # 只对非 padding 位置计算
    shift_teacher = teacher_logits[:, :-1, :].contiguous()
    shift_student = student_logits[:, :-1, :].contiguous()
    shift_mask = attention_mask[:, 1:].contiguous()

    # 展平
    B, T, V = shift_student.shape
    flat_teacher = shift_teacher.view(-1, V)[shift_mask.view(-1).bool()]
    flat_student = shift_student.view(-1, V)[shift_mask.view(-1).bool()]

    # 软标签 KL 散度
    soft_teacher = F.softmax(flat_teacher / temperature, dim=-1)
    log_soft_student = F.log_softmax(flat_student / temperature, dim=-1)
    kd_loss = F.kl_div(log_soft_student, soft_teacher, reduction="batchmean") * (temperature ** 2)

    # 自回归 CE loss（Student 预测下一个 token）
    labels = input_ids[:, 1:].contiguous().view(-1)[shift_mask.view(-1).bool()]
    ce_loss = F.cross_entropy(flat_student, labels)

    total_loss = alpha * kd_loss + (1 - alpha) * ce_loss
    return total_loss

def align_vocab_logits(
    teacher_logits: torch.Tensor,    # (B, T, V_teacher)
    student_logits: torch.Tensor,    # (B, T, V_student)
    teacher_tokenizer,
    student_tokenizer,
    shared_tokens: List[str],        # 共享 token 列表
) -> tuple:
    """
    对齐不同词表的 logits
    只计算共享 token 上的 KL 散度
    """
    # 获取共享 token 在各自词表中的索引
    teacher_ids = [teacher_tokenizer.vocab[t] for t in shared_tokens]
    student_ids = [student_tokenizer.vocab[t] for t in shared_tokens]

    # 提取共享 token 的 logits
    teacher_shared = teacher_logits[..., teacher_ids]   # (B, T, |shared|)
    student_shared = student_logits[..., student_ids]   # (B, T, |shared|)

    return teacher_shared, student_shared

class HiddenStateDistillation(nn.Module):
    """
    中间层隐藏状态蒸馏
    通过线性投影对齐不同维度的隐藏状态
    """
    def __init__(
        self,
        teacher_hidden_dim: int = 4096,   # Teacher 隐藏维度
        student_hidden_dim: int = 1024,   # Student 隐藏维度
        teacher_layers: List[int] = [3, 7, 11, 15, 19, 23],  # Teacher 层
        student_layers: List[int] = [1, 3, 5, 7, 9, 11],     # Student 层
        loss_type: str = "mse",           # mse / cos
    ):
        super().__init__()
        assert len(teacher_layers) == len(student_layers)

        # 投影矩阵：student_dim → teacher_dim
        self.projectors = nn.ModuleList([
            nn.Linear(student_hidden_dim, teacher_hidden_dim, bias=False)
            for _ in range(len(teacher_layers))
        ])
        self.loss_type = loss_type

    def forward(
        self,
        teacher_hidden_states: List[torch.Tensor],  # Teacher 各层隐藏状态
        student_hidden_states: List[torch.Tensor],  # Student 各层隐藏状态
    ) -> torch.Tensor:
        total_loss = 0

        for i, (t_layer, s_layer) in enumerate(zip(self.teacher_layers, self.student_layers)):
            t_hidden = teacher_hidden_states[t_layer]  # (B, T, teacher_dim)
            s_hidden = student_hidden_states[s_layer]  # (B, T, student_dim)

            # 投影 student 到 teacher 维度
            s_projected = self.projectors[i](s_hidden)  # (B, T, teacher_dim)

            if self.loss_type == "mse":
                loss = F.mse_loss(s_projected, t_hidden)
            elif self.loss_type == "cos":
                loss = 1 - F.cosine_similarity(s_projected, t_hidden, dim=-1).mean()
            total_loss += loss

        return total_loss / len(self.teacher_layers)

class AttentionDistillation(nn.Module):
    """
    注意力分布蒸馏：让 Student 模仿 Teacher 的注意力模式
    """
    def __init__(self, teacher_layers: List[int], student_layers: List[int]):
        super().__init__()
        self.teacher_layers = teacher_layers
        self.student_layers = student_layers

    def forward(
        self,
        teacher_attentions: List[torch.Tensor],  # List of (B, heads, T, T)
        student_attentions: List[torch.Tensor],
    ) -> torch.Tensor:
        total_loss = 0

        for t_layer, s_layer in zip(self.teacher_layers, self.student_layers):
            t_attn = teacher_attentions[t_layer]  # (B, H_t, T, T)
            s_attn = student_attentions[s_layer]  # (B, H_s, T, T)

            # 对齐 head 数量（取平均或投影）
            if t_attn.shape[1] != s_attn.shape[1]:
                t_attn = t_attn.mean(dim=1, keepdim=True)
                s_attn = s_attn.mean(dim=1, keepdim=True)

            # KL 散度
            loss = F.kl_div(
                s_attn.log(),
                t_attn,
                reduction="batchmean",
            )
            total_loss += loss

        return total_loss / len(self.teacher_layers)

from openai import OpenAI
import json

client = OpenAI()

def generate_distillation_data(
    seed_instructions: List[str],
    teacher_model: str = "gpt-4o",
    num_per_seed: int = 5,
    temperature: float = 0.8,
) -> List[dict]:
    """
    用 Teacher 模型生成蒸馏数据
    """
    distill_data = []

    for seed in seed_instructions:
        # 让 Teacher 生成变体 + 响应
        for _ in range(num_per_seed):
            # Step 1: 生成指令变体
            var_resp = client.chat.completions.create(
                model=teacher_model,
                messages=[
                    {"role": "system", "content": "基于以下指令，生成一个语义相似但措辞不同的变体指令。只返回变体指令本身。"},
                    {"role": "user", "content": seed}
                ],
                temperature=temperature,
            )
            instruction = var_resp.choices[0].message.content

            # Step 2: Teacher 生成响应
            resp = client.chat.completions.create(
                model=teacher_model,
                messages=[
                    {"role": "system", "content": "你是一个专业助手，请详细回答用户的问题。"},
                    {"role": "user", "content": instruction}
                ],
                temperature=temperature,
            )
            response = resp.choices[0].message.content

            distill_data.append({
                "instruction": instruction,
                "input": "",
                "output": response,
                "teacher_model": teacher_model,
            })

    return distill_data


# 批量生成
seeds = [
    "解释什么是梯度消失问题",
    "如何实现一个简单的 RAG 系统",
    "比较 LoRA 和 QLoRA 的区别",
    "什么是注意力机制",
    "解释 Transformer 架构",
]
data = generate_distillation_data(seeds, teacher_model="gpt-4o", num_per_seed=10)

with open("distill_data.json", "w", encoding="utf-8") as f:
    json.dump(data, f, ensure_ascii=False, indent=2)

def cot_distillation_data(
    questions: List[str],
    teacher_model: str = "gpt-4o",
) -> List[dict]:
    """生成带思维链的蒸馏数据"""
    data = []
    for q in questions:
        resp = client.chat.completions.create(
            model=teacher_model,
            messages=[
                {"role": "system", "content": """请按以下格式回答问题：

 reasoning:
<逐步推理过程>

answer:
<最终答案>"""},
                {"role": "user", "content": q}
            ],
            temperature=0.3,  # 较低温度保证推理质量
        )
        full_response = resp.choices[0].message.content

        # 解析推理和答案
        parts = full_response.split("answer:")
        reasoning = parts[0].replace("reasoning:", "").strip()
        answer = parts[1].strip() if len(parts) > 1 else ""

        data.append({
            "question": q,
            "reasoning": reasoning,
            "answer": answer,
            "full_response": full_response,
        })
    return data

def multi_turn_distillation(
    topic: str,
    teacher_model: str = "gpt-4o",
    num_turns: int = 5,
    student_model: str = "gpt-3.5-turbo",  # 模拟 Student
) -> List[dict]:
    """
    多轮对话蒸馏：Student 提问，Teacher 回答
    """
    conversation = []
    messages = [{"role": "system", "content": f"讨论主题：{topic}"}]

    for turn in range(num_turns):
        # Student 生成问题/回复
        student_resp = client.chat.completions.create(
            model=student_model,
            messages=messages + [{"role": "user", "content": "请继续对话。"}],
        )
        student_msg = student_resp.choices[0].message.content
        messages.append({"role": "assistant", "content": student_msg})

        # Teacher 给出更好的回复
        teacher_resp = client.chat.completions.create(
            model=teacher_model,
            messages=messages + [{"role": "user", "content": "请给出更好的回复。"}],
        )
        teacher_msg = teacher_resp.choices[0].message.content

        conversation.append({
            "turn": turn,
            "student_response": student_msg,
            "teacher_response": teacher_msg,
        })

        # 用 Teacher 的回复继续对话
        messages.append({"role": "assistant", "content": teacher_msg})

    return conversation

class MiniLMDistillation(nn.Module):
    """
    MiniLM 蒸馏：自注意力 + 值注意力
    """
    def __init__(self, temperature: float = 1.0):
        super().__init__()
        self.temp = temperature

    def forward(
        self,
        teacher_self_attns: torch.Tensor,    # (B, H_t, T, T)
        teacher_value_attns: torch.Tensor,   # (B, H_t, T, T)  softmax(VV^T)
        student_self_attns: torch.Tensor,
        student_value_attns: torch.Tensor,
    ) -> torch.Tensor:
        # 1. 自注意力蒸馏
        # 对齐 head 数量
        if teacher_self_attns.shape[1] != student_self_attns.shape[1]:
            # 多头平均
            teacher_self_attns = teacher_self_attns.mean(dim=1, keepdim=True)
            student_self_attns = student_self_attns.mean(dim=1, keepdim=True)

        sa_loss = F.kl_div(
            (student_self_attns / self.temp).log(),
            teacher_self_attns / self.temp,
            reduction="batchmean",
        ) * (self.temp ** 2)

        # 2. 值注意力蒸馏
        va_loss = F.kl_div(
            (student_value_attns / self.temp).log(),
            teacher_value_attns / self.temp,
            reduction="batchmean",
        ) * (self.temp ** 2)

        return sa_loss + va_loss

# 以 GPT-2 Medium (355M) → GPT-2 Small (124M) 为例
config = {
    "teacher_model": "gpt2-medium",      # 355M
    "student_model": "gpt2-small",       # 124M
    "distill_method": "logit+hidden",
    "temperature": 2.0,
    "alpha": 0.5,                         # KD loss 权重
    "beta": 0.3,                          # Hidden loss 权重
    "learning_rate": 5e-5,
    "batch_size": 32,
    "epochs": 3,
    "max_seq_len": 512,
    "warmup_ratio": 0.1,
    "weight_decay": 0.01,
}

def rlaif_pipeline(
    student_model,
    teacher_model,
    prompts: List[str],
    beta: float = 0.1,          # KL 惩罚系数
):
    """
    RLAIF: Teacher 作为奖励模型，Student 通过 RL 优化
    """
    import trl

    # Step 1: Teacher 生成偏好数据
    preference_data = []
    for prompt in prompts:
        # Student 生成两个回复
        resp_a = student_model.generate(prompt, temperature=0.8)
        resp_b = student_model.generate(prompt, temperature=0.8)

        # Teacher 判断哪个更好
        judge_prompt = f"""以下是对同一个问题的两个回复，请判断哪个更好。

问题：{prompt}

回复 A：{resp_a}
回复 B：{resp_b}

更好的回复是（A/B）："""
        teacher_judge = teacher_model.generate(judge_prompt)
        preferred = "A" if "A" in teacher_judge else "B"

        preference_data.append({
            "prompt": prompt,
            "chosen": resp_a if preferred == "A" else resp_b,
            "rejected": resp_b if preferred == "A" else resp_a,
        })

    # Step 2: DPO 训练
    dpo_trainer = trl.DPOTrainer(
        model=student_model,
        beta=beta,
        dataset=preference_data,
    )
    dpo_trainer.train()

有 Teacher 模型权重？
├── 是 → 白盒蒸馏
│   ├── 计算资源充足 → Logit + 中间层蒸馏
│   └── 计算资源有限 → 仅 Logit 蒸馏
└── 否 → 黑盒蒸馏
    ├── 需要推理能力 → CoT 蒸馏
    ├── 需要对话能力 → 多轮对话蒸馏
    └── 需要对齐人类偏好 → RLAIF

# 问题：alpha 过高导致 Student 过度模仿 Teacher 的错误
# 解决：alpha 0.5-0.7 通常最优，保留硬标签约束

# 问题：温度过高导致软标签过于平滑
# 解决：temperature 1-4 通常最优

# 问题：Teacher 用 GPT-4 tokenizer，Student 用 LLaMA tokenizer
# 解决：
# 1. 重新分词 Teacher 的训练数据
# 2. 或使用共享子集进行 Logit 蒸馏

# 问题：蒸馏后 Student 在某些任务上比 SFT 还差
# 原因：蒸馏数据分布偏移
# 解决：混合蒸馏数据 + 原始 SFT 数据（比例 7:3）