Tequila 三值量化训练项目 · 代码检视 2026-09-19（第 1 轮）

① 项目目标与预期

② 项目简介

③ Mermaid 训练流程图

④ 分模块核心代码片段

def replace_linear_layers(model, eps=1e-3, group_size=128):
    """替换 MLP 中的 nn.Linear 为 UltraQuantLinear"""
    replaced = 0
    # 注意: 需与 Qwen3.5 hybrid attention 结构核对, 确认 q/k/v/o_proj 是否已含于 linear_attn
    skip_keywords = ["linear_attn", "lm_head", "embed_tokens"]
    to_replace = []
    for full_name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            if any(kw in full_name for kw in skip_keywords):
                continue
            to_replace.append((full_name, module))

    for full_name, module in to_replace:
        parts = full_name.split(".")
        parent = model
        for p in parts[:-1]:
            parent = getattr(parent, p)
        child_name = parts[-1]

        quant_layer = UltraQuantLinear(
            module.in_features, module.out_features,
            bias=module.bias is not None, eps=eps, group_size=group_size
        )
        # 复用原权重张量, 省一份显存
        quant_layer.weight = module.weight
        if module.bias is not None:
            quant_layer.bias = module.bias
        setattr(parent, child_name, quant_layer)
        replaced += 1
    return replaced

def absmean(x):
    """每输出通道一个 scale: [out, in] -> [out, 1]"""
    scale = x.abs().mean([-1], keepdim=True)
    delta = scale / 2
    return scale, delta


class UltraQuantV2(torch.autograd.Function):
    """腾讯 Tequila 核心: dead weights 重新激活为 bias"""
    @staticmethod
    def forward(ctx, input, eps, scale_corr=False):
        original_shape = input.shape
        x = input.reshape(-1, input.shape[-1]) if input.dim() == 2 else input.reshape(input.shape[0], -1)

        scale, delta = absmean(x)
        A = torch.zeros_like(x)
        A[x >= delta] = 1
        A[x <= -delta] = -1
        A = A * scale                       # 活权重 = 符号 * scale

        B = torch.zeros_like(x)
        mask_B = A == 0                     # 死区掩码
        eps_t = torch.tensor(eps, device=x.device, dtype=x.dtype)
        B[mask_B] = eps_t * x[mask_B]       # 死权重 -> eps * 原值, 作为偏置分量

        A = A.reshape(original_shape)
        B = B.reshape(original_shape)
        mask_B = mask_B.reshape(original_shape)

        ctx.save_for_backward(mask_B, eps_t, scale)
        ctx.original_shape = original_shape
        ctx.scale_corr = scale_corr
        return A, B

    @staticmethod
    def backward(ctx, grad_A, grad_B):
        mask_B, eps_t, scale = ctx.saved_tensors
        original_shape = ctx.original_shape

        grad_A = grad_A.reshape(mask_B.shape)
        grad_B = grad_B.reshape(mask_B.shape)

        # 活权重走 STE, 死权重走直接梯度路径, 两者相加
        grad_output = grad_A.clone()
        grad_output[mask_B] = grad_output[mask_B] + eps_t * grad_B[mask_B]

        # 官方注释掉的 scale 修正: 默认关闭保持与官方一致
        if ctx.scale_corr:
            gscale = torch.ones_like(grad_output) * scale
            gscale[mask_B] = 1
            grad_output = grad_output * gscale

        grad_output = grad_output.reshape(original_shape)
        return grad_output, None, None

class UltraQuantLinear(nn.Linear):
    """output = X @ A + rowsum(B) + bias"""

    def __init__(self, in_features, out_features, bias=False, eps=1e-3, group_size=128):
        super().__init__(in_features, out_features, bias=bias)
        self.eps = eps
        self.group_size = group_size    # 接口预留, 当前 forward 未使用

    def forward(self, input_):
        A, B = UltraQuantV2.apply(self.weight, self.eps)
        A = A.to(input_.dtype)
        B = B.to(input_.dtype)
        # B.sum 等价 ones @ B.T, 免一次 matmul; 推理可离线 fuse
        out = F.linear(input_, A) + B.sum(dim=-1)
        if self.bias is not None:
            out = out + self.bias
        return out

