Mã nguồn thí nghiệm

"""Sinh các hình minh họa phần nền tảng. Mọi số liệu trên hình đều được TÍNH TRỰC TIẾP bởi code này
(trừ Hình 1 lấy số liệu từ Horowitz, ISSCC 2014)."""
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from scipy.stats import norm

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 10, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.dpi": 110, "savefig.dpi": 110,
                     "savefig.bbox": "tight"})
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "figs") + os.sep
os.makedirs(OUT, exist_ok=True)
C_SIGN, C_EXP, C_MAN, C_INT = "#d9534f", "#4a78c2", "#4caf6e", "#9e9e9e"
rng = np.random.default_rng(0)

# ---------------------------------------------------------------- Hình 1: năng lượng
ops = [("Cộng INT8", 0.03), ("Cộng INT32", 0.1), ("Cộng FP16", 0.4), ("Cộng FP32", 0.9),
       ("Nhân INT8", 0.2), ("Nhân INT32", 3.1), ("Nhân FP16", 1.1), ("Nhân FP32", 3.7),
       ("Đọc cache 8 KB (64 bit)", 10), ("Đọc cache 1 MB (64 bit)", 100),
       ("Đọc DRAM (64 bit)", 1300)]
colors = ["#4caf6e"] * 2 + ["#4a78c2"] * 2 + ["#4caf6e"] * 2 + ["#4a78c2"] * 2 + ["#e0923a"] * 3
fig, ax = plt.subplots(figsize=(8, 4.6))
y = np.arange(len(ops))[::-1]
ax.barh(y, [v for _, v in ops], color=colors)
ax.set_yticks(y, [n for n, _ in ops]); ax.set_xscale("log")
for yi, (n, v) in zip(y, ops):
    lab = "1300–2600" if "DRAM" in n else f"{v:g}"
    ax.text((2600 if "DRAM" in n else v) * 1.15, yi, lab + " pJ", va="center", fontsize=9)
ax.set_xlim(0.01, 2e4); ax.set_xlabel("Năng lượng mỗi phép toán (pJ, thang log)")
ax.set_title("Năng lượng phép toán và truy cập bộ nhớ (công nghệ 45 nm, 0.9 V)")
ax.errorbar([1950], [y[-1]], xerr=[[650], [650]], fmt="none", ecolor="k", capsize=3)
ax.legend(handles=[Rectangle((0, 0), 1, 1, color=c) for c in ["#4caf6e", "#4a78c2", "#e0923a"]],
          labels=["Số nguyên", "Dấu phẩy động", "Bộ nhớ"], loc="upper right", frameon=False)
fig.text(0.01, -0.02, "Nguồn số liệu: M. Horowitz, “Computing's energy problem”, ISSCC 2014.", fontsize=8)
fig.savefig(OUT + "fig01_energy.png"); plt.close(fig)

# ---------------------------------------------------------------- Hình 2: bố cục bit
fmts = [("FP32", [("S", 1, C_SIGN), ("Exponent (8)", 8, C_EXP), ("Mantissa (23)", 23, C_MAN)]),
        ("FP16", [("S", 1, C_SIGN), ("Exp (5)", 5, C_EXP), ("Mantissa (10)", 10, C_MAN)]),
        ("BF16", [("S", 1, C_SIGN), ("Exponent (8)", 8, C_EXP), ("Man (7)", 7, C_MAN)]),
        ("FP8 E4M3", [("S", 1, C_SIGN), ("E (4)", 4, C_EXP), ("M (3)", 3, C_MAN)]),
        ("FP8 E5M2", [("S", 1, C_SIGN), ("E (5)", 5, C_EXP), ("M (2)", 2, C_MAN)]),
        ("INT8", [("Số nguyên bù 2 (8)", 8, C_INT)]),
        ("INT4", [("INT (4)", 4, C_INT)])]
fig, ax = plt.subplots(figsize=(9, 3.8))
for r, (name, parts) in enumerate(fmts):
    x0, yy = 0, len(fmts) - 1 - r
    for lab, nb, c in parts:
        ax.add_patch(Rectangle((x0, yy + 0.1), nb, 0.8, facecolor=c, edgecolor="white", lw=1.5))
        if nb >= 2 or lab == "S":
            ax.text(x0 + nb / 2, yy + 0.5, lab, ha="center", va="center", color="white", fontsize=8.5,
                    fontweight="bold")
        x0 += nb
    ax.text(-0.6, yy + 0.5, name, ha="right", va="center", fontweight="bold")
    ax.text(x0 + 0.5, yy + 0.5, f"{x0} bit", va="center", fontsize=9, color="#555")
ax.set_xlim(-6, 36); ax.set_ylim(0, len(fmts)); ax.axis("off")
ax.set_title("Bố cục bit của các kiểu dữ liệu thường gặp (S = bit dấu)")
fig.savefig(OUT + "fig02_bit_layout.png"); plt.close(fig)

# ---------------------------------------------------------------- Hình 3: lưới giá trị FP8 vs INT8
def e4m3_values():
    vals = []
    for e in range(16):
        for m in range(8):
            if e == 15 and m == 7:      # mã NaN duy nhất của E4M3 (biến thể OCP "fn")
                continue
            vals.append((m / 8) * 2.0 ** (1 - 7) if e == 0 else (1 + m / 8) * 2.0 ** (e - 7))
    return np.array(vals)
v8 = e4m3_values()
assert np.isclose(v8.max(), 448.0)
s_int8 = 448 / 127
vi8 = np.arange(0, 128) * s_int8
fig, axes = plt.subplots(2, 1, figsize=(9, 3.6), sharex=True, gridspec_kw={"hspace": 0.6})
for ax, vals, title, c in [(axes[0], v8, "FP8 E4M3 (các giá trị dương biểu diễn được)", C_EXP),
                           (axes[1], vi8, f"INT8 đối xứng, cùng giá trị lớn nhất 448 (bước S = 448/127 ≈ {s_int8:.2f})", C_INT)]:
    vv = vals[vals <= 40]
    ax.vlines(vv, 0, 1, color=c, lw=1)
    ax.set_yticks([]); ax.set_title(title, fontsize=9.5, loc="left"); ax.spines["left"].set_visible(False)
axes[1].set_xlabel("Giá trị (chỉ hiển thị đoạn [0, 40])")
fig.suptitle("Lưới dấu phẩy động dày gần 0, thưa ở xa; lưới số nguyên cách đều", y=1.03)
fig.savefig(OUT + "fig03_fp_grid.png"); plt.close(fig)

# ---------------------------------------------------------------- Hình 4: bậc thang quantize
def quant_params(alpha, beta, qmin, qmax):
    S = (beta - alpha) / (qmax - qmin)
    Z = int(np.clip(np.round(qmin - alpha / S), qmin, qmax))
    return S, Z
def fake_quant(x, S, Z, qmin, qmax):
    q = np.clip(np.round(x / S) + Z, qmin, qmax)
    return S * (q - Z)
S, Z = quant_params(-1.0, 3.0, 0, 7)
x = np.linspace(-2, 4, 4000)
xh = fake_quant(x, S, Z, 0, 7)
lo, hi = S * (0 - Z), S * (7 - Z)
fig, axes = plt.subplots(2, 1, figsize=(8, 5.6), sharex=True, gridspec_kw={"height_ratios": [2, 1]})
ax = axes[0]
ax.plot(x, x, "--", color="#999", label="Lý tưởng: x̂ = x")
ax.plot(x, xh, color=C_EXP, lw=2, label="x̂ = S·(q − Z)")
for a in axes:
    a.axvspan(-2, lo, color=C_SIGN, alpha=0.08); a.axvspan(hi, 4, color=C_SIGN, alpha=0.08)
ax.text(-1.95, 2.9, "vùng clipping", color=C_SIGN, fontsize=9)
ax.text(hi + 0.05, -1.6, "vùng clipping", color=C_SIGN, fontsize=9)
ax.set_ylabel("x̂ (sau quantize → dequantize)"); ax.legend(loc="upper left", bbox_to_anchor=(0.2, 1), frameon=False)
ax.set_title(f"Quantize 3 bit không dấu trên [−1, 3]:  S = 4/7 ≈ {S:.4f},  Z = {Z}\n"
             f"Dải biểu diễn thực tế sau khi làm tròn Z: [{lo:.3f}, {hi:.3f}]")
err = xh - x
axes[1].plot(x, err, color=C_SIGN)
axes[1].axhline(S / 2, ls=":", color="#555"); axes[1].axhline(-S / 2, ls=":", color="#555")
axes[1].text(-1.9, S / 2 + 0.05, "±S/2", fontsize=9)
axes[1].set_ylim(-1.3, 1.3); axes[1].set_ylabel("Sai số x̂ − x"); axes[1].set_xlabel("x")
fig.savefig(OUT + "fig04_staircase.png"); plt.close(fig)

