このチュートリアルでは、Dory 1.0 4Bという、NVIDIAのNemotron Labsが開発した小さな推論モデルを使用します。このモデルの中間層は共有ブロックとなっており、各トークンに対して1回、2回、3回、または4回適用することができます。各追加のパスやループによって、トークンあたりの計算量が増加しますが、パラメータは一切追加されません。そのため、1つのチェックポイントでは、再学習を行わずに4種類の異なる計算バージョンで実行することができます。私たちはNVIDIAのモデルカードに記載されているllama.cppブランチを構築し、3.7 GBのGGUFファイルをダウンロードし、ファイルから直接アーキテクチャを読み取り、各ループがどれだけパラメータとキャッシュメモリを消費するかを予測します。その後、これらの予測値を実行中のサーバーと比較し、各ループ数におけるデコード速度を測定し、安価な最後のループのキャッシュがモデルの予測を変えるかどうかをテストし、GSM8K評価を実施して、ループ数が増えるにつれて精度、回答長さ、解決する問題あたりの時間がどのように変化するかを確認します。
コードをコピーimport os
import re
import sys
import json
import time
import math
import random
import signal
import shutil
import hashlib
import platform
import statistics
import traceback
import subprocess
import collections
import urllib.request
from pathlib import Path
RESULTS = {} # step -> one-line result, for the summary table
CLAIMS = [] # (step, claim, held)
CURRENT = ["1"]
def banner(title):
print("\n" + "=" * 78)
print(title)
print("=" * 78)
def section(name):
def wrap(fn):
def run(*a, **kw):
CURRENT[0] = name.split(".")[0]
banner(name)
t0 = time.time()
try:
out = fn(*a, **kw)
RESULTS[name] = out if isinstance(out, str) else "ok"
return out
except Exception as e:
RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
traceback.print_exc(limit=3)
return None
finally:
print(f"\n[{time.time() - t0:.1f} s]")
return run
return wrap
def claim(text, held):
CLAIMS.append((CURRENT[0], text, bool(held)))
return held
banner("1. Build the pinned Dory llama.cpp fork and record the environment")
FORK = "https://github.com/brianchmielNV/llama.cpp.git"
FORK_COMMIT = "98a74c97ea09ac72f6ceda1559808fc0a2f447ce" # the commit NVIDIA's model card pins; ggml-org PR #30306 is still a draft
GGUF_REPO = "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-GGUF"
GGUF_REV = "df53755cb8741843af8e019e088819fbef17e29b"
GGUF_FILE = "NVIDIA-Nemotron-Labs-Dory1.0-4B-Q5_K_M.gguf" # the card's commands call it dory-q5_k_m.gguf
GGUF_SIZE = 3_673_288_864
GGUF_SHA256 = "96080be49ffd67b916018b7bdef8c484762576fb52780e10ab98ac4978a660e0"
GSM8K_REV = "740312add88f781978c0658806c59bc2815b9866"
WORK = Path(os.environ.get("DORY_WORK") or ("/content/dory" if os.path.isdir("/content") else "dory_work")).resolve()
WORK.mkdir(parents=True, exist_ok=True)
(WORK / "logs").mkdir(exist_ok=True)
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "gguf==0.19.0", "huggingface_hub", "pandas", "pyarrow"], check=True)
SYSTEM = platform.system()
GPU = ""
if shutil.which("nvidia-smi"):
q = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,compute_cap", "--format=csv,noheader"],
capture_output=True, text=True)
GPU = q.stdout.strip().splitlines()[0] if q.returncode == 0 and q.stdout.strip() else ""
BACKEND = "metal" if SYSTEM == "Darwin" else ("cuda" if GPU else "cpu")
SRC = WORK / "llama.cpp-dory"
SERVER_BIN = SRC / "build" / "bin" / "llama-server"
free_gb = shutil.disk_usage(WORK).free / 1e9
print(f"{SYSTEM} {platform.machine()}, Python {platform.python_version()}, {os.cpu_count()} CPU cores, "
f"backend {BACKEND}{' (' + GPU + ')' if GPU else ''}, {free_gb:.1f} GB free disk")
if not SERVER_BIN.exists():
t0 = time.time()
if not (SRC / ".git").exists():
subprocess.run(["git", "init", "-q", str(SRC)], check=True)
subprocess.run(["git", "-C", str(SRC), "remote", "add", "origin", FORK], check=True)
subprocess.run(["git", "-C", str(SRC), "fetch", "-q", "--depth", "1", "origin", FORK_COMMIT], check=True)
subprocess.run(["git", "-C", str(SRC), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
flags = ["-DCMAKE_BUILD_TYPE=Release", "-DLLAMA_BUILD_SERVER=ON", "-DLLAMA_BUILD_TESTS=OFF",
"-DLLAMA_BUILD_EXAMPLES=OFF", "-DLLAMA_BUILD_APP=OFF", "-DLLAMA_BUILD_UI=OFF",
"-DLLAMA_USE_PREBUILT_UI=ON", "-DLLAMA_OPENSSL=OFF"]
if BACKEND == "metal": # the card's Apple-silicon flags
flags += ["-DCMAKE_OSX_ARCHITECTURES=arm64", "-DGGML_METAL=ON", "-DGGML_METAL_EMBED_LIBRARY=ON",
"-DGGML_BLAS=ON", "-DGGML_BLAS_VENDOR=Apple"]
elif BACKEND == "cuda": # kernels for this GPU only (a T4 is 7.5); flash-attention kernels are unused with --flash-attn off
cap = GPU.split(",")[-1].strip().replace(".", "")
flags += ["-DGGML_CUDA=ON", f"-DCMAKE_CUDA_ARCHITECTURES={cap}", "-DGGML_CUDA_FA=OFF", "-DGGML_NATIVE=OFF"]
env = dict(os.environ, HF_UI_VERSION="b11551")
with open(WORK / "build.log", "w") as log:
subprocess.run(["cmake", "-S", str(SRC), "-B", str(SRC / "build")] + flags, check=True, env=env, stdout=log, stderr=log)
subprocess.run(["cmake", "--build", str(SRC / "build"), "--parallel", str(os.cpu_count()), "--target", "llama-server"],
check=True, env=env, stdout=log, stderr=log)
print(f"built llama-server in {time.time() - t0:.0f} s")
head = subprocess.run(["git", "-C", str(SRC), "rev-parse", "HEAD"], capture_output=True, text=True).stdout.strip()
assert head == FORK_COMMIT, f"checked out {head}, expected {FORK_COMMIT}"
version = subprocess.run([str(SERVER_BIN), "--version"], capture_output=True, text=True)
print(f"fork commit {head[:12]}; {(version.stdout + version.stderr).strip().splitlines()[-1]}")
RESULTS["1. Build the pinned Dory llama.cpp fork and record the environment"] = f"fork {head[:12]}, backend {BACKEND}"
NVIDIAのモデルカードピンを利用するため、私たちはDoryブランチの正確なコミットを基にllama-serverを構築しています。なぜなら、llama.cppにおけるDoryのサポートはまだオープンなドラフトプルリクエストであり、リリースされたバイナリファイルではモデルを読み込むことができないからです。このセルはプラットフォームを検出し、Mac上ではカードのMetalフラグを渡し、Colab上では単独でGPU用のCUDAカーネルをコンパイルします。カードの–flash-attnが全く使用されないflash-attentionカーネルはスキップされ、これによりT4のビルドが管理可能になります。 その後、チェックアウトされたコミットがピンされるものであるかを確認し、したがって後のすべての数字はコードのある特定のバージョンに遡る。
コードをコピー@section("2. Download the Q5_K_M GGUF at a pinned revision and check its SHA-256")
def download_step():
global GGUF
from huggingface_hub import hf_hub_download
t0 = time.time()
GGUF = Path(hf_hub_download(GGUF_REPO, GGUF_FILE, revision=GGUF_REV, local_dir=str(WORK)))
h = hashlib.sha256()
with open(GGUF, "rb") as f:
for block in iter(lambda: f.read(1 << 24), b""):
h.update(block)
size, digest = GGUF.stat().st_size, h.hexdigest()
print(f"{GGUF.name}: {size:,} bytes, sha256 {digest[:16]}... ({time.time() - t0:.0f} s)")
assert size == GGUF_SIZE and digest == GGUF_SHA256, "the file does not match the card's size and checksum"
claim("the GGUF matches the size and SHA-256 published on the card", True)
return f"{size / 1e9:.2f} GB, checksum matches the card"
download_step()
Q5_K_M GGUFは、Hugging Faceの固定レポジトリリビュエーションから提供され、何かがそのファイルにアクセスする前に必ずハッシュ処理が行われます。カードのコマンドではファイルをdory-q5_k_m.ggufと呼んでいますが、レポジトリではNVIDIA-Nemotron-Labs-Dory1.0-4B-Q5_K_M.ggufと名前が付けられているため、カードのコマンドをそのままコピーすることはできません。実際の名前、サイズ、SHA-256値を一致させることで、間違ったファイルや破損したダウンロードの両方を防ぐことができます。 3.67GBのファイルは、カードに公表されたチェックサムと完全に一致しているため、その後のすべての測定値はNVIDIAが公表した同じバイト数を使用する。
コードをコピー@section("3. Read the architecture from the GGUF: three layer groups and a default of four loops")
def architecture_step():
global META, GROUPS, PATTERN
import gguf
reader = gguf.GGUFReader(str(GGUF))
META = {k: f.contents() for k, f in reader.fields.items() if k.startswith("dory.")}
for k in ("dory.block_count", "dory.input_layer_count", "dory.recurrent_layer_count", "dory.output_layer_count",
"dory.num_loops", "dory.recurrent_kv_cache_mode", "dory.attention.sliding_window", "dory.context_length"):
print(f"{k:34s} {META[k]}")
PATTERN = META["dory.layer_pattern"].split("-")[:-1] # one symbol per attention layer: S sliding window, * full
n_in, n_rec = META["dory.input_layer_count"], META["dory.recurrent_layer_count"]
print(f"\nlayer entries alternate attention and MLP: {len(PATTERN)} attention layers "
f"({PATTERN.count('S')} sliding-window, {PATTERN.count('*')} full attention)")
GROUPS = collections.defaultdict(lambda: {"params": 0, "bytes": 0})
dtypes = collections.Counter()
for t in reader.tensors:
m = re.match(r"blk\.(\d+)\.", t.name)
if m:
i = int(m.group(1))
g = "input block" if i < n_in else "recurrent block" if i < n_in + n_rec else "output block"
else:
g = "embedding" if t.name.startswith("token_embd") else "LM head"
GROUPS[g]["params"] += int(math.prod(int(x) for x in t.shape))
GROUPS[g]["bytes"] += int(t.n_bytes)
dtypes[t.tensor_type.name] += 1
print(f"\n{'group':16s} {'params':>9s} {'bytes':>9s}")
for g in ("embedding", "input block", "recurrent block", "output block", "LM head"):
print(f"{g:16s} {GROUPS[g]['params'] / 1e9:7.3f} B {GROUPS[g]['bytes'] / 1e9:6.3f} GB")
print(f"tensor types: {dict(dtypes)}")
claim("the GGUF stores 16 input, 40 recurrent and 16 output layer entries (8, 20 and 8 blocks)",
(n_in, n_rec, META["dory.output_layer_count"]) == (16, 40, 16))
claim("the GGUF's built-in loop count is 4, not the 2 the card recommends", META["dory.num_loops"] == 4)
total = sum(v["params"] for v in GROUPS.values())
return (f"{total / 1e9:.2f} B parameters, {GROUPS['recurrent block']['params'] / 1e9:.2f} B of them in the shared "
f"recurrent block; built-in num_loops = {META['dory.num_loops']}")
architecture_step()
GGUFヘッダーは、モデルを読み込む前に、その構造がどのようになっているかを示します。Doryの72層のエントリは、16層の入力グループ、40層の共有再帰ブロックグループ、そして16層の出力グループに分かれており、AttentionとMLPが交互に配置されています。また、3つのスライディングウィンドウ層ごとに1つの完全注意層が存在します。各グループによるテンソルサイズの合計を見ると、再帰ブロックがモデルの39.8億パラメータの1.84を占めています。 ヘッダーには罠もあります。内蔵されているループ数は4であり、カードでは2が推奨されています。そのため、オーバーロードなしでサーバーを起動すると、最も高い設定が無意識のうちに実行されます。
コードをコピーLOOPS = (1, 2, 3, 4)
CTX, UBATCH = 4096, 512
def predicted_kv(loops, mode, ctx=CTX, ubatch=UBATCH):
"""KV cache under the fork's allocation rule: per_loop keeps one cache per loop for every recurrent attention layer,
last_loop keeps one in all. Full-attention layers hold ctx cells; sliding-window layers hold window + ubatch, padded to 256."""
n_in, n_rec = META["dory.input_layer_count"] // 2, META["dory.recurrent_layer_count"] // 2
rec, fixed = PATTERN[n_in:n_in + n_rec], PATTERN[:n_in] + PATTERN[n_in + n_rec:]
copies = loops if mode == "per_loop" else 1
full = fixed.count("*") + copies * rec.count("*")
swa = fixed.count("S") + copies * rec.count("S")
swa_cells = min(ctx, -(-(META["dory.attention.sliding_window"] + ubatch) // 256) * 256)
per_cell = 2 * META["dory.attention.head_count_kv"] * META["dory.attention.key_length"] * 2 # K and V, 2 bytes each
return full, swa, (full * ctx + swa * swa_cells) * per_cell / 2**20
@section("4. Predict what each extra loop costs: parameters per token and KV-cache size")
def cost_model_step():
global PER_TOKEN
once = ("input block", "output block", "LM head") # the embedding is a one-row lookup, so it costs no matrix work
PER_TOKEN = {}
print(f"{'loops':>5s} {'params per token':>17s} {'vs 1 loop':>10s} {'weight bytes read':>18s} {'vs 1 loop':>10s}")
for n in LOOPS:
p = sum(GROUPS[g]["params"] for g in once) + n * GROUPS["recurrent block"]["params"]
b = sum(GROUPS[g]["bytes"] for g in once) + n * GROUPS["recurrent block"]["bytes"]
PER_TOKEN[n] = (p, b)
print(f"{n:5d} {p / 1e9:15.2f} B {p / PER_TOKEN[1][0]:9.2f}x {b / 1e9:16.2f} GB {b / PER_TOKEN[1][1]:9.2f}x")
print(f"\npredicted KV cache at --ctx-size {CTX}, f16 keys and values, micro-batch {UBATCH}:")
print(f"{'loops':>5s} {'per_loop: full + SWA layers':>28s} {'MiB':>7s} {'last_loop: full + SWA layers':>29s} {'MiB':>7s}")
for n in LOOPS:
a, b = predicted_kv(n, "per_loop"), predicted_kv(n, "last_loop")
print(f"{n:5d} {a[0]:>21d} + {a[1]:<4d} {a[2]:7.1f} {b[0]:>22d} + {b[1]:<4d} {b[2]:7.1f}")
full_ctx = META["dory.context_length"]
print(f"\nat the full {full_ctx:,}-token context: per_loop {predicted_kv(2, 'per_loop', full_ctx)[2] / 1024:.1f} GiB at 2 loops "
f"and {predicted_kv(4, 'per_loop', full_ctx)[2] / 1024:.1f} GiB at 4; last_loop "
f"{predicted_kv(4, 'last_loop', full_ctx)[2] / 1024:.1f} GiB at any loop count")
return (f"4 loops apply {PER_TOKEN[4][0] / PER_TOKEN[1][0]:.2f}x the parameters per token of 1 loop; predicted KV at 4 loops "
f"{predicted_kv(4, 'per_loop')[2]:.0f} MiB per_loop vs {predicted_kv(4, 'last_loop')[2]:.0f} MiB last_loop")
cost_model_step()
それらのグループサイズから、何も実行する前にループがどれだけのコストを要するかを予測できます。トークンは入力ブロックを通過し、各ループごとに再帰ブロックを1回、出力ブロック、そしてLMヘッドを通過します。したがって、4回のループを行うと、1回のループのパラメータの2.51倍の計算量がかかり、2.28倍のウェイトバイトが読み込まれます。これは、メモリに限られたデコーダにとってより良い指針です。 KVキャッシュはフォークの割り当て規則に従います。各ループごとに、各再帰的注意層の各ループに対して別々のキャッシュを保持し、4,096トークンセットコンテキストにおいて1ループ時の306MiBから4ループ時の816MiBに増加します。一方、last_loopは306MiBのままです。
コードをコピーPORT = 8090
KV_LINE = re.compile(r"llama_kv_cache: size =\s*([\d.]+) MiB \(\s*(\d+) cells,\s*(\d+) layers")
class DoryServer:
"""llama-server with one loop count and KV mode. Both are read when the model loads, so each setting is a fresh start."""
def __init__(self, loops=None, mode="per_loop", ngl=999, ubatch=UBATCH):
self.loops, self.mode, self.ngl, self.ubatch = loops, mode, ngl, ubatch
self.name = f"{'default' if loops is None else loops}-{mode}-ngl{ngl}-ub{ubatch}"
def __enter__(self):
cmd = [str(SERVER_BIN), "-m", str(GGUF), "--host", "127.0.0.1", "--port", str(PORT), "-ngl", str(self.ngl),
"--ctx-size", str(CTX), "--parallel", "1", "--batch-size", str(self.ubatch), "--ubatch-size", str(self.ubatch),
"--cache-type-k", "f16", "--cache-type-v", "f16", "--flash-attn", "off", "--no-context-shift",
"--jinja", "--reasoning", "on", "--temp", "1.0", "--top-p", "0.95", "--top-k", "0", "--min-p", "0",
"--repeat-penalty", "1"]
if self.loops is not None:
cmd += ["--override-kv", f"dory.num_loops=int:{self.loops},dory.recurrent_kv_cache_mode=str:{self.mode}"]
self.log = WORK / "logs" / f"server-{self.name}.log"
with open(self.log, "w") as f:
self.proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT, start_new_session=True)
t0 = time.time()
while True:
try:
with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r:
if r.status == 200:
break
except Exception:
pass
if self.proc.poll() is not None or time.time() - t0 > 600:
self.stop()
raise RuntimeError(f"llama-server did not start; see {self.log}")
time.sleep(0.5)
self.load_s = time.time() - t0
text = self.log.read_text(errors="ignore")
sizes = KV_LINE.findall(text) # the full-attention cache first, then the sliding-window cache
self.kv_mib = sum(float(s[0]) for s in sizes)
self.kv_layers = tuple(int(s[2]) for s in sizes)
m = re.search(r"offloaded (\d+)/(\d+) layers", text)
self.offloaded = f"{m.group(1)}/{m.group(2)}" if m else "n/a"
return self
def __exit__(self, *exc):
self.stop()
def stop(self):
if self.proc.poll() is None:
os.killpg(self.proc.pid, signal.SIGTERM)
try:
self.proc.wait(60)
except subprocess.TimeoutExpired:
os.killpg(self.proc.pid, signal.SIGKILL)
self.proc.wait()
def post(self, path, payload, timeout=1800):
req = urllib.request.Request(f"http://127.0.0.1:{PORT}{path}", data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.loads(r.read())
def split_reasoning(message):
"""The server returns the thinking in reasoning_content when it can parse it; otherwise it stays inside <think> tags."""
content, reasoning = message.get("content") or "", message.get("reasoning_content") or ""
m = re.search(r"<think>(.*?)</think>", content, re.S)
if m:
reasoning, content = reasoning + m.group(1), content.replace(m.group(0), "")
return reasoning.strip(), content.strip()
@section("5. Start the server, prove the loop override took effect, and catch the silent default")
def override_step():
rows = []
for loops, mode in ((2, "per_loop"), (2, "last_loop"), (None, "per_loop")):
with DoryServer(loops, mode) as s:
want = predicted_kv(META["dory.num_loops"] if loops is None else loops, mode)
rows.append((s.name, s.kv_layers, s.kv_mib, want, s.offloaded, s.load_s))
if loops == 2 and mode == "per_loop":
r = s.post("/v1/chat/completions", {"messages": [{"role": "user", "content": "What is 17 multiplied by 23?"}],
"max_tokens": 1024, "seed": 1, "cache_prompt": False})
reasoning, answer = split_reasoning(r["choices"][0]["message"])
print(f"2 loops, a first request: {r['usage']['completion_tokens']} tokens, "
f"{len(reasoning.split())} words of reasoning, answer: {answer[:120]!r}")
print(f" fields in the message: {sorted(r['choices'][0]['message'])}; "
f"timings: {sorted(r.get('timings', {}))[:6]}...\n")
print(f"{'server':28s} {'KV layers (full, SWA)':>22s} {'KV MiB':>8s} {'predicted':>16s} {'offloaded':>10s} {'load s':>7s}")
for name, layers, mib, want, off, load in rows:
print(f"{name:28s} {str(layers):>22s} {mib:8.1f} {str(want[:2]) + f' {want[2]:.1f}':>16s} {off:>10s} {load:7.1f}")
ok = all(layers == want[:2] for _, layers, _, want, _, _ in rows)
claim("the --override-kv loop count and cache mode are applied (KV layer counts match the prediction)", ok)
default = rows[-1]
claim("without an override the server runs the GGUF's 4 loops (the most expensive setting)",
default[1] == predicted_kv(4, "per_loop")[:2])
return f"override applied: {ok}; no override -> {default[1]} KV layers, {default[2]:.0f} MiB, i.e. 4 loops"
override_step()
ループ回数とキャッシュモードはモデルが読み込まれるときに読み取られるため、各設定には独自のサーバープロセスが必要である。DoryServerヘルパーは各設定ごとに1つを起動し、ヘルスチェックを待ち、サーバーのログからKVキャッシュ行を解析する。解析されたレイヤー数を予測値と比較することで、フラグのみを信頼するのではなく、オーバーライドが実際に効果を発揮したことを証明する。 また、オーバーライドなしでサーバーを1台起動し、ネガティブコントロールとして、サーバーがモデルの推論と回答をどのように分離するかを確認するためにチャットリクエストを1回送信する。
コードをコピーCONFIGS = [(1, "per_loop"), (2, "per_loop"), (2, "last_loop"), (3, "per_loop"), (3, "last_loop"), (4, "per_loop"), (4, "last_loop")]
SPEED_PROMPT = ("Explain, step by step, how a recurrent block that is applied several times per token changes the amount of "
"computation a language model spends on each token, and what that means for memory and latency.")
@section("6. Measure KV memory and decode speed for every loop count and cache mode")
def sweep_step():
global SPEED
SPEED = {}
print(f"{'loops':>5s} {'mode':10s} {'KV MiB':>8s} {'predicted':>9s} {'prefill tok/s':>14s} {'decode tok/s':>13s} "
f"{'decode vs 1 loop':>17s} {'params ratio':>13s}")
for loops, mode in CONFIGS:
with DoryServer(loops, mode) as s:
runs = [s.post("/completion", {"prompt": SPEED_PROMPT, "n_predict": 192, "temperature": 0, "ignore_eos": True,
"cache_prompt": False, "seed": 1})["timings"] for _ in range(3)]
dec = statistics.median(t["predicted_per_second"] for t in runs)
pre = statistics.median(t["prompt_per_second"] for t in runs)
SPEED[(loops, mode)] = {"kv_mib": s.kv_mib, "kv_layers": s.kv_layers, "decode": dec, "prefill": pre,
"decode_runs": [round(t["predicted_per_second"], 2) for t in runs]}
base = SPEED[(1, "per_loop")]["decode"]
print(f"{loops:5d} {mode:10s} {s.kv_mib:8.1f} {predicted_kv(loops, mode)[2]:9.1f} {pre:14.1f} {dec:13.1f} "
f"{base / dec:16.2f}x {PER_TOKEN[loops][0] / PER_TOKEN[1][0]:12.2f}x")
flat = len({round(SPEED[(n, 'last_loop')]['kv_mib'], 1) for n in (2, 3, 4)} | {round(SPEED[(1, 'per_loop')]['kv_mib'], 1)}) == 1
grows = SPEED[(1, "per_loop")]["kv_mib"] < SPEED[(2, "per_loop")]["kv_mib"] < SPEED[(3, "per_loop")]["kv_mib"] < SPEED[(4, "per_loop")]["kv_mib"]
claim("last_loop keeps the KV cache at the one-loop size for every loop count", flat)
claim("per_loop KV memory grows with every added loop", grows)
slow = SPEED[(1, "per_loop")]["decode"] / SPEED[(4, "per_loop")]["decode"]
json.dump({f"{k[0]}-{k[1]}": v for k, v in SPEED.items()}, open(WORK / "speed.json", "w"), indent=1)
return (f"KV {SPEED[(1, 'per_loop')]['kv_mib']:.0f} MiB at 1 loop; 4 loops: {SPEED[(4, 'per_loop')]['kv_mib']:.0f} MiB per_loop, "
f"{SPEED[(4, 'last_loop')]['kv_mib']:.0f} MiB last_loop; decode {slow:.2f}x slower at 4 loops than at 1")
sweep_step()
ヘルパーが設置された後、7種類の設定ごとにサーバーを起動し、キャッシュモードにおいて1回から4回のループを行い、予測の隣に測定されたKVメモリを記録する。デコード速度は、固定されたプロンプトにおけるサーバーの自身のタイミングから得られ、3回にわたって192個のトークンを無駄なく生成し、中央値を保持する。したがって、行間で変わる唯一のものはループ設定である。 1ループの速度を各行の速度で割ることで、測定された減速が得られる。この値を予測比率の隣に記すことで、デコードが計算制限によるものかメモリ制限によるものかを確認する。
コードをコピーdef top_probs(response):
"""Next-token candidates as {token: logprob}, from either of the server's two probability formats."""
first = response["completion_probabilities"][0]
if "top_logprobs" in first:
return {c["token"]: c["logprob"] for c in first["top_logprobs"]}
return {c["tok_str"]: math.log(max(c["prob"], 1e-30)) for c in first["probs"]}
GATE_PROMPTS = ["The capital of France is", "def fibonacci(n):\n ", "Water boils at a temperature of",
"In 1969, the first humans to walk on the Moon were", "The derivative of x squared is"]
@section("7. Correctness gates: the same seed gives the same answer, and the GPU agrees with the CPU")
def gates_step():
payload = {"messages": [{"role": "user", "content": "Name three prime numbers greater than 50."}],
"max_tokens": 512, "seed": 7, "cache_prompt": False}
probe = lambda s, p: top_probs(s.post("/completion", {"prompt": p, "n_predict": 1, "n_probs": 10, "temperature": 0,
"cache_prompt": False}))
with DoryServer(2, "per_loop") as s:
a, b = (s.post("/v1/chat/completions", payload) for _ in range(2))
same = a["choices"][0]["message"] == b["choices"][0]["message"]
accel = [probe(s, p) for p in GATE_PROMPTS]
print(f"two sampled chat answers with seed 7 at temperature 1.0: identical = {same}")
with DoryServer(2, "per_loop", ngl=0) as s:
cpu = [probe(s, p) for p in GATE_PROMPTS]
agree, gaps = 0, []
for p, x, y in zip(GATE_PROMPTS, accel, cpu):
tx, ty = max(x, key=x.get), max(y, key=y.get)
agree += tx == ty
gaps += [abs(x[t] - y[t]) for t in set(x) & set(y)]
print(f" {p[:40]!r:44s} {BACKEND}: {tx!r:12s} CPU: {ty!r:12s}")
print(f"{BACKEND} vs CPU: top-1 token agrees on {agree} of {len(GATE_PROMPTS)} prompts; largest log-probability gap "
f"over shared top-10 tokens {max(gaps):.3f}")
claim("a fixed seed reproduces a sampled answer exactly", same)
claim(f"the {BACKEND} build agrees with the CPU build on the top-1 token", agree >= len(GATE_PROMPTS) - 1)
return f"seeded answers identical: {same}; {BACKEND} vs CPU top-1 agreement {agree}/{len(GATE_PROMPTS)}"
gates_step()
設定間の比較は、測定値が一致し、加速されたビルドが参照CPUパスと同じモデルを計算する場合にのみ意味があります。カードのサンプリング温度で同じシードを持つ2つのチャットリクエストは、同一のメッセージを返さなければなりません。そうでないと、各質問に対する比較ではモデルの動作とサンプリングノイズが混ざり合ってしまいます。 その後、レイヤーがオフロードされていないCPUのみのサーバーを起動し、5つのプロンプトにおける最も可能性の高い次のトークンをGPUまたはMetalサーバーと比較する。なぜなら、Doryのサポートはすべてのバックエンドでアップストリームでレビューされていないからだ。
コードをコピーdef load_gsm8k():
import pandas as pd
from huggingface_hub import hf_hub_download
path = hf_hub_download("openai/gsm8k", "main/test-00000-of-00001.parquet", repo_type="dataset", revision=GSM8K_REV,
local_dir=str(WORK / "gsm8k"))
df = pd.read_parquet(path)
df["gold"] = df.answer.str.extract(r"####\s*([-\d,\.]+)")[0].str.replace(",", "").astype(float)
return df
def tv_distance(p, q):
"""Total-variation distance between two top-k next-token distributions, each renormalised over its own top k."""
zp, zq = sum(math.exp(v) for v in p.values()), sum(math.exp(v) for v in q.values())
return 0.5 * sum(abs(math.exp(p.get(t, -1e9)) / zp - math.exp(q.get(t, -1e9)) / zq) for t in set(p) | set(q))
@section("8. What an extra loop changes: next-token distributions on 50 prompts")
def mechanism_step():
global GSM8K
GSM8K = load_gsm8k()
prompts = list(GSM8K.sample(50, random_state=11).question)
configs = [(n, m, ub) for ub in (UBATCH, 16) for n in (1, 2, 4) for m in ("per_loop", "last_loop")]
dist, lengths = {}, []
for n, m, ub in configs:
with DoryServer(n, m, ubatch=ub) as s:
got = [s.post("/completion", {"prompt": p, "n_predict": 1, "n_probs": 20, "temperature": 0, "cache_prompt": False})
for p in prompts]
dist[(n, m, ub)] = [top_probs(r) for r in got]
lengths = lengths or [r["timings"]["prompt_n"] for r in got]
print(f"50 GSM8K questions as raw prompts, {min(lengths)} to {max(lengths)} tokens: one chunk at a {UBATCH}-token "
f"micro-batch, {math.ceil(min(lengths) / 16)} to {math.ceil(max(lengths) / 16)} chunks at 16\n")
print(f"{'comparison':52s} {'same top-1':>10s} {'mean TV':>8s} {'max |dlogp|':>12s}")
def compare(label, a, b):
x, y = dist[a], dist[b]
same = sum(max(p, key=p.get) == max(q, key=q.get) for p, q in zip(x, y))
tv = statistics.mean(tv_distance(p, q) for p, q in zip(x, y))
mx = max(abs(p[t] - q[t]) for p, q in zip(x, y) for t in set(p) & set(q))
print(f"{label:52s} {same:>7d}/50 {tv:8.4f} {mx:12.5f}")
return same, tv, mx
out = {}
out["modes_1"] = compare("1 loop: per_loop vs last_loop", (1, "per_loop", UBATCH), (1, "last_loop", UBATCH))
for n in (2, 4):
out[f"modes_{n}"] = compare(f"{n} loops: per_loop vs last_loop, one {UBATCH}-token chunk",
(n, "per_loop", UBATCH), (n, "last_loop", UBATCH))
out[f"chunks_{n}"] = compare(f"{n} loops: per_loop vs last_loop, 16-token chunks", (n, "per_loop", 16), (n, "last_loop", 16))
print()
out["1v2"] = compare("1 loop vs 2 loops (per_loop)", (1, "per_loop", UBATCH), (2, "per_loop", UBATCH))
out["2v4"] = compare("2 loops vs 4 loops (per_loop)", (2, "per_loop", UBATCH), (4, "per_loop", UBATCH))
out["1v4"] = compare("1 loop vs 4 loops (per_loop)", (1, "per_loop", UBATCH), (4, "per_loop", UBATCH))
claim("with the prompt in one chunk, last_loop gives the same next-token distribution as per_loop",
all(out[f"modes_{n}"][2] == 0 for n in (1, 2, 4)))
claim("split into small chunks, last_loop and per_loop disagree at 2 and 4 loops",
all(out[f"chunks_{n}"][1] > 0 for n in (2, 4)))
return (f"top-1 agreement 1 vs 4 loops {out['1v4'][0]}/50 (mean TV {out['1v4'][1]:.3f}); last_loop = per_loop in one chunk: "
f"{all(out[f'modes_{n}'][2] == 0 for n in (1, 2, 4))}")
mechanism_step()
次のトークン分布は、モデル内のループがほとんどコストをかけずにどのように変化するかを示しています。原始プロンプトとして使用された50件のGSM8K質問については、各設定下で上位20の次のトークン候補を求め、トップ1の一致率と総変動距離で比較しました。 フォークのコードは鋭い予測を示している。last_loopは各ループで再帰層のキャッシュを上書きするため、per_loopとは異なるのは、プロンプトの前の部分が最終ループのキーだけを保持している場合のみである。これは、プロンプト全体が一つのマイクロバッチに収まる場合には起こらない。16個のトークンのマイクロバッチで検証を行うことで、その点が正確に確認される。
コードをコピーEVAL_CONFIGS = [(1, "per_loop"), (2, "per_loop"), (2, "last_loop"), (4, "per_loop"), (4, "last_loop")]
N_QUESTIONS = 30
MAX_TOKENS = 2048
INSTRUCTION = "\n\nSolve the problem. End your reply with a last line of the form 'Final answer: <number>'."
def grade(text, gold):
m = re.findall(r"Final answer:\s*\$?\s*(-?[\d,]*\.?\d+)", text)
nums = m or re.findall(r"-?[\d,]*\.?\d+", text)
if not nums:
return False
try:
return abs(float(nums[-1].replace(",", "")) - gold) < 1e-6
except ValueError:
return False
@section("9. Paired GSM8K evaluation across loop counts and cache modes")
def eval_step():
global EVAL, SAMPLE
SAMPLE = GSM8K.sample(N_QUESTIONS, random_state=20261106).reset_index()
EVAL = {}
for loops, mode in EVAL_CONFIGS:
path = WORK / f"gsm8k-{loops}-{mode}.jsonl"
done = {json.loads(l)["i"]: json.loads(l) for l in open(path)} if path.exists() else {}
todo = [i for i in range(len(SAMPLE)) if i not in done]
if todo:
with DoryServer(loops, mode) as s, open(path, "a") as f:
for i in todo:
row = SAMPLE.iloc[i]
t0 = time.time()
r = s.post("/v1/chat/completions", {"messages": [{"role": "user", "content": row.question + INSTRUCTION}],
"max_tokens": MAX_TOKENS, "seed": 1000 + i, "cache_prompt": False})
reasoning, answer = split_reasoning(r["choices"][0]["message"])
rec = {"i": i, "correct": grade(answer, row.gold), "tokens": r["usage"]["completion_tokens"],
"finish": r["choices"][0]["finish_reason"], "wall_s": round(time.time() - t0, 2),
"decode_tps": round(r["timings"]["predicted_per_second"], 2), "answer": answer[-200:]}
f.write(json.dumps(rec) + "\n")
f.flush()
done[i] = rec
EVAL[(loops, mode)] = [done[i] for i in range(len(SAMPLE))]
rows = EVAL[(loops, mode)]
acc = sum(r["correct"] for r in rows)
print(f"{loops} loops {mode:9s}: {acc:2d}/{len(rows)} correct, median {statistics.median(r['tokens'] for r in rows):6.0f} "
f"tokens, {sum(r['finish'] == 'length' for r in rows)} cut off at {MAX_TOKENS}, "
f"median {statistics.median(r['wall_s'] for r in rows):5.1f} s per question")
return "; ".join(f"{n}/{m[:4]} {sum(r['correct'] for r in EVAL[(n, m)])}/{N_QUESTIONS}" for n, m in EVAL_CONFIGS)
eval_step()
ペア付きのGSM8K評価では、5つの設定下で同じセットされた質問が行われ、カードのサンプリング設定と各質問に対する固定シードが考慮される。そのため、各質問はすべてのループ数において同じ条件に適合する。各回答については、データセットの最終回答との正確性、生成されたトークン数（推論を含む）、2,048トークン制限に達したかどうか、およびウォールタイムが記録される。 結果は設定ごとに1つのJSONLファイルとして送信されるので、Colabセッションが中断されても停止した場所から再開できます。
コードをコピーdef mcnemar(a, b):
"""Exact two-sided McNemar p-value from paired booleans."""
only_a = sum(x and not y for x, y in zip(a, b))
only_b = sum(y and not x for x, y in zip(a, b))
n, k = only_a + only_b, min(only_a, only_b)
p = 1.0 if n == 0 else min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2**n)
return only_a, only_b, p
def wilson(k, n, z=1.96):
c = (k + z * z / 2) / (n + z * z)
h = z * math.sqrt(k * (n - k) / n + z * z / 4) / (n + z * z)
return 100 * (c - h), 100 * (c + h)
@section("10. Statistics: accuracy, tokens and time per solved question")
def stats_step():
global STATS
rng = random.Random(7)
ref = EVAL[(1, "per_loop")]
print(f"{'setting':18s} {'accuracy':>9s} {'95% CI':>13s} {'median tokens':>14s} {'tokens vs 1 loop [95% CI]':>27s} "
f"{'s per solved':>13s}")
STATS = {}
for key in EVAL_CONFIGS:
rows = EVAL[key]
k, n = sum(r["correct"] for r in rows), len(rows)
ratios = [r["tokens"] / q["tokens"] for r, q in zip(rows, ref)]
boots = sorted(statistics.median(rng.choices(ratios, k=len(ratios))) for _ in range(2000))
per_solved = sum(r["wall_s"] for r in rows) / max(k, 1)
STATS[key] = {"acc": k / n, "ratio": statistics.median(ratios), "ratio_ci": (boots[50], boots[1949]), "per_solved": per_solved,
"tokens": statistics.median(r["tokens"] for r in rows)}
lo, hi = wilson(k, n)
print(f"{key[0]} loops {key[1]:9s} {100 * k / n:8.1f}% {lo:5.1f}-{hi:5.1f}% {STATS[key]['tokens']:14.0f} "
f"{STATS[key]['ratio']:12.2f} [{boots[50]:.2f}, {boots[1949]:.2f}] {per_solved:12.1f}")
print("\npaired accuracy contrasts (McNemar exact; 'only A' = questions only the first setting solved)")
pairs = [((2, "per_loop"), (1, "per_loop")), ((4, "per_loop"), (1, "per_loop")), ((4, "per_loop"), (2, "per_loop")),
((2, "last_loop"), (2, "per_loop")), ((4, "last_loop"), (4, "per_loop"))]
for a, b in pairs:
oa, ob, p = mcnemar([r["correct"] for r in EVAL[a]], [r["correct"] for r in EVAL[b]])
print(f" {a[0]} {a[1]:9s} vs {b[0]} {b[1]:9s}: only A {oa:2d}, only B {ob:2d}, p = {p:.3f}")
STATS[(a, b)] = (oa, ob, p)
fewer = STATS[(4, "per_loop")]["ratio_ci"][1] < 1
claim("on GSM8K, 4 loops need fewer generated tokens than 1 loop (95% CI of the median ratio below 1)", fewer)
return (f"accuracy {', '.join(f'{k[0]}/{k[1][:4]} {100 * STATS[k]['acc']:.0f}%' for k in EVAL_CONFIGS)}; "
f"tokens at 4 loops {STATS[(4, 'per_loop')]['ratio']:.2f}x those at 1 loop")
stats_step()
30組のペア問題では、小さな精度の差を解決することはできないため、統計セルはサンプルが示すことも示せないことも明らかにするように構築されている。精度はウィルソン間隔と共に得られる。各対比では、1つの設定で解決できる問題に対して正確なマクネマー検定が用いられる。トークン数は、1回のループにおける各問題の比率の中央値と、ブートストラッピング間隔を使って比較される。これは、より多くのループで生成されるトークン数が少なくなるというカードの主張を検証するものである。 解決した問題あたりの時間は、スピード、長さ、正確性を合わせた数値であり、ユーザーが実際に支払う金額となります。
コードをコピー@section("11. The card's claims against this run")
def claims_step():
s, sp = STATS, SPEED
rows = [
("The GGUF runs 2 loops by default", f"built-in num_loops = {META['dory.num_loops']}", META["dory.num_loops"] == 2),
("Last-loop caching keeps the KV footprint at the one-loop size",
f"{sp[(1, 'per_loop')]['kv_mib']:.0f} MiB at 1 loop, {sp[(4, 'last_loop')]['kv_mib']:.0f} at 4 last_loop",
round(sp[(1, "per_loop")]["kv_mib"]) == round(sp[(4, "last_loop")]["kv_mib"])),
("Per-loop KV memory grows with the loop count",
" / ".join(f"{sp[(n, 'per_loop')]['kv_mib']:.0f}" for n in LOOPS) + " MiB",
sp[(4, "per_loop")]["kv_mib"] > sp[(1, "per_loop")]["kv_mib"]),
("Loop count trades compute per token without a model swap",
f"decode {sp[(1, 'per_loop')]['decode']:.1f} -> {sp[(4, 'per_loop')]['decode']:.1f} tok/s, same file",
sp[(4, "per_loop")]["decode"] < sp[(1, "per_loop")]["decode"]),
("More loops solve tasks with fewer generated tokens",
f"GSM8K median ratio at 4 loops {s[(4, 'per_loop')]['ratio']:.2f} [{s[(4, 'per_loop')]['ratio_ci'][0]:.2f}, "
f"{s[(4, 'per_loop')]['ratio_ci'][1]:.2f}]", s[(4, "per_loop")]["ratio_ci"][1] < 1),
("Last-loop accuracy stays close to per-loop",
f"4 loops: {100 * s[(4, 'last_loop')]['acc']:.0f}% vs {100 * s[(4, 'per_loop')]['acc']:.0f}%, "
f"McNemar p = {s[((4, 'last_loop'), (4, 'per_loop'))][2]:.2f}", s[((4, "last_loop"), (4, "per_loop"))][2] > 0.05),
]
print(f"{'claim (model card)':60s} {'this run':48s} verdict")
for text, found, ok in rows:
print(f"{text:60s} {found:48s} {'holds' if ok else 'does not hold'}")
claim(text, ok)
print(f"\nnot testable here: the benchmark table (SWE-bench, Terminal-Bench, GPQA ...), the 256K context, NVFP4 pre-training;"
f"\nwith {N_QUESTIONS} paired questions only large accuracy differences are resolvable")
return f"{sum(bool(ok) for *_, ok in rows)} of {len(rows)} card claims hold in this run"
claims_step()
モデルカードを再言することなく、この実行結果とその主張を照合する。各行は、カードの文とそれを検証する測定値を対応させる：キャッシュサイズに関する主張やループのデフォルトはトークン効率と最後のループのキャッシング精度であり、判断結果は手書きで入力されるのではなく、保存された結果から得られる。ノートブックが達成できない主張、例えばベンチマークテーブル、256Kのコンテキスト、NVFP4の事前学習などは、静かに除外されるのではなく、未検証として記載される。
コードをコピーbanner("SUMMARY")
for name, res in RESULTS.items():
print(f" {name:<100s} {res}")
print("\nchecks decided in this run:")
for step, text, held in CLAIMS:
print(f" [{'held' if held else 'did not hold'}] step {step}: {text}")
failed = [k for k, v in RESULTS.items() if str(v).startswith("SKIPPED")]
print(f"\n{sum(h for *_, h in CLAIMS)} of {len(CLAIMS)} checks held; steps failed or skipped: {failed or 'none'}")
print(f"llama.cpp fork {FORK_COMMIT[:12]}, GGUF @ {GGUF_REV[:7]}, backend {BACKEND}{' ' + GPU if GPU else ''}, "
f"Python {platform.python_version()}")
要約では、各ステップの一行目の結果と実行時に決定されたすべてのチェックが記載され、どれが成功し、どれが失敗したかが明記され、さらにフォークコミット、GGUFリビュウ、バックエンドも記載される。これにより、ノートブックは何が確認されたか、そしてどのハードウェアで行われたかを正確に記録する役割も果たす。
結論として、Doryはテスト時の計算をロード時の設定に変換する。3.7 GBのファイルが4つの異なるモデルで実行される。20ブロックの再帰型コアを通過する1回の処理から4回までである。GGUFヘッダーだけで、各設定がどれだけパラメータとキャッシュメモリを消費するかが予測される。ノートブックはこれらの予測をサーバーで確認し、速度を測定し、最後のループのキャッシュをコードからの正確な予測と比較して、精度と回答の長さをGSM8Kサンプルで評価する。 測定値が何であれ、2つの実用的な教訓があります。常にループ数を明示的に設定しなければならない。なぜならGGUFのデフォルトは4であるからです。また、forkコミットを固定する必要がある。なぜならllama.cppにおけるDoryサポートはまだオープンな草案段階にあるからです。
見てみてください 全コードはここに。また、私たちをフォローしても構いません Twitter そして、私たちの参加も忘れずに 150k+MB サブReddit そして、購読する 私たちのニュースレター。待って！Telegramを使っていますか？ 今なら、Telegramでも私たちと参加できます。
GitHubリポジトル、Hugging Faceページ、製品リリース、ウェビナーなどをプロモーションするために私たちと協力したいですか？私たちと連絡を取る
NVIDIA Dory 4Bをllama.cppで実行する方法：ループごとに手順のガイドという記事が、MarkTechPostに最初に掲載されました。