【案例】大健康行业智能问诊系统 · 底座选型与训练

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""选基座：随机初始化 vs 加载预训练权重，差别有多大。

这一步是整个项目最容易被轻轻放过、又最致命的一步。
同一份 config.json 有两种用法：

    GPT2LMHeadModel(config=cfg)              # 只借架构，权重随机
    GPT2LMHeadModel.from_pretrained(路径)     # 架构 + 预训练好的权重

两行代码长得像，打印出来的模型结构一模一样，参数量一模一样，
训练也照样能跑、loss 也照样会降——但前者等于从零开始学中文，
后者才是「微调」。用一万条医疗问答从零训一个语言模型，
学出来的只会是语法都不通的胡话。

纯标准库，可直接运行：用参数量和数据量算一笔账，说明为什么。
"""


def tokens_in_corpus(n_samples, avg_chars):
    """中文经 BertTokenizer 基本一字一 token，粗估即可。"""
    return n_samples * avg_chars


def chinchilla_tokens(n_params, ratio=20):
    """训练一个模型「应当」喂多少 token 的经验参考量级。

    这个比例来自缩放律研究的经验结论：参数量与训练 token 数大致同步增长，
    量级上每个参数配约 20 个 token。它不是精确公式，这里只用来看数量级差距。
    """
    return n_params * ratio


def fmt(n):
    for unit, base in (("万亿", 1e12), ("亿", 1e8), ("万", 1e4)):
        if n >= base:
            return "%.2f%s" % (n / base, unit)
    return "%d" % n


if __name__ == "__main__":
    N_PARAMS = 102_068_736      # 本项目使用的中文 GPT-2 规模
    N_SAMPLES = 20_000          # 手上的医疗问答条数
    AVG_CHARS = 109             # 实测的平均长度

    have = tokens_in_corpus(N_SAMPLES, AVG_CHARS)
    need = chinchilla_tokens(N_PARAMS)

    print("一、从零训练这个规模的模型，需要多少语料")
    print("  模型参数量            %s" % fmt(N_PARAMS))
    print("  从零训练的参考 token 量 %s" % fmt(need))
    print("  手上医疗语料的 token 量 %s" % fmt(have))
    print("  差距                  %.0f 倍" % (need / have))
    print("\n  → 差三个数量级。指望用这点数据从零学会中文，不现实。")

    print("\n二、两种写法的区别")
    rows = [
        ("模型结构", "完全相同", "完全相同"),
        ("参数量", "完全相同", "完全相同"),
        ("初始权重", "按正态分布随机采样", "预训练学到的值"),
        ("初始 loss", "≈ ln(词表大小)", "明显更低"),
        ("训练能不能跑", "能跑，loss 也会降", "能跑"),
        ("学到的东西", "从零学中文，数据远远不够", "把已有的中文能力调向医疗问答"),
        ("适用场景", "词表/架构大改，且有海量语料", "绝大多数业务场景"),
    ]
    print("  %-14s %-24s %s" % ("对比项", "随机初始化", "加载预训练权重"))
    for a, b, c in rows:
        print("  %-14s %-24s %s" % (a, b, c))

    import math
    print("\n三、开训第一步的 loss 能一眼认出用的是哪种")
    vocab = 21128
    print("  随机初始化时，模型对每个 token 等概率瞎猜，")
    print("  交叉熵 ≈ ln(%d) = %.2f" % (vocab, math.log(vocab)))
    print("  → 第一个 batch 打出 loss≈%.1f，说明权重是随机的；" % math.log(vocab))
    print("    打出 loss≈2~4，说明预训练权重确实加载上了。")
    print("  这是开训 30 秒内就能做的一次核对，别等训完才发现。")

    print("\n四、选基座的三条判据")
    for i, line in enumerate([
        "语言对不对：中文任务用中文语料训过的底座，别拿纯英文底座硬掰",
        "尺寸配不配：底座越大效果越好但显存越贵，先用小底座把流程跑通",
        "许可证准不准：商用项目要看清开源协议，这件事上线前一定会被问到",
    ], 1):
        print("  %d. %s" % (i, line))

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""token 准确率：位移对齐这一步，写错了分数会虚高。

自回归模型第 i 个位置的输出，预测的是第 i+1 个位置的 token。
所以拿 logits 和 labels 直接比是错位的，必须先对齐：

    logits = logits[:, :-1, :]   去掉最后一个位置的预测（它没有下文可对）
    labels = labels[:, 1:]       去掉第一个位置的标签（它没有上文来预测）

对齐之后再把 -100 的位置剔掉，剩下的才是「真正被预测、且该算分」的 token。

这里用纯 Python 列表把整个过程复现一遍，不依赖 torch，可直接运行。
"""