class SlimTokenDataset(Dataset):
    def __init__(self, chunks_dir, cache_path="cache/input_ids_only.pt", max_samples=None):
        if os.path.exists(cache_path):
            payload = torch.load(cache_path, map_location="cpu", weights_only=True)
            self.input_ids = payload["input_ids"]
            self.attention_mask = payload["attention_mask"]
        else:
            import glob
            chunk_files = sorted(glob.glob(os.path.join(chunks_dir, "chunk_*.pt")))
            ids_list, mask_list = [], []
            for cf in chunk_files:
                d = torch.load(cf, map_location="cpu", weights_only=True)
                ids_list.extend(item["input_ids"] for item in d)
                mask_list.extend(item["attention_mask"] for item in d)
                del d
            self.input_ids = torch.cat(ids_list, dim=0)
            self.attention_mask = torch.cat(mask_list, dim=0)
            # 原子写: 先写 .tmp 再 os.replace, 防中途被杀留下损坏缓存
            tmp_path = cache_path + ".tmp"
            torch.save({"input_ids": self.input_ids,
                        "attention_mask": self.attention_mask}, tmp_path)
            os.replace(tmp_path, cache_path)

        if max_samples:
            self.input_ids = self.input_ids[:max_samples]
            self.attention_mask = self.attention_mask[:max_samples]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, idx):
        return {"input_ids": self.input_ids[idx],
                "attention_mask": self.attention_mask[idx]}


def collate_fn(batch):
    return {
        "input_ids": torch.stack([b["input_ids"] for b in batch]),
        "attention_mask": torch.stack([b["attention_mask"] for b in batch]),
    }

# ── 优化器（8bit AdamW 省显存）──
import bitsandbytes as bnb
optimizer = bnb.optim.AdamW8bit(
    [p for p in model.parameters() if p.requires_grad],
    lr=args.lr, weight_decay=0.0
)

# ── 学习率调度: 按"优化器步数"配置 ──
opt_steps = args.max_steps // args.grad_accum
warmup_steps = int(opt_steps * args.warmup_ratio)
lr_scheduler = get_scheduler(
    "cosine", optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=opt_steps
)

# ── 训练循环核心 ──
while step < args.max_steps and not StopRequested:
    for batch in loader:
        steps_in_run = step - start_step
        # 启动/resume 前 30 步是显存峰值期, 不检查 STOP
        if step >= args.max_steps or (steps_in_run >= 30 and
                                     (StopRequested or check_stop_file())):
            break

        input_ids = batch["input_ids"].to("cuda", non_blocking=True)
        attention_mask = batch["attention_mask"].to("cuda", non_blocking=True)

        out = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = out.logits
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = input_ids[..., 1:].contiguous()
        loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=tokenizer.pad_token_id
        )

        loss_norm = loss / args.grad_accum
        loss_norm.backward()

        if (step + 1) % args.grad_accum == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            # 不能周期性 empty_cache: 会造成显存锯齿峰值误触监控

        step += 1
        if step % args.save_interval == 0:
            ckpt_path = os.path.join(args.output_dir, f"tequila_{step}.pt")
            save_checkpoint(model, optimizer, lr_scheduler, step, ckpt_path, log)

def save_checkpoint(model, optimizer, lr_scheduler, step, path, log_obj):
    t0 = time.time()
    model_state = {
        name: p.data.detach().cpu().clone()
        for name, p in model.named_parameters() if p.requires_grad
    }
    state = {
        "step": step,
        "model": model_state,
        "optimizer": optimizer.state_dict(),
        "scheduler": lr_scheduler.state_dict(),
    }
    # 原子写: 先 .tmp 再 replace, 防中断产生 0MB/截断文件
    torch.save(state, path + ".tmp")
    os.replace(path + ".tmp", path)

    # 滚动清理: 只保留最近 3 个, final/deploy 保留
    import glob as _glob
    olds = sorted(_glob.glob(os.path.join(os.path.dirname(path), 'tequila_*.pt')),
                  key=os.path.getmtime)
    for old in olds[:-3]:
        if 'final' not in old and 'deploy' not in old:
            try:
                os.remove(old)
            except OSError:
                pass
    log_obj.info(f"  Saved checkpoint ({time.time()-t0:.1f}s): {path}")


