大模型蒸馏实战

# -*- coding: utf-8 -*-
"""软标签蒸馏损失：温度缩放 softmax + KL 散度，逐步拆开算给你看。

刻意只用标准库（math），不依赖 numpy / torch，任何机器上都能直接跑。
真实工程里把下面的 softmax_t / kl_divergence 换成框架自带算子即可，
公式和温度的作用完全一样。

跑法：python3 kd_soft_label_loss.py
"""
import math

# 艾青自传那篇文章的分类 logits（教师模型 softmax 之前的原始输出）。
# 类别顺序固定，后面所有列表都按这个顺序。
LABELS = ["军事", "文学", "家居", "体育", "财经"]
TEACHER_LOGITS = [3.20, 4.45, -2.90, -6.85, -3.30]
STUDENT_LOGITS = [2.10, 3.00, -1.20, -3.40, -1.80]

# 硬标签：只有「文学」是 1，其余全是 0。信息量到此为止。
HARD_LABEL = [0, 1, 0, 0, 0]


def softmax_t(logits, temperature=1.0):
    """带温度的 softmax。

    T = 1   原始分布，差距被指数函数拉得很开
    T > 1   分布被「抹平」，小概率类别的信息被放大 —— 蒸馏就靠这个
    T < 1   分布更尖锐，趋近硬标签
    """
    scaled = [x / temperature for x in logits]
    m = max(scaled)                       # 减最大值，防指数溢出
    exps = [math.exp(x - m) for x in scaled]
    total = sum(exps)
    return [e / total for e in exps]


def kl_divergence(p, q, eps=1e-12):
    """KL(p || q)：用 q 去近似 p 时损失了多少信息。

    p 是教师分布（目标），q 是学生分布（被优化的一方）。
    两个分布的差异要用 KL 散度度量，不是 MSE —— 概率分布之间比的是信息量。
    """
    return sum(pi * math.log((pi + eps) / (qi + eps)) for pi, qi in zip(p, q))


def distill_loss(student_logits, teacher_logits, temperature=4.0):
    """软标签蒸馏损失。

    注意乘 T^2：温度把 logits 缩小了 T 倍，梯度跟着小 T^2 倍，
    乘回来才能让软标签项和硬标签项的梯度在一个量级上。
    """
    p = softmax_t(teacher_logits, temperature)
    q = softmax_t(student_logits, temperature)
    return kl_divergence(p, q) * temperature * temperature


def show(title, dist):
    print(title)
    for name, v in zip(LABELS, dist):
        bar = "#" * max(1, int(v * 60))
        print("  %-4s %10.6f  %s" % (name, v, bar))


if __name__ == "__main__":
    print("=" * 62)
    print("一、硬标签 vs 软标签：同一篇文章，信息量差多少")
    print("=" * 62)
    print("硬标签：", dict(zip(LABELS, HARD_LABEL)))
    show("教师软标签（T=1）：", softmax_t(TEACHER_LOGITS, 1.0))
    print("""
  硬标签只说了「这是文学」。
  软标签还说了「它有两成像军事，而且像军事的程度
  远高于像体育」—— 这正是艾青自传里那段与王震将军
  的往事留下的痕迹。类别之间的隐含关系，就藏在这些
  小数里。""")

    print()
    print("=" * 62)
    print("二、温度怎么影响分布")
    print("=" * 62)
    for t in (1.0, 2.0, 4.0, 8.0):
        d = softmax_t(TEACHER_LOGITS, t)
        print("T=%-4.1f  " % t + "  ".join(
            "%s=%.5f" % (n, v) for n, v in zip(LABELS, d)))
    print("""
  T 越大，分布越平，「家居 / 体育 / 财经」这些极小概率
  被放大到学生看得见的量级。蒸馏对温度非常敏感：
  T 太小等于在学硬标签，T 太大则噪声盖过信号。""")

    print()
    print("=" * 62)
    print("三、蒸馏损失")
    print("=" * 62)
    for t in (1.0, 2.0, 4.0, 8.0):
        print("T=%-4.1f  KL=%.6f  loss(KL x T^2)=%.6f"
              % (t,
                 kl_divergence(softmax_t(TEACHER_LOGITS, t),
                               softmax_t(STUDENT_LOGITS, t)),
                 distill_loss(STUDENT_LOGITS, TEACHER_LOGITS, t)))