# ---------------------------------------------------------------- Hình 5: đối xứng vs bất đối xứng trên dữ liệu sau ReLU
act = np.maximum(rng.normal(0.3, 1.0, 200_000), 0)
amax = act.max()
bits = 3
# đối xứng có dấu, dải hạn chế [-(2^{b-1}-1), 2^{b-1}-1]
qs = 2 ** (bits - 1) - 1
S_sym = amax / qs
lv_sym = np.arange(-qs, qs + 1) * S_sym
# bất đối xứng không dấu [0, 2^b - 1]
S_asym, Z_asym = quant_params(0.0, amax, 0, 2 ** bits - 1)
lv_asym = (np.arange(0, 2 ** bits) - Z_asym) * S_asym
def mse_sym(x, b):
    q = 2 ** (b - 1) - 1; s = np.abs(x).max() / q
    return np.mean((np.clip(np.round(x / s), -q, q) * s - x) ** 2)
def mse_asym(x, b):
    s, z = quant_params(x.min(), x.max(), 0, 2 ** b - 1)
    return np.mean((fake_quant(x, s, z, 0, 2 ** b - 1) - x) ** 2)
fig, axes = plt.subplots(1, 2, figsize=(10, 3.4), sharey=True)
for ax, lv, t in [(axes[0], lv_sym, f"Đối xứng ({bits} bit có dấu): {np.sum(lv_sym >= 0)}/{len(lv_sym)} mức dùng được"),
                  (axes[1], lv_asym, f"Bất đối xứng ({bits} bit không dấu): {len(lv_asym)}/{len(lv_asym)} mức dùng được")]:
    ax.hist(act, bins=120, range=(-amax, amax), color="#bbb", density=True)
    for l in lv:
        ax.axvline(l, color=C_SIGN if l < -1e-9 else C_EXP, lw=1.2, ls="--" if l < -1e-9 else "-")
    ax.set_title(t, fontsize=9.5); ax.set_xlabel("Giá trị activation sau ReLU"); ax.set_ylim(0, 0.6)
    ax.annotate(f"đỉnh tại 0 ({np.mean(act == 0):.0%} giá trị)\nbị cắt bớt cho dễ nhìn", xy=(0, 0.58), xytext=(-4.8, 0.45),
                fontsize=8, arrowprops=dict(arrowstyle="->", color="#555"))
axes[0].set_ylabel("Mật độ")
fig.suptitle("Mức lượng tử (đường dọc) đặt trên histogram activation; mức màu đỏ nét đứt bị lãng phí", y=1.04)
fig.savefig(OUT + "fig05_sym_asym.png"); plt.close(fig)
sym_asym_mse = {b: (mse_sym(act, b), mse_asym(act, b)) for b in (4, 8)}

# ---------------------------------------------------------------- Hình 6: đánh đổi rounding - clipping
w = rng.laplace(0, 1.0, 1_000_000)
bits = 4; q = 2 ** (bits - 1) - 1
cs = np.linspace(0.5, np.abs(w).max(), 400)
r_err, c_err = [], []
for c in cs:
    s = c / q
    inside = np.abs(w) <= c
    wq = np.clip(np.round(w / s), -q, q) * s
    e2 = (wq - w) ** 2
    r_err.append(e2[inside].sum() / w.size); c_err.append(e2[~inside].sum() / w.size)
r_err, c_err = np.array(r_err), np.array(c_err); tot = r_err + c_err
c_opt = cs[np.argmin(tot)]
fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(cs, r_err, label="Sai số làm tròn (rounding)", color=C_EXP)
ax.plot(cs, c_err, label="Sai số cắt (clipping)", color=C_SIGN)
ax.plot(cs, tot, label="Tổng MSE", color="k", lw=2)
ax.axvline(c_opt, ls="--", color="k"); ax.axvline(cs[-1], ls=":", color="#777")
ax.text(c_opt + 0.2, 2e-4, f"ngưỡng tối ưu c* ≈ {c_opt:.2f}\nMSE = {tot.min():.4f}", fontsize=9)
ax.text(cs[-1] - 0.2, 5e-3, f"min–max: c = max|w| = {cs[-1]:.2f}\nMSE = {tot[-1]:.4f}", ha="right", fontsize=9)
ax.set_yscale("log"); ax.set_xlabel("Ngưỡng cắt c  (dải quantize [−c, c])"); ax.set_ylabel("MSE (thang log)")
ax.set_title("INT4 đối xứng trên 10⁶ trọng số phân phối Laplace(0, 1)"); ax.legend(frameon=False)
fig.savefig(OUT + "fig07_round_clip.png"); plt.close(fig)
clip_stats = (c_opt, tot.min(), cs[-1], tot[-1])

# ---------------------------------------------------------------- Hình 7: SQNR theo số bit
def sqnr_db(x, xh): return 10 * np.log10(np.mean(x ** 2) / np.mean((x - xh) ** 2))
def q_sym(x, b, c):
    qq = 2 ** (b - 1) - 1; s = c / qq
    return np.clip(np.round(x / s), -qq, qq) * s
bs = np.arange(2, 11)
u = rng.uniform(-1, 1, 1_000_000); g = rng.normal(0, 1, 1_000_000)
s_uni = [sqnr_db(u, q_sym(u, b, 1.0)) for b in bs]
s_gmm = [sqnr_db(g, q_sym(g, b, np.abs(g).max())) for b in bs]
def best_clip(x, b):
    cands = np.linspace(0.5, np.abs(x).max(), 300)
    return max(sqnr_db(x, q_sym(x, b, c)) for c in cands)
g_small = g[:200_000]
s_gopt = [best_clip(g_small, b) for b in bs]
fig, ax = plt.subplots(figsize=(7.5, 4))
ax.plot(bs, 20 * np.log10(2 ** bs - 2), "k--", label="Lý thuyết cho phân phối đều: 20·log₁₀(2ᵇ−2) ≈ 6.02·b dB")
ax.plot(bs, s_uni, "o", color=C_MAN, label="Thực nghiệm: dữ liệu đều U(−1, 1)")
ax.plot(bs, s_gmm, "s-", color=C_SIGN, label="Gauss N(0,1), dải min–max")
ax.plot(bs, s_gopt, "^-", color=C_EXP, label="Gauss N(0,1), ngưỡng cắt tối ưu")
ax.set_xlabel("Số bit b"); ax.set_ylabel("SQNR (dB)"); ax.legend(frameon=False, fontsize=8.5)
ax.set_title("Mỗi bit thêm vào tăng SQNR khoảng 6 dB")
fig.savefig(OUT + "fig06_sqnr.png"); plt.close(fig)
sqnr_tab = list(zip(bs, s_uni, s_gmm, s_gopt))

# ---------------------------------------------------------------- Hình 8: STE
x = np.linspace(-1.6, 1.6, 3000)
qq = 3; s = 1 / qq
fq = np.clip(np.round(x / s), -qq, qq) * s
fig, axes = plt.subplots(1, 2, figsize=(10, 3.4))
axes[0].plot(x, x, "--", color="#999", label="x"); axes[0].plot(x, fq, color=C_EXP, lw=2, label="FakeQuant(x)")
axes[0].set_title("Lượt xuôi (forward): hàm bậc thang"); axes[0].legend(frameon=False); axes[0].set_xlabel("x")
axes[1].plot(x, np.zeros_like(x), color=C_SIGN, lw=2, label="Đạo hàm thật: 0 hầu khắp nơi\n(không xác định tại các bước nhảy)")
axes[1].plot(x, ((x >= -1) & (x <= 1)).astype(float) + 0.02, color=C_MAN, lw=2, label="STE: 1 trong [α, β], 0 ngoài")
axes[1].set_ylim(-0.3, 1.4); axes[1].set_title("Lượt ngược (backward): đạo hàm dùng để cập nhật")
axes[1].legend(frameon=False, fontsize=8.5, loc="upper right"); axes[1].set_xlabel("x")
fig.savefig(OUT + "fig12_ste.png"); plt.close(fig)

# ---------------------------------------------------------------- Hình 9: NF4 vs INT4
offset = 0.9677083
v1 = norm.ppf(np.linspace(offset, 0.5, 9)[:-1])
v3 = -norm.ppf(np.linspace(offset, 0.5, 8)[:-1])
nf4 = np.sort(np.concatenate([v1, [0.0], v3])); nf4 /= nf4.max()
int4 = np.arange(-7, 8) / 7.0
W = rng.normal(0, 1, (20_000, 64))
Wn = (W / np.abs(W).max(axis=1, keepdims=True)).ravel()      # chuẩn hóa absmax theo khối 64
def nearest(x, grid):
    idx = np.abs(x[:, None] - grid[None, :]).argmin(1); return grid[idx]