def load_checkpoint(model, optimizer, lr_scheduler, path, log_obj, load_optimizer=True):
    ckpt = torch.load(path, map_location="cpu")
    model_state = ckpt.get("model", {})
    n_loaded = 0
    for name, param in model.named_parameters():
        if name in model_state:
            param.data.copy_(model_state[name].to(param.device))
            n_loaded += 1
    if load_optimizer and "optimizer" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer"])
    if "scheduler" in ckpt:
        lr_scheduler.load_state_dict(ckpt["scheduler"])
    start_step = ckpt.get("step", 0)
    del ckpt, model_state
    gc.collect()
    log_obj.info(f"Resumed {n_loaded} tensors from step {start_step}: {path}")
    return start_step

while True:
    time.sleep(args.check_interval)

    ret = proc.poll()
    if ret is not None:
        exit_reason = f"training exited (code {ret}); 重启由日循环管理器负责"
        break

    gpus = get_gpu_usage()
    if not gpus:
        continue

    over = [g for g in gpus if g[1] > threshold_mb]

    # 熔断: 任一卡超阈值 -> 优雅停止 (启动豁免期内不告警)
    if over and not stop_sent and time.time() - launch_time > startup_grace:
        worst = max(over, key=lambda g: g[1])
        log.warning(f"VRAM ALERT on GPU {worst[0]}: {worst[1]/1024:.1f}GB")
        with open(STOP_FILE, 'w') as f:
            f.write(str(time.time()))
        stop_sent = True
        stop_sent_at = time.time()
        continue

    # STOP 已发送: 显存回落撤销, 超时强杀
    if stop_sent:
        peak_now = max((g[1] for g in gpus), default=0)
        if peak_now < threshold_mb * 0.95:
            log.info("VRAM back under threshold, canceling STOP")
            stop_sent = False
            if os.path.exists(STOP_FILE):
                os.remove(STOP_FILE)
        elif time.time() - stop_sent_at > args.grace_sec:
            log.warning("Training did not exit in grace, force killing...")
            kill_process_tree(proc.pid)
            break
        continue

def in_window(now=None):
    """训练窗口: 每天 06:00 -> 17:40"""
    now = now or datetime.now()
    if now.hour < 6:
        return False
    if now.hour > 17:
        return False
    if now.hour == 17 and now.minute >= 40:
        return False
    return True


def stop_training():
    with open(STOP_FILE, 'w') as f:
        f.write(str(time.time()))
    log.info("STOP written, waiting graceful exit (max 300s)")
    for _ in range(100):
        time.sleep(3)
        if proc is None or proc.poll() is not None:
            break
    if proc and proc.poll() is None:
        log.warning("still alive after 300s, force kill")
        subprocess.run(['taskkill', '/F', '/T', '/PID', str(proc.pid)],
                       capture_output=True)
    if os.path.exists(STOP_FILE):
        os.remove(STOP_FILE)


def completed():
    """完成判定: Done! 且步数 >= MAX_STEPS"""
    nl = newest('logs/train_tequila_2026*.log')
    if nl and os.path.exists(nl):
        with open(nl, encoding='utf-8', errors='ignore') as f:
            content = f.read()
        if 'Done!' in content:
            import re
            m = re.search(r'Done! (\d+) steps', content)
            done_steps = int(m.group(1)) if m else -1
            return done_steps >= MAX_STEPS
    return False