==============================================================
一、硬标签 vs 软标签：同一篇文章，信息量差多少
==============================================================
硬标签： {'军事': 0, '文学': 1, '家居': 0, '体育': 0, '财经': 0}
教师软标签（T=1）：
  军事     0.222512  #############
  文学     0.776644  ##############################################
  家居     0.000499  #
  体育     0.000010  #
  财经     0.000335  #

  硬标签只说了「这是文学」。
  软标签还说了「它有两成像军事，而且像军事的程度
  远高于像体育」—— 这正是艾青自传里那段与王震将军
  的往事留下的痕迹。类别之间的隐含关系，就藏在这些
  小数里。

==============================================================
二、温度怎么影响分布
==============================================================
T=1.0   军事=0.22251  文学=0.77664  家居=0.00050  体育=0.00001  财经=0.00033
T=2.0   军事=0.33773  文学=0.63096  家居=0.01599  体育=0.00222  财经=0.01310
T=4.0   军事=0.34935  文学=0.47751  家居=0.07603  体育=0.02832  财经=0.06879
T=8.0   军事=0.29726  文学=0.34753  家居=0.13867  体育=0.08464  财经=0.13191

  T 越大，分布越平，「家居 / 体育 / 财经」这些极小概率
  被放大到学生看得见的量级。蒸馏对温度非常敏感：
  T 太小等于在学硬标签，T 太大则噪声盖过信号。

==============================================================
三、蒸馏损失
==============================================================
T=1.0   KL=0.025434  loss(KL x T^2)=0.025434
T=2.0   KL=0.067847  loss(KL x T^2)=0.271388
T=4.0   KL=0.060104  loss(KL x T^2)=0.961661
T=8.0   KL=0.022986  loss(KL x T^2)=1.471111

# 魔搭社区（modelscope）下载基座模型，--local_dir 指定落点
modelscope download --model Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4 \
  --local_dir /workspace/deepseekDistllation/models/Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4

# 能科学上网的可以直接用 huggingface 的模型库。
# 下载完先确认目录里有权重文件、tokenizer 和 config，缺一个都起不来。

# ---------------------------------------------------------------------------
# 一、只起基座（对照组）
# ---------------------------------------------------------------------------
CUDA_VISIBLE_DEVICES=0 API_PORT=8000 llamafactory-cli api \
  --model_name_or_path /workspace/models/Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4

# ---------------------------------------------------------------------------
# 二、基座 + 一个 LoRA（实验组）
#     同一个基座换不同的 --adapter_name_or_path，就能切换到不同能力的模型
# ---------------------------------------------------------------------------
CUDA_VISIBLE_DEVICES=0 API_PORT=8000 llamafactory-cli api \
  --model_name_or_path /workspace/models/Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4 \
  --adapter_name_or_path /workspace/saves/lora/cot/checkpoint-1250 \
  --finetuning_type lora

# ---------------------------------------------------------------------------
# 三、换 vLLM 部署，并同时挂载多个 LoRA
#     调用时用 model 参数点名要哪一个；model 给基座名或空串就是纯基座推理
# ---------------------------------------------------------------------------
CUDA_VISIBLE_DEVICES=0 API_PORT=8000 vllm serve \
  /workspace/models/Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4 \
  --enable-lora \
  --lora-modules \
    cot=/workspace/saves/lora/cot/checkpoint-1250 \
    cat=/workspace/saves/lora/classify/checkpoint-120 \
    spo=/workspace/saves/lora/triple/checkpoint-375

# 服务起来后观察显卡：linux 用 watch -n 1 nvidia-smi

# -*- coding: utf-8 -*-
"""访问 LLaMA-Factory 或 vLLM 的 api server，并测输出速度。

两者都是 OpenAI 兼容接口，代码完全一样，区别只在 model 参数：
  LLaMA-Factory 的 api server  -> model 随便填（通常写 "test"）
  vLLM 且 --enable-lora        -> model 填 --lora-modules 里定义的名字，
                                  填空串或基座名就是用基座模型推理

运行前：
  export LLM_BASE_URL=http://localhost:8000/v1
  export LLM_API_KEY=local            # 本地服务不校验，但也别在源码里写字面量
  export LLM_MODEL=cot                # 想测基座就换成基座模型名
"""
import os
import time