sub = Wn[:400_000]
mse_nf4 = np.mean((nearest(sub, nf4) - sub) ** 2); mse_int4 = np.mean((nearest(sub, int4) - sub) ** 2)
fig, axes = plt.subplots(2, 1, figsize=(9, 4.6), sharex=True)
for ax, grid, t, c in [(axes[0], int4, f"INT4 đối xứng (15 mức cách đều) — MSE = {mse_int4:.2e}", C_INT),
                       (axes[1], nf4, f"NF4 (16 mức theo phân vị phân phối chuẩn) — MSE = {mse_nf4:.2e}", C_EXP)]:
    ax.hist(Wn, bins=200, color="#ddd", density=True)
    for l in grid: ax.axvline(l, color=c, lw=1.3)
    ax.set_title(t, fontsize=9.5, loc="left"); ax.set_yticks([])
axes[1].set_xlabel("Trọng số sau khi chia cho absmax của khối 64 phần tử (w ~ N(0,1))")
fig.savefig(OUT + "fig16_nf4.png"); plt.close(fig)

# ---------------------------------------------------------------- in số liệu dùng trong bài
print("S,Z (3bit,[-1,3]):", S, Z, "range", lo, hi)
print("sym/asym MSE on ReLU act:", {b: (f"{a:.3e}", f"{c:.3e}") for b, (a, c) in sym_asym_mse.items()})
print("clip stats c_opt, mse_opt, c_max, mse_max:", clip_stats)
print("SQNR table (b, uniform, gauss-minmax, gauss-opt):")
for r in sqnr_tab: print("  %d  %.2f  %.2f  %.2f" % r)
print("NF4 levels:", np.round(nf4, 4)); print("mse nf4 / int4:", mse_nf4, mse_int4, mse_int4 / mse_nf4)
print("E4M3 #pos finite values:", len(v8), "min subnormal", v8[v8 > 0].min(), "max", v8.max())

"""Sơ đồ khối (vẽ bằng matplotlib để không phụ thuộc công cụ ngoài)."""
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9.5, "savefig.dpi": 110, "savefig.bbox": "tight"})
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "figs") + os.sep
os.makedirs(OUT, exist_ok=True)

def box(ax, x, y, w, h, text, fc="#eef3fb", ec="#4a78c2", fs=9, bold=False):
    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle="round,pad=0.02,rounding_size=0.08",
                                fc=fc, ec=ec, lw=1.4))
    ax.text(x, y, text, ha="center", va="center", fontsize=fs, fontweight="bold" if bold else "normal")

def arrow(ax, p, q, text=None, color="#333", off=(0, 0.12), fs=8, ls="-"):
    ax.add_patch(FancyArrowPatch(p, q, arrowstyle="-|>", mutation_scale=12, color=color, lw=1.3, linestyle=ls))
    if text:
        ax.text((p[0] + q[0]) / 2 + off[0], (p[1] + q[1]) / 2 + off[1], text, ha="center", fontsize=fs, color=color)

# ------------------------------------------------------------ Hình 13: một lớp Linear/Conv chạy toàn số nguyên
fig, ax = plt.subplots(figsize=(12, 4.2)); ax.set_xlim(0, 12); ax.set_ylim(0, 4.4); ax.axis("off")
G, O, B = ("#eaf6ee", "#4caf6e"), ("#fdf1e4", "#e0923a"), ("#eef3fb", "#4a78c2")
box(ax, 1.0, 3.2, 1.6, 0.7, "q_x\nuint8", *G)
box(ax, 1.0, 1.6, 1.6, 0.7, "q_w\nint8 (Z_w = 0)", *G)
box(ax, 3.4, 2.4, 2.0, 1.0, "Nhân int8 × int8\ncộng dồn int32\nΣ q_w·(q_x − Z_x)", *B)
box(ax, 5.9, 2.4, 1.7, 0.8, "+ q_b\n(bias int32)", *B)
box(ax, 8.2, 2.4, 2.2, 1.0, "Requantize\n(acc · M₀) >> (31 + n)\n[số nguyên 64 bit]", *B)
box(ax, 10.6, 2.4, 1.6, 0.8, "+ Z_y, clamp\n[0, 255]", *B)
box(ax, 10.6, 0.8, 1.6, 0.7, "q_y  uint8\n→ lớp tiếp theo", *G)
arrow(ax, (1.8, 3.2), (2.4, 2.7)); arrow(ax, (1.8, 1.6), (2.4, 2.1))
arrow(ax, (4.4, 2.4), (5.05, 2.4), "int32"); arrow(ax, (6.75, 2.4), (7.1, 2.4), "int32")
arrow(ax, (9.3, 2.4), (9.8, 2.4)); arrow(ax, (10.6, 2.0), (10.6, 1.15))
box(ax, 5.9, 4.0, 9.0, 0.55, "Tính TRƯỚC khi convert (offline):   Z_x·Σ q_w ,   q_b = round(b / (S_w·S_x)) ,   M = S_w·S_x / S_y = M₀·2^−(31+n)", *O, fs=8.5)
ax.text(6, 0.25, "Toàn bộ đường đi chỉ dùng số nguyên → chạy được trên NPU / DSP / vi điều khiển không có FPU.",
        ha="center", fontsize=9, style="italic", color="#444")
fig.savefig(OUT + "fig09_integer_pipeline.png"); plt.close(fig)

# ------------------------------------------------------------ Hình 16: đồ thị có nút fake-quant (dùng cho PTQ mô phỏng & QAT)
fig, ax = plt.subplots(figsize=(12, 2.9)); ax.set_xlim(0, 12); ax.set_ylim(0, 3); ax.axis("off")
FQ = ("#fdecea", "#d9534f")
box(ax, 0.6, 1.3, 0.9, 0.6, "x (FP32)", *G)
box(ax, 1.9, 1.3, 1.1, 0.6, "FakeQuant\n(act, 8 bit)", *FQ, fs=8)
box(ax, 3.5, 2.5, 1.0, 0.5, "W (FP32)", *G); box(ax, 3.5, 1.75, 1.1, 0.5, "FakeQuant\n(weight)", *FQ, fs=8)
box(ax, 5.0, 1.3, 1.6, 0.7, "Conv + BN đã gộp\n(tính bằng FP32)", *B)
box(ax, 6.7, 1.3, 0.9, 0.6, "ReLU", *B)
box(ax, 8.1, 1.3, 1.1, 0.6, "FakeQuant\n(act)", *FQ, fs=8)
box(ax, 9.6, 1.3, 1.1, 0.6, "Lớp kế tiếp …", *B)
arrow(ax, (1.05, 1.3), (1.35, 1.3)); arrow(ax, (2.45, 1.3), (4.2, 1.3))
arrow(ax, (3.5, 2.25), (3.5, 2.0)); arrow(ax, (3.5, 1.5), (4.3, 1.45))
arrow(ax, (5.8, 1.3), (6.25, 1.3)); arrow(ax, (7.15, 1.3), (7.55, 1.3)); arrow(ax, (8.65, 1.3), (9.05, 1.3))
ax.text(6, 0.3, "FakeQuant(x) = S·(clamp(round(x/S) + Z, q_min, q_max) − Z): giá trị vẫn là FP32 nhưng chỉ nhận các mức lượng tử.\n"
        "Lượt ngược dùng STE. Khi convert, mỗi cặp FakeQuant được thay bằng phép quantize / dequantize thật.",
        ha="center", fontsize=8.5, color="#444")
fig.savefig(OUT + "fig10_fake_quant.png"); plt.close(fig)