def main():
    log.info("watchdog start: watching for training completion")
    train_log = newest('logs/train_tequila_2026*.log')
    mon_log = newest('logs/gpu_monitor_dual_*.log')
    alert_sent = False
    mon_base_len = os.path.getsize(mon_log) if mon_log else 0

    while True:
        time.sleep(60)
        # 新日志文件出现则切换（自愈重启会产生新文件）
        nl = newest('logs/train_tequila_2026*.log')
        if nl and nl != train_log:
            train_log = nl
            log.info(f"switched to new train log: {os.path.basename(nl)}")

        # 告警: 自愈预算耗尽（只检查启动后新增内容）
        if not alert_sent and mon_log and os.path.exists(mon_log):
            with open(mon_log, encoding='utf-8', errors='ignore') as f:
                content = f.read()
            new_part = content[mon_base_len:]
            if 'budget exhausted' in new_part:
                send_mail("[quanti_train] 告警: 训练自愈预算耗尽, 需人工介入",
                          "24L 训练自愈重启 10 次仍失败, 请检查 GPU 状态与日志。",
                          html=False)
                alert_sent = True
            mon_base_len = len(content)

        # 完成: 训练日志出现 Done! 且步数达到目标
        done = False
        if train_log and os.path.exists(train_log):
            with open(train_log, encoding='utf-8', errors='ignore') as f:
                content = f.read()
            if 'Done!' in content:
                import re
                m = re.search(r'Done! (\d+) steps', content)
                done_steps = int(m.group(1)) if m else -1
                done = done_steps >= 39900   # TODO: 应从配置动态解析
        if not done:
            continue

        log.info("training DONE detected, starting relay")
        time.sleep(120)  # 等 final 保存稳定

        # 1. 导出三值部署 + HF 目录
        run([PY, '-u', 'export_ternary_model.py',
             '--ckpt', final, '--out', deploy])
        run([PY, '-u', 'export_hf_from_deploy.py',
             '--deploy', deploy, '--out', hf_dir])

        # 2. lm-eval 同口径
        for tasks, lim, bs, tag in [(['hellaswag', 'arc_challenge'], '1000', '16', 'ha'),
                                    (['ceval-valid', 'mmlu'], '25', '8', 'cm')]:
            cmd = [PY, '-m', 'lm_eval', '--model', 'hf',
                   '--model_args', f'pretrained={hf_dir},dtype=bfloat16',
                   '--tasks', ','.join(tasks), '--limit', lim, '--num_fewshot', '0',
                   '--batch_size', bs, '--device', 'cuda:1', '--seed', '42',
                   '--log_samples', '--output_path', out_path]
            subprocess.run(cmd, capture_output=True, text=True,
                           timeout=14400, env=env)

        # 3. 报告邮件
        send_mail(f"[quanti_train-Tequila] 完成: avg {avg:.2f}%", body)
        return

# ── export_ternary_model.py 核心 ──
with torch.no_grad():
    for full_name, module in list(model.named_modules()):
        if isinstance(module, UltraQuantLinear):
            quant_names.add(full_name)
            w = module.weight.data
            scale, delta = absmean(w)
            A = torch.zeros_like(w)
            A[w >= delta] = 1
            A[w <= -delta] = -1
            codes = A.to(torch.int8)                        # {-1,0,+1}
            scale_fp16 = scale.squeeze(-1).to(torch.float16)  # [out]

            Bb = torch.zeros_like(w)
            mask = A == 0
            Bb[mask] = module.eps * w[mask]                 # 用 module.eps, 不硬编码
            # 与 eval 完全一致的 bf16 舍入顺序, 保证 bit 级一致
            b_out = Bb.sum(dim=-1).to(torch.bfloat16)
            if module.bias is not None:
                b_out = b_out + module.bias.data

            deploy[full_name + ".weight"] = {
                "codes": codes.cpu(),
                "scale": scale_fp16.cpu()
            }
            deploy[full_name + ".bias"] = b_out.cpu()
            q_numel += codes.numel()

# 非量化参数原生 bf16 存
for name, p in model.named_parameters():
    root = name.rsplit(".", 1)[0]
    if root in quant_names:
        continue
    deploy[name] = p.data.cpu().clone()
    bf16_numel += p.numel()


