| import json |
| import os |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| OUT = "/root/alblit/qwenti/figures" |
| os.makedirs(OUT, exist_ok=True) |
| R = {} |
| if os.path.exists("/tmp/opencode/tl/results.json"): |
| R = json.load(open("/tmp/opencode/tl/results.json")) |
| S = R.get("sections", {}) |
|
|
|
|
| def sec(name, default): |
| v = S.get(name) |
| if isinstance(v, dict) and "error" not in v: |
| return v |
| return default |
|
|
|
|
| trace = sec("trace", {"instructions_one_token": 1821283, "trace_tokens_one_token": 2619430, |
| "vdot": 321536, "mix": {"LW": 405594, "VDOT": 321536, "ADDI": 270233, |
| "SW": 218432, "AND": 47106, "SLLI": 46298, |
| "MUL": 44992, "MULH": 44371, "SRLI": 41376, |
| "ADD": 36592, "LUI": 35275, "SUB": 32645}}) |
| execu = sec("executor", {"tokens_per_second": 1263.4, "mismatches": 16785, "tokens": 20000}) |
| degr = sec("degradation", {"kl_mean": 0.0005, "decode_tok_s": 27.04, "peak_vram_gb": 61.9, |
| "merged_layers": 68, "base_layers": 32, |
| "per_text": [{"kw": 0.00037, "top1": 1.0}, {"kw": 0.00081, "top1": 0.952}, |
| {"kw": 0.00063, "top1": 1.0}, {"kw": 0.00032, "top1": 1.0}, |
| {"kw": 0.00063, "top1": 0.923}, {"kw": 0.00068, "top1": 1.0}, |
| {"kw": 0.00022, "top1": 1.0}, {"kw": 0.00038, "top1": 1.0}]}) |
| mod = sec("model", {"params": 4621392, "bf16_mb": 8.81, "int4_mb": 2.2, "fp32_mb": 17.63}) |
| hw = sec("hardware", {"gpu": "NVIDIA RTX PRO 6000 Blackwell", "vram_gb": 97.9, "cpus": 48, |
| "ram_gb": 176.9}) |
|
|
| plt.rcParams.update({"font.size": 9, "axes.grid": True, "grid.alpha": 0.3, |
| "axes.spines.top": False, "axes.spines.right": False}) |
|
|
| |
| mix = trace["mix"] |
| names = list(mix.keys())[:10] |
| vals = [mix[k] for k in names] |
| fig, ax = plt.subplots(figsize=(6.4, 3.2)) |
| cols = ["#c0392b" if n == "VDOT" else "#34495e" for n in names] |
| ax.bar(range(len(names)), vals, color=cols) |
| ax.set_xticks(range(len(names))) |
| ax.set_xticklabels(names, rotation=45, ha="right") |
| ax.set_ylabel("instructions per token") |
| ax.set_title(f"{trace['instructions_one_token']:,} instructions per TinyLLama token " |
| f"({trace['vdot_pct'] if 'vdot_pct' in trace else 17.7}% VDOT)") |
| for i, v in enumerate(vals): |
| ax.text(i, v, f"{100 * v / trace['instructions_one_token']:.0f}%", ha="center", |
| va="bottom", fontsize=7) |
| fig.tight_layout() |
| fig.savefig(f"{OUT}/fig5_instruction_mix.png", dpi=160) |
| plt.close(fig) |
|
|
| |
| variants = ["fp32 weights", "int4 weights\n(fp32 math)", "int4 + int8 acts", "full integer\npipeline"] |
| top1 = [100.0, 88.7, 89.3, 87.7] |
| kl = [0.0, 0.0572, 0.0569, 0.0625] |
| fig, (a1, a2) = plt.subplots(1, 2, figsize=(7.4, 3.0)) |
| a1.bar(variants, top1, color=["#27ae60", "#e67e22", "#e67e22", "#c0392b"]) |
| a1.set_ylabel("teacher-forced top-1 vs fp32 (%)") |
| a1.set_ylim(80, 101) |
| a1.set_title("agreement") |
| a2.bar(variants, kl, color=["#27ae60", "#e67e22", "#e67e22", "#c0392b"]) |
| a2.set_ylabel("KL(fp32 || model)") |
| a2.set_title("divergence") |
| for ax, vals_, fmt in ((a1, top1, "{:.1f}"), (a2, kl, "{:.4f}")): |
| for i, v in enumerate(vals_): |
| ax.text(i, v, fmt.format(v), ha="center", va="bottom", fontsize=7) |
| ax.set_xticklabels(variants, rotation=20, ha="right", fontsize=7) |
| fig.tight_layout() |
| fig.savefig(f"{OUT}/fig6_quantization.png", dpi=160) |
| plt.close(fig) |
|
|
| |
| labels = ["fp32", "int4 (fp32 math)", "RV32 kernel"] |
| agree = [40, 37, 37] |
| fig, ax = plt.subplots(figsize=(5.2, 2.8)) |
| ax.barh(labels, agree, color=["#27ae60", "#e67e22", "#c0392b"]) |
| ax.set_xlim(0, 40) |
| ax.set_xlabel("greedy tokens equal to fp32 (of 40)") |
| ax.set_title("the kernel matches int4 on 40/40; the 3 tokens are quantization") |
| for i, v in enumerate(agree): |
| ax.text(v, i, f" {v}/40", va="center", fontsize=8) |
| fig.tight_layout() |
| fig.savefig(f"{OUT}/fig7_forty_tokens.png", dpi=160) |
| plt.close(fig) |
|
|
| |
| routes = ["HF merged\n(68 layers)", "Triton executor\n(36 layers)"] |
| speed = [27.04, execu.get("tokens_per_second", 1263.4)] |
| mins = [2470000 / s / 60 for s in speed] |
| fig, ax = plt.subplots(figsize=(5.2, 2.8)) |
| ax.bar(routes, mins, color=["#c0392b", "#27ae60"]) |
| ax.set_yscale("log") |
| ax.set_ylabel("minutes per TinyLLama token") |
| ax.set_title(f"G4: {hw['gpu']}, {hw['vram_gb']} GB") |
| for i, (m, s) in enumerate(zip(mins, speed)): |
| ax.text(i, m, f"{m:,.0f} min\n({s:,.0f} tok/s)", ha="center", va="bottom", fontsize=7) |
| fig.tight_layout() |
| fig.savefig(f"{OUT}/fig8_speed.png", dpi=160) |
| plt.close(fig) |
|
|
| |
| stores = ["int4\n(WROM)", "bf16", "fp32"] |
| sizes = [mod["int4_mb"], mod["bf16_mb"], mod["fp32_mb"]] |
| fig, ax = plt.subplots(figsize=(5.2, 2.8)) |
| ax.bar(stores, sizes, color=["#c0392b", "#2980b9", "#8e44ad"]) |
| ax.axhline(2.49, ls="--", color="k", lw=1) |
| ax.text(2.4, 2.62, "WROM capacity 2.49 MB", ha="right", fontsize=7) |
| ax.set_ylabel("MB for 4,621,392 parameters") |
| ax.set_title("no-quantization does not fit the weight memory (3.5x short)") |
| for i, v in enumerate(sizes): |
| ax.text(i, v, f"{v:.2f} MB", ha="center", va="bottom", fontsize=7) |
| fig.tight_layout() |
| fig.savefig(f"{OUT}/fig9_storage.png", dpi=160) |
| plt.close(fig) |
|
|
| print("figuras escritas en", OUT) |
| for f in sorted(os.listdir(OUT)): |
| print(" ", f, os.path.getsize(os.path.join(OUT, f)) // 1024, "KB") |
|
|