# ------------------------------------------------------------ Hình 17: quy trình chọn phương pháp
fig, ax = plt.subplots(figsize=(11, 5.6)); ax.set_xlim(0, 11); ax.set_ylim(0, 6); ax.axis("off")
Y_ = ("#fff8e1", "#c9a227")
box(ax, 5.5, 5.5, 4.4, 0.6, "Mô hình FP32 đã huấn luyện + baseline accuracy", *G, bold=True)
box(ax, 5.5, 4.5, 4.6, 0.65, "Phần cứng đích hỗ trợ kiểu gì?\n(GPU: FP16/BF16/FP8/INT8; CPU/NPU/MCU: INT8)", *Y_, fs=8.5)
box(ax, 1.7, 3.4, 3.1, 0.7, "Chỉ cần nhẹ/nhanh vừa phải:\nFP16 / BF16 (gần như không mất acc)", *B, fs=8.5)
box(ax, 5.5, 3.4, 2.8, 0.7, "Weight-only INT8 / INT4\nhoặc dynamic quantization\n(không cần dữ liệu calibration)", *B, fs=8.5)
box(ax, 9.3, 3.4, 2.6, 0.7, "Static full-integer INT8\n+ calibration 100–1000 mẫu", *B, fs=8.5)
box(ax, 7.4, 2.1, 3.4, 0.75, "Acc giảm quá ngưỡng cho phép?\n→ per-channel, đổi calibration,\nCLE / bias correction / AdaRound", *Y_, fs=8.5)
box(ax, 7.4, 0.8, 3.4, 0.75, "Vẫn chưa đạt?\n→ mixed precision cho lớp nhạy cảm\n→ QAT (fine-tune có fake-quant)", *FQ, fs=8.5)
box(ax, 2.6, 1.4, 3.0, 0.8, "Đo lại TRÊN THIẾT BỊ ĐÍCH:\nkích thước, độ trễ, độ chính xác", *G, fs=8.5, bold=True)
arrow(ax, (5.5, 5.2), (5.5, 4.8))
arrow(ax, (4.2, 4.3), (2.4, 3.75)); arrow(ax, (5.5, 4.2), (5.5, 3.75)); arrow(ax, (6.8, 4.3), (8.7, 3.75))
arrow(ax, (9.3, 3.05), (8.4, 2.45)); arrow(ax, (5.9, 3.05), (6.6, 2.45))
arrow(ax, (7.4, 1.72), (7.4, 1.18), "có", off=(0.25, -0.05))
arrow(ax, (5.7, 2.1), (4.1, 1.6), "không", off=(0, 0.1))
arrow(ax, (1.6, 3.05), (2.2, 1.8)); arrow(ax, (5.7, 0.8), (4.1, 1.2))
fig.savefig(OUT + "fig17_decision.png"); plt.close(fig)
print("ok")

"""Thí nghiệm numpy: (A) lớp Linear chạy hoàn toàn bằng số nguyên, (B) GPTQ so với làm tròn thường (RTN),
(C) mô phỏng outlier activation trong LLM và cách SmoothQuant / LLM.int8() xử lý."""
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 10, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.dpi": 110, "savefig.dpi": 110, "savefig.bbox": "tight"})
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "figs") + os.sep
os.makedirs(OUT, exist_ok=True)
rng = np.random.default_rng(42)

# ============================================================ (A) Integer-only Linear
def affine_params(lo, hi, qmin=0, qmax=255):
    lo, hi = min(lo, 0.0), max(hi, 0.0)            # dải phải chứa 0 để 0 biểu diễn chính xác
    S = (hi - lo) / (qmax - qmin)
    Z = int(np.clip(round(qmin - lo / S), qmin, qmax))
    return S, Z

def quantize_multiplier(M):
    """Viết M (0 < M < 1) dưới dạng M0 * 2^(-31-shift), M0 là int32 trong [2^30, 2^31)."""
    assert 0 < M < 1
    shift = 0
    while M < 0.5:
        M *= 2; shift += 1
    M0 = int(round(M * (1 << 31)))
    if M0 == (1 << 31):                          # làm tròn lên chạm 2^31
        M0 //= 2; shift -= 1
    return M0, shift

def rounding_right_shift(x, n):                   # chia cho 2^n, làm tròn gần nhất (x là int64)
    return (x + (1 << (n - 1))) >> n

K, N_OUT, BATCH = 256, 64, 1000
X = rng.normal(0.8, 1.0, (BATCH, K)).astype(np.float32)          # activation lệch dương (ví dụ sau GELU/Add)
W = rng.normal(0, 0.05, (N_OUT, K)).astype(np.float32)
b = rng.normal(0, 0.1, N_OUT).astype(np.float32)
Y = X @ W.T + b                                                          # tham chiếu FP32

S_x, Z_x = affine_params(X.min(), X.max())                   # activation: uint8 bất đối xứng
S_w = np.abs(W).max() / 127; Z_w = 0                        # weight: int8 đối xứng
S_y, Z_y = affine_params(Y.min(), Y.max())                   # output: uint8 bất đối xứng

q_x = np.clip(np.round(X / S_x) + Z_x, 0, 255).astype(np.int32)
q_w = np.clip(np.round(W / S_w), -127, 127).astype(np.int32)
q_b = np.round(b / (S_w * S_x)).astype(np.int32)            # bias int32, scale S_w*S_x, zero-point 0

# Lõi int8 x int8 -> cộng dồn int32 ; số hạng Z_x * sum(q_w) tính trước lúc convert
acc = q_x @ q_w.T - Z_x * q_w.sum(axis=1) + q_b            # int32
M = S_w * S_x / S_y
M0, shift = quantize_multiplier(M)
q_y = Z_y + rounding_right_shift(acc.astype(np.int64) * M0, 31 + shift)
q_y = np.clip(q_y, 0, 255)
Y_int = S_y * (q_y - Z_y)                                    # dequantize chỉ để so sánh

# Đường "mô phỏng" (fake-quant bằng float) để kiểm tra đường số nguyên tính đúng
Y_sim = (S_x * (q_x - Z_x)) @ (S_w * q_w).T + S_w * S_x * q_b
q_y_sim = np.clip(np.round(Y_sim / S_y) + Z_y, 0, 255)
rel = lambda a, r: np.linalg.norm(a - r) / np.linalg.norm(r)
print("=== (A) Integer-only Linear ===")
print(f"S_x={S_x:.6f} Z_x={Z_x} | S_w={S_w:.6f} | S_y={S_y:.6f} Z_y={Z_y}")
print(f"M = {M:.10f} = {M0} * 2^-(31+{shift}) -> {M0 / 2 ** (31 + shift):.10f}")
print("max |q_y(integer) - q_y(float sim)| =", int(np.abs(q_y - q_y_sim).max()), "LSB;",
      "tỉ lệ khác nhau:", f"{np.mean(q_y != q_y_sim):.4%}")
print(f"sai số tương đối so với FP32: {rel(Y_int, Y):.4%}")
print("max |acc| =", int(np.abs(acc).max()), "; giới hạn int16 =", 2 ** 15 - 1,
      "; cận trên lý thuyết 255*127*K =", 255 * 127 * K)

# ============================================================ (B) GPTQ vs RTN
def quant_rtn_rows(W, scales, qmax=7):
    return np.clip(np.round(W / scales[:, None]), -qmax, qmax) * scales[:, None]

def gptq(W, H, scales, qmax=7, damp=0.01):
    """GPTQ tối giản (không lazy-batch): lượng tử từng cột, bù sai số vào các cột chưa lượng tử."""
    W = W.copy(); d = W.shape[1]
    H = H + damp * np.mean(np.diag(H)) * np.eye(d)
    Hinv = np.linalg.cholesky(np.linalg.inv(H)).T            # Cholesky trên (upper) của H^-1
    Q = np.zeros_like(W)
    for i in range(d):
        w = W[:, i]
        q = np.clip(np.round(w / scales), -qmax, qmax) * scales
        Q[:, i] = q
        err = (w - q) / Hinv[i, i]
        W[:, i + 1:] -= np.outer(err, Hinv[i, i + 1:])
    return Q

d_in, d_out, n_cal = 256, 128, 4096
A = rng.normal(0, 1, (d_in, d_in)) * (np.linspace(1, 0.05, d_in) ** 2)[None, :]   # trộn -> đặc trưng tương quan
Xc = (A @ rng.normal(0, 1, (d_in, n_cal))).astype(np.float64)
Xt = (A @ rng.normal(0, 1, (d_in, n_cal))).astype(np.float64)
Wg = rng.normal(0, 0.02, (d_out, d_in))
scales = np.abs(Wg).max(axis=1) / 7                           # INT4 đối xứng, scale theo hàng
H = 2 * Xc @ Xc.T / n_cal
Q_rtn = quant_rtn_rows(Wg, scales)
Q_gptq = gptq(Wg, H, scales)
out_err = lambda Q: np.linalg.norm((Wg - Q) @ Xt) ** 2 / np.linalg.norm(Wg @ Xt) ** 2
w_err = lambda Q: np.mean((Wg - Q) ** 2)
print("\n=== (B) GPTQ vs RTN, INT4 theo hàng ===")
print(f"RTN : sai số trọng số MSE={w_err(Q_rtn):.3e}, sai số output tương đối={out_err(Q_rtn):.4%}")
print(f"GPTQ: sai số trọng số MSE={w_err(Q_gptq):.3e}, sai số output tương đối={out_err(Q_gptq):.4%}")
gptq_res = (w_err(Q_rtn), out_err(Q_rtn), w_err(Q_gptq), out_err(Q_gptq))

# ============================================================ (C) Outlier activation, LLM.int8(), SmoothQuant
T, Kc, Mo = 2048, 512, 512
Xl = rng.normal(0, 1, (T, Kc))
outlier_ch = rng.choice(Kc, 6, replace=False)
Xl[:, outlier_ch] *= 60                                          # vài kênh có biên độ lớn gấp hàng chục lần
Wl = rng.normal(0, 0.02, (Mo, Kc))
Yl = Xl @ Wl.T

