diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fecc107ff..d52bbed27 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -402,6 +402,42 @@ jobs: tests.test_inefficiency.TinyEfficiencyTest.test_profile_phases_present_and_nonneg \ tests.test_inefficiency.TinyEfficiencyTest.test_disk_wait_not_dominant + # The token-exact parity gate from #7 (c/tests/test_cluster_sharding.py): the + # engine's local CPU run must already reproduce the transformers oracle + # token-exactly, and delegating the routed-expert shard to a cluster worker + # must not perturb a single position. It exercises the fmt=6 (E8/IQ3, + # rotation-bearing) path and its fmt=4 (grouped int4, no-rotation) control. + # + # Not in `make check`: the two fixtures are gitignored (regenerated, not + # committed -- see .gitignore) and need the pinned torch/transformers to + # build, the same rationale as inkling-oracle and efficiency above. The test + # skips cleanly when the fixtures are absent, so leaving fixture generation + # out of CI would turn the gate into a silent no-op on a clean host -- the + # worker path would never actually run. This job regenerates both fixtures + # first so the gate is exercised, not dormant. + # + # Socket safety: the test spawns a cluster worker subprocess and a + # coordinator, each binding a 127.0.0.1 port obtained from _free_port() + # (bind to port 0), so it is safe on a shared runner. + cluster-parity: + name: Cluster parity gate (fmt6/fmt4, token-exact) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: pip + cache-dependency-path: c/tools/oracle-requirements.txt + - name: Install torch (CPU) + transformers + run: pip install -r c/tools/oracle-requirements.txt + - name: Build colibri + run: make -C c colibri + - name: Generate the fmt6/fmt4 quantized fixtures + run: cd c && python3 tools/make_glm_oracle.py --fmt6 && python3 tools/make_glm_oracle.py --fmt4 + - name: Cluster parity gate (token-exact) + run: cd c && python3 -m unittest -v tests.test_cluster_sharding + engine-hip-syntax: name: HIP syntax check runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 5644439f9..982972c7f 100644 --- a/.gitignore +++ b/.gitignore @@ -82,6 +82,8 @@ c/glm_tiny/ c/glm_tiny_i2/ c/glm_tiny_i4/ c/glm_tiny_mix/ +c/glm_tiny_fmt6/ +c/glm_tiny_fmt4/ c/glm_bench_medium/ c/bench/ diff --git a/README.md b/README.md index 7afc1bef1..384ec3b63 100644 --- a/README.md +++ b/README.md @@ -162,8 +162,40 @@ the right experts get. It works because routing has measurable structure (see the [expert atlas](https://github.com/JustVugg/colibri/issues/175)) — and structure is cacheable. -The engine is a single C file (`c/glm.c`) plus small headers. No BLAS, no Python -at runtime, no GPU required. +The engine is a single C file (`c/colibri.c`) plus small headers. No BLAS, no +Python at runtime, no GPU required. + +### Local cluster mode + +The coordinator keeps token generation, routing, and KV state local while +disk-backed expert workers execute routed FFNs on other Macs. A layer's routed +batch-union is sent as one persistent TCP request, so a token does not incur one +round trip per expert. + +Start the optional registration service: + +```bash +./coli cluster coordinator --host 0.0.0.0 --port 8765 +``` + +On each worker, with the same converted model available locally: + +```bash +./coli cluster worker --model /nvme/glm52_i4 --port 9100 \ + --coordinator http://COORDINATOR:8765 --advertise-host WORKER_IP +``` + +Run the coordinator with discovery, or provide `--cluster-workers +HOST:PORT,...` for a static setup: + +```bash +./coli serve --model /nvme/glm52_i4 \ + --cluster-coordinator http://127.0.0.1:8765 +``` + +The transport is disabled unless workers are configured, so the existing +single-machine path remains unchanged. Dense-layer sharding and browser/WebGPU +workers are separate follow-up seams. ## How it works diff --git a/c/Makefile b/c/Makefile index 798d81fc8..3061bbb99 100644 --- a/c/Makefile +++ b/c/Makefile @@ -894,6 +894,9 @@ iobench$(EXE): iobench.c compat.h tests/test_serve_sentinel$(EXE): tests/test_serve_sentinel.c compat.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) +tests/test_cluster_protocol$(EXE): tests/test_cluster_protocol.c colibri.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h quant.h sample.h kv_persist.h telemetry.h route_trace.h + $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) + tests/test_ue8m0$(EXE): tests/test_ue8m0.c st.h json.h compat.h $(CC) $(CFLAGS) $< -o $@ $(LDFLAGS) @@ -1244,7 +1247,7 @@ install: colibri$(EXE) inkling$(EXE) kimi_k3$(EXE) olmoe$(EXE) \ $(INSTALL) -m 755 deepseek_v4$(EXE) $(DESTDIR)$(LIBEXECDIR)/deepseek_v4$(EXE); \ fi $(INSTALL) -m 644 family_registry.py resource_plan.py doctor.py autotune.py \ - openai_server.py v4_dsml.py version.py $(DESTDIR)$(LIBEXECDIR)/ + openai_server.py cluster.py v4_dsml.py version.py $(DESTDIR)$(LIBEXECDIR)/ $(INSTALL) -m 644 tools/*.py $(DESTDIR)$(LIBEXECDIR)/tools/ @# The dashboard is an optional build artifact (cd web && npm run build), so install @# it only when it exists. It goes NEXT TO openai_server.py, which probes ./web/dist. diff --git a/c/cluster.py b/c/cluster.py new file mode 100644 index 000000000..6f0e5a53c --- /dev/null +++ b/c/cluster.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Registration and discovery control plane for local expert workers.""" + +import argparse +import json +import os +import sys +import threading +import time +from http.server import ThreadingHTTPServer +from urllib.request import Request, urlopen + +import openai_server + + +PROTOCOL_VERSION = 1 + + +class ClusterRegistry: + def __init__(self, stale_after=30.0): + self.stale_after = float(stale_after) + self._nodes = {} + self._lock = threading.Lock() + + def register(self, node): + required = {"node_id", "host", "port", "role"} + missing = sorted(required - set(node)) + if missing: + raise ValueError("missing node fields: " + ", ".join(missing)) + port = int(node["port"]) + if not 1 <= port <= 65535: + raise ValueError("port must be between 1 and 65535") + role = str(node["role"]) + if role not in ("expert", "dense", "coordinator"): + raise ValueError("role must be expert, dense, or coordinator") + record = dict(node) + record.update(protocol_version=PROTOCOL_VERSION, port=port, last_seen=time.time()) + with self._lock: + self._nodes[str(node["node_id"])] = record + return record + + def heartbeat(self, node_id): + with self._lock: + node = self._nodes.get(str(node_id)) + if node is None: + raise KeyError(node_id) + node["last_seen"] = time.time() + return dict(node) + + def snapshot(self): + now = time.time() + with self._lock: + nodes = [dict(node) for node in self._nodes.values() + if now - node["last_seen"] <= self.stale_after] + nodes.sort(key=lambda node: (node["role"], node["node_id"])) + return {"protocol_version": PROTOCOL_VERSION, "nodes": nodes} + + def expert_endpoints(self): + return [f"{node['host']}:{node['port']}" + for node in self.snapshot()["nodes"] if node["role"] == "expert"] + + +class _Handler(openai_server.APIHandler): + # The original registry served HTTP/1.0: one request per connection, no + # keep-alive. Keep that wire shape; the shared APIHandler machinery (Host + # guard, body cap, deadline reader, connection bookkeeping) still applies. + protocol_version = "HTTP/1.0" + server_version = "colibri-cluster/1" + + def log_message(self, *_args): + return + + def _fail(self, error): + self.send_json(error.status, {"error": error.message}) + + def do_GET(self): # noqa: N802 - stdlib handler API + try: + self._check_host() + except openai_server.APIError as error: + self._fail(error) + return + if self.path in ("/health", "/v1/cluster/topology"): + self.send_json(200, self.server.registry.snapshot()) + return + self.send_json(404, {"error": "not found"}) + + def do_POST(self): # noqa: N802 - stdlib handler API + try: + self._check_host() + body = self.read_json() + if self.path == "/v1/cluster/register": + self.send_json(200, self.server.registry.register(body)) + elif self.path == "/v1/cluster/heartbeat": + self.send_json(200, self.server.registry.heartbeat(body["node_id"])) + else: + self.send_json(404, {"error": "not found"}) + except openai_server.APIError as error: + self._fail(error) + except (KeyError, ValueError, TypeError) as error: + self.send_json(400, {"error": str(error)}) + + +class ClusterServer(openai_server.APIServer): + def __init__(self, address, registry, allowed_hosts=()): + ThreadingHTTPServer.__init__(self, address, _Handler) + self.registry = registry + # Shared APIHandler plumbing reads these off the server. The registry has + # no API key or CORS surface, and its Host guard stays loopback + bind + # address (the same default openai_server enforces). A cross-host worker + # registers from a LAN IP, so the operator opts those hosts in via + # --allowed-host (the same #597 escape hatch coli serve exposes); the + # default stays loopback + bind address. + self.cors_origins = () + self.allowed_hosts = tuple( + h.strip().lower() for h in allowed_hosts if h and h.strip()) + # Connection-cap bookkeeping for the inherited APIServer.process_request / + # _release / close_request. There is no engine or model to carry, so + # initialise the tracking state directly rather than APIServer.__init__. + self._conn_lock = threading.Lock() + self._conn_live = 0 + self._conn_by_ip = {} + self._conn_owner = {} + + +def _endpoint(coordinator, path): + return coordinator.rstrip("/") + "/v1/cluster/" + path + + +def serve(host="127.0.0.1", port=8765, stale_after=30.0, allowed_hosts=()): + if allowed_hosts and "*" in allowed_hosts: + print("WARNING: --allowed-host '*' accepts ANY Host header " + "(DNS-rebinding guard disabled)", file=sys.stderr) + server = ClusterServer((host, port), ClusterRegistry(stale_after), allowed_hosts) + print(f"colibri cluster coordinator listening on http://{host}:{port}", flush=True) + try: + server.serve_forever() + finally: + server.server_close() + + +def register(coordinator, node): + request = Request(_endpoint(coordinator, "register"), + data=json.dumps(node).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urlopen(request, timeout=5) as response: + return json.load(response) + + +def heartbeat(coordinator, node_id): + request = Request(_endpoint(coordinator, "heartbeat"), + data=json.dumps({"node_id": node_id}).encode(), + headers={"Content-Type": "application/json"}, method="POST") + with urlopen(request, timeout=5) as response: + return json.load(response) + + +def discover_workers(coordinator): + with urlopen(_endpoint(coordinator, "topology"), timeout=5) as response: + topology = json.load(response) + return [f"{node['host']}:{int(node['port'])}" + for node in topology.get("nodes", []) if node.get("role") == "expert"] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8765) + parser.add_argument("--stale-after", type=float, default=30.0) + parser.add_argument("--allowed-host", action="append", + default=[h.strip() for h in os.environ.get("COLI_ALLOWED_HOSTS", "").split(",") if h.strip()], + help="additional Host header accepted by the DNS-rebinding guard; repeat as needed") + args = parser.parse_args() + serve(args.host, args.port, args.stale_after, args.allowed_host) + + +if __name__ == "__main__": + main() diff --git a/c/coli b/c/coli index e0c63d89c..697ca8b29 100755 --- a/c/coli +++ b/c/coli @@ -323,17 +323,32 @@ def engine_for(model): local = os.path.join(HERE, name + _EXE) return local if os.path.exists(local) else os.path.join(_LIBEXEC, name + _EXE) -def need_model(model, engine=None): +def require_model(model, file="tokenizer.json", engine=None, target="", build="coli build"): + """Shared model-directory / required-file / engine presence check. + + `need_model` and `need_worker_model` differ only in which file they require + (tokenizer.json vs config.json) and whether non-GLM families are allowed + (expert workers live in the GLM engine only). The exit wording and + the check order are shared so the two validators cannot drift apart. + """ if not model: sys.exit(f"{C.yel}no model directory given.{C.r}\n {NO_MODEL_HINT}") if not os.path.isdir(model): sys.exit(f"{C.yel}model not found:{C.r} {model}\n set COLI_MODEL or use --model") - if not os.path.exists(os.path.join(model,"tokenizer.json")): - sys.exit(f"{C.yel}tokenizer.json is missing from {model}{C.r}") + if file and not os.path.exists(os.path.join(model, file)): + sys.exit(f"{C.yel}{file} is missing from {model}{C.r}") + if engine and not os.path.exists(engine): + label = f"{target} " if target else "" + sys.exit(f"{C.yel}{label}engine is not built.{C.r} Run: {build}") + +def need_model(model, engine=None): + # Directory/file checks come FIRST: resolving the family reads config.json, + # and a missing model dir must say "model not found", not "unsupported model". + require_model(model, file="tokenizer.json") engine = engine or engine_for(model) - if not os.path.exists(engine): - target = resolve_model(model).descriptor.build_target - sys.exit(f"{C.yel}{target} engine is not built.{C.r} Run: make -C c {target}") + target = resolve_model(model).descriptor.build_target + require_model(model, file="tokenizer.json", engine=engine, target=target, + build=f"make -C c {target}") # One-shot runs keep the historic 1024. Interactive sessions get 16384, because # the browser and the TUI both offer a per-request limit and the server treats @@ -455,6 +470,17 @@ def dsv4_cuda_available(model=None): dll = os.path.join(os.path.dirname(eng), "coli_cuda_dsv4.dll") return os.path.exists(dll) +def need_worker_model(model): + require_model(model, file="config.json") + family = resolve_model(model).descriptor + if family.id != "glm": + sys.exit(f"{C.yel}cluster expert workers currently support the GLM engine only;{C.r} " + f"{family.display_name} models cannot serve experts") + engine = engine_for(model) + require_model(model, file="config.json", engine=engine, target=family.build_target, + build=f"make -C c {family.build_target}") + return engine + def cuda_binary(): if not os.path.exists(GLM): return False if sys.platform == "linux": @@ -555,6 +581,8 @@ def env_for(a): e.setdefault("PIPE", "1") e.setdefault("PILOT_REAL", "1") e["COLI_POLICY"]=a.policy + if getattr(a, "cluster_workers", None): e["CLUSTER_WORKERS"] = a.cluster_workers + if getattr(a, "cluster_coordinator", None): e["CLUSTER_COORDINATOR"] = a.cluster_coordinator if a.ram: e["RAM_GB"]=str(a.ram) e["NGEN"]=str(ngen_for(a)) if a.topp: e["TOPP"]=str(a.topp) @@ -1437,9 +1465,20 @@ def cmd_serve(a): model_id=a.model_id or family.default_model_id try: if a.temp is not None: os.environ["COLI_TEMP"] = str(a.temp) + env=_serve_engine_env(a,arch) + if a.cluster_coordinator and not a.cluster_workers: + from cluster import discover_workers + try: + workers=discover_workers(a.cluster_coordinator) + except OSError as error: + sys.exit(f"{C.yel}cannot discover cluster workers:{C.r} {error}") + if not workers: + sys.exit(f"{C.yel}cluster coordinator has no live expert workers{C.r}") + env["CLUSTER_WORKERS"] = ",".join(workers) + print(f" {C.dim}[CLUSTER] discovered {len(workers)} expert worker(s){C.r}", file=sys.stderr) openai_server.serve(a.model,a.host,a.port,model_id,a.api_key, a.cap,ngen_for(a,interactive=True,family=family),engine, - _serve_engine_env(a,arch),a.cors_origin, + env,a.cors_origin, a.max_queue,a.queue_timeout,a.kv_slots,allowed_hosts=a.allowed_host, family=family) finally: @@ -1476,6 +1515,34 @@ def _pid_alive(pid): finally: k32.CloseHandle(handle) +def cmd_cluster_coordinator(a): + from cluster import serve + serve(a.host, a.port, a.stale_after, a.allowed_host) + +def cmd_cluster_worker(a): + engine=need_worker_model(a.model) + coordinator=a.coordinator or os.environ.get("CLUSTER_COORDINATOR") + host=a.advertise_host or os.environ.get("CLUSTER_ADVERTISE_HOST", "127.0.0.1") + node_id=a.node_id or f"{host}:{a.port}" + stop_heartbeat=threading.Event() + if coordinator: + from cluster import heartbeat, register + register(coordinator, {"node_id":node_id, "host":host, "port":a.port, + "role":"expert", "layers":a.layers}) + def keep_registered(): + while not stop_heartbeat.wait(10): + try: heartbeat(coordinator, node_id) + except OSError: pass + threading.Thread(target=keep_registered, name="colibri-cluster-heartbeat", daemon=True).start() + e=env_for(a) + e.update({"EXPERT_WORKER":"1", "CLUSTER_WORKER_PORT":str(a.port), + "COLI_MMAP":os.environ.get("COLI_MMAP", "1")}) + print(f" {C.dim}[CLUSTER] expert worker · {host}:{a.port} · layers {a.layers}{C.r}") + try: + return subprocess.call([engine,str(a.cap),str(a.ebits),str(a.dbits)],env=e) + finally: + stop_heartbeat.set() + def cmd_stop(a): """Shut down a running `coli serve` AND its engine — one command, no pkill. The engine re-execs itself for OMP tuning, so its process is named `exe`, @@ -1677,6 +1744,10 @@ def main(): common.add_argument("--ngen", type=int, default=None) # rete di sicurezza: la fine vera la decidono gli stop token common.add_argument("--topp", type=float, default=0); common.add_argument("--topk", type=int, default=0) common.add_argument("--temp", type=float, default=None) # temperatura token (0=greedy, default 1.0+nucleus .95) + common.add_argument("--cluster-workers", default=os.environ.get("CLUSTER_WORKERS"), + help="comma-separated expert workers, host:port,...") + common.add_argument("--cluster-coordinator", default=os.environ.get("CLUSTER_COORDINATOR"), + help="control-plane URL used to discover expert workers") ap=argparse.ArgumentParser(prog="coli", parents=[common], description="colibri — run GLM-5.2 locally") ap.add_argument("--version", action="version", version=f"colibri {_version}") sub=ap.add_subparsers(dest="cmd") @@ -1725,6 +1796,20 @@ def main(): ps.add_argument("--allowed-host",action="append", default=[h.strip() for h in os.environ.get("COLI_ALLOWED_HOSTS","").split(",") if h.strip()], help="additional Host header accepted by the DNS-rebinding guard; repeat as needed") + pcluster=sub.add_parser("cluster", help="run the local-cluster control plane or an expert worker") + cluster_sub=pcluster.add_subparsers(dest="cluster_cmd") + pcoord=cluster_sub.add_parser("coordinator", help="serve worker registration and discovery") + pcoord.add_argument("--host",default="127.0.0.1"); pcoord.add_argument("--port",type=int,default=8765) + pcoord.add_argument("--stale-after",type=float,default=30.0) + pcoord.add_argument("--allowed-host",action="append", + default=[h.strip() for h in os.environ.get("COLI_ALLOWED_HOSTS","").split(",") if h.strip()], + help="additional Host header accepted by the DNS-rebinding guard; repeat as needed") + pworker=cluster_sub.add_parser("worker", parents=[common], help="serve disk-backed expert compute") + pworker.add_argument("--port",type=int,default=int(os.environ.get("CLUSTER_WORKER_PORT","9100"))) + pworker.add_argument("--coordinator",default=os.environ.get("CLUSTER_COORDINATOR")) + pworker.add_argument("--node-id",default=None); pworker.add_argument("--advertise-host",default=None) + pworker.add_argument("--layers",default="all",help="topology label, e.g. 0-37") + pworker.add_argument("--ebits",type=int,default=8); pworker.add_argument("--dbits",type=int,default=8) pst=sub.add_parser("stop", parents=[common], help="shut down a running coli serve and its engine") pst.add_argument("--port",type=int,default=8000); pst.add_argument("--dry-run",action="store_true") pw=sub.add_parser("web", parents=[common], help="serve + open the dashboard in a browser") @@ -1758,6 +1843,9 @@ def main(): "doctor":cmd_doctor,"tune":cmd_tune, "run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench, "convert":cmd_convert,"web":cmd_web}.get(a.cmd) + if a.cmd=="cluster": + if a.cluster_cmd=="coordinator": handler=cmd_cluster_coordinator + elif a.cluster_cmd=="worker": handler=cmd_cluster_worker if handler: try: sys.exit(handler(a) or 0) diff --git a/c/colibri.c b/c/colibri.c index 385a87d51..33b36bf5e 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -30,6 +30,10 @@ #include #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include /* select() serve-loop polling (#68); not on native MinGW */ +#include +#include +#include +#include #endif #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include @@ -128,6 +132,11 @@ static const float *g_pre_sh; /* routing precalcolata dalla GPU (Metal layer CB o device router CUDA, #431): * moe() la usa e salta la FASE A. NULL = router su CPU. */ static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; +#if !defined(_WIN32) +typedef struct { int fd; char host[128]; int port; } ClusterWorker; +static ClusterWorker g_cluster_workers[16]; +static int g_cluster_n; +#endif #ifdef __APPLE__ #include /* host_statistics64: MemAvailable di macOS */ #endif @@ -1236,12 +1245,14 @@ static inline float siluf(float x){ return x/(1.f+expf(-x)); } * SEMPRE — e' la forma CANONICA, pura funzione del fmt del down-tensor. Perche' resta * byte-identica al vecchio inline: (a) i fallback device-lost Vulkan ricaricano solo * expert che erano in registry (vk_reg_at), e vk_registry_fill ammette solo fmt 2/4/5 — - * li' d->fmt!=6 e la rotazione e' un no-op; (b) l'oracolo e ogni modello single-format - * (l'output di una conversione normale) non hanno mai un expert fmt=6 accanto a fmt 2/4/5. + * li' d->fmt!=6 e la rotazione e' un no-op; (b) l'oracolo e' single-format, e una + * conversione normale single-pass produce un modello single-format: le due copie che + * POSSONO vedere fmt=6 lo fanno solo su un container mixed-format, che il tooling + * normale non emette mai. * Le DUE copie che saltavano la rotazione e potevano davvero vedere fmt=6 — la CPU-share - * del blocco Vulkan e il fallback CUDA — ora seguono il canonico: su un container - * mixed-format (fmt per-tensor da qt_resolve_fmt, nessuna uniformita' tra expert, vedi il - * commento MB_BUILD) e' la correzione di un bug latente, NON un no-op. */ + * del blocco Vulkan e il fallback CUDA early-issued — ora seguono il canonico: su un + * container mixed-format (fmt per-tensor da qt_resolve_fmt, nessuna uniformita' tra + * expert, vedi il commento MB_BUILD) e' la correzione di un bug latente, NON un no-op. */ static void expert_ffn(float *hh, float *gg, float *uu, const float *xg, QT *g, QT *u, QT *d, int nr, int I){ expert_gate_up(gg,uu,xg,g,u,nr); for(int64_t z=0;z<(int64_t)nr*I;z++) gg[z]=siluf(gg[z])*uu[z]; @@ -2623,6 +2634,211 @@ static int expert_load(Model *m, int layer, int eid, ESlot *s, int fatal, int de return rc; } +#if !defined(_WIN32) +/* Expert-worker protocol. Headers use network-order u32 values; activation + * bytes remain raw little-endian f32, matching the native engine ABI. One + * request contains the routed batch-union for a layer. */ +#define COLI_CLUSTER_MAGIC "COLIEX01" +#define COLI_CLUSTER_VERSION 1u +static int cluster_io(int fd, void *buf, size_t n, int write_mode){ + char *p=(char*)buf; + while(n){ + ssize_t r=write_mode?send(fd,p,n,0):recv(fd,p,n,MSG_WAITALL); + if(r<=0){ if(r<0&&errno==EINTR) continue; return -1; } + p+=r; n-=(size_t)r; + } + return 0; +} +static int cluster_u32(int fd, uint32_t *v, int write_mode){ + uint32_t x=write_mode?htonl(*v):0; + if(cluster_io(fd,write_mode?(void*)&x:(void*)v,sizeof(x),write_mode)) return -1; + if(!write_mode) *v=ntohl(*v); + return 0; +} +static int cluster_connect_one(const char *spec, ClusterWorker *out){ + char copy[256]; strncpy(copy,spec,sizeof(copy)-1); copy[sizeof(copy)-1]=0; + char *colon=strrchr(copy,':'); if(!colon||colon==copy||!colon[1]) return -1; + *colon=0; int port=atoi(colon+1); if(port<1||port>65535) return -1; + char portbuf[16]; snprintf(portbuf,sizeof(portbuf),"%d",port); + struct addrinfo hint={0},*ai=NULL; hint.ai_socktype=SOCK_STREAM; + if(getaddrinfo(copy,portbuf,&hint,&ai)!=0) return -1; + int fd=-1; + for(struct addrinfo *p=ai;p;p=p->ai_next){ + fd=socket(p->ai_family,p->ai_socktype,p->ai_protocol); + if(fd<0) continue; + if(connect(fd,p->ai_addr,p->ai_addrlen)==0) break; + close(fd); fd=-1; + } + freeaddrinfo(ai); if(fd<0) return -1; + out->fd=fd; + size_t hostlen=strlen(copy); if(hostlen>=sizeof(out->host)) hostlen=sizeof(out->host)-1; + memcpy(out->host,copy,hostlen); out->host[hostlen]=0; + out->port=port; return 0; +} +static void cluster_close_all(void){ + for(int i=0;i=0) close(g_cluster_workers[i].fd); + g_cluster_n=0; +} +static void cluster_init(void){ + const char *list=getenv("CLUSTER_WORKERS"); if(!list||!*list) return; + char *copy=strdup(list),*save=NULL; + for(char *tok=strtok_r(copy,",",&save);tok&&g_cluster_n<16;tok=strtok_r(NULL,",",&save)){ + while(*tok==' '||*tok=='\t') tok++; + if(!cluster_connect_one(tok,&g_cluster_workers[g_cluster_n])) g_cluster_n++; + else fprintf(stderr,"[CLUSTER] cannot connect to expert worker %s\n",tok); + } + free(copy); + if(g_cluster_n<1){ fprintf(stderr,"[CLUSTER] no expert workers reachable\n"); exit(1); } + fprintf(stderr,"[CLUSTER] coordinator connected to %d expert worker(s)\n",g_cluster_n); +} +typedef struct { int eid,nr; int *rows; float *weights,*inputs; } ClusterItem; +static int cluster_item(const int *idxs,const float *ws,const int *keff,int K,int S, + int eid,ClusterItem *it,int D,const float *x){ + it->eid=eid; it->nr=0; + for(int s=0;snr++; break; } + if(!it->nr) return 0; + it->rows=malloc((size_t)it->nr*sizeof(int)); + it->weights=malloc((size_t)it->nr*sizeof(float)); + it->inputs=falloc((int64_t)it->nr*D); int r=0; + for(int s=0;srows[r]=s; it->weights[r]=ws[(int64_t)s*K+k]; + memcpy(it->inputs+(int64_t)r*D,x+(int64_t)s*D,(size_t)D*sizeof(float)); r++; break; + } + return 1; +} +static void cluster_item_free(ClusterItem *it){ free(it->rows); free(it->weights); free(it->inputs); memset(it,0,sizeof(*it)); } +static void cluster_moe_batch(Model *m,int layer,float *x,int S,float *out, + const int *idxs,const float *ws,const int *keff,int K, + const int *uniq,int base,int nb){ + int D=m->c.hidden; + for(int wi=0;wifd,(void*)COLI_CLUSTER_MAGIC,8,1)) goto fail; + v=COLI_CLUSTER_VERSION; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)layer; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)D; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)m->c.moe_inter; if(cluster_u32(w->fd,&v,1)) goto fail; + v=(uint32_t)n; if(cluster_u32(w->fd,&v,1)) goto fail; + for(int j=0;jfd,&v,1)) goto fail; + v=(uint32_t)items[j].nr; if(cluster_u32(w->fd,&v,1)) goto fail; + if(cluster_io(w->fd,items[j].inputs,(size_t)items[j].nr*D*sizeof(float),1)) goto fail; + } + if(cluster_io(w->fd,magic,8,0)||memcmp(magic,COLI_CLUSTER_MAGIC,8)) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=COLI_CLUSTER_VERSION) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=0) goto fail; + if(cluster_u32(w->fd,&v,0)||v!=(uint32_t)n) goto fail; + for(int j=0;jfd,&eid,0)||cluster_u32(w->fd,&nr,0) || + eid!=(uint32_t)items[j].eid || nr!=(uint32_t)items[j].nr) goto fail; + float *y=falloc((int64_t)nr*D); + if(cluster_io(w->fd,y,(size_t)nr*D*sizeof(float),0)){ free(y); goto fail; } + for(uint32_t r=0;rhost,w->port,layer); + exit(1); + } +} +typedef struct { int eid,nr; float *inputs; } ClusterRequestItem; +/* Mirror expert_load_impl's resolution so the worker can NAME the tensor it can't + * resolve instead of dying inside a fatal path (st_die_missing / st_read_* exit(1)). + * Returns 1 and fills `out` with the first missing expert tensor — quantized: a + * gate/up/down weight or its .qs sidecar; unquantized: the full weight — else 0. */ +static int worker_expert_missing_tensor(Model *m,int layer,int eid,char *out,size_t outsz){ + static const char *suf[3]={"gate_proj","up_proj","down_proj"}; + char nm[288], qn[320]; + snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,eid); + snprintf(qn,sizeof(qn),"%s.qs",nm); + if(!st_has(&m->S,qn)){ /* unquantized: full weights only */ + for(int k=0;k<3;k++){ + snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.%s.weight",layer,eid,suf[k]); + if(!st_has(&m->S,nm)){ snprintf(out,outsz,"%s",nm); return 1; } + } + return 0; + } + for(int k=0;k<3;k++){ /* quantized: weight + .qs sidecar */ + snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.%s.weight",layer,eid,suf[k]); + snprintf(qn,sizeof(qn),"%s.qs",nm); + if(!st_has(&m->S,nm)||!st_has(&m->S,qn)){ snprintf(out,outsz,"%s",nm); return 1; } + } + return 0; +} +static int cluster_worker_run(const char *snap,int port,int ebits,int dbits){ + Model m; memset(&m,0,sizeof(m)); m.ebits=ebits; m.dbits=dbits; load_cfg(&m.c,snap); + { const char *xd=getenv("COLI_MODEL_DIRS"); /* SPLIT: expert shards spread across N drives */ + st_init_multi(&m.S,snap,(xd&&*xd)?xd:NULL); } + int nr_layers=m.c.n_layers+1; ESlot *cache=calloc((size_t)nr_layers,sizeof(ESlot)); + for(int i=0;i=(uint32_t)nr_layers||n<1||n>64){ + close(cfd); cfd=-1; break; + } + ClusterRequestItem *items=calloc(n,sizeof(*items)); int bad=0; + for(uint32_t j=0;j=(uint32_t)m.c.n_experts||nr<1||nr>65536){bad=1;break;} + items[j].eid=(int)eid; items[j].nr=(int)nr; items[j].inputs=falloc((int64_t)nr*D); + if(cluster_io(cfd,items[j].inputs,(size_t)nr*D*sizeof(float),0)){bad=1;break;} + } + if(bad){ for(uint32_t j=0;jeid!=items[j].eid || !slot->slab){ + char miss[288]; + if(worker_expert_missing_tensor(&m,(int)layer,items[j].eid,miss,sizeof(miss))){ + fprintf(stderr,"[CLUSTER] worker cannot resolve expert tensor %s (layer %d, expert %d)\n", + miss,(int)layer,items[j].eid); + bad=1; break; + } + if(expert_load(&m,(int)layer,items[j].eid,slot,0,0)){bad=1;break;} + } + int rows=items[j].nr; float *g=falloc((int64_t)rows*I),*u=falloc((int64_t)rows*I),*y=falloc((int64_t)rows*D); + /* fmt=6: the gate/up input is per-item here (never reused after this + * expert), so rotate it in place — the same Q^T x that moe() applies + * once per layer via E8_XE. The down-input rotation lives in expert_ffn. */ + if(slot->g.fmt==6) e8_rot_rows(items[j].inputs,rows,D); + expert_ffn(y,g,u,items[j].inputs,&slot->g,&slot->u,&slot->d,rows,I); + v=(uint32_t)items[j].eid; if(cluster_u32(cfd,&v,1)){bad=1;free(g);free(u);free(y);break;} + v=(uint32_t)rows; if(cluster_u32(cfd,&v,1)||cluster_io(cfd,y,(size_t)rows*D*sizeof(float),1)){bad=1;free(g);free(u);free(y);break;} + free(g);free(u);free(y); + } + for(uint32_t j=0;j=0)close(cfd); + } + close(fd); return 0; +} +#endif + #ifdef __linux__ /* io_uring expert batches. One owner prepares all reads for a block, submits * them in one syscall, and reaps CQEs on demand. The kernel, rather than a set @@ -4428,6 +4644,12 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int int shared_on_gpu=0; (void)shared_on_gpu; /* set by the Metal path when Phase E was fused */ for(int base=0;base2?atoi(argv[2]):8; int dbits= argc>3?atoi(argv[3]):ebits; +#if !defined(_WIN32) + if(getenv("EXPERT_WORKER")){ + int port=getenv("CLUSTER_WORKER_PORT")?atoi(getenv("CLUSTER_WORKER_PORT")):9100; + if(port<1||port>65535){fprintf(stderr,"CLUSTER_WORKER_PORT must be 1..65535\n");return 2;} + return cluster_worker_run(snap,port,ebits,dbits); + } +#else + if(getenv("EXPERT_WORKER")){ + fprintf(stderr,"[CLUSTER] expert workers are not supported on Windows yet\n"); return 2; + } +#endif int kv_limit=(getenv("SERVE_BATCH")&&atoi(getenv("SERVE_BATCH")))?512:16; if(getenv("SERVE") && (kv_slot_count()<1 || kv_slot_count()>kv_limit)){ fprintf(stderr,"KV_SLOTS must be between 1 and %d\n",kv_limit); return 2; @@ -9649,6 +9882,12 @@ int main(int argc, char **argv){ fprintf(stderr,"METAL: fast SSD (%.1f GB/s) — page cache favored, expert cache minimal (cap 1); override with --cap\n", coli_ssd_gbs); printf("== GLM C engine (glm_moe_dsa), cache=%d experts/layer | experts@%d-bit dense@%d-bit | idot: " IDOT_KERNEL " ==\n", cap, ebits, dbits); g_mem_avail_boot = mem_available_gb(); +#if !defined(_WIN32) + if(getenv("CLUSTER_WORKERS") && *getenv("CLUSTER_WORKERS")){ + cluster_init(); + atexit(cluster_close_all); + } +#endif Model m; double t0=now_s(); model_init(&m,snap,cap,ebits,dbits); if(!g_direct_heat_explicit){ /* COLI_DISKCLASS_WINDOW default, needs m.c (topk/n_layers) */ /* CURRENT-STATE CALIBRATION: the "8" multiplier (recency window ~= the last 8 diff --git a/c/tests/test_cluster.py b/c/tests/test_cluster.py new file mode 100644 index 000000000..0f84f8654 --- /dev/null +++ b/c/tests/test_cluster.py @@ -0,0 +1,48 @@ +import json +import sys +import threading +import unittest +from http.client import HTTPConnection +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from cluster import ClusterRegistry, ClusterServer, PROTOCOL_VERSION + + +class ClusterRegistryTests(unittest.TestCase): + def test_registers_and_discovers_expert_nodes(self): + registry = ClusterRegistry() + registry.register({"node_id": "mac-a", "host": "10.0.0.2", "port": 9100, + "role": "expert", "layers": "all"}) + registry.register({"node_id": "mac-b", "host": "10.0.0.3", "port": 9101, + "role": "dense", "layers": "38-75"}) + self.assertEqual(registry.expert_endpoints(), ["10.0.0.2:9100"]) + self.assertEqual(registry.snapshot()["protocol_version"], PROTOCOL_VERSION) + + def test_http_topology_registration_and_heartbeat(self): + server = ClusterServer(("127.0.0.1", 0), ClusterRegistry()) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + conn = HTTPConnection(*server.server_address) + body = json.dumps({"node_id": "mac-a", "host": "127.0.0.1", + "port": 9100, "role": "expert"}) + conn.request("POST", "/v1/cluster/register", body, + {"Content-Type": "application/json"}) + self.assertEqual(conn.getresponse().status, 200) + conn.request("POST", "/v1/cluster/heartbeat", json.dumps({"node_id": "mac-a"}), + {"Content-Type": "application/json"}) + self.assertEqual(conn.getresponse().status, 200) + conn.request("GET", "/v1/cluster/topology") + response = conn.getresponse() + self.assertEqual(response.status, 200) + self.assertEqual(len(json.loads(response.read())["nodes"]), 1) + conn.close() + finally: + server.shutdown() + server.server_close() + + +if __name__ == "__main__": + unittest.main() diff --git a/c/tests/test_cluster_protocol.c b/c/tests/test_cluster_protocol.c new file mode 100644 index 000000000..1dc4bdf51 --- /dev/null +++ b/c/tests/test_cluster_protocol.c @@ -0,0 +1,97 @@ +/* COLIEX01 cluster wire contract: network-order headers, raw f32 activations, + * and one request/response over a socketpair. No model fixture required. */ +#define main coli_engine_main_unused +#include "../colibri.c" +#undef main + +#include + +#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) +typedef struct { int fd; int failed; } ClusterProtocolArgs; + +static void *cluster_protocol_worker(void *opaque) +{ + ClusterProtocolArgs *args = opaque; + int fd = args->fd; + char magic[8]; + uint32_t version, layer, D, I, n, eid, nr; + float input[6], output[6]; + + /* Request: magic(8) version layer D I n, then eid nr and nr*D inputs. */ + if (cluster_io(fd, magic, sizeof(magic), 0) || memcmp(magic, COLI_CLUSTER_MAGIC, 8) || + cluster_u32(fd, &version, 0) || version != COLI_CLUSTER_VERSION || + cluster_u32(fd, &layer, 0) || layer != 7 || + cluster_u32(fd, &D, 0) || D != 3 || + cluster_u32(fd, &I, 0) || I != 5 || + cluster_u32(fd, &n, 0) || n != 1 || + cluster_u32(fd, &eid, 0) || eid != 42 || + cluster_u32(fd, &nr, 0) || nr != 2 || + cluster_io(fd, input, sizeof(input), 0)) { + args->failed = 1; + return NULL; + } + for (int i = 0; i < 6; i++) output[i] = input[i] * 2.0f; + + /* Response: magic(8) version status(0) n, then eid nr and nr*D outputs. */ + version = COLI_CLUSTER_VERSION; + if (cluster_io(fd, (void *)COLI_CLUSTER_MAGIC, 8, 1) || + cluster_u32(fd, &version, 1) || + cluster_u32(fd, &(uint32_t){0}, 1) || + cluster_u32(fd, &n, 1) || + cluster_u32(fd, &eid, 1) || + cluster_u32(fd, &nr, 1) || + cluster_io(fd, output, sizeof(output), 1)) + args->failed = 1; + return NULL; +} + +static void test_wire_round_trip(void) +{ + int sockets[2]; + assert(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0); + ClusterProtocolArgs args = {sockets[1], 0}; + pthread_t thread; + assert(pthread_create(&thread, NULL, cluster_protocol_worker, &args) == 0); + + uint32_t value; + float input[6] = {1.0f, -2.0f, 0.5f, 3.0f, -4.0f, 0.25f}, output[6] = {0}; + + assert(cluster_io(sockets[0], (void *)COLI_CLUSTER_MAGIC, 8, 1) == 0); + value = COLI_CLUSTER_VERSION; assert(cluster_u32(sockets[0], &value, 1) == 0); + value = 7; assert(cluster_u32(sockets[0], &value, 1) == 0); /* layer */ + value = 3; assert(cluster_u32(sockets[0], &value, 1) == 0); /* D */ + value = 5; assert(cluster_u32(sockets[0], &value, 1) == 0); /* moe_inter */ + value = 1; assert(cluster_u32(sockets[0], &value, 1) == 0); /* n */ + value = 42; assert(cluster_u32(sockets[0], &value, 1) == 0); /* eid */ + value = 2; assert(cluster_u32(sockets[0], &value, 1) == 0); /* nr */ + assert(cluster_io(sockets[0], input, sizeof(input), 1) == 0); + + char magic[8]; + assert(cluster_io(sockets[0], magic, 8, 0) == 0); + assert(memcmp(magic, COLI_CLUSTER_MAGIC, 8) == 0); + value = 0; assert(cluster_u32(sockets[0], &value, 0) == 0 && value == COLI_CLUSTER_VERSION); + value = 1; assert(cluster_u32(sockets[0], &value, 0) == 0 && value == 0); /* status */ + value = 0; assert(cluster_u32(sockets[0], &value, 0) == 0 && value == 1); /* n */ + value = 0; assert(cluster_u32(sockets[0], &value, 0) == 0 && value == 42); /* eid */ + value = 0; assert(cluster_u32(sockets[0], &value, 0) == 0 && value == 2); /* nr */ + assert(cluster_io(sockets[0], output, sizeof(output), 0) == 0); + assert(output[0] == 2.0f && output[1] == -4.0f && output[2] == 1.0f && + output[3] == 6.0f && output[4] == -8.0f && output[5] == 0.5f); + + assert(pthread_join(thread, NULL) == 0); + assert(args.failed == 0); + close(sockets[0]); + close(sockets[1]); +} +#endif + +int main(void) +{ +#if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) + test_wire_round_trip(); + puts("cluster protocol tests: ok"); +#else + puts("cluster protocol tests: skipped on Windows"); +#endif + return 0; +} diff --git a/c/tests/test_cluster_sharding.py b/c/tests/test_cluster_sharding.py new file mode 100644 index 000000000..2b1118504 --- /dev/null +++ b/c/tests/test_cluster_sharding.py @@ -0,0 +1,212 @@ +"""Token-exact parity gate: local vs cluster-delegated expert sharding (#7). + +Mirrors test_dense_sharding.py (#550) but STRONGER. #550 asserted only that the +baseline and delegated runs produce *equal* teacher-forcing signatures. This gate +asserts BOTH (a) identical token signatures AND (b) ZERO `[ORACLE] mismatch` +lines in BOTH runs — i.e. the local CPU run must already reproduce the +transformers oracle token-exactly, and delegating the routed-expert shard to a +cluster worker must not perturb a single position. + +The engine routes ONLY routed experts to the worker: `cluster_moe_batch` shards +`uniq[]`, which `moe()` builds exclusively from the per-position routed expert +ids (`idxs[]`); the shared expert is computed locally on the coordinator in both +the baseline and delegated runs, so it is not part of the shard under test. + +Fixtures (ticket #3, generated by tools/make_glm_oracle.py — regenerated, not +committed): + glm_tiny_fmt6/ fmt=6 (E8/IQ3) routed experts — the ONLY format that carries + the activation rotation, hence the only one that exercises the + worker's `e8_rot_rows` pre-rotation (the fmt=6 bug under test). + glm_tiny_fmt4/ fmt=4 (grouped int4) routed experts — the no-rotation control. + +Regeneration (run from c/): + python3 tools/make_glm_oracle.py --fmt6 + python3 tools/make_glm_oracle.py --fmt4 + +Determinism: both the worker and the coordinator run the SAME c/colibri binary +(same ARCH) with the SAME numeric-path env block, so the comparison is +meaningful rather than accidental. +""" + +import os +import re +import socket +import subprocess +import time +import unittest +from pathlib import Path + + +HERE = Path(__file__).resolve().parent +C_DIR = HERE.parent +ENGINE = C_DIR / "colibri" +FMT6 = C_DIR / "glm_tiny_fmt6" +FMT4 = C_DIR / "glm_tiny_fmt4" + +# cap, ebits, dbits — same order the engine parses argv (c/colibri.c main()). +ENGINE_ARGS = ["64", "16", "16"] + +# Numeric-path determinism pinning, applied to BOTH the worker and the +# coordinator. The worker is always CPU (it returns from main() before GPU +# init), and both coordinator runs use the same binary, so ARCH matches. +_NUMERIC_ENV = { + "IDOT": "0", + "I4S": "1", + "COLI_NO_FUSED_PAIR": "1", + "SPEC_PIN": "0", + "XEXP": "0", + "I4_ACC512": "0", + "I3_AVX512": "0", + "DRAFT": "0", + "COLI_CUDA": "0", + "COLI_METAL": "0", + "COLI_VULKAN": "0", + "COLI_NO_OMP_TUNE": "1", +} + + +def _fixture_ok(d: Path) -> bool: + return ( + (d / "config.json").exists() + and (d / "model.safetensors").exists() + and (d / "ref_glm.json").exists() + ) + + +def _available() -> bool: + return ENGINE.exists() and _fixture_ok(FMT6) + + +def _skip_reason() -> str: + if not ENGINE.exists(): + return "colibri is not built (run: make colibri)" + return "glm_tiny_fmt6 fixture absent (run: python3 tools/make_glm_oracle.py --fmt6)" + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _tf_signature(result: subprocess.CompletedProcess[str]): + match = re.search( + r"PREFILL \(teacher-forcing\).*:\s+(\d+)/(\d+) positions", + result.stdout, + ) + mismatches = tuple( + line for line in result.stderr.splitlines() if line.startswith("[ORACLE] mismatch") + ) + return (match.groups() if match else None, mismatches) + + +@unittest.skipUnless(_available(), _skip_reason()) +class ClusterShardingParityTest(unittest.TestCase): + """Local CPU must equal cluster-delegated expert sharding, token-exact.""" + + def _run_parity(self, fixture_dir: Path): + port = _free_port() + worker_env = { + **os.environ, + "SNAP": str(fixture_dir), + "EXPERT_WORKER": "1", + "CLUSTER_WORKER_PORT": str(port), + **_NUMERIC_ENV, + } + worker = subprocess.Popen( + [str(ENGINE), *ENGINE_ARGS], + cwd=C_DIR, + env=worker_env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if worker.poll() is not None: + stderr = worker.stderr.read() if worker.stderr else "" + self.fail(f"cluster worker exited early: {stderr}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.1): + break + except OSError: + time.sleep(0.05) + else: + self.fail("cluster worker did not start listening") + + common_env = { + **os.environ, + "SNAP": str(fixture_dir), + "REF": str(fixture_dir / "ref_glm.json"), + "TF": "1", + "TEMP": "0", + **_NUMERIC_ENV, + } + baseline = subprocess.run( + [str(ENGINE), *ENGINE_ARGS], + cwd=C_DIR, + env=common_env, + capture_output=True, + text=True, + check=False, + ) + delegated = subprocess.run( + [str(ENGINE), *ENGINE_ARGS], + cwd=C_DIR, + env={**common_env, "CLUSTER_WORKERS": f"127.0.0.1:{port}"}, + capture_output=True, + text=True, + check=False, + ) + self.assertEqual(baseline.returncode, 0, baseline.stderr) + self.assertEqual(delegated.returncode, 0, delegated.stderr) + + base_sig = _tf_signature(baseline) + del_sig = _tf_signature(delegated) + self.assertIsNotNone( + base_sig[0], + f"baseline produced no teacher-forcing signature:\n" + f"stdout={baseline.stdout}\nstderr={baseline.stderr}", + ) + self.assertEqual( + base_sig[1], + (), + f"baseline mismatched the oracle ({base_sig[0]}): {base_sig[1]}", + ) + self.assertEqual( + del_sig[1], + (), + f"delegated mismatched the oracle ({del_sig[0]}): {del_sig[1]}", + ) + self.assertEqual( + base_sig, + del_sig, + "cluster delegation changed teacher-forced token predictions", + ) + finally: + worker.terminate() + try: + worker.wait(timeout=2) + except subprocess.TimeoutExpired: + worker.kill() + worker.wait() + for stream in (worker.stdout, worker.stderr): + if stream is not None: + stream.close() + + def test_fmt6_parity(self): + """fmt=6 (rotation-bearing) routed experts: worker must rotate the input.""" + self._run_parity(FMT6) + + @unittest.skipUnless( + _fixture_ok(FMT4), + "glm_tiny_fmt4 fixture absent (run: python3 tools/make_glm_oracle.py --fmt4)", + ) + def test_fmt4_parity(self): + """fmt=4 (no-rotation control) routed experts.""" + self._run_parity(FMT4) + + +if __name__ == "__main__": + unittest.main() diff --git a/c/tools/README.md b/c/tools/README.md index 01f08e597..659394ba7 100644 --- a/c/tools/README.md +++ b/c/tools/README.md @@ -19,3 +19,17 @@ Run them from `c/`, for example: python3 tools/convert_fp8_to_int4.py --selftest python3 tools/make_glm_bench_model.py --output /tmp/colibri-bench ``` + +`make_glm_oracle.py` also produces the quantized routed-expert fixtures for the +fmt=6 (E8/IQ3, rotation-bearing) and fmt=4 (grouped int4, no-rotation control) +parity gate (#3/#7). Only the routed experts are quantized; shared/dense/attn +stay f32, and the reference (`ref_glm.json` inside each fixture dir) is computed +from the dequantized weights so the engine reproduces it token-exactly: + +```sh +python3 tools/make_glm_oracle.py --fmt6 # -> glm_tiny_fmt6/ +python3 tools/make_glm_oracle.py --fmt4 # -> glm_tiny_fmt4/ +# verify the engine loads the formats directly (32/32 expected): +SNAP=./glm_tiny_fmt6 REF=./glm_tiny_fmt6/ref_glm.json TF=1 COLI_TEMP=0 ./colibri 64 16 16 +SNAP=./glm_tiny_fmt4 REF=./glm_tiny_fmt4/ref_glm.json TF=1 COLI_TEMP=0 ./colibri 64 16 16 +``` diff --git a/c/tools/make_glm_oracle.py b/c/tools/make_glm_oracle.py index 5e16f5168..b8c05b506 100644 --- a/c/tools/make_glm_oracle.py +++ b/c/tools/make_glm_oracle.py @@ -14,7 +14,23 @@ 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).""" +EN: ingests. Default: bf16 (original oracle unchanged). + +--fmt6 / --fmt4 (ticket #3): quantize ONLY the ROUTED experts to fmt=6 (E8/IQ3, the only +rotation-bearing format) or fmt=4 (grouped int4, the no-rotation control); shared/dense/attn +stay f32. This is the fixture that lets the token-exact parity gate (#7) exercise the E8 +activation pre-rotation (E8_XE / e8_rot_rows) — the default f32 oracle never touches it. + +The E8 super-block is 256 weights (98 bytes), so the routed-expert contraction dims must be +multiples of 256; these two flags therefore generate the model with hidden_size=256 and +moe_intermediate_size=256 (the f32 default keeps 128/32). Weights are packed from the +ORIGINAL weights and the ref is computed from the DEQUANTIZED weights (round-trip through the +same quantizer), so the engine's decode reproduces the reference token-exactly. + +Regeneration (run from c/): + python3 tools/make_glm_oracle.py --fmt6 # -> glm_tiny_fmt6/ (model.safetensors, config.json, ref_glm.json) + python3 tools/make_glm_oracle.py --fmt4 # -> glm_tiny_fmt4/ + # verify: SNAP=./glm_tiny_fmt6 REF=./glm_tiny_fmt6/ref_glm.json TF=1 COLI_TEMP=0 ./colibri 64 16 16""" import json, sys, argparse from pathlib import Path @@ -63,20 +79,78 @@ def _tf_version_tuple(): from glm_fp8_emit import (fp8_block_quantize, fp8_block_dequantize, keep_f32, save_fp8_safetensors, unfuse_experts) +# Codec per le fixture quantizzate (fmt=6/fmt=4). Importati a livello di modulo: sono +# numpy-only (niente torch/transformers) e safetensors e' gia' richiesto da OGNI ramo di +# save (default bf16 incluso), quindi non toccano il path default dependency-free. +# EN: codecs for the quantized fixtures. Imported at module scope: numpy-only (no +# torch/transformers) and safetensors is already required by every save path, so they +# leave the dependency-free default path untouched. +import numpy as np +import iq3_pack +from convert_fp8_to_int4 import quant_int4_grouped +from convert_fmt4_to_fmt2 import dequant_fmt4 +from safetensors.torch import save_file + +FMT6 = 6 # container E8/IQ3 con rotazione (98B per super-blocco di 256 pesi) +FMT4 = 4 # int4 grouped, braccio di controllo senza rotazione +FMT6_SCALE_TAG = 6.0 # il singolo float del companion .qs che il motore legge come tag fmt=6 +GROUPED_INT4_BITS = 4 # bit depth del braccio int4 grouped +GROUPED_INT4_GROUP = 64 # group size (input dim per scala) del braccio int4 grouped +# EN: FMT6 = rotation-bearing E8/IQ3 container; FMT4 = grouped int4 control; FMT6_SCALE_TAG +# EN: = single .qs float the engine reads to detect fmt=6; grouped int4 uses 4 bits / gs=64. + +def quantize_routed(w, fmt): + """Quantizza una matrice di pesi routed [O, I] nel container fmt=6 (E8/IQ3) o fmt=4 + (grouped int4). Ritorna (packed_or_q, scale_or_none): per fmt=6 i byte E8 impacchettati + (scale None, il codec non ha scale esterne); per fmt=4 i nibble U8 appiattiti + le scale + f32. UNICO encoder condiviso da round-trip e save, cosi' i due path non possono divergere. + EN: quantize routed weight matrix w [O,I] to fmt=6 (E8/IQ3) or fmt=4 (grouped int4). + Returns (packed_or_q, scale_or_none): fmt=6 -> packed E8 bytes (scale None); fmt=4 -> flat + U8 nibbles + f32 scales. Single shared encoder for round-trip and save, so they can't drift.""" + if fmt == FMT6: + return iq3_pack.encode(iq3_pack.rotate_rows(w)), None + q, s = quant_int4_grouped(w, GROUPED_INT4_BITS, GROUPED_INT4_GROUP) + return q, s + ap = argparse.ArgumentParser() -ap.add_argument("--fp8", action="store_true", +_quant = ap.add_mutually_exclusive_group() +_quant.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") +_quant.add_argument("--fmt6", action="store_true", + help="quantizza SOLO gli expert routed in fmt=6 (E8/IQ3, #452) e lascia il " + "resto (shared/dense/attn) f32; oracolo per esercitare la rotazione E8. " + "EN: quantize ONLY routed experts to fmt=6 (E8/IQ3); everything else f32") +_quant.add_argument("--fmt4", action="store_true", + help="come --fmt6 ma fmt=4 (int4 grouped gs=64, nessuna rotazione): il " + "braccio di controllo senza rotazione. " + "EN: like --fmt6 but fmt=4 (grouped int4, no rotation) — the control arm") args = ap.parse_args() +fmt = FMT6 if args.fmt6 else FMT4 if args.fmt4 else 0 + torch.manual_seed(1234) +# E8/IQ3 (fmt=6) e il suo controllo int4-grouped (fmt=4) richiedono che le +# dimensioni di contrazione degli expert routed siano multiple di 256: il codec E8 +# impacchetta 256 pesi per super-blocco (98 byte) e quant_e8 rifiuta qualsiasi altra +# I (iq3_pack.encode: `K % 256 == 0`). La config tiny di default (hidden=128, +# moe_inter=32) NON e' quantizzabile in E8, quindi le fixture quantizzate usano +# hidden=256 / moe_inter=256. Il default f32 (nessuna flag) resta INVARIATO a 128/32. +# EN: fmt=6's E8 super-block (256 weights/98B) forces routed-expert contraction dims +# to be multiples of 256; the quantized fixtures use hidden=256/moe_inter=256 while +# the f32 default is unchanged at 128/32. +if fmt: + _HIDDEN, _MOE_INTER = 256, 256 +else: + _HIDDEN, _MOE_INTER = 128, 32 + cfg = GlmMoeDsaConfig( vocab_size=256, - hidden_size=128, + hidden_size=_HIDDEN, intermediate_size=64, # MLP densa (primi 3 layer) - moe_intermediate_size=32, # expert + moe_intermediate_size=_MOE_INTER, # expert num_hidden_layers=5, # 3 densi + 2 sparse first_k_dense_replace=3, num_attention_heads=4, @@ -128,6 +202,38 @@ def _tf_version_tuple(): q, s = fp8_block_quantize(p) p.copy_(fp8_block_dequantize(q, s)) +# --fmt6/--fmt4: round-trip degli expert routed FUSI (gate_up_proj [E,2M,I] e +# down_proj [E,I,M]) attraverso la quantizzazione PRIMA di calcolare il riferimento, +# cosi' ref_glm.json riflette ESATTAMENTE i pesi che il motore decodifichera' dal +# container quantizzato (stesso pattern del ramo --fp8). La quantizzazione opera per +# riga sull'ultima dim (= la dim di contrazione), quindi quantizzare il fuso 3-D +# equivale a quantizzare gli expert 2-D non fusi: le righe sono identiche. +# EN: round-trip the FUSED routed experts through the quantizer before computing the +# ref (same pattern as --fp8), so the ref matches exactly what the engine decodes. +# Quantization is per-row on the last (contraction) dim, so fused and unfused encodings +# are identical. +if fmt: + # Cattura i pesi ORIGINALI (pre-round-trip) PRIMA di toccare il modello: servono + # per il SAVE — i pesi impacchettati devono venire dai pesi originali, non da quelli + # gia' dequantizzati (encode(decode(x)) != encode(x): il codec E8 non e' idempotente + # byte-per-byte). state_dict() ritorna viste che aliasano i parametri, quindi clone + # esplicito. EN: clone the original weights before the round-trip; the SAVE path must + # pack the ORIGINAL weights (encode∘decode is not idempotent), and state_dict() + # aliases the live parameters. + sd_orig = {k: v.detach().clone() for k, v in model.state_dict().items()} + with torch.no_grad(): + for n, p in model.named_parameters(): + if ".mlp.experts." not in n or not n.endswith(("gate_up_proj", "down_proj")): + continue + shp = tuple(p.shape); I = shp[-1] + w = p.detach().float().numpy().reshape(-1, I) + packed, scale = quantize_routed(w, fmt) + if fmt == FMT6: + w_eff = iq3_pack.unrotate_rows(iq3_pack.decode(packed, I)) + else: + w_eff = dequant_fmt4(packed.reshape(-1, (I + 1) // 2), scale, w.shape[0], I) + p.copy_(torch.from_numpy(w_eff.reshape(shp))) + 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)}") @@ -154,15 +260,46 @@ def _tf_version_tuple(): sd = model.state_dict() unfuse_experts(sd) -Path("glm_tiny").mkdir(parents=True, exist_ok=True) # safetensors/json won't create the dir themselves -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") +if fmt: + # Salva SOLO gli expert routed quantizzati (fmt=6 o fmt=4 + .qs); shared/dense/attn/ + # norme restano f32. I pesi impacchettati provengono dai pesi ORIGINALI (sd_orig, + # non fusi), non da quelli gia' dequantizzati, cosi' il motore decodifica ESATTAMENTE + # i pesi che il round-trip sopra ha usato per il riferimento. EN: quantize ONLY the + # routed experts (fmt=6 or fmt=4 + .qs); everything else stays f32. Packed weights + # come from the ORIGINAL (unfused) weights, not the dequantized ones, so the engine + # decodes exactly the weights the round-trip used for the reference. + outdir = "glm_tiny_fmt6" if fmt == FMT6 else "glm_tiny_fmt4" + Path(outdir).mkdir(parents=True, exist_ok=True) + sd = unfuse_experts(sd_orig) + out = {} + for name, t in sd.items(): + if ".mlp.experts." in name and name.endswith( + (".gate_proj.weight", ".up_proj.weight", ".down_proj.weight")): + w = t.detach().float().numpy() + packed, scale = quantize_routed(w, fmt) + if fmt == FMT6: + out[name] = torch.from_numpy(np.ascontiguousarray(packed.reshape(-1))) + out[name + ".qs"] = torch.tensor([FMT6_SCALE_TAG], dtype=torch.float32) + else: + out[name] = torch.from_numpy(packed) + out[name + ".qs"] = torch.from_numpy(scale) + else: + out[name] = t.detach().contiguous() + save_file(out, f"{outdir}/model.safetensors") + json.dump(cfg.to_dict(), open(f"{outdir}/config.json", "w")) + json.dump({"prompt_ids": prompt, "full_ids": full, "tf_pred": tf_pred}, + open(f"{outdir}/ref_glm.json", "w")) + print(f"saved: {outdir}/ (routed experts fmt={fmt}, shared/dense/attn f32) + " + f"{outdir}/ref_glm.json") 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("saved: glm_tiny/ (weights + config) and ref_glm.json" - + (" [fp8]" if args.fp8 else "")) + Path("glm_tiny").mkdir(parents=True, exist_ok=True) # safetensors/json won't create the dir themselves + 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: + 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("saved: glm_tiny/ (weights + config) and ref_glm.json" + + (" [fp8]" if args.fp8 else ""))