from openai import OpenAI

# 同一道题跑基座和蒸馏后模型，差异最容易看出来：
# 基座会给一段四平八稳的议论，蒸馏过的会先在 <think> 里推一遍再作答。
DEFAULT_QUESTION = (
    "对于「初三女生在搀扶跌倒老奶奶后反被冤枉，但仍选择资助她千元」的新闻事件，"
    "你有什么看法？"
)


def make_client():
    base_url = os.environ.get("LLM_BASE_URL", "http://localhost:8000/v1")
    api_key = os.environ.get("LLM_API_KEY")
    if not api_key:
        raise RuntimeError("环境变量 LLM_API_KEY 没设置（本地服务可随便给一个占位值）")
    return OpenAI(api_key=api_key, base_url=base_url)


def call_server(question, model=None, max_tokens=20000):
    client = make_client()
    model = model or os.environ.get("LLM_MODEL", "test")
    result = client.chat.completions.create(
        messages=[{"role": "user", "content": question}],
        model=model,
        max_tokens=max_tokens,
    )
    return result.choices[0].message.content


def bench(question=DEFAULT_QUESTION, model=None):
    """打印输出长度、耗时、速度三个数，这就是部署优化的验收指标。"""
    start = time.time()
    text = call_server(question, model=model)
    cost = time.time() - start
    print("输出长度：%d 字  耗时：%.2f 秒  输出速度：%.2f 字/秒"
          % (len(text), cost, len(text) / cost))
    return text, cost


def compare(models, question=DEFAULT_QUESTION):
    """依次测多个 model（多个 LoRA 或基座），横向比速度。

    vLLM 能同时挂载多个 LoRA，切换只是换一个 model 名字，
    不用重启服务、不用再占一份显存 —— 对比测试因此变得很便宜。
    """
    rows = []
    for name in models:
        print("---- model = %s ----" % name)
        text, cost = bench(question, model=name)
        rows.append((name, len(text), cost, len(text) / cost))
        print(text[:200] + " ……")
        print()
    print("%-12s %10s %10s %12s" % ("model", "输出字数", "耗时(秒)", "速度(字/秒)"))
    for name, n, cost, speed in rows:
        print("%-12s %10d %10.2f %12.2f" % (name, n, cost, speed))


if __name__ == "__main__":
    bench()

# -*- coding: utf-8 -*-
"""大模型 + 传统方法混合策略：AC 自动机全量扫描，只把歧义片段送大模型。

场景：2.86 亿篇文章里识别车系。全量喂大模型算得起吗？算不起。
做法：先用 AC 自动机把所有车系字符串一次扫出来，
      没出现车系的直接丢，无歧义的直接记账，
      只有「常用词歧义」和「英文数字串歧义」的片段才取上下文送大模型判定。

AC 自动机用标准库手写（无第三方依赖），本地可直接跑：
  python3 hybrid_pipeline_demo.py
"""
from collections import deque

# ---------------------------------------------------------------------------
# 一、AC 自动机（Aho-Corasick）：一次扫描同时匹配上万个关键词
# ---------------------------------------------------------------------------