def q_sym(x, S, qmax=127):
    return np.clip(np.round(x / S), -qmax, qmax) * S
def w8_per_channel(W):                                           # per-output-channel
    return q_sym(W, np.abs(W).max(axis=1, keepdims=True) / 127)
def a8_per_tensor(X):
    return q_sym(X, np.abs(X).max() / 127)
def a8_per_token(X):
    return q_sym(X, np.abs(X).max(axis=1, keepdims=True) / 127)

res = {}
res["W8A8, activation per-tensor"] = rel(a8_per_tensor(Xl) @ w8_per_channel(Wl).T, Yl)
res["W8A8, activation per-token"] = rel(a8_per_token(Xl) @ w8_per_channel(Wl).T, Yl)
# LLM.int8(): cột outlier (|x| > 6) tính FP16, phần còn lại int8 per-token x per-channel
mask = (np.abs(Xl) > 6.0).any(axis=0)
X_reg = Xl.copy(); X_reg[:, mask] = 0
W_reg = Wl.copy(); W_reg[:, mask] = 0
Y_mix = a8_per_token(X_reg) @ w8_per_channel(W_reg).T + Xl[:, mask] @ Wl[:, mask].T
res["LLM.int8() (tách outlier)"] = rel(Y_mix, Yl)
# SmoothQuant alpha = 0.5
alpha = 0.5
s = np.abs(Xl).max(axis=0) ** alpha / np.abs(Wl).max(axis=0) ** (1 - alpha)
Xs, Ws = Xl / s, Wl * s
assert np.allclose(Xs @ Ws.T, Yl)                                 # phép biến đổi tương đương toán học
res["SmoothQuant (α=0.5) + per-tensor"] = rel(a8_per_tensor(Xs) @ w8_per_channel(Ws).T, Yl)
print("\n=== (C) Outlier activation ===")
print("số cột được LLM.int8 tách:", int(mask.sum()), "; kênh outlier thật:", sorted(outlier_ch.tolist()))
for k, v in res.items(): print(f"{k:38s}: {v:.4%}")

fig, axes = plt.subplots(1, 3, figsize=(13, 3.6), gridspec_kw={"width_ratios": [1.2, 1.2, 1.3]})
ch = np.arange(Kc)
axes[0].vlines(ch, 0, np.abs(Xl).max(0), color="#d9534f", lw=1.2)
axes[0].set_title("max|X| theo kênh — TRƯỚC khi làm mịn"); axes[0].set_xlabel("Kênh (chiều ẩn)")
axes[1].vlines(ch, 0, np.abs(Xs).max(0), color="#4a78c2", lw=1.2, label="max|X/s| (activation)")
axes[1].vlines(ch + 0.5, 0, -np.abs(Ws).max(0), color="#4caf6e", lw=1.2, label="−max|W·s| (trọng số, vẽ ngược xuống)")
axes[1].axhline(0, color="k", lw=0.5)
axes[1].set_title("SAU SmoothQuant (α = 0.5)"); axes[1].set_xlabel("Kênh (chiều ẩn)"); axes[1].legend(frameon=False, fontsize=8)
names = list(res.keys()); vals = [res[n] * 100 for n in names]
axes[2].barh(range(len(names))[::-1], vals, color=["#d9534f", "#e0923a", "#4caf6e", "#4a78c2"])
axes[2].set_yticks(range(len(names))[::-1], [n.replace(", ", "\n").replace(" + ", "\n+ ") for n in names], fontsize=8.5)
axes[2].set_xscale("log"); axes[2].set_xlabel("Sai số output tương đối (%, thang log)")
for i, v in zip(range(len(names))[::-1], vals): axes[2].text(v * 1.1, i, f"{v:.2f}%", va="center", fontsize=8.5)
axes[2].set_xlim(0.1, max(vals) * 6); axes[2].set_title("W8A8 trên dữ liệu có outlier")
fig.suptitle("Mô phỏng: 6/512 kênh activation có biên độ gấp 60 lần (dữ liệu tổng hợp)", y=1.03)
fig.savefig(OUT + "fig14_llm_outliers.png"); plt.close(fig)

# Hình GPTQ
fig, axes = plt.subplots(1, 2, figsize=(9, 3.2))
lab = ["RTN", "GPTQ"]
axes[0].bar(lab, [gptq_res[0], gptq_res[2]], color=["#9e9e9e", "#4a78c2"]); axes[0].set_title("Sai số TRỌNG SỐ (MSE)")
axes[1].bar(lab, [gptq_res[1] * 100, gptq_res[3] * 100], color=["#9e9e9e", "#4a78c2"]); axes[1].set_title("Sai số OUTPUT tương đối (%) trên dữ liệu kiểm tra")
for ax in axes: ax.ticklabel_format(axis="y", style="sci", scilimits=(-3, 3))
fig.suptitle("INT4 cho một lớp Linear 256→128, dữ liệu đầu vào có tương quan (tổng hợp)", y=1.04)
fig.savefig(OUT + "fig15_gptq.png"); plt.close(fig)

"""Thí nghiệm PyTorch THẬT trên bộ dữ liệu chữ số viết tay 8x8 (sklearn digits, 1797 ảnh, có sẵn offline).
Mọi con số trong bài phần thực nghiệm được in ra từ file này."""
import copy, os, time, warnings
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

warnings.filterwarnings("ignore")
torch.set_num_threads(1)
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 10, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.dpi": 110, "savefig.dpi": 110, "savefig.bbox": "tight"})
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "figs") + os.sep
os.makedirs(OUT, exist_ok=True)

# ------------------------------------------------------------------ dữ liệu
d = load_digits()
X = (d.images / 16.0).astype(np.float32)[:, None]            # (1797, 1, 8, 8), giá trị trong [0, 1]
y = d.target.astype(np.int64)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)
Xtr, Xte, ytr, yte = map(torch.from_numpy, (Xtr, Xte, ytr, yte))
Xcal = Xtr[:256]                                              # tập calibration: 256 ảnh từ tập train

def accuracy(model, X=Xte, y=yte):
    model.eval()
    with torch.no_grad():
        return (model(X).argmax(1) == y).float().mean().item()

# ------------------------------------------------------------------ mô hình float (có BatchNorm, có depthwise)
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1, self.bn1 = nn.Conv2d(1, 16, 3, padding=1, bias=False), nn.BatchNorm2d(16)
        self.dw, self.bn2 = nn.Conv2d(16, 16, 3, padding=1, groups=16, bias=False), nn.BatchNorm2d(16)
        self.pw, self.bn3 = nn.Conv2d(16, 32, 1, bias=False), nn.BatchNorm2d(32)
        self.conv4, self.bn4 = nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32)
        self.fc = nn.Linear(32, 10)
    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.dw(x)))
        x = F.max_pool2d(F.relu(self.bn3(self.pw(x))), 2)
        x = F.relu(self.bn4(self.conv4(x)))
        return self.fc(torch.flatten(F.adaptive_avg_pool2d(x, 1), 1))

def train(model, epochs=40, lr=3e-3, seed=0, wd=0.0):
    g = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
    for _ in range(epochs):
        model.train()
        perm = torch.randperm(len(Xtr), generator=g)
        for i in range(0, len(Xtr), 64):
            idx = perm[i:i + 64]
            loss = F.cross_entropy(model(Xtr[idx]), ytr[idx])
            opt.zero_grad(); loss.backward(); opt.step()
        sched.step()
    return model.eval()

torch.manual_seed(0)
fp_model = train(Net())
acc_fp = accuracy(fp_model)
print(f"[FP32] test accuracy = {acc_fp:.4f}  (n_test = {len(yte)})")

# ------------------------------------------------------------------ (1) BN folding
def fold_bn(conv, bn):
    std = torch.sqrt(bn.running_var + bn.eps)
    w = conv.weight * (bn.weight / std).reshape(-1, 1, 1, 1)
    b0 = conv.bias if conv.bias is not None else torch.zeros_like(bn.running_mean)
    b = bn.weight * (b0 - bn.running_mean) / std + bn.bias
    fused = nn.Conv2d(conv.in_channels, conv.out_channels, conv.kernel_size, conv.stride, conv.padding,
                      groups=conv.groups, bias=True)
    fused.weight.data, fused.bias.data = w.detach().clone(), b.detach().clone()
    return fused

