File size: 5,503 Bytes
df41178 | 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 | import os
"""実験23-C: **層ごとに違う設定**でモデル全体を量子化する。
quantize2.py（全層一律）との違いは1点だけ: 各層の k/cb/G/rank を
experiments/23_alloc/results/recipe.json から引く。平均ビット数は据え置きのまま、
痛い層に厚く・楽な層に薄く配る（レート歪み配分。allocate.py が解いた結果）。
環境変数 RECIPE でレシピを差し替えられる。レシピに無い層は従来の一律設定に戻す。
--- 以下は quantize2.py から引き継いだ堅牢化 ---
前版の失敗から直したこと:
1. チェックポイントを**層ごとの個別ファイル**に。全体の書き直しをしないので
処理時間が層数に依存しない（前版は10層ごとに8GBを書き直していた）。
2. 量子化した重みを**RAMに溜めない**。使い終わったら解放する
（前版はスワップを13GB使い切って停止した）。
3. float16で保存。容量半分。
4. I/Oエラーを**再試行**する（exFAT経由で一度失敗している）。
5. ログをディスクに残す。
6. 既存の層はスキップ。何度でも中断・再開できる。
"""
import json, os, sys, time
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.join(ROOT, "lib"))
import numpy as np
import gguf, ggufwrite, forward, recipe as RP
MODEL = (os.environ.get("MODEL_BLOB") or "/path/to/localai/ollama-models/blobs/"
"sha256-4a188102020e9c9530b687fd6400f775c45e90a0d7baafe65bd0a36963fbb7ba")
NAME = os.environ.get("NAME", "alloc")
CKPT = os.path.join(ROOT, "data", "ckpt", NAME)
LOG = os.path.join(ROOT, "data", "ckpt", NAME + ".log")
NTOK = int(os.environ.get("NTOK", "1024"))
RANK = int(os.environ.get("RANK", "32"))
K = int(os.environ.get("K", "8"))
CB = int(os.environ.get("CB", "8"))
GRP = int(os.environ.get("GRP", "256"))
os.makedirs(CKPT, exist_ok=True)
RECIPE = os.environ.get("RECIPE", os.path.join(
ROOT, "experiments/23_alloc/results/recipe.json"))
PLAN = json.load(open(RECIPE)) if os.path.exists(RECIPE) else {}
def say(msg):
line = f"[{time.strftime('%H:%M:%S')}] {msg}"
print(line, flush=True)
with open(LOG, "a") as f:
f.write(line + "\n")
def retry(fn, what, n=4):
"""外部SSDのI/Oは時々こける。数回粘ってから諦める。"""
for i in range(n):
try:
return fn()
except OSError as e:
say(f" I/O失敗({what}) {i+1}/{n}回目: {e}")
time.sleep(5 * (i + 1))
raise OSError(f"{what} が {n} 回失敗")
toks = json.load(open(os.path.join(ROOT, "data/calib/tokens_big.json")))[:NTOK]
m = forward.Model(MODEL, in_memory=True)
targets = {n for n, (d, t, o) in m.r.tensors.items()
if len(d) == 2 and n.startswith("blk.") and n.endswith(".weight")
and ("attn" in n or "ffn" in n)}
have = {f[:-4] for f in os.listdir(CKPT) if f.endswith(".npy")}
say(f"対象{len(targets)}層 / 済み{len(have)}層 / 残り{len(targets - have)}層 "
f"(較正{len(toks)}トークン)")
say(f"レシピ: {RECIPE} → {len(PLAN)}層ぶんの割り当てあり"
if PLAN else f"レシピ無し。一律 k={K} コード{2**CB} G={GRP} ランク{RANK}")
st = {"n": len(have), "t0": time.time(), "bits": 0.0, "w": 0}
def quantize(name, W, h):
p = os.path.join(CKPT, name + ".npy")
if os.path.exists(p):
return retry(lambda: np.load(p).astype(np.float32), f"読込 {name}")
pl = PLAN.get(name, {})
k = int(pl.get("k", K)); cb = int(pl.get("cb", CB))
G = int(pl.get("G", GRP)); rank = int(pl.get("rank", RANK))
X = np.ascontiguousarray(h.T.astype(np.float32))
inn = W.shape[1]
t0 = time.time()
try:
if X.shape[1] >= inn:
Wh, bpw = RP.fit(W, X, rank=rank, k=k, cb_bits=cb, G=G)
else:
# 標本不足で完全なヘッセ行列を使うと壊れる（序盤の失敗の教訓）。対角のみに落とす。
imp = np.sqrt((X ** 2).mean(1))
imp /= np.exp(np.log(imp + 1e-30).mean())
s = imp ** 0.75
Wh, bpw = RP.fit(W * s[None, :], np.eye(inn, dtype=np.float32),
rank=rank, k=k, cb_bits=cb, G=G)
Wh = Wh / s[None, :]
except Exception as e:
say(f" !! {name} 失敗({type(e).__name__}: {e}) → 元の重みのまま")
Wh, bpw = W.astype(np.float32), 16.0
retry(lambda: np.save(p, Wh.astype(np.float16)), f"保存 {name}")
st["n"] += 1; st["bits"] += bpw * W.size; st["w"] += W.size
say(f" [{st['n']:3d}/{len(targets)}] {name:<30} k={k:>2} cb={cb:>2} {bpw:.3f}bit "
f"{time.time()-t0:.0f}秒 累計{(time.time()-st['t0'])/60:.0f}分")
del X
return Wh.astype(np.float32)
forward.run(m, toks, quantizer=(targets, quantize))
if st["w"]:
say(f"今回量子化した分の平均 {st['bits']/st['w']:.4f} bit/重み")
say("全層完了。GGUFを組み立てる。")
# 組み立ては層を1枚ずつ読みながら行う（全部を同時に持たない）
class Lazy(dict):
def __contains__(self, k):
return os.path.exists(os.path.join(CKPT, k + ".npy"))
def __getitem__(self, k):
return np.load(os.path.join(CKPT, k + ".npy")).astype(np.float32)
OUT = os.path.join(ROOT, "data/models", NAME + ".gguf")
ggufwrite.Copier(MODEL).write(OUT, Lazy())
say(f"書き出し: {OUT} {os.path.getsize(OUT)/1e9:.2f}GB")
|