IGNORE = -100


def argmax(row):
    """手写 argmax，避免依赖任何第三方库。"""
    best_i, best_v = 0, row[0]
    for i, v in enumerate(row):
        if v > best_v:
            best_i, best_v = i, v
    return best_i


def calc_accuracy(logits, labels, ignore_index=IGNORE, shift=True):
    """返回 (正确数, 总数)。

    logits: [batch, seq_len, vocab]
    labels: [batch, seq_len]
    shift:  是否做位移对齐。设成 False 是为了演示不对齐会得到什么分数。
    """
    correct = total = 0

    for b in range(len(labels)):
        if shift:
            preds_src = logits[b][:-1]
            gold = labels[b][1:]
        else:
            preds_src = logits[b]
            gold = labels[b]

        for step, row in enumerate(preds_src):
            target = gold[step]
            if target == ignore_index:      # 提示词和 padding 不计分
                continue
            total += 1
            if argmax(row) == target:
                correct += 1

    return correct, total


def one_hot_logits(token_ids, vocab, peak=5.0):
    """造一组「预测完全正确」的 logits，用来做对照实验。"""
    out = []
    for tid in token_ids:
        row = [0.0] * vocab
        row[tid % vocab] = peak
        out.append(row)
    return out


if __name__ == "__main__":
    VOCAB = 10

    # 一条样本：前 3 个位置是提示词（屏蔽），后面是答案
    labels = [[-100, -100, -100, 4, 5, 6, 7]]

    # 造一个「完美模型」：位置 i 的输出正好等于位置 i+1 的标签
    perfect = [one_hot_logits([0, 0, 4, 5, 6, 7, 0], VOCAB)]

    print("一、对齐与不对齐，同一份预测两个分数")
    c1, t1 = calc_accuracy(perfect, labels, shift=True)
    c2, t2 = calc_accuracy(perfect, labels, shift=False)
    print("  做位移对齐：  %d/%d = %.1f%%" % (c1, t1, 100.0 * c1 / t1))
    print("  不做对齐：    %d/%d = %.1f%%" % (c2, t2, 100.0 * c2 / t2))
    assert c1 == t1, "完美预测在对齐后应当是 100%"
    print("  → 同一个模型，少写一行位移，分数从 100%% 掉到 %.0f%%。"
          % (100.0 * c2 / t2))
    print("    这个错误的方向不固定，有时虚高有时虚低，但一定是错的。")

    print("\n二、-100 到底挡掉了多少位置")
    flat = [x for row in labels for x in row]
    n_ignore = sum(1 for x in flat if x == IGNORE)
    print("  标签总长 %d，其中 -100 有 %d 个" % (len(flat), n_ignore))
    print("  对齐后参与计分的位置：%d 个" % t1)
    print("  → 如果把 -100 也算进分母，分数会被大量提示词位置稀释。")
    print("    这也是同一份模型在不同仓库里分数对不上的常见原因。")

    print("\n三、一个更接近真实的例子：预测对一半")
    half = [one_hot_logits([0, 0, 4, 5, 9, 9, 0], VOCAB)]   # 后两个猜错
    c3, t3 = calc_accuracy(half, labels)
    print("  正确 %d / 总数 %d = %.1f%%" % (c3, t3, 100.0 * c3 / t3))
    assert (c3, t3) == (2, 4)
    print("  断言通过：4 个待预测位置里对了 2 个")

    print("\n四、这个指标能说明什么、不能说明什么")
    print("  能说明：训练有没有在推进（它和 loss 应当同向变化）")
    print("  不能说明：答复对不对。token 级全对只意味着复述得一模一样，")
    print("           而同一个意思换种说法，token 准确率可能很低却完全正确。")
    print("  所以它只配当训练过程的监控指标，不能当验收标准。")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""梯度累积、预热、裁剪三件事在同一个循环里怎么配合。