class FoldedNet(nn.Module):
    """Mạng sau khi gộp BN; có chỗ cắm fake-quant cho weight và activation."""
    def __init__(self, net):
        super().__init__()
        self.convs = nn.ModuleList([fold_bn(net.conv1, net.bn1), fold_bn(net.dw, net.bn2),
                                    fold_bn(net.pw, net.bn3), fold_bn(net.conv4, net.bn4)])
        self.fc = copy.deepcopy(net.fc)
        self.wq = None            # hàm fake-quant weight: (tên lớp, W) -> W_hat
        self.aq = None            # danh sách fake-quant activation: aq[i](x)
    def _w(self, i, w): return w if self.wq is None else self.wq(i, w)
    def _a(self, i, x): return x if self.aq is None else self.aq[i](x)
    def forward(self, x):
        x = self._a(0, x)                                     # lượng tử input
        for i, c in enumerate(self.convs):
            x = F.relu(F.conv2d(x, self._w(i, c.weight), c.bias, c.stride, c.padding, groups=c.groups))
            if i == 2: x = F.max_pool2d(x, 2)
            x = self._a(i + 1, x)
        x = torch.flatten(F.adaptive_avg_pool2d(x, 1), 1)
        return F.linear(x, self._w(4, self.fc.weight), self.fc.bias)

folded = FoldedNet(fp_model).eval()
with torch.no_grad():
    diff = (folded(Xte) - fp_model(Xte)).abs().max().item()
print(f"[BN folding] max |logit_folded - logit_goc| = {diff:.2e}")

# ------------------------------------------------------------------ (2) per-tensor vs per-channel
def fq_sym(w, bits, per_channel):
    qmax = 2 ** (bits - 1) - 1
    if per_channel:
        amax = w.abs().reshape(w.shape[0], -1).max(1).values.clamp_min(1e-8)
        amax = amax.reshape(-1, *([1] * (w.dim() - 1)))
    else:
        amax = w.abs().max().clamp_min(1e-8)
    s = amax / qmax
    q = torch.clamp(torch.round(w / s), -qmax, qmax) * s
    return w + (q - w).detach()                               # STE

dw_w = folded.convs[1].weight.detach()
dw_ranges = dw_w.abs().reshape(16, -1).max(1).values.numpy()
res_gran = {}
for bits in (8, 4):
    for pc in (False, True):
        folded.wq = lambda i, w, b=bits, p=pc: fq_sym(w, b, p)
        res_gran[(bits, pc)] = (accuracy(folded),
                                 torch.mean((fq_sym(dw_w, bits, pc) - dw_w) ** 2).item())
folded.wq = None
print("[Granularity] (bit, per_channel) -> (acc chỉ quantize weight, MSE lớp depthwise):")
for k, v in res_gran.items(): print("   ", k, f"acc={v[0]:.4f}  mse_dw={v[1]:.3e}")

fig, axes = plt.subplots(1, 2, figsize=(10, 3.4))
axes[0].bar(range(16), dw_ranges, color="#4a78c2")
axes[0].axhline(dw_ranges.max(), ls="--", color="#d9534f")
axes[0].text(15.5, dw_ranges.max() * 1.02, "dải per-tensor dùng chung", ha="right", color="#d9534f", fontsize=9)
axes[0].set_xlabel("Kênh của lớp depthwise (sau khi gộp BN)"); axes[0].set_ylabel("max|w| của kênh")
axes[0].set_title(f"Lớp depthwise: max|w| lớn nhất / nhỏ nhất ≈ {dw_ranges.max() / dw_ranges.min():.1f} lần", fontsize=10)
lab = ["INT8\nper-tensor", "INT8\nper-channel", "INT4\nper-tensor", "INT4\nper-channel"]
keys = [(8, False), (8, True), (4, False), (4, True)]
vals = [res_gran[k][0] * 100 for k in keys]
bars = axes[1].bar(lab, vals, color=["#9e9e9e", "#4a78c2", "#9e9e9e", "#4a78c2"])
axes[1].axhline(acc_fp * 100, ls="--", color="k"); axes[1].text(3.4, acc_fp * 100 + 0.5, "FP32", ha="right")
for bb, v in zip(bars, vals): axes[1].text(bb.get_x() + bb.get_width() / 2, v + 0.5, f"{v:.1f}", ha="center", fontsize=9)
axes[1].set_ylim(min(vals) - 8, 102); axes[1].set_ylabel("Độ chính xác test (%)")
axes[1].set_title("Chỉ lượng tử trọng số (activation giữ FP32)", fontsize=10)
fig.savefig(OUT + "fig08_per_channel.png"); plt.close(fig)

# ------------------------------------------------------------------ (3) calibration activation
def collect_acts(model, X):
    acts = [[] for _ in range(5)]
    def mk(i):
        def f(x): acts[i].append(x.detach().flatten()); return x
        return f
    model.aq = [mk(i) for i in range(5)]
    with torch.no_grad(): model(X)
    model.aq = None
    return [torch.cat(a) for a in acts]

def fq_asym(x, lo, hi, bits):
    lo, hi = min(lo, 0.0), max(hi, 0.0)
    qmax = 2 ** bits - 1
    s = max(hi - lo, 1e-8) / qmax
    z = float(np.clip(round(-lo / s), 0, qmax))
    q = torch.clamp(torch.round(x / s) + z, 0, qmax)
    xq = (q - z) * s
    return x + (xq - x).detach()

def calib_minmax(a, bits): return a.min().item(), a.max().item()
def calib_percentile(a, bits, p=99.99):
    return 0.0, torch.quantile(a[torch.randperm(len(a))[:1_000_000]], p / 100).item()
def calib_mse(a, bits):
    amax = a.max().item(); best = (None, float("inf"))
    for t in np.linspace(amax / 50, amax, 200):
        e = torch.mean((fq_asym(a, 0.0, t, bits) - a) ** 2).item()
        if e < best[1]: best = (t, e)
    return 0.0, best[0]
def calib_kl(a, bits, nbins=2048):
    """Thuật toán entropy calibration kiểu TensorRT (Migacz, 2017), cho activation không âm."""
    a = a.numpy(); a = a[a > 0]
    hist, edges = np.histogram(a, bins=nbins, range=(0, a.max()))
    nlev = 2 ** bits - 1
    best_kl, best_i = np.inf, nbins
    for i in range(nlev, nbins + 1):
        p = hist[:i].astype(np.float64).copy(); p[-1] += hist[i:].sum()        # dồn phần bị cắt vào bin cuối
        chunks = np.array_split(hist[:i].astype(np.float64), nlev)
        q = np.concatenate([np.where(c > 0, c.sum() / max((c > 0).sum(), 1), 0) for c in chunks])
        pm, qm = p / p.sum(), q / max(q.sum(), 1e-12)
        mask = pm > 0
        if np.any(qm[mask] == 0): continue
        kl = np.sum(pm[mask] * np.log(pm[mask] / qm[mask]))
        if kl < best_kl: best_kl, best_i = kl, i
    return 0.0, float(edges[best_i])

acts_cal = collect_acts(folded, Xcal)
methods = {"min–max": calib_minmax, "percentile 99.99%": calib_percentile, "MSE": calib_mse, "KL (entropy)": calib_kl}

def ptq_eval(model, bits_w, bits_a, method, keep_ends_8bit=True):
    ranges = []
    for i, a in enumerate(acts_cal):
        b = 8 if (i == 0 and keep_ends_8bit) else bits_a
        ranges.append(calib_minmax(a, b) if i == 0 else methods[method](a, b))
    model.aq = [lambda x, r=r, b=(8 if (i == 0 and keep_ends_8bit) else bits_a): fq_asym(x, r[0], r[1], b)
                for i, r in enumerate(ranges)]
    model.wq = lambda i, w: fq_sym(w, 8 if (keep_ends_8bit and i in (0, 4)) else bits_w, True)
    acc = accuracy(model); model.aq = model.wq = None
    return acc, ranges

calib_res = {}
for bits in (8, 4):
    for m in methods:
        calib_res[(bits, m)] = ptq_eval(folded, bits, bits, m)
print("[Calibration] W{b}A{b} PTQ, weight per-channel, lớp đầu/cuối 8 bit:")
for (b, m), (acc, rg) in calib_res.items():
    print(f"   {b} bit | {m:18s}: acc={acc:.4f} | ngưỡng lớp pw (act #3) = {rg[3][1]:.3f}")

a3 = acts_cal[3].numpy(); a3p = a3[a3 > 0]
fig, ax = plt.subplots(figsize=(8.5, 3.8))
ax.hist(a3p, bins=300, color="#ccc", log=True)
cols = {"min–max": "#d9534f", "percentile 99.99%": "#e0923a", "MSE": "#4a78c2", "KL (entropy)": "#4caf6e"}
for m, c in cols.items():
    t = calib_res[(4, m)][1][3][1]
    ax.axvline(t, color=c, lw=2, label=f"{m}: ngưỡng = {t:.2f}, acc W4A4 = {calib_res[(4, m)][0] * 100:.1f}%")
ax.set_xlabel("Giá trị activation dương sau lớp pointwise + ReLU (256 ảnh calibration)")
ax.set_ylabel("Số lượng (thang log)"); ax.legend(frameon=False, fontsize=8.5)
ax.set_title(f"Ngưỡng cắt do các phương pháp calibration chọn (4 bit); {np.mean(a3 == 0):.0%} giá trị bằng 0 không vẽ")
fig.savefig(OUT + "fig11_calibration.png"); plt.close(fig)

