半速的张量核 — NVIDIA 消费显卡的 AI 算力限制，和 gpu-unlock 怎么解开它

# bench_tensor.py：gpu-unlock 解锁前后测速（稠密 GEMM + BF16 小模型训练），自动记录 SM 频率和功耗
# 用法：python bench_tensor.py --tag before|after --out bench-<tag>.json
import argparse, json, platform, statistics, subprocess, threading, time
import torch

def smi_sampler(stop, rows):
    q = "clocks.sm,power.draw,temperature.gpu"
    while not stop.is_set():
        out = subprocess.run(["nvidia-smi", f"--query-gpu={q}", "--format=csv,noheader,nounits"],
                             capture_output=True, text=True).stdout.strip()
        try:
            sm, pw, t = [float(x) for x in out.split(",")]
            rows.append((sm, pw, t))
        except ValueError:
            pass
        stop.wait(0.5)

def timed(fn, iters):
    torch.cuda.synchronize()
    s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
    s.record()
    for _ in range(iters):
        fn()
    e.record(); torch.cuda.synchronize()
    return s.elapsed_time(e) / 1000 / iters

def gemm_case(name, n, warm_s, samples, iters):
    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
    if name == "int8":
        a = torch.randint(-64, 64, (n, n), dtype=torch.int8, device="cuda")
        b = torch.randint(-64, 64, (n, n), dtype=torch.int8, device="cuda").t().contiguous().t()
        fn = lambda: torch._int_mm(a, b)
    else:
        dt = {"fp16": torch.float16, "bf16": torch.bfloat16, "tf32": torch.float32, "fp32": torch.float32}[name]
        torch.backends.cuda.matmul.allow_tf32 = name == "tf32"
        a = torch.randn(n, n, dtype=dt, device="cuda"); b = torch.randn(n, n, dtype=dt, device="cuda")
        fn = lambda: a @ b
    end = time.time() + warm_s
    while time.time() < end:
        fn()
    torch.cuda.synchronize()
    stop, rows = threading.Event(), []
    th = threading.Thread(target=smi_sampler, args=(stop, rows)); th.start()
    secs = [timed(fn, iters) for _ in range(samples)]
    stop.set(); th.join()
    med = statistics.median(secs)
    tflops = 2 * n ** 3 / med / 1e12
    return {"case": name, "n": n, "tflops": round(tflops, 2), "ms": round(med * 1000, 3),
            "sm_mhz": round(statistics.median(r[0] for r in rows)) if rows else None,
            "power_w": round(statistics.median(r[1] for r in rows), 1) if rows else None,
            "temp_c": round(statistics.median(r[2] for r in rows)) if rows else None}

def train_case(warm_steps, steps):
    torch.manual_seed(0)
    d, layers, seq, batch = 1024, 6, 512, 16
    layer = torch.nn.TransformerEncoderLayer(d, 16, 4 * d, batch_first=True, dropout=0.0, norm_first=True)
    model = torch.nn.TransformerEncoder(layer, layers, enable_nested_tensor=False).cuda()
    head = torch.nn.Linear(d, d).cuda()
    opt = torch.optim.AdamW(list(model.parameters()) + list(head.parameters()), lr=1e-4)
    x = torch.randn(batch, seq, d, device="cuda"); y = torch.randn(batch, seq, d, device="cuda")
    def step():
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = torch.nn.functional.mse_loss(head(model(x)), y)
        loss.backward(); opt.step(); opt.zero_grad(set_to_none=True)
        return loss
    for _ in range(warm_steps):
        step()
    torch.cuda.synchronize()
    stop, rows = threading.Event(), []
    th = threading.Thread(target=smi_sampler, args=(stop, rows)); th.start()
    t0 = time.time()
    for _ in range(steps):
        loss = step()
    torch.cuda.synchronize(); dt = time.time() - t0
    stop.set(); th.join()
    return {"case": "train_bf16_transformer", "layers": layers, "d_model": d, "seq": seq, "batch": batch,
            "step_ms": round(dt / steps * 1000, 1), "tokens_per_s": round(batch * seq * steps / dt),
            "final_loss": round(loss.item(), 6),
            "sm_mhz": round(statistics.median(r[0] for r in rows)) if rows else None,
            "power_w": round(statistics.median(r[1] for r in rows), 1) if rows else None}

def env_info():
    def sh(cmd):
        return subprocess.run(cmd, shell=True, capture_output=True, text=True).stdout.strip()
    return {"gpu": torch.cuda.get_device_name(0), "torch": torch.__version__, "cuda": torch.version.cuda,
            "driver": sh("nvidia-smi --query-gpu=driver_version --format=csv,noheader"),
            "kernel": platform.release(), "cmdline_has_gpu_unlock": "gpu_unlock" in open("/proc/cmdline").read(),
            "gpu_unlock_loaded": bool(sh("lsmod | grep -w gpu_unlock")),
            "ac_online": sh("cat /sys/class/power_supply/AC*/online 2>/dev/null"),
            "time": time.strftime("%Y-%m-%d %H:%M:%S %z")}

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--warm", type=float, default=8); ap.add_argument("--samples", type=int, default=5)
    a = ap.parse_args()
    res = {"tag": a.tag, "env": env_info(), "gemm": [], "train": None}
    for n in (4096, 8192):
        for case in ("fp16", "bf16", "tf32", "fp32", "int8"):
            r = gemm_case(case, n, a.warm, a.samples, iters=10 if n == 4096 else 4)
            print(json.dumps(r, ensure_ascii=False), flush=True)
            res["gemm"].append(r)
    res["train"] = train_case(warm_steps=10, steps=30)
    print(json.dumps(res["train"], ensure_ascii=False), flush=True)
    json.dump(res, open(a.out, "w"), ensure_ascii=False, indent=2)