# ── export_hf_from_deploy.py 核心: 重建并保存 HF 目录 ──
with torch.no_grad():
    for full_name, v in deploy.items():
        if isinstance(v, dict) and not full_name.endswith(".weight"):
            continue
        parts = full_name.split(".")
        parent = model
        for p in parts[:-1]:
            parent = getattr(parent, p)
        leaf = parts[-1]
        if isinstance(v, dict):
            codes = v["codes"]
            scale = v["scale"]
            w = (codes.float() * scale.float().unsqueeze(1)).to(torch.bfloat16)
            bias_key = full_name.replace(".weight", ".bias")
            b = deploy[bias_key].to(torch.bfloat16)
            parent.weight.data = w
            if parent.bias is None:
                parent.bias = nn.Parameter(b)
            else:
                parent.bias.data = b
        else:
            getattr(parent, leaf).data.copy_(v.to(torch.bfloat16))

model.save_pretrained(args.out, safe_serialization=True)
tok.save_pretrained(args.out)

TARGET_TOKENS = 68_000_000
SEQ = 512
OUT = 'cache/corpus_wiki64m_seq512.pt'

def clean(text):
    t = HEAD_JUNK.sub('', text)
    t = '\n'.join(line.strip() for line in t.splitlines() if line.strip())
    return t.strip()

def main():
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    from datasets import load_dataset
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained('./models/original', use_fast=True,
                                        trust_remote_code=True)
    eos = tok.eos_token_id
    ds = load_dataset('wikimedia/wikipedia', '20231101.zh',
                      split='train', streaming=True)

    ids_chunks, n_tokens, n_articles = [], 0, 0
    batch_texts, BATCH = [], 256

    def flush_batch():
        nonlocal n_tokens
        enc = tok(batch_texts, add_special_tokens=False)['input_ids']
        for ids in enc:
            if len(ids) < 10:
                continue
            arr = torch.tensor(ids + [eos], dtype=torch.int32)
            ids_chunks.append(arr)
            n_tokens += arr.numel()

    for row in ds:
        text = clean(row['text'])
        if len(text) < 200:
            continue
        if sum('\u4e00' <= c <= '\u9fff' for c in text[:500]) < 100:
            continue
        batch_texts.append(text[:20000])
        if len(batch_texts) >= BATCH:
            flush_batch()
            batch_texts = []
        n_articles += 1
        if n_tokens >= TARGET_TOKENS:
            break
    if batch_texts:
        flush_batch()

    all_ids = torch.cat(ids_chunks)
    n_seqs = all_ids.numel() // SEQ
    packed = all_ids[:n_seqs * SEQ].reshape(n_seqs, SEQ).long()
    attn = torch.ones_like(packed)
    tmp = OUT + '.tmp'
    torch.save({'input_ids': packed, 'attention_mask': attn}, tmp)
    os.replace(tmp, OUT)

def row_hashes(ids: torch.Tensor):
    """每行哈希: sha1 前 8 字节 -> int64"""
    b = ids.contiguous().view(torch.uint8).numpy().tobytes()
    # 注意: ids 为 int64, 每 token 8 字节; 此处若写 *4 会导致行切片错位
    row_bytes = ids.shape[1] * ids.element_size()
    return [int.from_bytes(hashlib.sha1(b[i*row_bytes:(i+1)*row_bytes]).digest()[:8], 'big')
            for i in range(ids.shape[0])]


def main():
    d = torch.load(CORPUS, map_location='cpu', weights_only=True)
    ids = d['input_ids']

    # 1. 行级精确去重
    hs = row_hashes(ids)
    seen, dup_idx = set(), []
    for i, h in enumerate(hs):
        if h in seen:
            dup_idx.append(i)
        else:
            seen.add(h)
    if dup_idx:
        keep = torch.ones(ids.shape[0], dtype=torch.bool)
        keep[torch.tensor(dup_idx)] = False
        ids = ids[keep]

    # 2. 语料 shingle 哈希集
    flat = ids.reshape(-1)
    win, stride = 20, 64
    sh = flat.unfold(0, win, stride)
    sh_bytes = sh.contiguous().view(torch.uint8).numpy()
    corp_set = set()
    for i in range(sh_bytes.shape[0]):
        corp_set.add(int.from_bytes(hashlib.sha1(sh_bytes[i]).digest()[:8], 'big'))

    # 3. 评测题污染检查
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained('./models/original', use_fast=True,
                                        trust_remote_code=True)
    bench = json.load(open(BENCH, encoding='utf-8'))
    total_hits = 0
    for name, samps in bench.items():
        hits = 0
        for s in samps:
            q_ids = tok(s['q'], add_special_tokens=False)['input_ids']
            for h in token_shingle_hashes(torch.tensor(q_ids, dtype=torch.int64)):
                if h in corp_set:
                    hits += 1
                    break
        total_hits += hits

    # 4. 保存去重后语料
    tmp = OUT + '.tmp'
    torch.save({'input_ids': ids, 'attention_mask': torch.ones_like(ids)}, tmp)
    os.replace(tmp, OUT)