# ------------------------------------------------------------------ (4) PTQ vs QAT theo số bit
class EMARange:
    def __init__(self, lo, hi, m=0.9): self.lo, self.hi, self.m = lo, hi, m
    def update(self, x):
        self.lo = self.m * self.lo + (1 - self.m) * min(x.min().item(), 0.0)
        self.hi = self.m * self.hi + (1 - self.m) * x.max().item()

def qat(fmodel, bits, epochs=15, lr=5e-4, seed=0):
    model = copy.deepcopy(fmodel)
    _, ranges = ptq_eval(model, bits, bits, "MSE")                 # khởi tạo dải bằng PTQ
    emas = [EMARange(lo, hi) for lo, hi in ranges]
    def mk(i):
        b = 8 if i == 0 else bits
        def f(x):
            if model.training: emas[i].update(x.detach())
            return fq_asym(x, emas[i].lo, emas[i].hi, b)
        return f
    model.aq = [mk(i) for i in range(5)]
    model.wq = lambda i, w: fq_sym(w, 8 if i in (0, 4) else bits, True)
    g = torch.Generator().manual_seed(seed)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        model.train()
        perm = torch.randperm(len(Xtr), generator=g)
        for i in range(0, len(Xtr), 64):
            idx = perm[i:i + 64]
            loss = F.cross_entropy(model(Xtr[idx]), ytr[idx])
            opt.zero_grad(); loss.backward(); opt.step()
    acc = accuracy(model); model.aq = model.wq = None
    return acc

bit_list = [8, 6, 4, 3, 2]
ptq_minmax = [ptq_eval(folded, b, b, "min–max")[0] for b in bit_list]
ptq_mse = [ptq_eval(folded, b, b, "MSE")[0] for b in bit_list]
qat_acc = [np.mean([qat(folded, b, seed=s) for s in range(3)]) for b in bit_list]
print("[PTQ vs QAT] bits:", bit_list)
print("   PTQ min-max:", [f"{a:.4f}" for a in ptq_minmax])
print("   PTQ MSE    :", [f"{a:.4f}" for a in ptq_mse])
print("   QAT (TB 3 seed):", [f"{a:.4f}" for a in qat_acc])

fig, ax = plt.subplots(figsize=(7.5, 4))
ax.axhline(acc_fp * 100, ls="--", color="k", label=f"FP32 ({acc_fp * 100:.1f}%)")
ax.plot(bit_list, np.array(ptq_minmax) * 100, "s-", color="#d9534f", label="PTQ, calibration min–max")
ax.plot(bit_list, np.array(ptq_mse) * 100, "o-", color="#e0923a", label="PTQ, calibration MSE")
ax.plot(bit_list, np.array(qat_acc) * 100, "^-", color="#4a78c2", label="QAT (15 epoch, trung bình 3 lần chạy)")
ax.set_xticks(bit_list, [f"W{b}A{b}" for b in bit_list]); ax.invert_xaxis()
ax.set_ylabel("Độ chính xác test (%)"); ax.legend(frameon=False, fontsize=9)
ax.set_title("PTQ và QAT theo độ rộng bit (lớp đầu/cuối giữ 8 bit)")
fig.savefig(OUT + "fig13_ptq_qat.png"); plt.close(fig)

# ------------------------------------------------------------------ (5) Backend int8 thật của PyTorch
import torch.ao.quantization as tq

def size_kb(m):
    torch.save(m.state_dict(), "/tmp/_m.pt"); s = os.path.getsize("/tmp/_m.pt") / 1024; os.remove("/tmp/_m.pt"); return s

class QNet(nn.Module):
    """Cùng kiến trúc Net nhưng viết theo yêu cầu của eager-mode quantization."""
    def __init__(self):
        super().__init__()
        self.quant = tq.QuantStub()
        self.conv1, self.bn1, self.relu1 = nn.Conv2d(1, 16, 3, padding=1, bias=False), nn.BatchNorm2d(16), nn.ReLU()
        self.dw, self.bn2, self.relu2 = nn.Conv2d(16, 16, 3, padding=1, groups=16, bias=False), nn.BatchNorm2d(16), nn.ReLU()
        self.pw, self.bn3, self.relu3 = nn.Conv2d(16, 32, 1, bias=False), nn.BatchNorm2d(32), nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.conv4, self.bn4, self.relu4 = nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU()
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(32, 10)
        self.dequant = tq.DeQuantStub()
    def forward(self, x):
        x = self.quant(x)
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.dw(x)))
        x = self.pool(self.relu3(self.bn3(self.pw(x))))
        x = self.relu4(self.bn4(self.conv4(x)))
        x = self.fc(torch.flatten(self.gap(x), 1))
        return self.dequant(x)

qm = QNet(); qm.load_state_dict(fp_model.state_dict(), strict=False); qm.eval()
assert accuracy(qm) == acc_fp
torch.backends.quantized.engine = "x86"
qm.qconfig = tq.get_default_qconfig("x86")
tq.fuse_modules(qm, [["conv1", "bn1", "relu1"], ["dw", "bn2", "relu2"], ["pw", "bn3", "relu3"],
                     ["conv4", "bn4", "relu4"]], inplace=True)
tq.prepare(qm, inplace=True)
with torch.no_grad(): qm(Xcal)
tq.convert(qm, inplace=True)
acc_eager = accuracy(qm)
print(f"[torch.ao eager static int8] acc = {acc_eager:.4f}; size FP32 = {size_kb(fp_model):.1f} KB, INT8 = {size_kb(qm):.1f} KB")
print("   ", qm.conv1)

# PT2E với torchao
try:
    from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
    import torchao.quantization.pt2e.quantizer.x86_inductor_quantizer as xiq
    m = copy.deepcopy(fp_model).eval()
    ep = torch.export.export(m, (Xcal[:4],), dynamic_shapes={"x": {0: torch.export.Dim("batch")}}).module()
    quantizer = xiq.X86InductorQuantizer().set_global(xiq.get_default_x86_inductor_quantization_config())
    pm = prepare_pt2e(ep, quantizer)
    with torch.no_grad(): pm(Xcal)
    cm = convert_pt2e(pm)
    with torch.no_grad():
        acc_pt2e = (cm(Xte).argmax(1) == yte).float().mean().item()
    print(f"[torchao PT2E + X86InductorQuantizer] acc = {acc_pt2e:.4f}")
except Exception as e:
    print("[PT2E] lỗi:", repr(e)[:300])

# torchao quantize_ (weight-only / dynamic) trên một MLP lớn hơn + đo độ trễ
from torchao.quantization import quantize_, Int8WeightOnlyConfig, Int8DynamicActivationInt8WeightConfig
def mlp(): return nn.Sequential(nn.Linear(1024, 4096), nn.ReLU(), nn.Linear(4096, 4096), nn.ReLU(), nn.Linear(4096, 1024)).eval()
torch.manual_seed(0); base = mlp(); xin = torch.randn(16, 1024)
def lat(m, n=30):
    with torch.no_grad():
        for _ in range(5): m(xin)
        t = time.perf_counter()
        for _ in range(n): m(xin)
    return (time.perf_counter() - t) / n * 1000
with torch.no_grad(): ref = base(xin)
rows = [("FP32", base)]
dyn_ao = tq.quantize_dynamic(copy.deepcopy(base), {nn.Linear}, dtype=torch.qint8); rows.append(("torch.ao quantize_dynamic (int8)", dyn_ao))
for name, cfg in [("torchao Int8WeightOnly", Int8WeightOnlyConfig()),
                  ("torchao Int8DynamicActivationInt8Weight", Int8DynamicActivationInt8WeightConfig())]:
    try:
        mm = copy.deepcopy(base); quantize_(mm, cfg); rows.append((name, mm))
    except Exception as e:
        print(name, "lỗi:", repr(e)[:200])
print("[MLP 1024-4096-4096-1024, batch 16, 1 luồng CPU]")
for name, m in rows:
    try:
        with torch.no_grad(): o = m(xin)
        err = ((o - ref).norm() / ref.norm()).item()
        print(f"   {name:42s} size={size_kb(m) / 1024:6.1f} MB  latency={lat(m):7.2f} ms  rel_err={err:.2e}")
    except Exception as e:
        print(f"   {name}: lỗi khi chạy {repr(e)[:200]}")

