🧊 动态混合精度量化 · 流程 & 源码

def collect_target_linears(model):
    """
    收集可量化线性层，跳过 vision 塔与 lm_head（保留 FP16 稳定输出）。
    """
    targets = []                                # 存储 (层名, 模块) 的列表
    for name, module in model.named_modules():  # 遍历模型所有模块
        if not isinstance(module, nn.Linear):   # 只关注线性层
            continue
        nl = name.lower()
        if "vision" in nl or "lm_head" in nl:   # 跳过视觉塔和输出头
            continue
        targets.append((name, module))          # 加入目标列表
    return targets
            

@torch.no_grad()
def compute_sensitivity(model, calib_batches, group_size=128, device="cuda"):
    targets = collect_target_linears(model)      # 获取所有目标层
    err_keys = [2, 3, 4, 8]                      # 需要评估的比特数

    # ---------- CPU 预计算：各层 4 种比特的量化结果 ----------
    pre = {}
    for name, module in targets:
        W = module.weight.detach().float().cpu()    # 权重移到 CPU
        pre[id(module)] = {
            "nparams": W.numel(),                    # 该层参数量
            "packed": {b: quantize_weight(W, b, group_size) for b in err_keys},
        }

    # 统计量：累计误差和分母 ||Y||²
    stats = {id(m): {f"err{b}": 0.0 for b in err_keys} | {"y2": 0.0}
             for _, m in targets}
    name_of = {id(m): n for n, m in targets}

    def make_hook(mid):
        layer = pre[mid]   # 该层的预量化结果
        def hook(module, inp, out):
            Wf = module.weight.detach().float().to(device)  # 原始权重到 GPU
            A = inp[0].detach().reshape(-1, Wf.shape[1]).to(device).float()
            ns = stats[mid]
            Y = (Wf @ A.T).T                    # 原始输出
            ns["y2"] += float((Y * Y).sum())   # 累积分母
            for b in err_keys:                  # 对每个比特位宽
                qw, sc, zp = layer["packed"][b]
                Wq = dequant_weight(qw.to(device), sc.to(device), zp.to(device),
                                    b, group_size, Wf.shape[1])
                Yq = (Wq @ A.T).T               # 量化后的输出
                ns[f"err{b}"] += float(((Y - Yq) ** 2).sum())  # 累加平方误差
                del qw, sc, zp, Wq
            del Wf, A, Y, Yq
        return hook

    # 注册前向 hook
    handles = [m.register_forward_hook(make_hook(id(m))) for _, m in targets]
    model.eval()
    for batch in calib_batches:                # 遍历校准数据
        batch = {k: v.to(device) for k, v in batch.items()}
        model(**batch)
    for h in handles:                         # 移除 hook
        h.remove()

    # 汇总相对误差
    sens = {}
    for _, m in targets:
        mid = id(m)
        s = stats[mid]
        y2 = s["y2"] + 1e-12
        sens[name_of[mid]] = {f"err{b}": s[f"err{b}"] / y2 for b in err_keys} | {
            "nparams": pre[mid]["nparams"]}
    return sens
            

def allocate_bits(sens, avg_bits=4.0, floor=2):
    nxt_lower = {8: 4, 4: 3, 3: 2}         # 降级路径
    nparams = {n: sens[n]["nparams"] for n in sens}
    total = sum(nparams.values()) + 1e-12
    prec = {n: 8 for n in sens}            # 初始全 8bit

    def weighted(p):
        return sum(p[n] * nparams[n] for n in p) / total

    # ---------- Pass 1: 贪心降级直到达标 ----------
    guard = 0
    while weighted(prec) > avg_bits + 1e-9 and guard < 10 ** 6:
        guard += 1
        best = None
        for n in sens:
            c = prec[n]
            if c == 2 or c not in nxt_lower:
                continue
            nx = nxt_lower[c]
            dec = sens[n][f"err{c}"] - sens[n][f"err{nx}"]   # 降档带来的误差增量
            bits_saved = c - nx
            cost = (dec * nparams[n]) / bits_saved          # 每释放 1 bit 的代价
            if best is None or cost < best[0]:
                best = (cost, n, nx)
        if best is None:
            break
        _, n, nx = best
        prec[n] = nx

    # ---------- Pass 2: 强制下限（单调上调，必然终止） ----------
    if floor and floor > 2:
        bump = {2: 3, 3: 4}
        changed = True
        while changed:
            changed = False
            for n in sens:
                if prec[n] < floor and prec[n] in bump:
                    prec[n] = bump[prec[n]]
                    changed = True
    return prec
            

def apply_quantization(model, prec, group_size=128):
    replaced = {}
    for name, module in model.named_modules():
        # 只替换 4/8 bit 的线性层
        if isinstance(module, nn.Linear) and name in prec and prec[name] in (4, 8):
            ql = QuantLinear.from_linear(module, prec[name], group_size)
            # 处理嵌套路径，找到父模块并替换属性
            if "." in name:
                parent_name, attr = name.rsplit(".", 1)
                parent = model.get_submodule(parent_name)
            else:
                parent, attr = model, name
            setattr(parent, attr, ql)          # 原地替换为量化层
            replaced[name] = prec[name]        # 记录已替换的层
    return replaced
            

def save_quantized(model, out_dir, prec, group_size, meta=None):
    from safetensors.torch import save_file
    os.makedirs(out_dir, exist_ok=True)
    tensors, keys = {}, {}
    for name, module in model.named_modules():
        if isinstance(module, QuantLinear):
            key = name.replace(".", "__")    # 替换 '.' 为 '__' 作为 safetensors key
            tensors[f"{key}.qweight"] = module.qweight
            tensors[f"{key}.scale"] = module.scale
            tensors[f"{key}.zp"] = module.zp
            keys[name] = {"bits": module.bits, "group_size": module.group_size,
                          "qweight": f"{key}.qweight", "scale": f"{key}.scale", "zp": f"{key}.zp"}
    save_file(tensors, os.path.join(out_dir, "model_quant.safetensors"))
    out = {"group_size": group_size, "precision_map": prec, "quant_keys": keys, "meta": meta or {}}
    with open(os.path.join(out_dir, "quant_config.json"), "w", encoding="utf-8") as f:
        json.dump(out, f, ensure_ascii=False, indent=2)
    return out_dir

def load_quantized_into(model, out_dir):
    from safetensors.torch import load_file
    cfg = json.load(open(os.path.join(out_dir, "quant_config.json"), encoding="utf-8"))
    sd = load_file(os.path.join(out_dir, "model_quant.safetensors"))
    gs = cfg["group_size"]
    replaced = 0
    for name, info in cfg["quant_keys"].items():
        if not isinstance(model.get_submodule(name), nn.Linear):
            continue
        bits = info["bits"]
        qw, sc, zp = sd[info["qweight"]], sd[info["scale"]], sd[info["zp"]]
        if "." in name:
            parent_name, attr = name.rsplit(".", 1)
            parent = model.get_submodule(parent_name)
        else:
            parent, attr = model, name
        bias = getattr(parent, attr).bias   # 保留原偏置
        ql = QuantLinear(bits, gs, qw, sc, zp, bias)
        setattr(parent, attr, ql)           # 替换为量化层
        replaced += 1
    return replaced