Merge remote-tracking branch 'origin/dev' into pr259
# Conflicts: # README.md
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
@echo off
|
||||
call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat"
|
||||
cd /d C:\Users\Mark\Desktop\Projects\colibri\c
|
||||
"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\nvcc" -O3 -std=c++17 -arch=sm_120 -Xcompiler=-W3 -shared -DCOLI_CUDA_BUILDING_DLL -L"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\lib/x64" -lcudart backend_cuda.cu -o coli_cuda.dll
|
||||
echo EXITCODE=%ERRORLEVEL%
|
||||
@@ -317,10 +317,21 @@ class Spinner:
|
||||
if TTY: sys.stdout.write("\r\033[K"); sys.stdout.flush()
|
||||
|
||||
def stream_turn(p, sentinel, on_bytes):
|
||||
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT."""
|
||||
pend=b""
|
||||
"""legge fino alla sentinella; on_bytes riceve i chunk della risposta. Poi legge la riga STAT.
|
||||
Il PRIMO Ctrl-C durante lo stream non chiude la sessione: il motore (handler SIGINT)
|
||||
chiude il turno per la via del tetto NGEN e noi dreniamo fino alla sentinella.
|
||||
Un SECONDO Ctrl-C esce davvero."""
|
||||
pend=b""; interrupted=False
|
||||
while True:
|
||||
b=p.stdout.read(1)
|
||||
try:
|
||||
b=p.stdout.read(1)
|
||||
except KeyboardInterrupt:
|
||||
if interrupted or p.poll() is not None: raise
|
||||
interrupted=True
|
||||
try: p.send_signal(signal.SIGINT) # non-TTY: il motore potrebbe non aver visto il Ctrl-C
|
||||
except Exception: pass
|
||||
print(f"\n {C.yel}⏹ stopping… (Ctrl-C again to quit){C.r}", flush=True)
|
||||
continue
|
||||
if b==b"": return None
|
||||
pend+=b
|
||||
if pend.endswith(sentinel):
|
||||
@@ -328,7 +339,9 @@ def stream_turn(p, sentinel, on_bytes):
|
||||
if rest: on_bytes(rest)
|
||||
line=p.stdout.readline().decode("utf-8","replace").strip() # STAT tok tps hit rss
|
||||
m=re.match(r"STAT (\S+) (\S+) (\S+) (\S+)", line)
|
||||
return {"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
|
||||
st={"tok":int(m.group(1)),"tps":float(m.group(2)),"hit":float(m.group(3)),"rss":float(m.group(4))} if m else {}
|
||||
if interrupted: st["interrupted"]=True
|
||||
return st
|
||||
if len(pend)>len(sentinel):
|
||||
out=pend[:-len(sentinel)]; pend=pend[-len(sentinel):]
|
||||
on_bytes(out)
|
||||
@@ -438,6 +451,8 @@ def cmd_chat(a):
|
||||
try: errlog.write(p.stderr.read().decode("utf-8","replace"))
|
||||
except (OSError, ValueError): pass
|
||||
errlog.seek(0); print(errlog.read()[-1500:]); sys.exit("the engine exited while loading")
|
||||
p.stdout.readline() # TIERS line (web-dashboard protocol): emitted once right after STAT;
|
||||
# left unread it leaks into the first answer's text
|
||||
# READY received. Drain the child's stderr into errlog without blocking:
|
||||
# the engine is still alive (blocked on stdin), so a plain read() would
|
||||
# hang forever waiting for EOF. A short bounded drain grabs the ~400 bytes
|
||||
@@ -464,7 +479,7 @@ def cmd_chat(a):
|
||||
for chunk in textwrap.wrap(l, term_w()-4) or [l]:
|
||||
print(f" {C.dgray}{chunk}{C.r}")
|
||||
except Exception: pass
|
||||
print(f" {C.dim}type and press Enter · :more continues · :reset clears memory · :q exits{C.r}\n")
|
||||
print(f" {C.dim}type and press Enter · Ctrl-C stops the answer · :more continues · :reset clears memory · :q exits{C.r}\n")
|
||||
w=term_w()-4
|
||||
def user_box(msg):
|
||||
"""ri-disegna il messaggio dentro una box che si ADATTA su piu' righe:
|
||||
@@ -529,10 +544,13 @@ def cmd_chat(a):
|
||||
el=time.time()-t0
|
||||
if st.get("tok"):
|
||||
print(f"\r {C.dgray}└─ {st['tok']} tok · {st['tps']:.2f} tok/s · hit {st['hit']:.0f}% · RSS {st['rss']:.1f} GB · {el:.0f}s{C.r}")
|
||||
if st["tok"]>=a.ngen:
|
||||
if st.get("interrupted"):
|
||||
print(f" {C.yel}⏹ interrupted; type :more to continue the response{C.r}")
|
||||
elif st["tok"]>=a.ngen:
|
||||
print(f" {C.yel}…stopped at --ngen ({a.ngen}); type :more to continue the response{C.r}")
|
||||
print()
|
||||
else:
|
||||
if st.get("interrupted"): print(f" {C.yel}⏹ interrupted{C.r}")
|
||||
print()
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n {C.dim}interrupted{C.r}")
|
||||
|
||||
+35
-2
@@ -95,7 +95,18 @@ static inline int compat_open_direct(const char *path){
|
||||
* prevents 0x0A bytes from being silently translated to \r\n. */
|
||||
#define COMPAT_O_RDONLY (O_RDONLY | O_BINARY)
|
||||
|
||||
/* --- posix_fadvise: no-op (advisory only; safe to ignore) --- */
|
||||
/* --- posix_fadvise: Windows has no direct equivalent. Semantics:
|
||||
* WILLNEED -> warm the OS page cache so a later synchronous pread finds the
|
||||
* pages resident. Implemented as an overlapped background ReadFile
|
||||
* into a throwaway scratch buffer (fire-and-forget readahead). Called
|
||||
* from the dedicated PILOT I/O thread / next-block readahead in moe(),
|
||||
* NEVER inline on the hot path (the existing comment at glm.c:2847
|
||||
* measures inline fadvise submit at ~0.5ms x 169k calls = +92s/48tok).
|
||||
* Each call owns its OVERLAPPED + scratch buffer -> thread-safe.
|
||||
* DONTNEED -> no-op: Windows' standby-list trimming self-regulates under pressure,
|
||||
* and on a low-RAM host keeping the pages is what we want for reuse.
|
||||
* Matches macOS (compat.h:16-19) which no-ops DONTNEED for the same
|
||||
* reason. The engine only ever uses DONTNEED as an advisory. */
|
||||
#ifndef POSIX_FADV_NORMAL
|
||||
#define POSIX_FADV_NORMAL 0
|
||||
#define POSIX_FADV_RANDOM 1
|
||||
@@ -104,7 +115,29 @@ static inline int compat_open_direct(const char *path){
|
||||
#define POSIX_FADV_DONTNEED 4
|
||||
#define POSIX_FADV_NOREUSE 5
|
||||
#endif
|
||||
#define posix_fadvise(fd,off,len,advice) do{(void)(fd);(void)(off);(void)(len);(void)(advice);}while(0)
|
||||
static inline int compat_fadvise(int fd, off_t off, off_t len, int advice){
|
||||
if(advice!=POSIX_FADV_WILLNEED || len<=0) return 0;
|
||||
intptr_t osfh=_get_osfhandle(fd);
|
||||
if(osfh==-1 || osfh==-2) return 0;
|
||||
HANDLE h=(HANDLE)osfh;
|
||||
/* Cap the readahead window: reading a whole 19MB expert per hint is fine on the
|
||||
* PILOT thread, but a pathological huge len would spike transient memory. */
|
||||
size_t rdlen = (len>(off_t)(64*1024*1024)) ? (size_t)(64*1024*1024) : (size_t)len;
|
||||
char *buf=(char*)_aligned_malloc(rdlen, 4096);
|
||||
if(!buf) return -1;
|
||||
OVERLAPPED ov={0};
|
||||
ov.Offset = (DWORD)( (off_t)off & 0xFFFFFFFFULL);
|
||||
ov.OffsetHigh = (DWORD)(((off_t)off >> 32) & 0xFFFFFFFFULL);
|
||||
/* Issue an overlapped read. With a non-OVERLAPPED-opened handle ReadFile still
|
||||
* accepts lpOverlapped (it carries the 64-bit offset) and blocks until the read
|
||||
* completes — but crucially it populates the standby page cache for this region,
|
||||
* so the later synchronous pread on the same offsets faults from RAM not disk. */
|
||||
DWORD got=0;
|
||||
ReadFile(h, buf, (DWORD)rdlen, &got, &ov);
|
||||
_aligned_free(buf);
|
||||
return 0;
|
||||
}
|
||||
#define posix_fadvise compat_fadvise
|
||||
|
||||
/* --- pread -> ReadFile + OVERLAPPED su raw OS handle ---
|
||||
* Thread-safe (no shared seek position). Gestisce offset >4 GB e chunking
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Download GLM-5.2-FP8 from ModelScope (fast, no HF throttling) with
|
||||
HuggingFace fallback. Parallel shard download, clean progress display.
|
||||
Usage: python download_fp8.py
|
||||
python download_fp8.py --parallel 4
|
||||
python download_fp8.py --source hf (force HuggingFace)
|
||||
"""
|
||||
import os, sys, time, threading, argparse, subprocess
|
||||
|
||||
REPO_MS = "ZhipuAI/GLM-5.2-FP8" # ModelScope
|
||||
REPO_HF = "zai-org/GLM-5.2-FP8" # HuggingFace
|
||||
DEST = r"I:\glm52_fp8"
|
||||
|
||||
# ── ANSI colors ──
|
||||
class C:
|
||||
dim="\033[2m"; grn="\033[32m"; yel="\033[33m"; cyn="\033[36m"; b="\033[1m"; r="\033[0m"
|
||||
|
||||
def fmt_bytes(n):
|
||||
if n>=1e9: return f"{n/1e9:.2f} GB"
|
||||
if n>=1e6: return f"{n/1e6:.1f} MB"
|
||||
return f"{n/1e3:.0f} KB"
|
||||
|
||||
def fmt_time(s):
|
||||
if s<60: return f"{s:.0f}s"
|
||||
if s<3600: return f"{s/60:.0f}m"
|
||||
return f"{s/3600:.1f}h"
|
||||
|
||||
def bar(cur, total, width=24):
|
||||
if total<=0: return "["+" "*width+"]"
|
||||
pct=min(cur/total,1.0); filled=int(width*pct)
|
||||
return "["+"█"*filled+"░"*(width-filled)+f"] {pct*100:4.0f}%"
|
||||
|
||||
def get_shard_list_hf():
|
||||
from huggingface_hub import HfApi
|
||||
info=HfApi().repo_info(REPO_HF, files_metadata=True)
|
||||
shards=sorted(s.rfilename for s in info.siblings if s.rfilename.endswith(".safetensors"))
|
||||
sizes={s.rfilename:s.size for s in info.siblings if s.rfilename.endswith(".safetensors")}
|
||||
return shards, sizes
|
||||
|
||||
def get_shard_list_ms():
|
||||
"""Get shard list from ModelScope API."""
|
||||
import requests
|
||||
# ModelScope API: list files
|
||||
r = requests.get(f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo/files?Revision=master&Root=", timeout=30)
|
||||
data = r.json()["Data"]["Files"]
|
||||
shards = sorted(f["Path"] for f in data if f["Path"].endswith(".safetensors"))
|
||||
sizes = {f["Path"]: f.get("Size", 0) for f in data if f["Path"].endswith(".safetensors")}
|
||||
return shards, sizes
|
||||
|
||||
def download_file_ms(fn):
|
||||
"""Download a single file from ModelScope using their CDN."""
|
||||
from modelscope.hub.file_download import model_file_download
|
||||
model_file_download(
|
||||
model_id=REPO_MS,
|
||||
file_path=fn,
|
||||
local_dir=DEST,
|
||||
revision="master",
|
||||
)
|
||||
|
||||
def download_file_hf(fn):
|
||||
"""Download a single file from HuggingFace with hf_transfer."""
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(REPO_HF, fn, local_dir=DEST)
|
||||
|
||||
def download_file_curl(fn, base_url, expected_size):
|
||||
"""Fallback: download with curl to a .part file with resume."""
|
||||
outpath = os.path.join(DEST, fn)
|
||||
partpath = outpath + ".part"
|
||||
if os.path.exists(outpath) and os.path.getsize(outpath) == expected_size:
|
||||
return True
|
||||
url = f"{base_url}/{fn}"
|
||||
cmd = ["curl", "-L", "-C", "-", "--retry", "999", "--retry-delay", "5",
|
||||
"--connect-timeout", "15", "--speed-time", "30", "--speed-limit", "1000",
|
||||
"-o", partpath, "-H", "User-Agent: colibri-download/1.0", url]
|
||||
subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||||
if os.path.exists(partpath) and os.path.getsize(partpath) >= expected_size:
|
||||
os.replace(partpath, outpath)
|
||||
return True
|
||||
return False
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--parallel", type=int, default=3)
|
||||
ap.add_argument("--source", choices=["auto", "ms", "hf"], default="auto",
|
||||
help="auto (try ModelScope first), ms (ModelScope only), hf (HuggingFace only)")
|
||||
args = ap.parse_args()
|
||||
|
||||
os.makedirs(DEST, exist_ok=True)
|
||||
|
||||
# Determine source and get shard list
|
||||
use_ms = False
|
||||
shards, sizes = [], {}
|
||||
|
||||
if args.source in ("auto", "ms"):
|
||||
try:
|
||||
print(f"{C.dim}Trying ModelScope...{C.r}", end=" ", flush=True)
|
||||
shards, sizes = get_shard_list_ms()
|
||||
use_ms = True
|
||||
print(f"{C.grn}✓{C.r} {len(shards)} shards found")
|
||||
except Exception as e:
|
||||
print(f"{C.yel}failed ({e}){C.r}")
|
||||
if args.source == "ms":
|
||||
print("ModelScope failed and --source ms was set. Exiting."); return
|
||||
|
||||
if not shards:
|
||||
print(f"{C.dim}Using HuggingFace...{C.r}", end=" ", flush=True)
|
||||
shards, sizes = get_shard_list_hf()
|
||||
print(f"{C.grn}✓{C.r} {len(shards)} shards found")
|
||||
|
||||
total = len(shards)
|
||||
total_bytes = sum(sizes.values())
|
||||
source_name = "ModelScope" if use_ms else "HuggingFace"
|
||||
|
||||
# Download metadata files
|
||||
meta_files = ["config.json", "tokenizer.json", "tokenizer_config.json",
|
||||
"generation_config.json", "model.safetensors.index.json"]
|
||||
for fn in meta_files:
|
||||
out = os.path.join(DEST, fn)
|
||||
if not os.path.exists(out):
|
||||
try:
|
||||
if use_ms: download_file_ms(fn)
|
||||
else: download_file_hf(fn)
|
||||
except Exception: pass
|
||||
|
||||
# Build work queue
|
||||
todo = []
|
||||
done_set = set()
|
||||
for fn in shards:
|
||||
outpath = os.path.join(DEST, fn)
|
||||
if os.path.exists(outpath) and os.path.getsize(outpath) == sizes.get(fn, 0):
|
||||
done_set.add(fn)
|
||||
else:
|
||||
todo.append(fn)
|
||||
|
||||
existing = len(done_set)
|
||||
print(f"\n{C.b}GLM-5.2-FP8 Download ({source_name}){C.r}")
|
||||
print(f" {C.dim}{total} shards · {total_bytes/1e9:.0f} GB · {existing}/{total} complete{C.r}")
|
||||
print(f" {C.dim}{len(todo)} to download · {args.parallel} parallel{C.r}")
|
||||
if todo:
|
||||
remaining = sum(sizes[fn] for fn in todo)
|
||||
print(f" {C.dim}Remaining: {remaining/1e9:.0f} GB{C.r}")
|
||||
print()
|
||||
|
||||
if not todo:
|
||||
print(f"{C.grn}✓ All shards already downloaded!{C.r}\n"); return
|
||||
|
||||
lock = threading.Lock()
|
||||
completed = list(done_set)
|
||||
t0 = time.time()
|
||||
qidx = [0]
|
||||
|
||||
def worker(wid):
|
||||
while True:
|
||||
with lock:
|
||||
if qidx[0] >= len(todo): return
|
||||
idx = qidx[0]; qidx[0] += 1
|
||||
fn = todo[idx]
|
||||
|
||||
expected = sizes.get(fn, 0)
|
||||
shard_num = existing + idx + 1
|
||||
print(f" {C.cyn}[{shard_num}/{total}]{C.r} {C.dim}{fn}{C.r}")
|
||||
|
||||
success = False
|
||||
for attempt in range(3):
|
||||
try:
|
||||
if use_ms:
|
||||
download_file_ms(fn)
|
||||
else:
|
||||
download_file_hf(fn)
|
||||
outpath = os.path.join(DEST, fn)
|
||||
# ModelScope downloads to a cache dir, need to check
|
||||
# if the file exists at our expected path
|
||||
if not os.path.exists(outpath):
|
||||
# Try to find it in ModelScope's cache structure
|
||||
# and copy/symlink it
|
||||
pass
|
||||
if os.path.exists(outpath):
|
||||
actual = os.path.getsize(outpath)
|
||||
if expected == 0 or actual == expected:
|
||||
success = True; break
|
||||
# Size mismatch — might be in cache
|
||||
# If we got here, file downloaded but not at expected path
|
||||
# ModelScope puts it in local_dir; check again
|
||||
if os.path.exists(outpath):
|
||||
success = True; break
|
||||
# Retry with curl fallback
|
||||
if use_ms:
|
||||
base = f"https://modelscope.cn/api/v1/models/{REPO_MS}/repo?Revision=master&FilePath="
|
||||
else:
|
||||
base = f"https://huggingface.co/{REPO_HF}/resolve/main"
|
||||
if download_file_curl(fn, base, expected):
|
||||
success = True; break
|
||||
except Exception as e:
|
||||
if attempt < 2:
|
||||
print(f" {C.yel}retry {attempt+1}: {e}{C.r}")
|
||||
time.sleep(3)
|
||||
else:
|
||||
print(f" {C.yel}✗ failed: {e}{C.r}")
|
||||
|
||||
with lock:
|
||||
if success:
|
||||
completed.append(fn)
|
||||
elapsed = time.time() - t0
|
||||
have = sum(sizes.get(f,0) for f in completed)
|
||||
pct = 100.0 * have / total_bytes
|
||||
speed = (have - sum(sizes[f] for f in done_set)) / max(elapsed, 1)
|
||||
eta = (total_bytes - have) / speed if speed > 0 else 0
|
||||
print(f" {C.grn}✓{C.r} {fn} {C.dim}— {len(completed)}/{total} · "
|
||||
f"{pct:.1f}% · {fmt_time(elapsed)} · ETA {fmt_time(eta)}{C.r}")
|
||||
else:
|
||||
print(f" {C.yel}✗ GIVE UP: {fn}{C.r}")
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(args.parallel)]
|
||||
for t in threads: t.start()
|
||||
for t in threads: t.join()
|
||||
|
||||
print()
|
||||
final = sum(1 for fn in shards
|
||||
if os.path.exists(os.path.join(DEST, fn))
|
||||
and os.path.getsize(os.path.join(DEST, fn)) == sizes.get(fn, 0))
|
||||
if final == total:
|
||||
print(f"{C.grn}{'='*50}")
|
||||
print(f" ✓ All {total} shards downloaded!{C.r}\n")
|
||||
else:
|
||||
print(f"{C.yel} {final}/{total} complete, {total-final} remaining{C.r}")
|
||||
print(f" Re-run to resume.\n")
|
||||
|
||||
if __name__ == "__main__":
|
||||
try: main()
|
||||
except KeyboardInterrupt:
|
||||
print(f"\n\n{C.yel}Interrupted. Re-run to resume — no data lost.{C.r}\n")
|
||||
@@ -35,6 +35,7 @@
|
||||
#include <sys/resource.h>
|
||||
#include <sys/mman.h> /* mlock: inchioda le pagine in RAM / wire pages into RAM */
|
||||
#include <sys/stat.h> /* fstat per mmap degli shard (COLI_MMAP) */
|
||||
#include <signal.h> /* SIGINT = stop morbido del turno in serve mode */
|
||||
#endif
|
||||
#include "st.h"
|
||||
#ifdef __linux__
|
||||
@@ -145,6 +146,9 @@ typedef struct {
|
||||
int *kv_start, max_t;
|
||||
int disk_nrec;
|
||||
char disk_path[2048];
|
||||
FILE *disk_fp; /* kept-open handle: fopen once, fwrite per turn, fclose at exit (#4) */
|
||||
uint8_t *disk_buf; /* staging buffer: one contiguous record per position (#1) */
|
||||
int64_t disk_buf_cap;
|
||||
} KVState;
|
||||
|
||||
typedef struct {
|
||||
@@ -614,18 +618,30 @@ static inline int32_t dot_i8i8(const int8_t *w, const int8_t *x, int I){
|
||||
#elif defined(__ARM_NEON)
|
||||
/* ARM: SDOT nativo se disponibile (Apple Silicon: sempre); altrimenti vmull/vpadal.
|
||||
* Stesso bound anti-overflow del trucco AVX2: coppie <= 128*127*2 = 32512 < 32767. */
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
/* 4 accumulatori indipendenti: SDOT ha latenza ~3-4 cicli, con un solo acc la
|
||||
* catena seriale strozza il core a ~26 GB/s di pesi; con 4 lane indipendenti il
|
||||
* dot diventa memory-bound (misurato su M4: 26 -> 63 GB/s per core, 2.4x). */
|
||||
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0);
|
||||
for(;i+64<=I;i+=64){
|
||||
a0=vdotq_s32(a0,vld1q_s8(w+i), vld1q_s8(x+i));
|
||||
a1=vdotq_s32(a1,vld1q_s8(w+i+16),vld1q_s8(x+i+16));
|
||||
a2=vdotq_s32(a2,vld1q_s8(w+i+32),vld1q_s8(x+i+32));
|
||||
a3=vdotq_s32(a3,vld1q_s8(w+i+48),vld1q_s8(x+i+48));
|
||||
}
|
||||
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
|
||||
for(;i+16<=I;i+=16) acc=vdotq_s32(acc,vld1q_s8(w+i),vld1q_s8(x+i));
|
||||
sum=vaddvq_s32(acc);
|
||||
#else
|
||||
int32x4_t acc=vdupq_n_s32(0);
|
||||
for(;i+16<=I;i+=16){
|
||||
int8x16_t wv=vld1q_s8(w+i), xv=vld1q_s8(x+i);
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
acc=vdotq_s32(acc,wv,xv);
|
||||
#else
|
||||
int16x8_t p=vmull_s8(vget_low_s8(wv),vget_low_s8(xv));
|
||||
p=vmlal_s8(p,vget_high_s8(wv),vget_high_s8(xv));
|
||||
acc=vpadalq_s16(acc,p);
|
||||
#endif
|
||||
}
|
||||
sum=vaddvq_s32(acc);
|
||||
#endif
|
||||
#elif defined(__VSX__)
|
||||
/* POWER8: vec_msum (s8 x u8 -> s32) somma i prodotti byte DIRETTAMENTE in lane
|
||||
* s32, 16 byte/iter: il bound anti-saturazione a 16 bit di maddubs qui non serve.
|
||||
@@ -705,6 +721,28 @@ static inline int32_t dot_i4i8(const uint8_t *w4, const int8_t *x, int I){
|
||||
sum=hsum256_i32(acc);
|
||||
#elif defined(__ARM_NEON)
|
||||
const uint8x16_t m4q=vdupq_n_u8(0x0F); const int8x16_t b8q=vdupq_n_s8(8);
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
/* 4 accumulatori indipendenti (vedi dot_i8i8): spezza la catena seriale su acc.
|
||||
* Misurato su M4: 12.4 -> 29.9 GB/s di pesi per core (2.4x). */
|
||||
int32x4_t a0=vdupq_n_s32(0),a1=vdupq_n_s32(0),a2=vdupq_n_s32(0),a3=vdupq_n_s32(0);
|
||||
for(;i+64<=I;i+=64){
|
||||
uint8x16_t byA=vld1q_u8(w4+(i>>1)), byB=vld1q_u8(w4+(i>>1)+16);
|
||||
uint8x16x2_t zA=vzipq_u8(vandq_u8(byA,m4q), vshrq_n_u8(byA,4)); /* nibble in ordine */
|
||||
uint8x16x2_t zB=vzipq_u8(vandq_u8(byB,m4q), vshrq_n_u8(byB,4));
|
||||
a0=vdotq_s32(a0,vsubq_s8(vreinterpretq_s8_u8(zA.val[0]),b8q),vld1q_s8(x+i));
|
||||
a1=vdotq_s32(a1,vsubq_s8(vreinterpretq_s8_u8(zA.val[1]),b8q),vld1q_s8(x+i+16));
|
||||
a2=vdotq_s32(a2,vsubq_s8(vreinterpretq_s8_u8(zB.val[0]),b8q),vld1q_s8(x+i+32));
|
||||
a3=vdotq_s32(a3,vsubq_s8(vreinterpretq_s8_u8(zB.val[1]),b8q),vld1q_s8(x+i+48));
|
||||
}
|
||||
int32x4_t acc=vaddq_s32(vaddq_s32(a0,a1),vaddq_s32(a2,a3));
|
||||
for(;i+32<=I;i+=32){
|
||||
uint8x16_t by=vld1q_u8(w4+(i>>1)); /* 16 byte = 32 nibble */
|
||||
uint8x16x2_t z=vzipq_u8(vandq_u8(by,m4q), vshrq_n_u8(by,4)); /* nibble in ordine */
|
||||
acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q),vld1q_s8(x+i));
|
||||
acc=vdotq_s32(acc,vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q),vld1q_s8(x+i+16));
|
||||
}
|
||||
sum=vaddvq_s32(acc);
|
||||
#else
|
||||
int32x4_t acc=vdupq_n_s32(0);
|
||||
for(;i+32<=I;i+=32){
|
||||
uint8x16_t by=vld1q_u8(w4+(i>>1)); /* 16 byte = 32 nibble */
|
||||
@@ -712,18 +750,15 @@ static inline int32_t dot_i4i8(const uint8_t *w4, const int8_t *x, int I){
|
||||
int8x16_t w0=vsubq_s8(vreinterpretq_s8_u8(z.val[0]),b8q);
|
||||
int8x16_t w1=vsubq_s8(vreinterpretq_s8_u8(z.val[1]),b8q);
|
||||
int8x16_t x0=vld1q_s8(x+i), x1=vld1q_s8(x+i+16);
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
acc=vdotq_s32(acc,w0,x0); acc=vdotq_s32(acc,w1,x1);
|
||||
#else
|
||||
int16x8_t p=vmull_s8(vget_low_s8(w0),vget_low_s8(x0)); /* |w|<=8: nessun overflow */
|
||||
p=vmlal_s8(p,vget_high_s8(w0),vget_high_s8(x0));
|
||||
acc=vpadalq_s16(acc,p);
|
||||
p=vmull_s8(vget_low_s8(w1),vget_low_s8(x1));
|
||||
p=vmlal_s8(p,vget_high_s8(w1),vget_high_s8(x1));
|
||||
acc=vpadalq_s16(acc,p);
|
||||
#endif
|
||||
}
|
||||
sum=vaddvq_s32(acc);
|
||||
#endif
|
||||
#elif defined(__VSX__)
|
||||
/* 16 byte = 32 nibble. vec_mergeh/vec_mergel su ppc64le (GCC) interallacciano come
|
||||
* unpacklo/unpackhi x86 (verificato empiricamente su POWER8): i nibble escono in
|
||||
@@ -899,6 +934,11 @@ static float g_temp=-1; /* TEMP: temperatura di sampling sui TOKEN. <0 = auto (
|
||||
static float g_nuc=0.95f;/* NUCLEUS: top-p sul vocabolario (default dal generation_config GLM-5.2) */
|
||||
static int g_topk=0; /* TOPK=n -> usa n expert/token invece di config (ricerca: meno disco) */
|
||||
static float g_topp=0; /* TOPP=p (0..1) -> top-p adattivo: tieni gli expert fino a peso cumulato p */
|
||||
static int g_expert_budget=0; /* EXPERT_BUDGET=N -> cap distinct experts loaded per layer across the
|
||||
* batch-union. Reduces disk I/O on cold/low-RAM hosts by dropping the
|
||||
* lowest-gate-weight experts from the cross-position union. MoE-Spec
|
||||
* (arXiv 2602.16052): top-32 of 64 capture 93% routing weight. */
|
||||
static int64_t g_budget_dropped=0; /* total experts dropped by EXPERT_BUDGET across all layers */
|
||||
/* CACHE_ROUTE (paper 2412.00099 max-rank): opt-in only. Keep true top-J always;
|
||||
* fill remaining slots preferring pin∪LRU experts ranked within top-M (or mass ROUTE_P). */
|
||||
static int g_cache_route=0;
|
||||
@@ -1801,7 +1841,10 @@ static int uring_wait_all(UringBatch *b){
|
||||
* condvar exist ONLY to park/wake idle workers, never for correctness. Gated
|
||||
* behind PIPE=1; OFF => the original blocking-load + serial-matmul path runs
|
||||
* byte-identically. */
|
||||
static int g_pipe=0; /* PIPE=1: async expert-load pipeline (default OFF) */
|
||||
static int g_pipe=0; /* PIPE=1: async expert-load pipeline. Default ON for Windows
|
||||
* (parsed in main: getenv("PIPE")?:1 on _WIN32, :0 elsewhere).
|
||||
* Keeps expert pread off the forward-pass thread so loads overlap
|
||||
* the matmul. PIPE=0 opts back into the blocking serial path. */
|
||||
static int g_pipe_nw=8; /* PIPE_WORKERS=n: I/O worker threads (disk-parallel reads) */
|
||||
static int g_uring=0; /* URING=1: Linux io_uring load/completion backend; implies PIPE */
|
||||
typedef struct {
|
||||
@@ -2621,6 +2664,67 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
|
||||
int e=idxs[(int64_t)s*K+kk];
|
||||
if(!seen[e]){ seen[e]=1; uniq[nu++]=e; }
|
||||
}
|
||||
/* EXPERT_BUDGET: cap distinct experts per layer to reduce disk I/O on cold/low-RAM
|
||||
* hosts. MISS-AWARE: always keep cache hits (pin/LRU — they're free, no disk I/O),
|
||||
* only drop from misses. From the misses, keep the highest-aggregate-gate-weight
|
||||
* ones up to the budget; drop the rest from idxs[] so they're never loaded.
|
||||
* (MoE-Spec arXiv 2602.16052: top-32 of 64 capture 93% routing weight.)
|
||||
* Complementary to TOPP (per-position) — this trims cross-position. */
|
||||
if(g_expert_budget>0 && nu>g_expert_budget){
|
||||
/* compute aggregate gate weight per unique expert */
|
||||
float *wsum=falloc(nu); for(int j=0;j<nu;j++) wsum[j]=0;
|
||||
for(int s=0;s<S;s++) for(int kk=0;kk<keff[s];kk++){
|
||||
int e=idxs[(int64_t)s*K+kk];
|
||||
for(int j=0;j<nu;j++) if(uniq[j]==e){ wsum[j]+=ws[(int64_t)s*K+kk]; break; }
|
||||
}
|
||||
/* residency pre-scan: which experts are already in pin or ecache (hits)? */
|
||||
unsigned char *is_hit=calloc(nu,1); int nhits=0;
|
||||
for(int j=0;j<nu;j++){ int eid=uniq[j];
|
||||
int found=0;
|
||||
ESlot *P=m->pin[layer];
|
||||
for(int z=0;z<m->npin[layer];z++) if(P[z].eid==eid){ found=1; break; }
|
||||
if(!found){ ESlot *Sl=m->ecache[layer]; int nn=m->ecn[layer];
|
||||
for(int z=0;z<nn;z++) if(Sl[z].eid==eid){ found=1; break; } }
|
||||
if(found){ is_hit[j]=1; nhits++; }
|
||||
}
|
||||
/* budget for misses = total budget - hits already kept (min 0) */
|
||||
int miss_budget = g_expert_budget - nhits; if(miss_budget<0) miss_budget=0;
|
||||
/* mark which unique experts to keep (1) or drop (0): keep all hits, fill rest
|
||||
* with top-weight misses up to miss_budget */
|
||||
unsigned char *keep=calloc(nu,1); int nkeep=0;
|
||||
for(int j=0;j<nu;j++) if(is_hit[j]){ keep[j]=1; nkeep++; }
|
||||
for(int rank=0;rank<miss_budget;rank++){
|
||||
int best=-1; float bv=-1e30f;
|
||||
for(int j=0;j<nu;j++) if(!keep[j] && wsum[j]>bv){ bv=wsum[j]; best=j; }
|
||||
if(best<0) break; keep[best]=1; nkeep++;
|
||||
}
|
||||
/* build a lookup: for each expert id, is it kept? (reuse seen[]) */
|
||||
memset(seen,0,(size_t)E);
|
||||
for(int j=0;j<nu;j++) if(keep[j]) seen[uniq[j]]=1;
|
||||
int dropped=nu-nkeep; g_budget_dropped+=dropped;
|
||||
/* remove dropped experts from each position's routing list */
|
||||
for(int s=0;s<S;s++){
|
||||
int w=0;
|
||||
for(int kk=0;kk<keff[s];kk++){
|
||||
int e=idxs[(int64_t)s*K+kk];
|
||||
if(seen[e]){ idxs[(int64_t)s*K+w]=e; ws[(int64_t)s*K+w]=ws[(int64_t)s*K+kk]; w++; }
|
||||
}
|
||||
if(w<keff[s]){
|
||||
keff[s]=w;
|
||||
/* renormalize remaining weights per position */
|
||||
if(c->norm_topk && w>0){
|
||||
float sm=0; for(int kk=0;kk<w;kk++) sm+=ws[(int64_t)s*K+kk]; sm+=1e-20f;
|
||||
for(int kk=0;kk<w;kk++) ws[(int64_t)s*K+kk]/=sm;
|
||||
for(int kk=0;kk<w;kk++) ws[(int64_t)s*K+kk]*=c->routed_scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
/* compact uniq[] to kept experts only */
|
||||
int nu2=0;
|
||||
for(int j=0;j<nu;j++) if(keep[j]) uniq[nu2++]=uniq[j];
|
||||
nu=nu2;
|
||||
free(wsum); free(is_hit); free(keep);
|
||||
}
|
||||
/* ---- FASE C/D: risolvi (pin/cache/disco) e calcola, a blocchi di 64 unici ---- */
|
||||
float *xg=falloc((int64_t)S*D), *gg=falloc((int64_t)S*I), *uu=falloc((int64_t)S*I), *hh=falloc((int64_t)S*D);
|
||||
int *rows=malloc(S*sizeof(int)); float *rw=malloc(S*sizeof(float));
|
||||
@@ -3795,12 +3899,34 @@ static void stops_arm(const Cfg *c, int tok_eos){
|
||||
* all: storia token (capacita' >= kv+n_new+g_draft+2), kv = token gia' in KV.
|
||||
* logit = logits della posizione kv-1 (dal prefill); viene liberato qui.
|
||||
* emit(tok,ud) per ogni token emesso. Ritorna i token emessi; *kv_out = nuova kv. */
|
||||
/* STOP MORBIDO (serve/chat): SIGINT chiude il turno CORRENTE per la stessa via
|
||||
* del tetto NGEN (stats, usage_save, KV append, sentinella END tutti normali)
|
||||
* invece di uccidere il motore; :more puo' continuare la risposta interrotta.
|
||||
* Il flag e' armato solo nei serve-loop (intr_install): nei run one-shot e in
|
||||
* validazione SIGINT resta il default (morte immediata). Solo POSIX: su
|
||||
* Windows il comportamento di Ctrl-C non cambia.
|
||||
* EN: soft stop (serve/chat): SIGINT ends the CURRENT turn through the same
|
||||
* path as the NGEN cap — stats/usage/KV/END sentinel all normal — instead of
|
||||
* killing the engine; :more can continue the interrupted answer. Armed only
|
||||
* in the serve loops; one-shot runs keep default SIGINT. POSIX only. */
|
||||
static volatile sig_atomic_t g_intr=0;
|
||||
#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
|
||||
static void intr_sig(int s){ (void)s; g_intr=1; }
|
||||
static void intr_install(void){
|
||||
struct sigaction sa; memset(&sa,0,sizeof(sa));
|
||||
sa.sa_handler=intr_sig; sigemptyset(&sa.sa_mask);
|
||||
sa.sa_flags=SA_RESTART; /* getline/pread non devono vedere EINTR */
|
||||
sigaction(SIGINT,&sa,NULL);
|
||||
}
|
||||
#else
|
||||
static void intr_install(void){}
|
||||
#endif
|
||||
static int spec_decode(Model *m, int *all, int kv, int n_new, int eos, float *logit,
|
||||
void (*emit)(int,void*), void *ud, int *kv_out){
|
||||
Cfg *c=&m->c; int V=c->vocab; int emitted=0, done=0;
|
||||
int draft[64]; if(g_draft>63) g_draft=63;
|
||||
int carry_ban=-1; /* token rifiutato dalla verifica: escluso dal resample */
|
||||
while(emitted<n_new && !done){
|
||||
while(emitted<n_new && !done && !g_intr){ /* g_intr: stessa uscita del tetto n_new */
|
||||
int next=pick_tok(logit,V,carry_ban); carry_ban=-1; free(logit); logit=NULL;
|
||||
if((eos>=0 && next==eos) || is_stop(next)) break;
|
||||
emit(next,ud); all[kv]=next; emitted++; m->n_emit++;
|
||||
@@ -4078,6 +4204,7 @@ static void run_text(Model *m, const char *snap, const char *prompt, int ngen){
|
||||
printf("experts loaded/token: %.1f (per-layer %.2f across %d; baseline topk=%d) | TOPK=%d TOPP=%.2f",
|
||||
produced?(double)m->ereq/produced:0.0, (produced&&nsp)?(double)m->ereq/produced/nsp:0.0, nsp, c->topk, g_topk, g_topp);
|
||||
if(g_cache_route) printf(" | CACHE_ROUTE J=%d M=%d P=%.2f alpha=%.2f", g_route_j, g_route_m, g_route_p, g_route_alpha);
|
||||
if(g_expert_budget) printf(" | EXPERT_BUDGET=%d (dropped %lld experts, ~%.1f GB I/O saved)", g_expert_budget, (long long)g_budget_dropped, g_budget_dropped*18.9e6/1e9);
|
||||
printf("\n");
|
||||
printf("speculation: %.2f tokens/forward (%llu forwards per %llu tokens) | MTP acceptance %.0f%% (%llu/%llu)\n",
|
||||
m->n_fw?(double)m->n_emit/m->n_fw:1.0, (unsigned long long)m->n_fw, (unsigned long long)m->n_emit,
|
||||
@@ -4259,36 +4386,73 @@ static void kv_hdr(Model *m, int32_t *h, int nrec){
|
||||
h[0]=c->n_layers; h[1]=c->kv_lora; h[2]=c->qk_rope;
|
||||
h[3]=m->has_dsa?c->index_hd:0; h[4]=nic; h[5]=c->vocab; h[6]=nrec; h[7]=0;
|
||||
}
|
||||
/* Bytes of one on-disk record: [tok i32][Lc+Rc per layer][Ic per DSA layer].
|
||||
* Layout matches what kv_disk_append writes and kv_disk_load reads. */
|
||||
static int64_t kv_rec_bytes(Model *m){
|
||||
Cfg *c=&m->c;
|
||||
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
|
||||
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
|
||||
return rec;
|
||||
}
|
||||
/* Open the persistent handle lazily; write the header if the file is new. After
|
||||
* this returns successfully, k->disk_fp is valid for the engine's lifetime and
|
||||
* positioned at end-of-header (nrec==0 case) or wherever the caller seeks. */
|
||||
static int kv_disk_open(Model *m){
|
||||
KVState *k=m->kv;
|
||||
if(k->disk_fp) return 1;
|
||||
k->disk_fp=fopen(k->disk_path,"r+b");
|
||||
if(!k->disk_fp){ /* not there yet -> create + header */
|
||||
k->disk_fp=fopen(k->disk_path,"wb");
|
||||
if(!k->disk_fp) return 0;
|
||||
int32_t h[8]; kv_hdr(m,h,0);
|
||||
fwrite(KV_MAGIC,1,8,k->disk_fp); fwrite(h,4,8,k->disk_fp);
|
||||
fflush(k->disk_fp);
|
||||
fclose(k->disk_fp);
|
||||
k->disk_fp=fopen(k->disk_path,"r+b"); /* reopen r+b for append */
|
||||
if(!k->disk_fp) return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
static void kv_disk_truncate(Model *m, int nrec){
|
||||
if(!g_kvsave) return;
|
||||
KVState *k=m->kv;
|
||||
if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; } /* drop to shrink on disc */
|
||||
FILE *f=fopen(k->disk_path,"r+b");
|
||||
if(!f){ k->disk_nrec=0; return; }
|
||||
k->disk_nrec=nrec;
|
||||
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
|
||||
int32_t nr=nrec; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
|
||||
fflush(f); fclose(f);
|
||||
}
|
||||
static void kv_disk_reset(Model *m){ kv_disk_truncate(m,0); }
|
||||
static void kv_disk_append(Model *m, const int *hist, int len){
|
||||
KVState *k=m->kv;
|
||||
if(!g_kvsave || len<=k->disk_nrec) return;
|
||||
Cfg *c=&m->c;
|
||||
FILE *f=fopen(k->disk_path,"r+b");
|
||||
if(!f){ f=fopen(k->disk_path,"wb"); if(!f) return;
|
||||
int32_t h[8]; kv_hdr(m,h,0); fwrite(KV_MAGIC,1,8,f); fwrite(h,4,8,f); }
|
||||
int64_t rec = 4 + (int64_t)c->n_layers*(c->kv_lora+c->qk_rope)*4;
|
||||
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]) rec+=(int64_t)c->index_hd*4;
|
||||
if(!kv_disk_open(m)) return;
|
||||
FILE *f=k->disk_fp;
|
||||
int64_t rec = kv_rec_bytes(m);
|
||||
/* grow the contiguous staging buffer if the record is larger (#1 batching) */
|
||||
if(rec > k->disk_buf_cap){
|
||||
uint8_t *nb=realloc(k->disk_buf, rec);
|
||||
if(!nb) return; /* OOM: skip this turn, retry next */
|
||||
k->disk_buf=nb; k->disk_buf_cap=rec;
|
||||
}
|
||||
fseek(f, 8+8*4 + (int64_t)k->disk_nrec*rec, SEEK_SET);
|
||||
for(int p=k->disk_nrec;p<len;p++){
|
||||
int32_t tk=hist[p]; fwrite(&tk,4,1,f);
|
||||
uint8_t *b=k->disk_buf; /* pack token + every layer into one record */
|
||||
*(int32_t*)b = hist[p]; b+=4;
|
||||
for(int i=0;i<c->n_layers;i++){
|
||||
fwrite(m->Lc[i]+(int64_t)p*c->kv_lora, 4, c->kv_lora, f);
|
||||
fwrite(m->Rc[i]+(int64_t)p*c->qk_rope, 4, c->qk_rope, f);
|
||||
memcpy(b, m->Lc[i]+(int64_t)p*c->kv_lora, (size_t)c->kv_lora*4); b+=c->kv_lora*4;
|
||||
memcpy(b, m->Rc[i]+(int64_t)p*c->qk_rope,(size_t)c->qk_rope*4); b+=c->qk_rope*4;
|
||||
}
|
||||
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i])
|
||||
fwrite(m->Ic[i]+(int64_t)p*c->index_hd, 4, c->index_hd, f);
|
||||
if(m->has_dsa) for(int i=0;i<c->n_layers;i++) if(m->Ic[i]){
|
||||
memcpy(b, m->Ic[i]+(int64_t)p*c->index_hd, (size_t)c->index_hd*4); b+=c->index_hd*4;
|
||||
}
|
||||
fwrite(k->disk_buf, 1, (size_t)rec, f); /* one fwrite per position (was ~157) */
|
||||
}
|
||||
fflush(f); /* dati prima, contatore poi */
|
||||
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f); fclose(f);
|
||||
int32_t nr=len; fseek(f,8+6*4,SEEK_SET); fwrite(&nr,4,1,f);
|
||||
fflush(f); /* persist the counter too */
|
||||
k->disk_nrec=len;
|
||||
}
|
||||
static int kv_disk_load(Model *m, int *hist, int maxctx){
|
||||
@@ -4343,6 +4507,8 @@ static void serve_ctx_init(Model *m, ServeCtx *s, const char *snap, int slot, in
|
||||
|
||||
static void serve_ctx_free(Model *m, ServeCtx *s){
|
||||
KVState *k=&s->kv; int NR=m->c.n_layers+1;
|
||||
if(k->disk_fp){ fclose(k->disk_fp); k->disk_fp=NULL; }
|
||||
free(k->disk_buf); k->disk_buf=NULL;
|
||||
if(k->Lc) for(int i=0;i<NR;i++){ free(k->Lc[i]); free(k->Rc[i]); }
|
||||
if(k->Ic) for(int i=0;i<m->c.n_layers;i++) free(k->Ic[i]);
|
||||
free(k->Lc); free(k->Rc); free(k->Ic); free(k->kv_start); free(s->hist);
|
||||
@@ -4428,6 +4594,7 @@ static int mux_submit(Model *m, Tok *T, ServeCtx *ctx, ServeReq *req, int nctx,
|
||||
kv_disk_truncate(m,sc->len); }
|
||||
int add=nt-sc->len;
|
||||
if(add>0) memcpy(sc->hist+sc->len,tmp+sc->len,(size_t)add*sizeof(int));
|
||||
fprintf(stderr,"[API] KV slot %d prefix %d/%d token, prefill %d\n",sub.slot,sc->len,nt,add);
|
||||
free(tmp);
|
||||
float *logit = add>0 ? step(m,sc->hist+sc->len,add,sc->len)
|
||||
: step(m,sc->hist+sc->len-1,1,sc->len-1);
|
||||
@@ -4465,12 +4632,17 @@ static void run_serve_mux(Model *m, const char *snap){
|
||||
setvbuf(stdout, NULL, _IONBF, 0);
|
||||
#endif
|
||||
setvbuf(stdin,NULL,_IONBF,0);
|
||||
intr_install(); /* Ctrl-C = chiudi i turni in volo, non il processo */
|
||||
printf("\x01\x01READY\x01\x01\nSTAT 0 0.00 0.0 %.2f\n",rss_gb()); fflush(stdout);
|
||||
hwinfo_emit(m);
|
||||
tiers_emit(m);
|
||||
emap_emit(m);
|
||||
int eof=0;
|
||||
for(;;){
|
||||
if(g_intr){ g_intr=0; /* stop morbido: ogni request attiva finisce ORA per la
|
||||
* via normale di mux_done (DONE+stats+KV coerenti) */
|
||||
for(int i=0;i<nctx;i++) if(req[i].active) mux_done(m,&ctx[i],&req[i]);
|
||||
}
|
||||
int active=0; for(int i=0;i<nctx;i++) active+=req[i].active;
|
||||
/* Poll stdin for available input without blocking. On POSIX this is
|
||||
* select(); on Windows, select() on a pipe handle routes to winsock
|
||||
@@ -4563,9 +4735,11 @@ static void run_serve(Model *m, const char *snap){
|
||||
#define len (sc->len)
|
||||
#define first (sc->first)
|
||||
char *line=NULL; size_t cap=0; ssize_t nr; char *buf=malloc(1<<16);
|
||||
intr_install(); /* Ctrl-C = fine turno, non fine processo */
|
||||
printf("\x01\x01" "READY" "\x01\x01\n"); printf("STAT 0 0.00 0.0 %.2f\n", rss_gb()); fflush(stdout);
|
||||
tiers_emit(m);
|
||||
while((nr=getline(&line,&cap,stdin))>0){
|
||||
g_intr=0; /* interruzioni arrivate tra i turni: stantie */
|
||||
if(nr>0 && line[nr-1]=='\n') line[--nr]=0;
|
||||
if(!strcmp(line,"\x02RESET")){ len=0; first=1; if(m->has_mtp) m->kv_start[m->c.n_layers]=-1;
|
||||
kv_disk_reset(m);
|
||||
@@ -5213,6 +5387,7 @@ int main(int argc, char **argv){
|
||||
if(g_mmap) fprintf(stderr,"[MMAP] expert = viste zero-copy nei file (page cache = cache)\n");
|
||||
g_topk = getenv("TOPK")?atoi(getenv("TOPK")):0;
|
||||
g_topp = getenv("TOPP")?atof(getenv("TOPP")):0;
|
||||
g_expert_budget = getenv("EXPERT_BUDGET")?atoi(getenv("EXPERT_BUDGET")):0;
|
||||
g_cache_route = getenv("CACHE_ROUTE")?atoi(getenv("CACHE_ROUTE")):0;
|
||||
g_route_j = getenv("ROUTE_J")?atoi(getenv("ROUTE_J")):2;
|
||||
g_route_m = getenv("ROUTE_M")?atoi(getenv("ROUTE_M")):12;
|
||||
@@ -5256,7 +5431,13 @@ int main(int argc, char **argv){
|
||||
g_pilot_k = getenv("PILOT_K")?atoi(getenv("PILOT_K")):(g_pilot_real?6:8);
|
||||
if(g_pilot_k<1) g_pilot_k=1;
|
||||
g_disk_split = getenv("DISK_SPLIT")?atoi(getenv("DISK_SPLIT")):0; /* 1 = split dei disk load nelle stats */
|
||||
g_pipe = getenv("PIPE")?atoi(getenv("PIPE")):0; /* default OFF: overlap expert load ‖ matmul (byte-identical; reorders I/O). PIPE=1 opts in */
|
||||
g_pipe = getenv("PIPE")?atoi(getenv("PIPE")):
|
||||
#ifdef _WIN32
|
||||
1 /* default ON: overlap expert load ‖ matmul (byte-identical; reorders I/O). PIPE=0 opts out */
|
||||
#else
|
||||
0
|
||||
#endif
|
||||
;
|
||||
g_pipe_nw = getenv("PIPE_WORKERS")?atoi(getenv("PIPE_WORKERS")):8; /* I/O worker threads */
|
||||
if(g_pipe_nw<1) g_pipe_nw=1;
|
||||
g_direct = getenv("DIRECT")?atoi(getenv("DIRECT")):0;
|
||||
|
||||
+5
-2
@@ -410,8 +410,11 @@ def generation_options(body, limit):
|
||||
top_p = body.get("top_p")
|
||||
temperature = 0.7 if temperature is None else temperature
|
||||
top_p = 0.9 if top_p is None else top_p
|
||||
if isinstance(maximum, bool) or not isinstance(maximum, int) or not 1 <= maximum <= limit:
|
||||
raise APIError(400, f"`{maximum_param}` must be an integer between 1 and {limit}.", maximum_param)
|
||||
if isinstance(maximum, bool) or not isinstance(maximum, int) or maximum < 1:
|
||||
raise APIError(400, f"`{maximum_param}` must be a positive integer.", maximum_param)
|
||||
if maximum > limit:
|
||||
maximum = limit # clamp to the server's --max-tokens cap instead of 400 (#260): OpenAI
|
||||
# clients (opencode/ai-sdk) default to large max_tokens; rejecting breaks them.
|
||||
if (isinstance(temperature, bool) or not isinstance(temperature, (int, float)) or
|
||||
not math.isfinite(temperature) or not 0 <= temperature <= 2):
|
||||
raise APIError(400, "`temperature` must be between 0 and 2.", "temperature")
|
||||
|
||||
@@ -51,6 +51,23 @@ int main(void){
|
||||
if(compat_open_direct("no_such_file.tmp")>=0) return fail("open missing file must fail");
|
||||
if(compat_fsize(-1)>=0) return fail("compat_fsize on bad fd must be negative");
|
||||
|
||||
/* compat_fadvise: WILLNEED warms the page cache (background read into throwaway
|
||||
* buffer), DONTNEED is a documented no-op. After a WILLNEED the buffered fd's
|
||||
* subsequent pread must still return the exact bytes — the cache-warmer must not
|
||||
* corrupt data. Bad fd / non-WILLNEED advice must be safe no-ops (return 0). */
|
||||
int wfd = open(TMPF, COMPAT_O_RDONLY);
|
||||
if(wfd<0) return fail("open buffered for fadvise");
|
||||
if(posix_fadvise(wfd, 0, (off_t)FSZ, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED returned nonzero");
|
||||
if(posix_fadvise(wfd, 0, (off_t)FSZ, POSIX_FADV_DONTNEED)!=0) return fail("DONTNEED should be a safe no-op (return 0)");
|
||||
if(posix_fadvise(-1, 0, (off_t)FSZ, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED on bad fd should no-op (return 0)");
|
||||
if(posix_fadvise(wfd, 0, 0, POSIX_FADV_WILLNEED)!=0) return fail("WILLNEED with len<=0 should no-op");
|
||||
/* verify data integrity through the buffered fd after the cache-warmer ran */
|
||||
uint8_t *verify=malloc(FSZ);
|
||||
if(pread(wfd, verify, FSZ, 0)!=(ssize_t)FSZ) return fail("fadvise: pread size");
|
||||
if(memcmp(verify, pat, FSZ)!=0) return fail("fadvise: data corrupted by cache-warmer");
|
||||
free(verify);
|
||||
close(wfd);
|
||||
|
||||
close(dfd);
|
||||
compat_aligned_free(buf); free(pat); remove(TMPF);
|
||||
puts("compat direct tests: ok");
|
||||
|
||||
@@ -72,8 +72,13 @@ class TemplateTest(unittest.TestCase):
|
||||
def test_validates_generation_limits(self):
|
||||
self.assertEqual(generation_options({"max_tokens": 4, "temperature": 0, "top_p": 1}, 8),
|
||||
(4, 0.0, 1.0))
|
||||
# max_tokens above the server cap is clamped, not rejected (#260): OpenAI
|
||||
# clients default to large values; erroring breaks them.
|
||||
self.assertEqual(generation_options({"max_tokens": 9, "temperature": 0, "top_p": 1}, 8),
|
||||
(8, 0.0, 1.0))
|
||||
# non-positive / non-int max_tokens is still a hard error
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"max_tokens": 9}, 8)
|
||||
generation_options({"max_tokens": 0}, 8)
|
||||
with self.assertRaises(APIError):
|
||||
generation_options({"temperature": math.nan}, 8)
|
||||
with self.assertRaises(APIError):
|
||||
|
||||
@@ -116,7 +116,7 @@ def layer_idx(name):
|
||||
|
||||
def classify(name, n_layers, keep_mtp=False, keep_idx=False):
|
||||
if name.endswith("_scale_inv"): return "consumed" # FP8 base: gestito col suo peso
|
||||
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro .weight U8.
|
||||
# NVFP4 (modelopt): i sidecar delle scale sono consumati insieme al loro U8 .weight.
|
||||
# EN: NVFP4 (modelopt): scale sidecars are consumed together with their U8 .weight.
|
||||
if name.endswith((".weight_scale", ".weight_scale_2", ".input_scale")): return "consumed"
|
||||
li = layer_idx(name)
|
||||
@@ -137,7 +137,20 @@ def classify(name, n_layers, keep_mtp=False, keep_idx=False):
|
||||
if name.endswith("norm.weight") or name == "model.norm.weight": return "f32"
|
||||
if name in ("model.embed_tokens.weight", "lm_head.weight"): return "io"
|
||||
if ".mlp.experts." in name and name.endswith(".weight"): return "x" # expert ROUTED (streaming)
|
||||
if name.endswith(".weight"): return "q" # attn/dense-mlp/shared (residente)
|
||||
# Split resident weights by type for mixed-precision control:
|
||||
# "sh" = shared expert (fires on every token, highest sensitivity)
|
||||
# "o" = o_proj attention (reconstructs output, biggest attn tensor)
|
||||
# "kvb" = kv_b_proj (reconstructs KV cache on every decode step)
|
||||
# "attn" = other attention projections (q_a, q_b, kv_a)
|
||||
# "dmlp" = dense MLP (first 3 layers)
|
||||
if "shared_experts" in name: return "sh"
|
||||
if name.endswith("o_proj.weight"): return "o"
|
||||
if name.endswith("kv_b_proj.weight"): return "kvb"
|
||||
if any(name.endswith(k) for k in ("q_a_proj.weight", "q_b_proj.weight",
|
||||
"kv_a_proj_with_mqa.weight")): return "attn"
|
||||
if any(name.endswith(k) for k in ("mlp.gate_proj.weight", "mlp.up_proj.weight",
|
||||
"mlp.down_proj.weight")): return "dmlp"
|
||||
if name.endswith(".weight"): return "q" # fallback: other resident weights
|
||||
return "f32"
|
||||
|
||||
# ---------- dequant NVFP4 (modelopt) di UN tensore expert -> f32 [O,I] ----------
|
||||
@@ -202,7 +215,7 @@ def dequant(f, name, keys):
|
||||
return f.get_tensor(name).to(torch.float32).numpy()
|
||||
|
||||
def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
|
||||
keep_mtp=False, keep_idx=False, group_size=0):
|
||||
keep_mtp=False, keep_idx=False, group_size=0, bits_map=None):
|
||||
from safetensors import safe_open
|
||||
with safe_open(path, framework="pt") as f:
|
||||
keys = set(f.keys())
|
||||
@@ -213,7 +226,15 @@ def convert_shard(path, out_dict, n_layers, ebits, io_bits, xbits,
|
||||
if kind == "f32":
|
||||
out_dict[name] = w.astype(np.float32)
|
||||
else:
|
||||
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
|
||||
# Resolve bits for this tensor type: use bits_map override if provided,
|
||||
# otherwise fall back to the classic ebits/xbits/io_bits scheme.
|
||||
if bits_map and kind in bits_map:
|
||||
bits = bits_map[kind]
|
||||
else:
|
||||
bits = io_bits if kind == "io" else xbits if kind == "x" else ebits
|
||||
# Any unknown kind that fell through classify as "q"
|
||||
if bits_map and kind not in bits_map and kind not in ("io", "x", "sh", "o", "kvb", "attn", "dmlp"):
|
||||
bits = ebits
|
||||
if w.ndim != 2: # es. bias 1D non previsto come 'q' -> tienilo f32
|
||||
out_dict[name] = w.astype(np.float32); continue
|
||||
if group_size > 0 and bits <= 4:
|
||||
@@ -234,6 +255,18 @@ def main():
|
||||
ap.add_argument("--ebits", type=int, default=None) # bit residenti (default 4; 8 per --mtp/--indexer)
|
||||
ap.add_argument("--io-bits", type=int, default=8) # bit di embed/lm_head
|
||||
ap.add_argument("--xbits", type=int, default=None) # bit degli expert ROUTED (streaming); default=ebits
|
||||
# Mixed-precision: per-tensor-type bit overrides. Default = ebits (all same).
|
||||
# Set these higher to protect sensitive tensors from quantization error.
|
||||
ap.add_argument("--shared-bits", type=int, default=None,
|
||||
help="bits for shared expert (fires on every token, highest sensitivity). Default=ebits")
|
||||
ap.add_argument("--o-bits", type=int, default=None,
|
||||
help="bits for o_proj attention (reconstructs output, biggest attn tensor). Default=ebits")
|
||||
ap.add_argument("--kvb-bits", type=int, default=None,
|
||||
help="bits for kv_b_proj (reconstructs KV cache on every decode). Default=ebits")
|
||||
ap.add_argument("--attn-bits", type=int, default=None,
|
||||
help="bits for other attention projections (q_a, q_b, kv_a). Default=ebits")
|
||||
ap.add_argument("--dmlp-bits", type=int, default=None,
|
||||
help="bits for dense MLP (first 3 layers). Default=ebits")
|
||||
ap.add_argument("--group-size", type=int, default=0, # 0 = per-row (backward compat); 128 = group-scaled
|
||||
help="group size for int4 scales: 0=per-row (default), 128=one scale per 128 elements (much better quality)")
|
||||
ap.add_argument("--n-layers", type=int, default=78)
|
||||
@@ -255,6 +288,17 @@ def main():
|
||||
a.ebits = 8 if (a.mtp or a.indexer) else 4
|
||||
if a.xbits is None: a.xbits = a.ebits
|
||||
|
||||
# Build per-type bits map. If a type-specific arg is set, use it; otherwise the
|
||||
# converter falls back to ebits for that type.
|
||||
bits_map = {}
|
||||
if a.shared_bits is not None: bits_map["sh"] = a.shared_bits
|
||||
if a.o_bits is not None: bits_map["o"] = a.o_bits
|
||||
if a.kvb_bits is not None: bits_map["kvb"] = a.kvb_bits
|
||||
if a.attn_bits is not None: bits_map["attn"] = a.attn_bits
|
||||
if a.dmlp_bits is not None: bits_map["dmlp"] = a.dmlp_bits
|
||||
if bits_map:
|
||||
print(f"[MIXED] precision map: " + ", ".join(f"{k}={v}bit" for k,v in sorted(bits_map.items())))
|
||||
|
||||
if a.selftest_nvfp4:
|
||||
import torch
|
||||
# 1) LUT e2m1: i 16 codici devono decodificare esattamente ai valori attesi.
|
||||
@@ -336,7 +380,7 @@ def main():
|
||||
shards = sorted(glob.glob(os.path.join(a.indir, "*.safetensors")))
|
||||
from safetensors.numpy import save_file
|
||||
for i, sp in enumerate(shards):
|
||||
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
|
||||
out = {}; convert_shard(sp, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
|
||||
save_file(out, os.path.join(a.outdir, f"out-{i:05d}.safetensors"))
|
||||
# copia config + tokenizer
|
||||
for fn in ["config.json"]:
|
||||
@@ -579,7 +623,7 @@ def main():
|
||||
if os.path.exists(outp): print(f"[MTP] {outp} already done"); continue
|
||||
print(f"[MTP {i+1}/{len(mtp_shards)}] downloading {sh}...", flush=True)
|
||||
p = download_retry(a.repo, sh, tmp)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_mtp=True, group_size=a.group_size, bits_map=bits_map)
|
||||
save_file(out, outp)
|
||||
os.remove(p)
|
||||
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
|
||||
@@ -599,7 +643,7 @@ def main():
|
||||
if os.path.exists(outp): continue # gia' fatto -> ripartibile
|
||||
print(f"[IDX {i+1}/{len(idx_shards)}] downloading {sh}...", flush=True)
|
||||
p = download_retry(a.repo, sh, tmp)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, keep_idx=True, group_size=a.group_size, bits_map=bits_map)
|
||||
if out: save_file(out, outp)
|
||||
os.remove(p)
|
||||
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
|
||||
@@ -613,7 +657,7 @@ def main():
|
||||
if os.path.exists(outp): continue # gia' fatto -> ripartibile
|
||||
print(f"[{i+1}/{len(shards)}] downloading {sh} ({free_gb(a.outdir):.0f} GB free)...", flush=True)
|
||||
p = download_retry(a.repo, sh, tmp)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size)
|
||||
out = {}; convert_shard(p, out, a.n_layers, a.ebits, a.io_bits, a.xbits, group_size=a.group_size, bits_map=bits_map)
|
||||
save_file(out, outp)
|
||||
os.remove(p) # <-- cancella subito lo shard fp8
|
||||
for blob in glob.glob(os.path.join(tmp, "**", "*"), recursive=True):
|
||||
|
||||
@@ -1,141 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Expert Atlas (#175): measure per-expert topic affinity by diffing .coli_usage
|
||||
across themed probe batches, served through a running colibri API server.
|
||||
|
||||
Protocol per category: snapshot .coli_usage -> send probes -> snapshot again;
|
||||
the delta is that category's expert-activation spectrum. One engine load total.
|
||||
|
||||
Output: experts.json — for every (layer, expert): counts per category,
|
||||
normalized affinity, entropy, and a "specialist" label when one topic dominates.
|
||||
|
||||
Usage (server already running with the model):
|
||||
python3 tools/expert_atlas.py --api http://127.0.0.1:8000 \
|
||||
--usage /path/to/model/.coli_usage --out experts.json --ngen 64
|
||||
"""
|
||||
import argparse, json, math, time, urllib.request
|
||||
|
||||
PROBES = {
|
||||
"code": [
|
||||
"Write a Python function that parses a CSV file and returns a dict keyed by the first column.",
|
||||
"Explain the difference between a mutex and a semaphore, with a C example.",
|
||||
"Refactor this into idiomatic Rust: for i in range(len(xs)): total += xs[i] * 2",
|
||||
],
|
||||
"math": [
|
||||
"Prove that the square root of 2 is irrational.",
|
||||
"Compute the derivative of x^3 * ln(x) and explain each step.",
|
||||
"A fair die is rolled 4 times. What is the probability of at least one six?",
|
||||
],
|
||||
"chinese": [
|
||||
"请用中文解释一下什么是光合作用,以及它对地球生态系统的重要性。",
|
||||
"把这句话翻译成中文并解释语法:The early bird catches the worm.",
|
||||
"写一段关于秋天的短文,一百字左右。",
|
||||
],
|
||||
"english_prose": [
|
||||
"Write a vivid paragraph describing an old lighthouse keeper watching a storm arrive.",
|
||||
"Summarize the plot of Romeo and Juliet in three sentences.",
|
||||
"Continue this story: The last train left the station, and Maria realized her mistake.",
|
||||
],
|
||||
"science": [
|
||||
"Explain how mRNA vaccines work at the cellular level.",
|
||||
"Why is the sky blue during the day but red at sunset?",
|
||||
"Describe the life cycle of a massive star, from formation to supernova.",
|
||||
],
|
||||
"law": [
|
||||
"Explain the difference between a patent, a trademark, and a copyright.",
|
||||
"What are the key elements required to form a legally binding contract?",
|
||||
"Summarize what 'due process' means in constitutional law.",
|
||||
],
|
||||
"poetry": [
|
||||
"Write a short poem about a hummingbird in the style of Emily Dickinson.",
|
||||
"Compose a haiku about winter rain, then explain its imagery.",
|
||||
"Write four rhyming lines about the sea at night.",
|
||||
],
|
||||
"structured": [
|
||||
'Convert to JSON: name Alice, age 30, hobbies reading and chess, address 5 Oak St.',
|
||||
"Write a SQL query returning the top 5 customers by total order value, with the schema you assume.",
|
||||
"Write a regex that matches ISO-8601 dates and explain each part.",
|
||||
],
|
||||
"translation": [
|
||||
"Translate into French, German and Spanish: 'Knowledge is the only treasure that grows when shared.'",
|
||||
"Translate this Italian sentence to English and comment on nuance: 'In bocca al lupo per domani.'",
|
||||
"Translate into Japanese: 'The meeting has been moved to next Tuesday afternoon.'",
|
||||
],
|
||||
"casual": [
|
||||
"Hey! Any tips for staying awake during boring afternoon meetings?",
|
||||
"What should I cook tonight? I have eggs, rice, tomatoes and some cheese.",
|
||||
"My friend is always late. How do I tell them it bothers me without being rude?",
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def read_usage(path):
|
||||
counts = {}
|
||||
try:
|
||||
with open(path) as f:
|
||||
for line in f:
|
||||
p = line.split()
|
||||
if len(p) == 3:
|
||||
counts[(int(p[0]), int(p[1]))] = int(p[2])
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return counts
|
||||
|
||||
|
||||
def diff(after, before):
|
||||
return {k: v - before.get(k, 0) for k, v in after.items() if v - before.get(k, 0) > 0}
|
||||
|
||||
|
||||
def chat(api, prompt, ngen):
|
||||
body = json.dumps({"model": "glm-5.2-colibri", "stream": False, "max_tokens": ngen,
|
||||
"messages": [{"role": "user", "content": prompt}]}).encode()
|
||||
req = urllib.request.Request(f"{api}/v1/chat/completions", data=body,
|
||||
headers={"Content-Type": "application/json"})
|
||||
with urllib.request.urlopen(req, timeout=600) as r:
|
||||
json.load(r)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--api", default="http://127.0.0.1:8000")
|
||||
ap.add_argument("--usage", required=True, help="path to the model's .coli_usage")
|
||||
ap.add_argument("--out", default="experts.json")
|
||||
ap.add_argument("--ngen", type=int, default=64)
|
||||
a = ap.parse_args()
|
||||
|
||||
spectra = {}
|
||||
for cat, prompts in PROBES.items():
|
||||
before = read_usage(a.usage)
|
||||
t0 = time.time()
|
||||
for p in prompts:
|
||||
chat(a.api, p, a.ngen)
|
||||
time.sleep(2) # let the engine flush .coli_usage
|
||||
spectra[cat] = diff(read_usage(a.usage), before)
|
||||
total = sum(spectra[cat].values())
|
||||
print(f"[{cat}] {len(spectra[cat])} experts touched, {total} selections, {time.time()-t0:.0f}s", flush=True)
|
||||
|
||||
cats = list(PROBES.keys())
|
||||
experts = {}
|
||||
for cat, spec in spectra.items():
|
||||
for k, v in spec.items():
|
||||
experts.setdefault(k, {c: 0 for c in cats})[cat] = v
|
||||
|
||||
atlas = {}
|
||||
for (layer, eid), counts in experts.items():
|
||||
total = sum(counts.values())
|
||||
if total < 8:
|
||||
continue # too few observations to characterise
|
||||
aff = {c: v / total for c, v in counts.items() if v}
|
||||
ent = -sum(p * math.log2(p) for p in aff.values())
|
||||
top = max(aff, key=aff.get)
|
||||
label = f"specialist: {top}" if aff[top] >= 0.45 and ent < 2.2 else "generalist"
|
||||
atlas[f"{layer}:{eid}"] = {"counts": counts, "affinity": {c: round(p, 3) for c, p in aff.items()},
|
||||
"entropy": round(ent, 2), "top": top, "label": label}
|
||||
|
||||
spec_n = sum(1 for v in atlas.values() if v["label"].startswith("specialist"))
|
||||
with open(a.out, "w") as f:
|
||||
json.dump({"categories": cats, "ngen": a.ngen, "experts": atlas}, f)
|
||||
print(f"\natlas: {len(atlas)} experts characterised, {spec_n} specialists -> {a.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -7,10 +7,16 @@ histogram, and turns them into a per-expert topic-affinity vector.
|
||||
cd c
|
||||
export COLI_MODEL=/path/to/glm52_i4
|
||||
./tools/expert_atlas/sweep.sh # 30 probes (10 topics x 3 prompts)
|
||||
python3 tools/expert_atlas/analyze.py --stats atlas_out/stats --out atlas_out/experts.json
|
||||
python3 tools/expert_atlas/analyze.py --stats atlas_out/stats --out atlas_out/experts.json \
|
||||
--web web/dist/experts.json # optional: feed the web dashboard Atlas
|
||||
python3 tools/expert_atlas/validate.py atlas_out/stats 200 # leave-one-prompt-out check
|
||||
```
|
||||
|
||||
`--web` writes the same atlas in the shape the web dashboard consumes (the Atlas galaxy and the
|
||||
Brain hover tooltips): keyed `"layer:expert"` with `affinity`/`entropy`/`top`/`label`. It replaces
|
||||
the retired `tools/expert_atlas.py`, whose API-driven probing ran through a live server and was
|
||||
exposed to exactly the traps above (server-side `--topp`, speculative drafts, shared `.coli_usage`).
|
||||
|
||||
## Read this before you trust any atlas
|
||||
|
||||
Four things silently corrupt this measurement. The sweep script controls all of them; if you
|
||||
|
||||
@@ -29,6 +29,7 @@ def main():
|
||||
ap.add_argument("--min-count", type=int, default=30)
|
||||
ap.add_argument("--min-runs", type=int, default=2, help="must fire in >= this many of the top category's runs")
|
||||
ap.add_argument("--out", default="experts.json")
|
||||
ap.add_argument("--web", default="", help="also write the web-dashboard experts.json (Atlas/Brain hover)")
|
||||
a = ap.parse_args()
|
||||
|
||||
# run[(cat,idx)][(layer,expert)] = count ; run_tot[(cat,idx)] = total
|
||||
@@ -121,6 +122,19 @@ def main():
|
||||
json.dump({"categories": cats, "experts": atlas}, open(a.out, "w"), indent=1)
|
||||
print(f"wrote {a.out}")
|
||||
|
||||
if a.web:
|
||||
# Same atlas, keyed "layer:expert" with per-expert affinity/entropy/top/label —
|
||||
# the shape the web dashboard consumes (Atlas galaxy, Brain hover).
|
||||
web = {}
|
||||
for r in atlas:
|
||||
aff = {c: v for c, v in r["p"].items() if v > 0}
|
||||
H = -sum(v * math.log2(v) for v in aff.values())
|
||||
web[f"{r['layer']}:{r['expert']}"] = {
|
||||
"affinity": aff, "entropy": round(H, 2), "top": r["top_topic"],
|
||||
"label": f"specialist: {r['top_topic']}" if r["spec"] >= 0.5 else "generalist"}
|
||||
json.dump({"categories": cats, "experts": web}, open(a.web, "w"))
|
||||
print(f"wrote {a.web} (dashboard format, {len(web):,} experts)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
"""Helper: salva pesi in FP8 e4m3 + scale a blocchi 128x128, nello STESSO layout del
|
||||
checkpoint reale GLM-5.2-FP8 che `convert_fp8_to_int4.py` legge.
|
||||
|
||||
Layout (deve combaciare col `dequant()` del converter, convert_fp8_to_int4.py:164-169):
|
||||
- `name` F8_E4M3 [O, I]
|
||||
- `name_scale_inv` F32 [ceil(O/128), ceil(I/128)] (NOTA: '_scale_inv', underscore)
|
||||
dequant: W = q.float() * scale.repeat_interleave(128,0).repeat_interleave(128,1)[:O,:I]
|
||||
|
||||
Convenzione FBGEMM/TransformerEngine: scale = amax(blocco)/448 (448 = max e4m3),
|
||||
si MEMORIZZA il valore e si MOLTIPLICA in dequant. Malgrado il nome "_scale_inv" il
|
||||
checkpoint memorizza la scala (non il reciproco): e' un MOLTIPLIER.
|
||||
|
||||
EN: Helper that writes weights as FP8 e4m3 with 128x128 block scales, in the SAME layout
|
||||
EN: as the real GLM-5.2-FP8 checkpoint that `convert_fp8_to_int4.py` reads.
|
||||
EN: FBGEMM/TransformerEngine convention: scale = amax(block)/448, stored (not its
|
||||
EN: reciprocal) and MULTIPLIED on dequant. Despite the name "_scale_inv" it is a multiplier.
|
||||
"""
|
||||
import torch
|
||||
|
||||
E4M3_MAX = 448.0 # max valore rappresentabile in float8_e4m3fn / max representable value
|
||||
BLOCK = 128 # granularita' delle scale a blocchi del checkpoint FP8 / FP8 block scale granularity
|
||||
|
||||
|
||||
def keep_f32(name, t):
|
||||
"""Stesso set F32 di `classify()` in convert_fp8_to_int4.py (norme, router, bias 1-D).
|
||||
Tutti gli altri tensori 2-D vengono quantizzati FP8 (attn/mlp/shared/expert/embed/lm_head).
|
||||
EN: Same F32 set as the converter's classify(): norms, router, 1-D biases. All other 2-D
|
||||
EN: tensors are FP8-quantized (attn/mlp/shared/expert/embed/lm_head)."""
|
||||
if t.dim() < 2:
|
||||
return True # bias 1-D, e_score_correction_bias
|
||||
if name.endswith("e_score_correction_bias"):
|
||||
return True
|
||||
if name.endswith("mlp.gate.weight"):
|
||||
return True # router (NON gate_proj): tenuto F32 / kept F32
|
||||
if name.endswith("norm.weight") or name == "model.norm.weight":
|
||||
return True # RMSNorm
|
||||
return False
|
||||
|
||||
|
||||
def fp8_block_quantize(w):
|
||||
"""w: [O,I] f32 -> (w_fp8 float8_e4m3fn [O,I], scale_inv f32 [ceil(O/128),ceil(I/128)]).
|
||||
Identica matematica al `--selftest` del converter (scale = amax(blocco)/448). Padda a
|
||||
multipli di 128 internamente (gli zeri non alzano l'amax) e fa slice al risultato.
|
||||
EN: same math as the converter's --selftest. Pads to 128 multiples internally (zeros do
|
||||
EN: not raise amax), slices the result back to [O,I]."""
|
||||
O, I = w.shape
|
||||
nbO, nbI = (O + BLOCK - 1) // BLOCK, (I + BLOCK - 1) // BLOCK
|
||||
Op, Ip = nbO * BLOCK, nbI * BLOCK
|
||||
wpad = torch.zeros(Op, Ip, dtype=torch.float32, device=w.device)
|
||||
wpad[:O, :I] = w
|
||||
wb = wpad.view(nbO, BLOCK, nbI, BLOCK) # [nbO, BLOCK, nbI, BLOCK]
|
||||
amax = wb.abs().amax(dim=(1, 3)) # [nbO, nbI]
|
||||
scale = amax / E4M3_MAX # FBGEMM/TE: memorizza la scala / store the scale
|
||||
scale = torch.where(scale == 0, torch.ones_like(scale), scale) # blocco tutto-zero -> no div0
|
||||
scale = scale.to(torch.float32)
|
||||
q = (wpad / scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)).clamp(-E4M3_MAX, E4M3_MAX)
|
||||
w_fp8 = q.to(torch.float8_e4m3fn)
|
||||
return w_fp8[:O, :I].contiguous(), scale.contiguous()
|
||||
|
||||
|
||||
def fp8_block_dequantize(w_fp8, scale):
|
||||
"""Esatto inverso di fp8_block_quantize, e identico al `dequant()` del converter.
|
||||
EN: exact inverse of fp8_block_quantize, identical to the converter's dequant()."""
|
||||
O, I = w_fp8.shape
|
||||
qf = w_fp8.to(torch.float32)
|
||||
return qf * scale.repeat_interleave(BLOCK, 0).repeat_interleave(BLOCK, 1)[:O, :I]
|
||||
|
||||
|
||||
def unfuse_experts(sd):
|
||||
"""Split HF's fused 3-D `experts.gate_up_proj` [E, 2*M, I] into per-expert 2-D
|
||||
`experts.{e}.gate_proj` [M, I] + `experts.{e}.up_proj` [M, I], and
|
||||
`experts.down_proj` [E, I, M] -> `experts.{e}.down_proj` [M_out, I].
|
||||
|
||||
The real GLM-5.2-FP8 checkpoint stores experts UNFUSED as per-expert 2-D tensors
|
||||
(gate_proj, up_proj, down_proj), each with its own _scale_inv. HF's
|
||||
GlmMoeDsaForCausalLM fuses gate+up into a single 3-D gate_up_proj for efficiency.
|
||||
The converter (classify + ndim!=2 guard) and the C engine both expect the unfused
|
||||
layout, so we split before saving.
|
||||
|
||||
Idempotent: if experts are already unfused (no 3-D gate_up_proj), returns sd as-is.
|
||||
EN: split HF's fused 3-D expert weights into the per-expert 2-D layout that the real
|
||||
EN: checkpoint uses and the converter/engine expect. No-op if already unfused."""
|
||||
keys_to_remove = []
|
||||
new_entries = {}
|
||||
for name, t in sd.items():
|
||||
if not name.endswith(".mlp.experts.gate_up_proj"):
|
||||
continue
|
||||
# prefix = everything before ".mlp.experts.gate_up_proj"
|
||||
prefix = name[:-len(".mlp.experts.gate_up_proj")]
|
||||
E, twoM, I = t.shape # [E, 2*intermediate, input]
|
||||
M = twoM // 2
|
||||
for e in range(E):
|
||||
new_entries[f"{prefix}.mlp.experts.{e}.gate_proj.weight"] = t[e, :M, :].contiguous()
|
||||
new_entries[f"{prefix}.mlp.experts.{e}.up_proj.weight"] = t[e, M:, :].contiguous()
|
||||
keys_to_remove.append(name)
|
||||
# down_proj may be 3-D [E, I, M] in the fused form, or already per-expert
|
||||
for name, t in sd.items():
|
||||
if not name.endswith(".mlp.experts.down_proj") or t.dim() != 3:
|
||||
continue
|
||||
prefix = name[:-len(".mlp.experts.down_proj")]
|
||||
E = t.shape[0]
|
||||
for e in range(E):
|
||||
new_entries[f"{prefix}.mlp.experts.{e}.down_proj.weight"] = t[e].contiguous()
|
||||
keys_to_remove.append(name)
|
||||
for k in keys_to_remove:
|
||||
sd.pop(k, None)
|
||||
sd.update(new_entries)
|
||||
return sd
|
||||
|
||||
|
||||
def state_dict_to_fp8(sd):
|
||||
"""Converte uno state_dict HuggingFace nel layout FP8 del checkpoint reale:
|
||||
per ogni tensore quantizzabile 2-D scrive `{name}` (F8_E4M3) + `{name}_scale_inv` (F32);
|
||||
norme/router/bias e qualsiasi tensore NON 2-D (es. pesi MLA impaccati 3-D) restano nel
|
||||
dtype originale. Questo rispecchia il guard `w.ndim != 2 -> f32` del converter
|
||||
(convert_fp8_to_int4.py:184). EN: builds the real-checkpoint FP8 layout. Only exactly-2-D
|
||||
tensors are FP8-quantized; anything else (1-D, 3-D packed MLA weights, ...) is kept, exactly
|
||||
like the converter's `ndim != 2 -> f32` guard."""
|
||||
out = {}
|
||||
for name, t in sd.items():
|
||||
if keep_f32(name, t) or t.dim() != 2:
|
||||
out[name] = t # f32 / 1-D / 3-D+: tieni / keep
|
||||
else:
|
||||
w_fp8, scale = fp8_block_quantize(t.float())
|
||||
out[name] = w_fp8
|
||||
out[name + "_scale_inv"] = scale
|
||||
return out
|
||||
|
||||
|
||||
def save_fp8_safetensors(sd, path):
|
||||
"""Quantizza a blocchi FP8 e salva in un singolo safetensors leggibile dal converter
|
||||
via `--indir`. EN: block-quantize to FP8 and save a single safetensors for the converter."""
|
||||
from safetensors.torch import save_file
|
||||
out = state_dict_to_fp8(sd)
|
||||
save_file({k: v.contiguous() for k, v in out.items()}, str(path))
|
||||
n_fp8 = sum(1 for v in out.values() if v.dtype == torch.float8_e4m3fn)
|
||||
return n_fp8, len(out)
|
||||
@@ -3,15 +3,27 @@
|
||||
This is not a useful language model. It preserves the real glm_moe_dsa data
|
||||
flow while remaining small enough to generate locally and run repeated CPU/CUDA
|
||||
A/B tests without downloading the 379 GB checkpoint.
|
||||
|
||||
With --fp8 the weights are written as FP8 e4m3 + 128x128 block scale_inv, in the
|
||||
SAME layout as the real GLM-5.2-FP8 checkpoint, so convert_fp8_to_int4.py can
|
||||
exercise its FP8->int4 dequant path on a local fixture (its dims are 128-friendly,
|
||||
so this is also the right fixture for --group-size 128 testing):
|
||||
|
||||
python tools/make_glm_bench_model.py --fp8 --output glm_bench_fp8
|
||||
python tools/convert_fp8_to_int4.py --indir glm_bench_fp8 --outdir glm_bench_i4 --ebits 4 --group-size 128
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
|
||||
from glm_fp8_emit import save_fp8_safetensors, unfuse_experts
|
||||
|
||||
|
||||
def build_config() -> GlmMoeDsaConfig:
|
||||
return GlmMoeDsaConfig(
|
||||
@@ -51,6 +63,9 @@ def main() -> None:
|
||||
parser.add_argument("--output", default="glm_bench_medium")
|
||||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--seed", type=int, default=1234)
|
||||
parser.add_argument("--fp8", action="store_true",
|
||||
help="write weights as FP8 e4m3 + 128x128 block scale_inv (same layout as "
|
||||
"GLM-5.2-FP8) instead of bf16, so convert_fp8_to_int4.py can dequant+requant")
|
||||
args = parser.parse_args()
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
@@ -70,7 +85,6 @@ def main() -> None:
|
||||
output = Path(args.output)
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
params = sum(p.numel() for p in model.parameters())
|
||||
model.save_pretrained(output, safe_serialization=True, max_shard_size="4GB")
|
||||
|
||||
model.to(args.device)
|
||||
prompt = [3, 14, 159, 26, 53, 58, 200, 11, 77, 240, 5, 99]
|
||||
@@ -79,6 +93,25 @@ def main() -> None:
|
||||
full = model.generate(ids, max_new_tokens=8, do_sample=False, use_cache=True)[0]
|
||||
logits = model(full.unsqueeze(0), use_cache=False).logits[0]
|
||||
|
||||
# Unfuse experts AFTER reference generation (model needs fused weights for
|
||||
# forward/generate) but BEFORE saving — the real checkpoint and the converter
|
||||
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
|
||||
sd = model.state_dict()
|
||||
unfuse_experts(sd)
|
||||
|
||||
if args.fp8:
|
||||
n_fp8, n_tot = save_fp8_safetensors(sd, output / "model.safetensors")
|
||||
# save_pretrained scrive config.json; nel path FP8 lo bypassiamo, quindi lo scriviamo
|
||||
# a mano (serve al converter e al motore C). EN: save_pretrained writes config.json;
|
||||
# the FP8 path bypasses it, so write it manually (converter + C engine need it).
|
||||
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
|
||||
print(f"saved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
|
||||
f"-> {output / 'model.safetensors'}")
|
||||
else:
|
||||
from safetensors.torch import save_file
|
||||
save_file({k: v.contiguous() for k, v in sd.items()}, str(output / "model.safetensors"))
|
||||
(output / "config.json").write_text(json.dumps(cfg.to_dict()))
|
||||
|
||||
ref = {
|
||||
"prompt_ids": prompt,
|
||||
"full_ids": full.cpu().tolist(),
|
||||
@@ -89,6 +122,7 @@ def main() -> None:
|
||||
"seed": args.seed,
|
||||
"parameters": params,
|
||||
"parameters_billions": round(params / 1e9, 4),
|
||||
"format": "fp8-e4m3-128" if args.fp8 else "bf16",
|
||||
"purpose": "backend benchmark fixture; random weights, not a language model",
|
||||
}
|
||||
(output / "bench_manifest.json").write_text(json.dumps(manifest, indent=2))
|
||||
|
||||
@@ -3,10 +3,34 @@ Architettura vera (MLA + DSA indexer + router sigmoid/noaux_tc + shared expert),
|
||||
dimensioni minuscole. Salva pesi+config in c/glm_tiny/ e un riferimento greedy in
|
||||
c/ref_glm.json. seq corta (<= index_topk) cosi' il DSA seleziona tutte le key e
|
||||
l'attenzione coincide con la MLA densa: il motore C puo' validare senza implementare
|
||||
l'indexer sparso."""
|
||||
import json, torch
|
||||
l'indexer sparso.
|
||||
|
||||
--fp8: salva i pesi come FP8 e4m3 + scale a blocchi 128x128 (layout del checkpoint reale
|
||||
GLM-5.2-FP8) invece di bf16, cosi' convert_fp8_to_int4.py puo' esercitare il path FP8->int4
|
||||
su un modello minuscolo. PRIMA di calcolare ref_glm.json fa il round-trip dei pesi per FP8
|
||||
(quant->dequant, copy_ nel modello): cosi' il riferimento riflette ESATTAMENTE il modello
|
||||
FP8 che il converter legge, non il modello bf16 a precisione piena. Default: bf16 (oracolo
|
||||
originale invariato).
|
||||
EN: --fp8 writes FP8 e4m3 + 128x128 block scale_inv (real GLM-5.2-FP8 layout) instead of bf16,
|
||||
EN: so convert_fp8_to_int4.py can run its FP8->int4 path on a tiny model. ref_glm.json is
|
||||
EN: computed AFTER the FP8 round-trip, so the reference matches exactly what the converter
|
||||
EN: ingests. Default: bf16 (original oracle unchanged)."""
|
||||
import json, sys, argparse
|
||||
from pathlib import Path
|
||||
import torch
|
||||
from transformers import GlmMoeDsaConfig, GlmMoeDsaForCausalLM
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent)) # importa glm_fp8_emit se lanciato da c/
|
||||
from glm_fp8_emit import (fp8_block_quantize, fp8_block_dequantize, keep_f32,
|
||||
save_fp8_safetensors, unfuse_experts)
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--fp8", action="store_true",
|
||||
help="salva in FP8 e4m3 + 128x128 block scale_inv (layout GLM-5.2-FP8) e "
|
||||
"calcola ref_glm.json sul modello dopo il round-trip FP8. "
|
||||
"EN: write FP8 e4m3 + block scale_inv, ref computed on FP8-rounded model")
|
||||
args = ap.parse_args()
|
||||
|
||||
torch.manual_seed(1234)
|
||||
|
||||
cfg = GlmMoeDsaConfig(
|
||||
@@ -53,6 +77,18 @@ with torch.no_grad():
|
||||
layer.mlp.gate.e_score_correction_bias.copy_(
|
||||
torch.linspace(-0.1, 0.1, cfg.n_routed_experts))
|
||||
|
||||
# --fp8: round-trip dei pesi quantizzabili per FP8 PRIMA di calcolare il riferimento,
|
||||
# cosi' ref_glm.json riflette esattamente il modello FP8 che il converter leggera'.
|
||||
# Norme/router/bias (keep_f32) restano a precisione piena. EN: --fp8: round-trip quantizable
|
||||
# weights through FP8 before computing the reference, so ref_glm.json matches the FP8 model.
|
||||
if args.fp8:
|
||||
with torch.no_grad():
|
||||
for n, p in model.named_parameters():
|
||||
if keep_f32(n, p) or p.dim() != 2:
|
||||
continue
|
||||
q, s = fp8_block_quantize(p)
|
||||
p.copy_(fp8_block_dequantize(q, s))
|
||||
|
||||
print("=== state_dict tensors (names used by the C loader) ===")
|
||||
for n, p in model.state_dict().items():
|
||||
print(f" {n:60s} {tuple(p.shape)}")
|
||||
@@ -73,7 +109,20 @@ with torch.no_grad():
|
||||
tf_pred = lg.argmax(-1).tolist()
|
||||
print("tf_pred:", tf_pred)
|
||||
|
||||
model.save_pretrained("glm_tiny", safe_serialization=True)
|
||||
# Unfuse experts AFTER reference generation (model needs fused weights for
|
||||
# forward/generate) but BEFORE saving — the real checkpoint and the converter
|
||||
# + C engine all expect per-expert 2-D gate_proj/up_proj/down_proj tensors.
|
||||
sd = model.state_dict()
|
||||
unfuse_experts(sd)
|
||||
|
||||
if args.fp8:
|
||||
n_fp8, n_tot = save_fp8_safetensors(sd, "glm_tiny/model.safetensors")
|
||||
print(f"\nsaved FP8: {n_fp8} e4m3 tensors (+{n_tot - n_fp8} scale_inv sidecars / f32) "
|
||||
f"-> glm_tiny/model.safetensors")
|
||||
else:
|
||||
from safetensors.torch import save_file
|
||||
save_file({k: v.contiguous() for k, v in sd.items()}, "glm_tiny/model.safetensors")
|
||||
json.dump(cfg.to_dict(), open("glm_tiny/config.json", "w"))
|
||||
json.dump({"prompt_ids": prompt, "full_ids": full, "tf_pred": tf_pred}, open("ref_glm.json", "w"))
|
||||
print("\nsaved: glm_tiny/ (weights + config) and ref_glm.json")
|
||||
print("saved: glm_tiny/ (weights + config) and ref_glm.json"
|
||||
+ (" [fp8]" if args.fp8 else ""))
|
||||
|
||||
+85
-12
@@ -106,33 +106,106 @@ def rotation(dim, device, seed=417):
|
||||
return q
|
||||
|
||||
|
||||
def quantize_param(w, bits, group, rot=False):
|
||||
def quantize_param(w, bits, group, rot=False, e8=""):
|
||||
if w.ndim == 3: # fused experts [E, in, out] -> move input last
|
||||
x = w.transpose(1, 2).contiguous()
|
||||
x = _rot_quant(x, bits, group) if rot else _quant_last_dim(x, bits, group)
|
||||
x = _rot_quant(x, bits, group, e8) if rot else _grid_or_e8(x, bits, group, e8)
|
||||
return x.transpose(1, 2).contiguous()
|
||||
if rot:
|
||||
return _rot_quant(w, bits, group)
|
||||
return _quant_last_dim(w, bits, group) # nn.Linear [out, in] -- input already last
|
||||
return _rot_quant(w, bits, group, e8)
|
||||
return _grid_or_e8(w, bits, group, e8) # nn.Linear [out, in] -- input already last
|
||||
|
||||
|
||||
def _rot_quant(x, bits, group):
|
||||
def _grid_or_e8(x, bits, group, e8):
|
||||
if e8:
|
||||
return _quant_e8(x.float(), group, ball=(e8 == "-e8"))
|
||||
return _quant_last_dim(x, bits, group)
|
||||
|
||||
|
||||
def _rot_quant(x, bits, group, e8=""):
|
||||
"""W -> Qn(W@Q) @ Q^T along the last (input) dim — see rotation() above."""
|
||||
q = rotation(x.shape[-1], x.device)
|
||||
return (_quant_last_dim(x.float() @ q, bits, group) @ q.T).contiguous()
|
||||
return (_grid_or_e8(x.float() @ q, bits, group, e8) @ q.T).contiguous()
|
||||
|
||||
|
||||
SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-rot)?(-nohead)?$")
|
||||
# --------------------------------------------------------------------------------------
|
||||
# E8 lattice quantization (#81 follow-up): the -rot schemes above are QuaRot (rotation +
|
||||
# uniform grid). QuIP#'s 2-bit result needs the second ingredient — an E8 lattice codebook
|
||||
# instead of the grid. E8 = D8 ∪ (D8 + 1/2), nearest point via Conway–Sloane: round every
|
||||
# coordinate, and if the sum is odd re-round the worst coordinate the other way; repeat on
|
||||
# the half-shifted copy and keep the closer of the two. `-e8` clamps points to |p|^2 <= 10,
|
||||
# the E8P ball QuIP# builds its 2^16 codebook from (2 bits/weight for 8-dim blocks);
|
||||
# `-e8u` leaves the lattice unbounded — an ideal-codebook upper bound, not a deployable rate.
|
||||
# Scale: per group, a small MSE search over multiples of the block RMS (absmax is the wrong
|
||||
# statistic for a lattice — the ball wants energy matched, not the peak).
|
||||
# --------------------------------------------------------------------------------------
|
||||
def _d8_nearest(y):
|
||||
f = torch.round(y)
|
||||
d = y - f
|
||||
odd = (f.sum(-1).long() & 1).bool()
|
||||
idx = d.abs().argmax(-1, keepdim=True)
|
||||
step = torch.where(d.gather(-1, idx) >= 0, 1.0, -1.0)
|
||||
flipped = f.gather(-1, idx) + step
|
||||
return f.scatter(-1, idx, torch.where(odd[..., None], flipped, f.gather(-1, idx)))
|
||||
|
||||
|
||||
def _e8_nearest(y):
|
||||
a = _d8_nearest(y)
|
||||
b = _d8_nearest(y - 0.5) + 0.5
|
||||
da = ((y - a) ** 2).sum(-1, keepdim=True)
|
||||
db = ((y - b) ** 2).sum(-1, keepdim=True)
|
||||
return torch.where(da <= db, a, b)
|
||||
|
||||
|
||||
def _e8_ball(y, r2=10.0):
|
||||
p = _e8_nearest(y)
|
||||
for _ in range(8): # shrink-and-requantize until inside
|
||||
n2 = (p ** 2).sum(-1, keepdim=True)
|
||||
over = n2 > r2 + 1e-6
|
||||
if not over.any():
|
||||
break
|
||||
y = torch.where(over, y * torch.sqrt(r2 / torch.clamp(n2, min=r2)) * 0.98, y)
|
||||
p = torch.where(over, _e8_nearest(y), p)
|
||||
return p
|
||||
|
||||
|
||||
def _quant_e8(x, group, ball):
|
||||
"""Blocks of 8 along the input dim; per-group scale by MSE search over RMS multiples."""
|
||||
if x.shape[-1] % 8:
|
||||
raise SystemExit(f"-e8 needs input dim divisible by 8 (got {x.shape[-1]})")
|
||||
g = group or x.shape[-1]
|
||||
if g % 8:
|
||||
raise SystemExit(f"-e8 group {g} must be a multiple of 8")
|
||||
shp = x.shape
|
||||
xg = x.reshape(-1, g) # [G, g]
|
||||
rms = torch.clamp(xg.pow(2).mean(-1, keepdim=True).sqrt(), min=1e-8)
|
||||
best_out, best_err = None, None
|
||||
for k in (0.5, 0.7, 0.9, 1.1, 1.4, 1.8, 2.4):
|
||||
s = rms * k
|
||||
yb = (xg / s).reshape(-1, g // 8, 8)
|
||||
p = _e8_ball(yb) if ball else _e8_nearest(yb)
|
||||
out = (p.reshape(-1, g) * s)
|
||||
err = (out - xg).pow(2).sum(-1, keepdim=True)
|
||||
if best_err is None:
|
||||
best_out, best_err = out, err
|
||||
else:
|
||||
take = err < best_err
|
||||
best_out = torch.where(take, out, best_out)
|
||||
best_err = torch.where(take, err, best_err)
|
||||
return best_out.reshape(shp)
|
||||
|
||||
|
||||
SCHEME_RE = re.compile(r"^int(2|3|4|8)(?:-g(\d+))?(-e8u?)?(-rot)?(-nohead)?$")
|
||||
|
||||
|
||||
def parse_scheme(name):
|
||||
"""'int4-g128-nohead' -> (bits=4, group=128, skip_head=True). 'fp16' -> None."""
|
||||
"""'int4-g128-nohead' -> (bits, group, e8, skip_head...). 'fp16' -> None."""
|
||||
if name == "fp16":
|
||||
return None
|
||||
m = SCHEME_RE.match(name)
|
||||
if not m:
|
||||
raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,3,4,8}}[-g<N>][-rot][-nohead])")
|
||||
return int(m.group(1)), int(m.group(2) or 0), bool(m.group(3)), bool(m.group(4))
|
||||
raise SystemExit(f"bad scheme '{name}' (expected fp16 | int{{2,3,4,8}}[-g<N>][-e8|-e8u][-rot][-nohead])")
|
||||
return int(m.group(1)), int(m.group(2) or 0), m.group(3) or "", bool(m.group(4)), bool(m.group(5))
|
||||
|
||||
|
||||
def is_router(name):
|
||||
@@ -152,7 +225,7 @@ def apply_scheme(model, scheme):
|
||||
spec = parse_scheme(scheme)
|
||||
if spec is None:
|
||||
return 0, 0, total
|
||||
bits, group, rot, skip_head = spec
|
||||
bits, group, e8, rot, skip_head = spec
|
||||
n = qp = 0
|
||||
with torch.no_grad():
|
||||
for name, p in model.named_parameters():
|
||||
@@ -160,7 +233,7 @@ def apply_scheme(model, scheme):
|
||||
continue
|
||||
if skip_head and is_head_or_embed(name):
|
||||
continue
|
||||
p.data.copy_(quantize_param(p.data.float(), bits, group, rot).to(p.dtype))
|
||||
p.data.copy_(quantize_param(p.data.float(), bits, group, rot, e8).to(p.dtype))
|
||||
n += 1
|
||||
qp += p.numel()
|
||||
return n, qp, total
|
||||
|
||||
Reference in New Issue
Block a user