def wilson_ci(p, n, z=1.96):
    """Wilson 95% 置信区间"""
    if n == 0:
        return (0.0, 0.0)
    den = 1 + z**2 / n
    center = (p + z**2 / (2 * n)) / den
    half = z * math.sqrt(p * (1 - p) / n + z**2 / (4 * n**2)) / den
    return (center - half, center + half)


def mcnemar_p(b, c):
    """McNemar 精确检验(双侧)"""
    n = b + c
    if n == 0:
        return 1.0
    k = min(b, c)
    return min(1.0, sum(math.comb(n, i) for i in range(0, k + 1)) / 2**(n - 1))


def main():
    table = {m: load_results(m) for m in MODELS}
    report = {}
    for m in MODELS:
        row = {}
        for task in HEADER_TASKS:
            if task == 'hellaswag':
                entry = table[m].get('hellaswag')
                p = entry['acc_norm'] if entry else None
                n = 1000
            elif task == 'arc_challenge':
                entry = table[m].get('arc_challenge')
                p = entry['acc'] if entry else None
                n = 1000
            else:
                subs = {t: v for t, v in table[m].items()
                        if t.startswith(task) and t != task}
                p = (sum(v['acc'] for v in subs.values()) / len(subs)) if subs else None
                n = sum(25 for _ in subs)
            row[task] = {'acc': p, 'n': n,
                         'ci': wilson_ci(p, n) if p is not None else None}
        report[m] = row

    # McNemar 逐题配对
    mcn = {}
    for task in ['hellaswag', 'arc_challenge', 'mmlu', 'ceval']:
        s = {m: load_samples(m, task) for m in MODELS}
        for a, b in [('12l', '24l'), ('bf16', '24l'), ('bf16', '12l'),
                     ('ptq', '24l'), ('ptq', '12l')]:
            common = set(s[a].keys()) & set(s[b].keys())
            b10 = sum(1 for k in common if s[a][k] and not s[b][k])
            b01 = sum(1 for k in common if s[b][k] and not s[a][k])
            p = mcnemar_p(b10, b01)
            mcn[f"{task}:{a}_vs_{b}"] = {'n': len(common), '10': b10, '01': b01, 'p': p}

    json.dump({'table': report, 'mcnemar': mcn},
              open('cache/p02_final_report.json', 'w', encoding='utf-8'),
              ensure_ascii=False, indent=1)

def ste_ref(w, x, scale_corr):
    """参考实现: 硬量化常数 + 恒等 STE 残差 + B 可微路径"""
    with torch.no_grad():
        scale, delta = absmean(w)
        hard = torch.zeros_like(w)
        hard[w >= delta] = 1
        hard[w <= -delta] = -1
        mask = hard == 0
        scale_vec = scale.squeeze(-1)
    # 每位置 STE 缩放: 活=scale(ModeB) 或 1(ModeA), 死=1
    per_pos = scale_vec.unsqueeze(1) if scale_corr else torch.ones_like(w)
    per_pos = torch.where(mask, torch.ones_like(per_pos), per_pos)
    A_const = (hard * per_pos).detach()
    A_ref = A_const + (w - w.detach()) * per_pos
    B_ref = torch.where(mask, EPS * w, torch.zeros_like(w))
    return (F.linear(x, A_ref) + B_ref.sum(dim=-1)).sum()