class AhoCorasick:
    def __init__(self):
        self.next = [{}]          # next[state][char] -> state
        self.fail = [0]
        self.out = [[]]           # 该状态命中的关键词列表

    def add(self, word):
        state = 0
        for ch in word:
            nxt = self.next[state].get(ch)
            if nxt is None:
                self.next.append({})
                self.fail.append(0)
                self.out.append([])
                nxt = len(self.next) - 1
                self.next[state][ch] = nxt
            state = nxt
        self.out[state].append(word)

    def build(self):
        """BFS 构造 fail 指针：匹配失败时跳到「最长的、也是某个模式前缀的后缀」。"""
        q = deque()
        for ch, s in self.next[0].items():
            self.fail[s] = 0
            q.append(s)
        while q:
            r = q.popleft()
            for ch, s in self.next[r].items():
                q.append(s)
                f = self.fail[r]
                while f and ch not in self.next[f]:
                    f = self.fail[f]
                f = self.next[f].get(ch, 0)
                self.fail[s] = 0 if f == s else f
                # 把 fail 指向状态的命中词并进来，一次匹配能报出所有重叠关键词
                self.out[s] += self.out[self.fail[s]]

    def search(self, text):
        """返回 [(起点, 终点, 关键词), ...]，一遍扫描，复杂度与关键词数量无关。"""
        hits = []
        state = 0
        for i, ch in enumerate(text):
            while state and ch not in self.next[state]:
                state = self.fail[state]
            state = self.next[state].get(ch, 0)
            for word in self.out[state]:
                hits.append((i - len(word) + 1, i + 1, word))
        return hits


# ---------------------------------------------------------------------------
# 二、车系词典：三类，只有后两类需要大模型
# ---------------------------------------------------------------------------

UNAMBIGUOUS = ["帕萨特", "传祺", "马自达", "凯迪拉克"]   # 无歧义，直接记账
COMMON_WORD = ["几何", "哪吒", "指挥官"]                 # 常用词歧义
ALNUM_STR = ["XC60", "SL350", "XC360"]                   # 英文数字串歧义
AMBIGUOUS = set(COMMON_WORD + ALNUM_STR)

CONTEXT = 20          # 歧义片段前后各取 20 个字
BATCH = 20            # 20 条数据打包进一条 prompt

JUDGE_PROMPT_HEAD = (
    "你是一位汽车行业的专家，需要你确定一组文本片段的特定字符串是不是代表一个车系，"
    "如果是车系，输出是；否则输出否。不允许输出其他任何字符。"
)


def build_automaton():
    ac = AhoCorasick()
    for w in UNAMBIGUOUS + COMMON_WORD + ALNUM_STR:
        ac.add(w)
    ac.build()
    return ac


def scan_article(ac, doc_id, text):
    """返回 (直接记账的命中, 需要送大模型的片段)。"""
    direct, pending = [], []
    for start, end, word in ac.search(text):
        if word in AMBIGUOUS:
            left = max(0, start - CONTEXT)
            right = min(len(text), end + CONTEXT)
            pending.append([text[left:right], word])
        else:
            direct.append((doc_id, word, start))
    return direct, pending


def make_batches(pending, batch=BATCH):
    """每 batch 条打包成一条 prompt，这一步直接把请求数除以 20。"""
    for i in range(0, len(pending), batch):
        chunk = pending[i:i + batch]
        yield JUDGE_PROMPT_HEAD + "\n现在真实的输入是：\n" + repr(chunk), chunk


# ---------------------------------------------------------------------------
# 三、成本估算：把账算给人看
# ---------------------------------------------------------------------------

TOTAL_DOCS = 2.86e8           # 文章总量
AVG_CHARS = 1000              # 平均字数
RATIO_AMBIGUOUS = 0.0254      # 有潜在歧义车系的文章占比
AVG_AMB_PER_DOC = 2.5         # 这类文章里平均每篇的潜在歧义车系个数
PER_RECORD_CHARS = 60         # 一条歧义记录（4.6 字车系 + 前后 20 字 + 结构符号）
TMPL_CHARS = 305              # 提示词模板长度
OUT_CHARS_PER_REQ = 20        # 每次请求输出 20 个「是/否」

PLAN1_PROMPT_CHARS = 91       # 方案一每篇文章附带的提示词长度
CHARS_PER_TOKEN = 0.6         # 中文字数换算成 token 的经验系数
PRICE_IN = 4 / 1e6            # 元 / token
PRICE_OUT = 16 / 1e6


def cost(chars_in, chars_out):
    return (chars_in * CHARS_PER_TOKEN * PRICE_IN,
            chars_out * CHARS_PER_TOKEN * PRICE_OUT)