这三样都作用在「参数更新」这一步附近，顺序写错了不会报错，
但会让训练效果说不清道不明地变差。正确顺序只有一种：

    前向 → 算 loss → loss / accum → 反向 → 【攒够了才】裁剪 → step → 调度 → 清零

最常见的两个错位：
  ① 每个 batch 都裁剪：累积到一半就把梯度削了，累积出来的方向失真
  ② 每个 batch 都 scheduler.step()：学习率按 batch 数推进，
     提前跑完整条曲线，后半程学习率一直贴着 0

纯标准库，把这个循环的时序完整模拟一遍，可直接运行。
"""
import math


class FakeScheduler:
    """线性预热 + 线性衰减，只记步数不碰真实参数。"""

    def __init__(self, base_lr, warmup, total):
        self.base_lr, self.warmup, self.total = base_lr, warmup, total
        self.step_count = 0

    def step(self):
        self.step_count += 1

    def get_lr(self):
        s = self.step_count
        if s < self.warmup:
            return self.base_lr * s / max(1, self.warmup)
        progress = (s - self.warmup) / max(1, self.total - self.warmup)
        return self.base_lr * max(0.0, 1.0 - progress)


def grad_norm(grads):
    """所有梯度拼成一个大向量之后的 L2 范数——裁剪判断的就是它。"""
    return math.sqrt(sum(g * g for g in grads))


def clip(grads, max_norm):
    """返回 (裁剪后的梯度, 缩放系数)。

    注意裁剪是「整体等比缩放」，不是逐个截断。
    所以它只改梯度的长度，不改方向。
    """
    total = grad_norm(grads)
    if total <= max_norm:
        return grads, 1.0
    scale = max_norm / (total + 1e-6)
    return [g * scale for g in grads], scale


def simulate(n_batches, accum, base_lr=2e-4, max_norm=1.0, warmup=3):
    """走一遍训练循环，记录每次真正更新时的状态。"""
    total_updates = n_batches // accum
    sched = FakeScheduler(base_lr, warmup, total_updates)

    buffer = [0.0, 0.0]          # 假装模型只有两个参数
    log = []

    for step in range(1, n_batches + 1):
        # 造一个有大有小的梯度，第 7 个 batch 故意来一个异常大的
        raw = [0.3 * step % 1.1, 0.2 * step % 0.9]
        if step == 7:
            raw = [8.0, 6.0]     # 异常样本

        loss_scale = 1.0 / accum         # ← 关键：loss 先除以累积步数
        buffer = [b + g * loss_scale for b, g in zip(buffer, raw)]

        if step % accum == 0:            # 攒够了才动手
            clipped, scale = clip(buffer, max_norm)
            sched.step()
            log.append({
                "batch": step,
                "update": sched.step_count,
                "norm_before": grad_norm(buffer),
                "scale": scale,
                "lr": sched.get_lr(),
            })
            buffer = [0.0, 0.0]          # 清零，必须在 step 之后
    return log, total_updates


if __name__ == "__main__":
    N_BATCHES, ACCUM = 24, 4

    log, total = simulate(N_BATCHES, ACCUM)
    print("一、正确时序：%d 个 batch，累积 %d 步 → %d 次参数更新"
          % (N_BATCHES, ACCUM, total))
    print("  %-8s %-8s %-12s %-8s %-10s"
          % ("batch", "update", "裁剪前范数", "缩放", "学习率"))
    for r in log:
        print("  %-8d %-8d %-12.4f %-8.3f %-10.3e"
              % (r["batch"], r["update"], r["norm_before"], r["scale"], r["lr"]))

    clipped_rows = [r for r in log if r["scale"] < 1.0]
    print("\n  发生裁剪的更新：%d 次（异常梯度被按比例压回 max_norm）"
          % len(clipped_rows))
    assert clipped_rows, "第 7 个 batch 的异常梯度应当触发裁剪"

    print("\n二、学习率曲线：先升后降")
    lrs = [r["lr"] for r in log]
    peak = lrs.index(max(lrs))
    print("  第 %d 次更新到达峰值 %.3e，之后单调下降"
          % (peak + 1, max(lrs)))
    assert lrs[-1] < max(lrs), "最后应当已经衰减下来"

    print("\n三、把 scheduler.step() 错放到每个 batch 会怎样")
    wrong = FakeScheduler(2e-4, warmup=3, total=total)
    for _ in range(N_BATCHES):
        wrong.step()
    print("  调度器被推进了 %d 次，而总步数只规划了 %d 次"
          % (wrong.step_count, total))
    print("  此时学习率 = %.3e（曲线早就跑完，后半程一直贴着 0）"
          % wrong.get_lr())

    print("\n四、等效批大小对照")
    print("  单卡批 4 × 累积 %d = 等效批 %d" % (ACCUM, 4 * ACCUM))
    print("  显存按「单卡批 4」算，而模型感觉到的是「批 %d」——" % (4 * ACCUM))
    print("  这正是小显存上训大批次的唯一办法。")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""医疗问诊模型的完整训练脚本。

流程：加载数据 → 加载底座 → 逐轮训练 → 逐轮评估 → 保存最好的那一轮。
每一步都带一个能当场核对的输出，因为训练出问题时不会报错，
只会让最终效果不对，事后很难倒查。
"""
import json
import os
from datetime import datetime

