File size: 5,544 Bytes
2049fc8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | 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})
# fig 1: instruction mix
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)
# fig 2: quantization cost
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)
# fig 3: 40-token agreement
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)
# fig 4: speed
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)
# fig 5: storage
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")
|