def estimate():
    # 方案一：所有文章全量喂大模型
    in1 = TOTAL_DOCS * (AVG_CHARS + PLAN1_PROMPT_CHARS)
    out1 = TOTAL_DOCS * (2 + (2 + 3.8) * 0.31)
    c1_in, c1_out = cost(in1, out1)

    # 方案二：只送 2.54% 文章里的歧义片段，且 20 条打一包
    amb_docs = TOTAL_DOCS * RATIO_AMBIGUOUS
    amb_items = amb_docs * AVG_AMB_PER_DOC
    requests = amb_items / BATCH
    in2 = requests * (TMPL_CHARS + PER_RECORD_CHARS * BATCH)
    out2 = requests * OUT_CHARS_PER_REQ
    c2_in, c2_out = cost(in2, out2)

    print("%-12s %14s %14s %14s" % ("", "输入字数(亿)", "输出字数(亿)", "总成本(元)"))
    print("%-12s %14.2f %14.4f %14.2f"
          % ("方案一", in1 / 1e8, out1 / 1e8, c1_in + c1_out))
    print("%-12s %14.2f %14.4f %14.2f"
          % ("方案二", in2 / 1e8, out2 / 1e8, c2_in + c2_out))
    print("%-12s %13.3f%% %13.3f%% %13.3f%%"
          % ("方案二/方案一", in2 / in1 * 100, out2 / out1 * 100,
             (c2_in + c2_out) / (c1_in + c1_out) * 100))
    print()
    print("需要大模型判定的文章：%.2f 万篇" % (amb_docs / 1e4))
    print("潜在歧义车系总数：    %.2f 万个" % (amb_items / 1e4))
    print("打包后的请求次数：    %.0f 次" % requests)


if __name__ == "__main__":
    ac = build_automaton()
    docs = [
        (1, "今天天气不错，和朋友去公园散步，没聊车。"),
        (2, "这款帕萨特的后备箱空间很大，传祺也不差。"),
        (3, "今年哪吒大火，票房达到了惊人的120亿美元，很多人都去看了。"),
        (4, "大部分朋友对此有一个误解，其实哪吒的目标从来都不是商务用车，"
            "而更偏向于表达用户的个性，XC60 的车主也有类似看法。"),
    ]

    print("=" * 62)
    print("一、AC 自动机扫描结果")
    print("=" * 62)
    all_pending = []
    for doc_id, text in docs:
        direct, pending = scan_article(ac, doc_id, text)
        if not direct and not pending:
            print("文章 %d：无车系，丢弃" % doc_id)
            continue
        for d in direct:
            print("文章 %d：无歧义命中 %s（位置 %d），直接记录" % d)
        for frag, word in pending:
            print("文章 %d：潜在歧义 %s，取上下文送大模型 -> %r" % (doc_id, word, frag))
        all_pending += pending

    print()
    print("=" * 62)
    print("二、打包后的 prompt（20 条一批，这里只有 %d 条）" % len(all_pending))
    print("=" * 62)
    for prompt, chunk in make_batches(all_pending):
        print(prompt[:300] + (" ……" if len(prompt) > 300 else ""))
        print("本批 %d 条，期望输出 %d 个字" % (len(chunk), len(chunk)))

    print()
    print("=" * 62)
    print("三、2.86 亿篇文章的成本估算")
    print("=" * 62)
    estimate()

==============================================================
一、AC 自动机扫描结果
==============================================================
文章 1：无车系，丢弃
文章 2：无歧义命中 帕萨特（位置 2），直接记录
文章 2：无歧义命中 传祺（位置 14），直接记录
文章 3：潜在歧义 哪吒，取上下文送大模型 -> '今年哪吒大火，票房达到了惊人的120亿美元，很多'
文章 4：潜在歧义 哪吒，取上下文送大模型 -> '大部分朋友对此有一个误解，其实哪吒的目标从来都不是商务用车，而更偏向于表达'
文章 4：潜在歧义 XC60，取上下文送大模型 -> '不是商务用车，而更偏向于表达用户的个性，XC60 的车主也有类似看法。'