import torch
from torch.optim import AdamW
from tqdm import tqdm
from transformers import (AutoTokenizer, GPT2LMHeadModel,
                          get_linear_schedule_with_warmup)

from dataset_and_collate import get_dataloaders


class Config:
    """所有能调的东西集中在这里，不散落在代码各处。"""
    # 底座：走环境变量，方便在本地目录与 Hub 名字之间切换
    base_model = os.environ.get("BASE_MODEL", "uer/gpt2-chinese-cluecorpussmall")
    vocab_path = os.environ.get("VOCAB_PATH", "vocab/vocab.txt")

    train_path = os.environ.get("TRAIN_PATH", "data/medical_train.jsonl")
    valid_path = os.environ.get("VALID_PATH", "data/medical_valid.jsonl")

    batch_size = 4
    grad_accum = 8          # 等效批大小 4 × 8 = 32
    epochs = 4
    lr = 2.6e-5
    eps = 1e-9
    max_grad_norm = 1.0
    warmup_ratio = 0.03

    log_every = 20
    save_dir = os.environ.get("SAVE_DIR", "checkpoints/best_model")
    device = "cuda" if torch.cuda.is_available() else "cpu"


def calc_accuracy(logits, labels, ignore_index=-100):
    """token 准确率。位移对齐这一步不能省。"""
    logits = logits[:, :-1, :]        # 去掉最后一个位置的预测
    labels = labels[:, 1:]            # 去掉第一个位置的标签

    preds = logits.argmax(dim=-1).reshape(-1)
    labels = labels.reshape(-1)

    not_pad = labels.ne(ignore_index)             # 提示词与 padding 不计分
    correct = preds.eq(labels).masked_select(not_pad).sum()
    return correct.item(), not_pad.sum().item()


def build_model(cfg):
    """加载底座并核对词表大小。

    词表对不上是这一步最容易踩的坑：分词器编出来的 id 超过模型词表范围，
    会在 embedding 查表时抛越界，而报错信息完全看不出是词表问题。
    """
    model = GPT2LMHeadModel.from_pretrained(cfg.base_model)
    tokenizer = AutoTokenizer.from_pretrained(cfg.vocab_path)

    assert tokenizer.vocab_size == model.config.vocab_size, (
        "分词器词表 %d 与模型词表 %d 不一致"
        % (tokenizer.vocab_size, model.config.vocab_size))

    n_param = sum(p.numel() for p in model.parameters())
    print("✅ 底座已加载，参数量 %d，词表 %d" % (n_param, model.config.vocab_size))
    return model.to(cfg.device), tokenizer


@torch.no_grad()
def evaluate(model, loader, cfg):
    """验证集上跑一遍，返回 (平均 loss, token 准确率)。"""
    model.eval()
    total_loss = correct = total = 0

    for batch in tqdm(loader, desc="评估", leave=False):
        batch = {k: v.to(cfg.device) for k, v in batch.items()}
        out = model(**batch)
        total_loss += out.loss.item()

        c, t = calc_accuracy(out.logits, batch["labels"])
        correct += c
        total += t

    return total_loss / max(1, len(loader)), correct / max(1, total)


