| import json |
| import os |
| import subprocess |
| import sys |
| import time |
| import traceback |
|
|
| os.environ.setdefault("CHIPFORMER_WROM_HEADS", "320") |
| R = {"sections": {}} |
|
|
|
|
| def sh(cmd, tail=600): |
| t = time.time() |
| p = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True) |
| out = ((p.stdout or "") + (p.stderr or ""))[-tail:] |
| print(f"[{time.time() - t:.0f}s rc={p.returncode}] {out}", flush=True) |
| return p.returncode, out |
|
|
|
|
| def safe(name): |
| def deco(fn): |
| try: |
| R["sections"][name] = fn() |
| print(f"OK {name}: {json.dumps(R['sections'][name])[:300]}", flush=True) |
| except Exception as e: |
| R["sections"][name] = {"error": f"{type(e).__name__}: {e}", |
| "tb": traceback.format_exc()[-800:]} |
| print(f"FAIL {name}: {type(e).__name__}: {e}", flush=True) |
| return fn |
| return deco |
|
|
|
|
| @safe("hardware") |
| def hardware(): |
| import torch |
| g = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", |
| "--format=csv,noheader"], capture_output=True, text=True).stdout.strip() |
| name, mem, drv = [s.strip() for s in g.split(",")] |
| return {"gpu": name, "vram_gb": round(int(mem.split()[0]) / 1024, 1), "driver": drv, |
| "torch": torch.__version__, "cuda": torch.version.cuda, |
| "capability": ".".join(map(str, torch.cuda.get_device_capability())), |
| "cpus": os.cpu_count(), |
| "ram_gb": round(int(open("/proc/meminfo").readline().split()[1]) / 1048576, 1), |
| "python": sys.version.split()[0]} |
|
|
|
|
| A = R["sections"] |
|
|
|
|
| @safe("setup") |
| def setup(): |
| sh("apt-get -qq install -y clang lld >/dev/null 2>&1; clang --version | head -1") |
| rc1, _ = sh("cd /content && rm -rf qwenputer && git clone -q https://github.com/divagr18/qwenputer " |
| "&& cd qwenputer && git checkout -q 2717232") |
| rc2, _ = sh("cd /content && rm -rf bundle && mkdir bundle && unzip -q colab_bundle.zip -d bundle " |
| "&& cp -r bundle/colab/patches/* qwenputer/ && cp bundle/colab/tl/* tl/") |
| rc3, out = sh("cd /content/tl && python3 export_full.py 2>&1 | tail -1") |
| rc4, _ = sh("cd /content/tl && clang --target=riscv32-unknown-elf -march=rv32im_zba -mabi=ilp32 " |
| "-O2 -DNTOK=40 -nostdlib -nostartfiles -static -fuse-ld=lld " |
| "-T /content/qwenputer/experimental/riscv/guests/stories260k/vd.ld " |
| "-I /content/qwenputer/experimental/riscv/chipformer/rv/env -I . " |
| "-o tl40.elf /content/qwenputer/experimental/riscv/chipformer/rv/env/crt0.S " |
| "tl.c vd2.S wrom_tl.c minic.c 2>&1 | grep ' error' | head -2") |
| return {"export": out.strip(), "rc": [rc1, rc2, rc3, rc4]} |
|
|
|
|
| @safe("model") |
| def model(): |
| from huggingface_hub import snapshot_download |
| p = snapshot_download("Maykeye/TinyLLama-v0", |
| allow_patterns=["model.safetensors", "config.json", "tokenizer.json"]) |
| for f in ("model.safetensors", "config.json", "tokenizer.json"): |
| subprocess.run(["cp", f"{p}/{f}", f"/content/tl/{f}"], check=True) |
| from safetensors import safe_open |
| n = 0 |
| with safe_open("/content/tl/model.safetensors", "pt") as f: |
| for k in f.keys(): |
| m = 1 |
| for d in f.get_slice(k).get_shape(): |
| m *= d |
| n += m |
| return {"params": n, "bf16_mb": round(n * 2 / 2 ** 20, 2), "int4_mb": round(n * 0.5 / 2 ** 20, 2), |
| "fp32_mb": round(n * 4 / 2 ** 20, 2)} |
|
|
|
|
| @safe("trace") |
| def trace(): |
| import collections |
| from chipformer.rv import elf, ref |
| from chipformer.rv import spec as S |
| m = ref.Machine.from_program(elf.program("/content/tl/tl40.elf")) |
| mix = collections.Counter() |
| toks = 0 |
| t0 = time.time() |
| while not m.halted and m.ram[0xE000 // 4 + 100] == 0: |
| w = m.rom[(m.pc - S.ROM_BASE) >> 2] if m.n != S.N_MAX else None |
| mix[ref.decode(w).op if w is not None else "R"] += 1 |
| toks += len(m.record()) |
| tot = sum(mix.values()) |
| return {"instructions_one_token": tot, "trace_tokens_one_token": toks, |
| "cpu_seconds": round(time.time() - t0, 1), |
| "tokens_per_instruction": round(toks / tot, 3), |
| "mix": {k: v for k, v in mix.most_common(12)}, |
| "vdot": mix["VDOT"], "vdot_pct": round(100 * mix["VDOT"] / tot, 1), |
| "overhead_pct": round(100 * (tot - mix["VDOT"]) / tot, 1)} |
|
|
|
|
| @safe("executor") |
| def executor(): |
| import numpy as np |
| import torch |
| from chipformer.rv import core, elf, fastexec, ref |
| from chipformer.rv import spec as S |
| c = core.build() |
| prog = elf.program("/content/tl/tl40.elf") |
| want = ref.Machine.from_program(prog).run(10 ** 8) |
| out = {"circuit": c.stats(), "reference_tokens": len(want)} |
| core.snapshot(c, ref.Machine.from_program(prog)) |
| N = 200000 |
| fr = fastexec.FreeRun(c, 1, "cuda", False, v2=True) |
| fr.load(0, S.tok("GO")) |
| t0 = time.time() |
| got = fr.run(N)[0] |
| torch.cuda.synchronize() |
| dt = time.time() - t0 |
| exp = np.array(want[1:N + 1]) |
| bad = np.nonzero(got != exp)[0] |
| out.update(tokens=N, seconds=round(dt, 1), tokens_per_second=round(N / dt, 1), |
| mismatches=int(len(bad)), |
| first_mismatch=None if not len(bad) else int(bad[0]), |
| minutes_per_model_token=round(2470000 / (N / dt) / 60, 1)) |
| return out |
|
|
|
|
| @safe("merge") |
| def merge(): |
| from huggingface_hub import snapshot_download |
| p = snapshot_download("Qwen/Qwen3.5-4B", allow_patterns=["*.safetensors", "*.json", "*.jinja"]) |
| rc, out = sh(f"cd /content/qwenputer/experimental/riscv && CHIPFORMER_WROM_HEADS=320 " |
| f"python3 -m chipformer.rv.merge {p} --out /content/merged 2>&1 | tail -4", 1500) |
| size = subprocess.run(["bash", "-lc", "du -sm /content/merged | cut -f1"], |
| capture_output=True, text=True).stdout.strip() |
| import re |
| layers = re.findall(r"'layers': (\d+)", out) |
| return {"rc": rc, "size_mb": int(size or 0), "circuit_layers": int(layers[0]) if layers else None, |
| "tail": out.strip()[-300:]} |
|
|
|
|
| @safe("degradation") |
| def degradation(): |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| tok = AutoTokenizer.from_pretrained("/content/merged") |
| base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16, |
| device_map="cuda").eval() |
| merged = AutoModelForCausalLM.from_pretrained("/content/merged", dtype=torch.bfloat16, |
| device_map="cuda").eval() |
|
|
| def logsm(z): |
| z = z - z.max(-1, keepdim=True).values |
| return z - torch.log(torch.exp(z).sum(-1, keepdim=True)) |
|
|
| texts = ["The capital of France is Paris, and the river that runs through it is the Seine.", |
| "In 1969, humans first landed on the Moon, a milestone of the space race.", |
| "A neural network learns by adjusting its weights to reduce a loss function.", |
| "The recipe calls for two cups of flour, one egg, and a pinch of salt.", |
| "Photosynthesis converts light energy into chemical energy inside chloroplasts.", |
| "The committee approved the budget after a long debate about public transport.", |
| "Mount Everest rises 8,849 metres above sea level on the border of Nepal and China.", |
| "She opened the book, read the first page twice, and put it down again."] |
| rows, kls, pos = [], [], 0 |
| with torch.no_grad(): |
| for t in texts: |
| ids = tok(t, return_tensors="pt").to("cuda") |
| a, b = base(**ids).logits[0].float(), merged(**ids).logits[0].float() |
| n = min(a.shape[-1], b.shape[-1]) |
| a, b = a[..., :n], b[..., :n] |
| k = (torch.exp(logsm(a)) * (logsm(a) - logsm(b))).sum(-1).mean().item() |
| agree = (a.argmax(-1) == b.argmax(-1)).float().mean().item() |
| rows.append({"kw": round(k, 6), "top1": round(agree, 3), "positions": int(a.shape[0])}) |
| kls.append(k) |
| pos += int(a.shape[0]) |
| ids = tok("The capital of France is", return_tensors="pt").to("cuda") |
| with torch.no_grad(): |
| out = merged.generate(**ids, max_new_tokens=20, do_sample=False) |
| text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True) |
| merged.generate(**ids, max_new_tokens=9, do_sample=False) |
| torch.cuda.synchronize() |
| t0 = time.time() |
| merged.generate(**ids, max_new_tokens=32, do_sample=False) |
| torch.cuda.synchronize() |
| dt = time.time() - t0 |
| return {"per_text": rows, "kl_mean": round(sum(kls) / len(kls), 6), "positions": pos, |
| "sample_continuation": text, "decode_tok_s": round(32 / dt, 2), |
| "peak_vram_gb": round(torch.cuda.max_memory_allocated() / 2 ** 30, 1), |
| "base_layers": len(base.model.layers), "merged_layers": len(merged.model.layers)} |
|
|
|
|
| json.dump(R, open("/content/results.json", "w"), indent=1) |
| print("RESULTS " + json.dumps(R), flush=True) |
|
|