==============================================================
二、打包后的 prompt（20 条一批，这里只有 3 条）
==============================================================
你是一位汽车行业的专家，需要你确定一组文本片段的特定字符串是不是代表一个车系，如果是车系，输出是；否则输出否。不允许输出其他任何字符。
现在真实的输入是：
[['今年哪吒大火，票房达到了惊人的120亿美元，很多', '哪吒'], ['大部分朋友对此有一个误解，其实哪吒的目标从来都不是商务用车，而更偏向于表达', '哪吒'], ['不是商务用车，而更偏向于表达用户的个性，XC60 的车主也有类似看法。', 'XC60']]
本批 3 条，期望输出 3 个字

==============================================================
三、2.86 亿篇文章的成本估算
==============================================================
                    输入字数(亿)        输出字数(亿)         总成本(元)
方案一                 3120.26        10.8623      759290.19
方案二                   13.67         0.1816        3454.22
方案二/方案一              0.438%         1.672%         0.455%

需要大模型判定的文章：726.44 万篇
潜在歧义车系总数：    1816.10 万个
打包后的请求次数：    908050 次

#!/bin/bash
# =============================================================================
# 学生模型 LoRA 训练骨架（LLaMA-Factory 命令行）
# 把 TODO 换成自己的路径和参数即可。webui 上点出来的那一次训练，
# 日志里会打印等价的命令行，可以直接抄回这里做成可复现的脚本。
# =============================================================================

set -euo pipefail

# ---- TODO: 三个路径，改这里 -------------------------------------------------
BASE_MODEL="TODO/models/Qwen/Qwen2.5-3B-Instruct-GPTQ-Int4"   # 基座模型目录
DATA_DIR="TODO/data/chatData"                                  # 放 dataset_info.json 的目录
OUTPUT_DIR="TODO/models/lora/qwen2.5-3b-cot"                   # 训练产物落点
DATASET_NAME="chat-train"                                      # dataset_info.json 里的键名

export CUDA_VISIBLE_DEVICES=0

llamafactory-cli train \
  `# ---- 决定路径：改一个字都会指向另一次实验 ----` \
  --model_name_or_path "$BASE_MODEL" \
  --dataset_dir "$DATA_DIR" \
  --dataset "$DATASET_NAME" \
  --output_dir "$OUTPUT_DIR" \
  --template qwen \
  `# ---- 决定性质：这几项定了这次训练到底在干什么 ----` \
  --stage sft \
  --do_train True \
  --finetuning_type lora \
  --quantization_bit 4 \
  --quantization_method bitsandbytes \
  --double_quantization True \
  `# ---- 决定效果：反复调的就是这一组 ----` \
  --learning_rate 1e-4 \
  --cutoff_len 6000 \
  --num_train_epochs 1.0 \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 2 \
  --warmup_steps 10 \
  --lr_scheduler_type cosine \
  --lora_rank 8 \
  --lora_alpha 16 \
  --lora_dropout 0.1 \
  --lora_target all \
  `# ---- 工程项：一般照抄 ----` \
  --preprocessing_num_workers 16 \
  --max_samples 100000 \
  --max_grad_norm 1.0 \
  --logging_steps 5 \
  --save_steps 100 \
  --packing False \
  --report_to none \
  --bf16 True \
  --plot_loss True \
  --trust_remote_code True \
  --optim adamw_torch \
  --ddp_timeout 180000000 \
  --include_num_input_tokens_seen True

# =============================================================================
# 训练完拿 checkpoint 起服务验证：
#
#   基座（对照组）
#     API_PORT=8000 llamafactory-cli api --model_name_or_path "$BASE_MODEL"
#
#   基座 + LoRA（实验组）
#     API_PORT=8000 llamafactory-cli api \
#       --model_name_or_path "$BASE_MODEL" \
#       --adapter_name_or_path "$OUTPUT_DIR/checkpoint-TODO" \
#       --finetuning_type lora
#
#   换 vLLM 提速，并同时挂多个 LoRA
#     CUDA_VISIBLE_DEVICES=0 API_PORT=8000 vllm serve "$BASE_MODEL" \
#       --enable-lora \
#       --lora-modules cot="$OUTPUT_DIR/checkpoint-TODO"
#
# 然后 python3 call_api_server.py 测速度、看输出差异。
# =============================================================================