def train_one_epoch(model, loader, optimizer, scheduler, cfg, epoch):
    model.train()
    started = datetime.now()
    total_loss = correct = total = 0

    bar = tqdm(loader, desc="第 %d 轮" % (epoch + 1))
    for step, batch in enumerate(bar):
        batch = {k: v.to(cfg.device) for k, v in batch.items()}

        # 传了 labels，模型内部就会自己算交叉熵，并自动做位移对齐
        out = model(**batch)
        loss = out.loss
        total_loss += loss.item()

        # 先按累积步数缩放，再反向。不缩放的话等效批大了、梯度也大了一倍
        (loss / cfg.grad_accum).backward()

        c, t = calc_accuracy(out.logits, batch["labels"])
        correct += c
        total += t

        # 攒够一个累积周期才真正更新
        if (step + 1) % cfg.grad_accum == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

        if (step + 1) % cfg.log_every == 0:
            bar.set_postfix(loss=total_loss / (step + 1),
                            acc=correct / max(1, total),
                            lr=scheduler.get_last_lr()[0])

    print("第 %d 轮训练完成，平均 loss %.4f，token 准确率 %.4f，耗时 %s"
          % (epoch + 1, total_loss / max(1, len(loader)),
             correct / max(1, total), datetime.now() - started))
    return total_loss / max(1, len(loader)), correct / max(1, total)


def main():
    cfg = Config()
    print("设备：%s" % cfg.device)

    train_loader, valid_loader, _tok = get_dataloaders(
        cfg.train_path, cfg.valid_path, cfg.batch_size)
    print("✅ 训练 %d 批 / 验证 %d 批" % (len(train_loader), len(valid_loader)))

    model, tokenizer = build_model(cfg)

    optimizer = AdamW(model.parameters(), lr=cfg.lr, eps=cfg.eps)

    # 总步数按「参数更新次数」算，不是按 batch 数。
    # 这里除以累积步数是整个调度器配置里最容易写错的一处。
    total_steps = len(train_loader) // cfg.grad_accum * cfg.epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=int(total_steps * cfg.warmup_ratio),
        num_training_steps=total_steps)
    print("✅ 总参数更新次数 %d，其中预热 %d 次"
          % (total_steps, int(total_steps * cfg.warmup_ratio)))

    history, best_loss = [], float("inf")

    for epoch in range(cfg.epochs):
        tr_loss, tr_acc = train_one_epoch(
            model, train_loader, optimizer, scheduler, cfg, epoch)
        va_loss, va_acc = evaluate(model, valid_loader, cfg)

        print("第 %d 轮评估：loss %.4f，token 准确率 %.4f"
              % (epoch + 1, va_loss, va_acc))
        history.append({"epoch": epoch + 1, "train_loss": tr_loss,
                        "train_acc": tr_acc, "valid_loss": va_loss,
                        "valid_acc": va_acc})

        # 只保留验证 loss 最低的那一轮，而不是最后一轮。
        # 最后一轮往往已经开始过拟合了。
        if va_loss < best_loss:
            best_loss = va_loss
            os.makedirs(cfg.save_dir, exist_ok=True)
            model.save_pretrained(cfg.save_dir)
            tokenizer.save_pretrained(cfg.save_dir)
            print("✅ 第 %d 轮是目前最好的，已保存到 %s" % (epoch + 1, cfg.save_dir))

    with open(os.path.join(cfg.save_dir, "history.json"), "w", encoding="utf-8") as f:
        json.dump(history, f, ensure_ascii=False, indent=2)

    print("\n训练结束。逐轮记录：")
    for h in history:
        print("  轮次 %d  训练 loss %.4f  验证 loss %.4f  验证准确率 %.4f"
              % (h["epoch"], h["train_loss"], h["valid_loss"], h["valid_acc"]))
    print("\n判读方法：训练 loss 降而验证 loss 开始升的那一轮，就是过拟合的起点。")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""同一个任务换成 LoRA 训法，用来和全量微调那份逐行对照。