def run_case(corr, seeds=(42, 1234, 20260916)):
    max_rel, min_cos, fd_max_rel = 0.0, 2.0, 0.0
    for seed in seeds:
        torch.manual_seed(seed)
        w = (torch.randn(OUT, IN, device=DEV) * 0.05).requires_grad_(True)
        x = torch.randn(B_, IN, device=DEV)

        # 自定义 backward
        w1 = w.detach().clone().requires_grad_(True)
        A, B = UltraQuantV2.apply(w1, EPS, corr)
        (F.linear(x, A) + B.sum(dim=-1)).sum().backward()
        g_custom = w1.grad.detach()

        # 参考实现
        w2 = w.detach().clone().requires_grad_(True)
        ste_ref(w2, x, corr).backward()
        g_ref = w2.grad.detach()

        rel = (g_custom - g_ref).norm() / g_ref.norm().clamp_min(1e-12)
        cos = torch.nn.functional.cosine_similarity(
            g_custom.flatten(), g_ref.flatten(), dim=0)
        max_rel = max(max_rel, rel.item())
        min_cos = min(min_cos, cos.item())

        # B 路径有限差分(跳过量化边界附近的死权重)
        with torch.no_grad():
            s, d = absmean(w)
            Am = torch.zeros_like(w)
            Am[w >= d] = 1
            Am[w <= -d] = -1
            Am = Am * s
            dead = (Am == 0).nonzero()
        w3 = w.detach().clone().requires_grad_(True)
        (_, B3) = UltraQuantV2.apply(w3, EPS, corr)
        B3.sum(dim=-1).sum().backward()
        h = 1e-4
        safe = [ij for ij in dead.tolist()
                if abs(abs(w[ij[0], ij[1]].item()) - d[ij[0], 0].item()) > 10 * h]
        for o, j in safe[::max(1, len(safe) // 24)]:
            wp = w.detach().clone(); wp[o, j] += h
            wm = w.detach().clone(); wm[o, j] -= h
            num = (loss_B_only(wp) - loss_B_only(wm)) / (2 * h)
            ana = w3.grad[o, j]
            if abs(num) > 1e-9:
                fd_max_rel = max(fd_max_rel, (abs(num - ana) / abs(num)).item())
    return max_rel, min_cos, fd_max_rel

def hadamard_128(device):
    H = torch.linalg.hadamard(128).float()          # [128,128], 元素 ±1
    return (H / (128 ** 0.5)).to(device)            # 归一化正交


def ternary_quant_error(w):
    """当前方案(与 UltraQuantV2 同口径): 返回相对误差"""
    scale, delta = absmean(w)
    A = torch.zeros_like(w)
    A[w >= delta] = 1
    A[w <= -delta] = -1
    Q = A * scale
    return ((w - Q).norm() / w.norm().clamp_min(1e-12)).item()


def ternary_quant_error_hadamard(w, H128):
    """分块哈达玛旋转版: 每 128 列一组旋转后量化, 误差在旋转空间度量"""
    out_features, in_features = w.shape
    n_blocks = in_features // 128
    err_sq, norm_sq = 0.0, 0.0
    for b in range(n_blocks):
        blk = w[:, b*128:(b+1)*128] @ H128          # 旋转
        scale, delta = absmean(blk)
        A = torch.zeros_like(blk)
        A[blk >= delta] = 1
        A[blk <= -delta] = -1
        Q = A * scale
        err_sq += ((blk - Q).norm() ** 2).item()
        norm_sq += (blk.norm() ** 2).item()
    if in_features % 128:                            # 余数列不旋转
        rest = w[:, n_blocks*128:]
        scale, delta = absmean(rest)
        A = torch.zeros_like(rest)
        A[rest >= delta] = 1
        A[rest <= -delta] = -1
        Q = A * scale
        err_sq += ((rest - Q).norm() ** 2).item()
        norm_sq += (rest.norm() ** 2).item()
    return (err_sq ** 0.5) / (norm_sq ** 0.5)

⑤ 核心超参配置表

⑥ 项目依赖清单