NVIDIA Dory 4B 模型 llama.cpp 逐循环运行指南发布

本教程使用 Dory 1.0 4B——这是 NVIDIA Nemotron Labs 推出的一个小型推理模型，其中间层构成一个共享块，对每个 token 可以应用一次、两次、三次或四次。每多跑一遍（即一个 loop），每个 token 就会消耗更多算力，却不增加任何参数，因此同一个 checkpoint 可以在四种不同的算力预算下运行，无需重新训练。我们构建 NVIDIA 模型卡所指定的 llama.cpp 分支，下载 3.7 GB 的 GGUF 文件，直接从文件中读取架构，以预测每个 loop 在参数和缓存内存上应付出的代价。随后，我们将这些预测与运行中的服务器进行对照，测量每种 loop 数量下的解码速度，测试更省成本的 last-loop 缓存是否会改变模型的预测结果，并运行一次配对 GSM8K 评估，观察准确率、答案长度和每道题耗时随 loop 增加而如何变化。

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]}")

（原文较长，此处译出前 6000 字，完整内容请点「打开原文」。）