和 train_medical_gpt2.py 的差别只有三处，其余一模一样：
  ① 加载底座之后多一步 get_peft_model
  ② 学习率要调大一到两个数量级
  ③ 保存的是旁支而不是整个模型

把两份并排看，就能明白 PEFT 在工程上到底改变了什么：
训练循环没变，变的是「谁在被更新」和「产物有多大」。
"""
import os

import torch
from peft import LoraConfig, TaskType, get_peft_model
from torch.optim import AdamW
from transformers import (AutoTokenizer, GPT2LMHeadModel,
                          get_linear_schedule_with_warmup)

from dataset_and_collate import get_dataloaders
from train_medical_gpt2 import Config, evaluate, train_one_epoch


class LoraTrainConfig(Config):
    # 只训旁支，主干冻着，步子可以迈大得多。
    # 全量微调用 2.6e-5，这里用 2e-4，差了近十倍——
    # 照搬全量微调的学习率是 LoRA 效果差的头号原因。
    lr = 2e-4
    epochs = 4
    save_dir = os.environ.get("SAVE_DIR", "checkpoints/lora-adapter")


def build_lora_model(cfg):
    """加载底座并挂上旁支。"""
    model = GPT2LMHeadModel.from_pretrained(cfg.base_model)
    tokenizer = AutoTokenizer.from_pretrained(cfg.vocab_path)

    lora_cfg = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=16,           # 惯例设成 r 的两倍，缩放系数 alpha/r = 2
        lora_dropout=0.05,
        bias="none",
        # GPT-2 把 Q/K/V 合并成一个 c_attn，所以只写这一个名字。
        # 换成 Qwen / LLaMA 要改成 q_proj / k_proj / v_proj / o_proj。
        target_modules=["c_attn"],
    )
    model = get_peft_model(model, lora_cfg)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print("✅ 可训练参数 %d / %d = %.4f%%"
          % (trainable, total, 100.0 * trainable / total))

    # 占比应当落在零点几个百分点。接近 100% 说明 target_modules 名字写错了，
    # 旁支根本没挂上去，这一轮训的其实是全量微调。
    assert trainable / total < 0.05, "旁支没挂上，检查 target_modules"

    return model.to(cfg.device), tokenizer


def main():
    cfg = LoraTrainConfig()
    train_loader, valid_loader, _ = get_dataloaders(
        cfg.train_path, cfg.valid_path, cfg.batch_size)

    model, tokenizer = build_lora_model(cfg)

    # 只把需要梯度的参数交给优化器。
    # 把全部参数一股脑传进去也不会报错，但优化器会为冻结参数也分配状态，
    # LoRA 省下来的显存就这么白白还回去了。
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = AdamW(params, lr=cfg.lr, eps=cfg.eps)

    total_steps = len(train_loader) // cfg.grad_accum * cfg.epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=int(total_steps * cfg.warmup_ratio),
        num_training_steps=total_steps)

    best_loss = float("inf")
    for epoch in range(cfg.epochs):
        train_one_epoch(model, train_loader, optimizer, scheduler, cfg, epoch)
        va_loss, va_acc = evaluate(model, valid_loader, cfg)
        print("第 %d 轮评估：loss %.4f，token 准确率 %.4f"
              % (epoch + 1, va_loss, va_acc))

        if va_loss < best_loss:
            best_loss = va_loss
            os.makedirs(cfg.save_dir, exist_ok=True)
            # 这里存下来的只有旁支，几百 KB 量级，不是几百 MB 的整模型
            model.save_pretrained(cfg.save_dir)
            tokenizer.save_pretrained(cfg.save_dir)

    size = sum(os.path.getsize(os.path.join(cfg.save_dir, f))
               for f in os.listdir(cfg.save_dir))
    print("\n✅ 旁支已保存，目录总大小 %.2f MB" % (size / 1024 / 1024))
    print("部署时两种选择：")
    print("  ① 底座 + 旁支分开加载：一个底座可以挂多个任务的旁支，随时切换")
    print("  ② merge_and_unload() 合并回权重：推理零额外开销，但只对应这一个任务")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""可复制改造的训练循环骨架。

把 TODO 填完即可开训。刻意保留了四个「开训前必须看到」的证据点，
因为训练过程出错几乎从不报错，只会让结果变差：

  证据点 1  第一个 batch 的 loss ≈ ln(词表大小)？→ 底座没加载上
  证据点 2  可训练参数占比是不是预期的量级？
  证据点 3  总更新次数与预热次数是否合理？
  证据点 4  第一轮结束后 loss 有没有真的往下走？
"""
import math
import os