"""QAT eager mode của torch.ao: fine-tune từ trọng số FP32 đã huấn luyện rồi convert sang INT8.
Chạy: python code/torch_qat_eager.py (tự huấn luyện lại mô hình FP32 trước)."""
import os
import warnings; warnings.filterwarnings("ignore")
src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'torch_experiments.py')).read()
exec(src.split('# ------------------------------------------------------------------ (1) BN folding')[0])
exec('import torch.ao.quantization as tq\n' + 'class QNet' + src.split('class QNet')[1].split('qm = QNet()')[0])
import torch.ao.nn.intrinsic.qat as nniqat
model = QNet(); model.load_state_dict(fp_model.state_dict(), strict=False)   # bắt đầu từ trọng số FP32 đã huấn luyện
model.train()
torch.backends.quantized.engine = "x86"
model.qconfig = tq.get_default_qat_qconfig("x86")
tq.fuse_modules_qat(model, [["conv1","bn1","relu1"],["dw","bn2","relu2"],["pw","bn3","relu3"],["conv4","bn4","relu4"]], inplace=True)
tq.prepare_qat(model, inplace=True)
opt = torch.optim.Adam(model.parameters(), lr=5e-4)
g = torch.Generator().manual_seed(0)
for ep in range(10):
    if ep == 6: model.apply(tq.disable_observer)          # cố định dải lượng tử
    if ep == 7: model.apply(nniqat.freeze_bn_stats)       # cố định thống kê BN
    model.train()
    perm = torch.randperm(len(Xtr), generator=g)
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i+64]
        loss = F.cross_entropy(model(Xtr[idx]), ytr[idx]); opt.zero_grad(); loss.backward(); opt.step()
model.eval()
qmodel = tq.convert(model)
print("QAT eager int8 acc:", accuracy(qmodel))

"""Phân tích độ nhạy theo lớp: lượng tử từng lớp một xuống 3 bit (trọng số + activation đầu ra), các lớp khác giữ FP32."""
import os
src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "torch_experiments.py")).read()
exec(src.split("# ------------------------------------------------------------------ (4) PTQ vs QAT")[0])

names = ["conv1 (3x3)", "dw (depthwise 3x3)", "pw (pointwise 1x1)", "conv4 (3x3)", "fc (Linear)"]
_, ranges = ptq_eval(folded, 8, 8, "MSE")               # dải activation tham chiếu
acts_ranges3 = [calib_mse(a, 3) for a in acts_cal]       # ngưỡng MSE cho 3 bit
BITS = 3
print(f"FP32: {acc_fp:.4f}")
for li, name in enumerate(names):
    folded.wq = lambda i, w, li=li: fq_sym(w, BITS, True) if i == li else w
    # activation đầu ra của lớp li là điểm aq[li+1] (với fc: không có activation sau, chỉ lượng tử trọng số)
    def mk(i, li=li):
        if i == li + 1 and i < 5:
            lo, hi = acts_ranges3[i]
            return lambda x: fq_asym(x, lo, hi, BITS)
        return lambda x: x
    folded.aq = [mk(i) for i in range(5)]
    acc = accuracy(folded)
    folded.wq = folded.aq = None
    print(f"{name:22s}: acc = {acc:.4f}  (giảm {100 * (acc_fp - acc):.2f} điểm %)")

"""Bài 8 — quét tham số α của SmoothQuant và đo sai số đầu ra.

Script này dựng lại đúng thí nghiệm (C) của `numpy_experiments.py`: cùng hạt giống 42
và cùng thứ tự rút số ngẫu nhiên, nên các mốc đã in trong giáo trình (per-tensor
7,7435% và SmoothQuant α=0,5 cho 1,4430%) được tái lập y hệt.

Chỉ cần NumPy. Nếu có matplotlib thì script vẽ thêm đường cong ra figs/fig18_alpha.png.

    python code/sweep_alpha.py
"""

import numpy as np

rng = np.random.default_rng(42)
rel = lambda a, r: np.linalg.norm(a - r) / np.linalg.norm(r)

# --- Rút lại đúng các mẫu của phần (A) và (B) để trạng thái RNG khớp với ------
# --- numpy_experiments.py khi bước vào phần (C). Không dùng tới kết quả. ------
K, N_OUT, BATCH = 256, 64, 1000
rng.normal(0.8, 1.0, (BATCH, K)).astype(np.float32)
rng.normal(0, 0.05, (N_OUT, K)).astype(np.float32)
rng.normal(0, 0.1, N_OUT).astype(np.float32)

d_in, d_out, n_cal = 256, 128, 4096
A = rng.normal(0, 1, (d_in, d_in)) * (np.linspace(1, 0.05, d_in) ** 2)[None, :]
(A @ rng.normal(0, 1, (d_in, n_cal))).astype(np.float64)
(A @ rng.normal(0, 1, (d_in, n_cal))).astype(np.float64)
rng.normal(0, 0.02, (d_out, d_in))

# --- Phần (C): activation có outlier theo kênh -------------------------------
T, Kc, Mo = 2048, 512, 512
Xl = rng.normal(0, 1, (T, Kc))
outlier_ch = rng.choice(Kc, 6, replace=False)
Xl[:, outlier_ch] *= 60
Wl = rng.normal(0, 0.02, (Mo, Kc))
Yl = Xl @ Wl.T


def q_sym(x, S, qmax=127):
    return np.clip(np.round(x / S), -qmax, qmax) * S


def w8_per_channel(W):                      # trọng số: mỗi kênh ĐẦU RA một scale
    return q_sym(W, np.abs(W).max(axis=1, keepdims=True) / 127)


def a8_per_tensor(X):                       # activation: một scale cho cả tensor
    return q_sym(X, np.abs(X).max() / 127)


ax = np.abs(Xl).max(axis=0)                 # biên độ activation theo kênh đầu vào
aw = np.abs(Wl).max(axis=0)                 # biên độ trọng số theo kênh đầu vào

spread = lambda v: v.max() / v.min()
base = rel(a8_per_tensor(Xl) @ w8_per_channel(Wl).T, Yl)

print("=== Trước biến đổi ===")
print(f"activation: chênh lệch giữa các kênh = {spread(ax):6.1f} lần")
print(f"trọng số  : chênh lệch giữa các cột  = {spread(aw):6.2f} lần")
print(f"sai số đầu ra W8A8 khi KHÔNG biến đổi = {base:.4%}\n")

print("=== Quét alpha ===")
print(f"{'alpha':>6} {'sai số đầu ra':>14} {'chênh kênh X':>14} {'chênh cột W':>13}")
rows = []
for alpha in np.round(np.arange(0.0, 1.0001, 0.05), 2):
    s = ax ** alpha / aw ** (1 - alpha)
    Xs, Ws = Xl / s, Wl * s
    assert np.allclose(Xs @ Ws.T, Yl)       # biến đổi tương đương về mặt toán học
    err = rel(a8_per_tensor(Xs) @ w8_per_channel(Ws).T, Yl)
    nx, nw = np.abs(Xs).max(axis=0), np.abs(Ws).max(axis=0)
    rows.append((alpha, err, spread(nx), spread(nw)))
    if abs(round(alpha, 2) * 10 % 1) < 1e-9:
        print(f"{alpha:6.2f} {err:13.4%} {spread(nx):13.1f}x {spread(nw):12.1f}x")

best = min(rows, key=lambda r: r[1])
a0 = rows[0]
print(f"\ncực tiểu trên lưới: alpha = {best[0]:.2f}, sai số = {best[1]:.4%}")
print(f"alpha = 0 cho {a0[1]:.4%}, tức TỆ HƠN cả khi không biến đổi ({base:.4%}) "
      f"— vì chia cho s = 1/max|W| làm chênh lệch giữa các kênh activation "
      f"tăng từ {spread(ax):.1f} lên {a0[2]:.1f} lần.")

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    al = [r[0] for r in rows]
    er = [r[1] * 100 for r in rows]
    fig, ax_ = plt.subplots(figsize=(6.4, 3.6))
    ax_.plot(al, er, "-o", ms=3.5, color="#1f6f68", label="SmoothQuant + activation per-tensor")
    ax_.axhline(base * 100, ls="--", lw=1.2, color="#b0413e", label=f"không biến đổi ({base:.2%})")
    ax_.axvline(best[0], ls=":", lw=1.2, color="#666")
    ax_.annotate(f"α* ≈ {best[0]:.2f}", (best[0], best[1] * 100),
                 textcoords="offset points", xytext=(8, 10), fontsize=9)
    ax_.set_xlabel("α (mức chuyển độ khó từ activation sang trọng số)")
    ax_.set_ylabel("Sai số đầu ra tương đối (%)")
    ax_.set_title("Bài 8 — sai số W8A8 theo α của SmoothQuant")
    ax_.legend(frameon=False, fontsize=8)
    fig.tight_layout()
    fig.savefig("figs/fig18_alpha.png", dpi=150)
    print("\nĐã lưu figs/fig18_alpha.png")
except ImportError:
    print("\n(Không có matplotlib — bỏ qua phần vẽ hình.)")