import torch
from torch.optim import AdamW
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          get_linear_schedule_with_warmup)

BASE_MODEL = os.environ.get("BASE_MODEL", "TODO-填底座")
SAVE_DIR = os.environ.get("SAVE_DIR", "checkpoints/best_model")

BATCH_SIZE = 4
GRAD_ACCUM = 8
EPOCHS = 3
LR = 2e-5              # TODO: 全量微调用 1e-5~5e-5；LoRA 用 1e-4~3e-4
MAX_GRAD_NORM = 1.0
WARMUP_RATIO = 0.03


def sanity_check_first_batch(model, batch, vocab_size):
    """证据点 1：用第一个 batch 的 loss 判断底座到底有没有加载上。"""
    model.eval()
    with torch.no_grad():
        loss = model(**batch).loss.item()

    expected_random = math.log(vocab_size)
    print("首个 batch loss = %.4f（随机初始化的参考值 ln(%d) = %.2f）"
          % (loss, vocab_size, expected_random))

    if loss > expected_random * 0.9:
        print("⚠ loss 贴近随机初始化的水平，底座很可能没加载上——")
        print("  检查是不是误用了 Model(config=...) 而不是 from_pretrained(...)")
    else:
        print("✅ loss 明显低于随机水平，预训练权重已生效")
    model.train()
    return loss


def report_trainable(model):
    """证据点 2：可训练参数占比。"""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print("可训练参数 %d / %d = %.4f%%" % (trainable, total, 100.0 * trainable / total))
    return trainable, total


def build_scheduler(optimizer, n_batches):
    """证据点 3：总步数一定要除以累积步数。"""
    total_steps = n_batches // GRAD_ACCUM * EPOCHS
    warmup = max(10, int(total_steps * WARMUP_RATIO))
    print("总参数更新 %d 次，其中预热 %d 次（占 %.1f%%）"
          % (total_steps, warmup, 100.0 * warmup / max(1, total_steps)))
    if warmup > total_steps * 0.3:
        print("⚠ 预热占比过高，学习率可能全程都没进入衰减阶段")
    return get_linear_schedule_with_warmup(optimizer, warmup, total_steps)


def train(model, train_loader, valid_loader, device):
    optimizer = AdamW([p for p in model.parameters() if p.requires_grad], lr=LR)
    scheduler = build_scheduler(optimizer, len(train_loader))

    best_loss = float("inf")
    for epoch in range(EPOCHS):
        model.train()
        running = 0.0

        for step, batch in enumerate(train_loader):
            batch = {k: v.to(device) for k, v in batch.items()}
            loss = model(**batch).loss
            running += loss.item()

            (loss / GRAD_ACCUM).backward()

            if (step + 1) % GRAD_ACCUM == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

        train_loss = running / max(1, len(train_loader))

        # 证据点 4
        model.eval()
        with torch.no_grad():
            valid_loss = sum(
                model(**{k: v.to(device) for k, v in b.items()}).loss.item()
                for b in valid_loader) / max(1, len(valid_loader))

        print("轮次 %d  训练 loss %.4f  验证 loss %.4f"
              % (epoch + 1, train_loss, valid_loss))

        if valid_loss < best_loss:
            best_loss = valid_loss
            os.makedirs(SAVE_DIR, exist_ok=True)
            model.save_pretrained(SAVE_DIR)
            print("  ✅ 已保存（当前最好）")
        else:
            print("  验证 loss 没有改善，这一轮不保存——过拟合可能已经开始")

    return best_loss


def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL).to(device)

    # TODO: 换成你自己的 dataloader
    raise SystemExit(
        "把 train_loader / valid_loader 接上之后，依次调用：\n"
        "  sanity_check_first_batch(model, next(iter(train_loader)), "
        "model.config.vocab_size)\n"
        "  report_trainable(model)\n"
        "  train(model, train_loader, valid_loader, device)")


if __name__ == "__main__":
    main()