From 5a7c81baaeb83a8a84ccc9f34b9d5c45d66e2df9 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 20:42:20 -0400 Subject: [PATCH 1/3] qwen3-8b: register variant + sharded safetensors loader Adds qwen3-8b to the variant registry (Qwen3-8B dir name, GGUF name). The factory now derives Config from config.json on disk, so untied lm_head (tie_word_embeddings=false on 8B) flows through naturally. Wires sk_qwen_load_safetensors_index from libsk.dylib so the python loader picks the sharded entry point when model.safetensors.index.json is present; falls back to the single-file path otherwise. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/models/qwen/__init__.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/SuperKittens/models/qwen/__init__.py b/SuperKittens/models/qwen/__init__.py index 44b220c..57e94aa 100644 --- a/SuperKittens/models/qwen/__init__.py +++ b/SuperKittens/models/qwen/__init__.py @@ -29,9 +29,11 @@ def _build_cfg_from_snapshot(snap: Path, **overrides) -> Config: _VARIANT_TO_DIR = { "qwen3-0.6b": "Qwen3-0.6B", + "qwen3-8b": "Qwen3-8B", } _VARIANT_TO_GGUF = { "qwen3-0.6b": "Qwen3-0.6B-Q8_0.gguf", + "qwen3-8b": "Qwen3-8B-Q8_0.gguf", } @@ -39,7 +41,7 @@ def _from_pretrained(variant: str = "qwen3-0.6b", quant: str | None = None, snapshot: str | None = None, gguf_path: str | None = None, **cfg_overrides) -> Qwen: spec = variant - dir_name = _VARIANT_TO_DIR.get(spec.lower(), "Qwen3-0.6B") + dir_name = _VARIANT_TO_DIR.get(spec.lower(), spec) sk_root = Path(__file__).resolve().parents[3] snap = Path(snapshot) if snapshot else (sk_root / "SuperKittens" / "model_weights" / dir_name) if not snap.exists(): @@ -59,12 +61,19 @@ def _from_pretrained(variant: str = "qwen3-0.6b", quant: str | None = None, from .qwen import _load idx_path = snap / "model.safetensors.index.json" single = snap / "model.safetensors" - target = idx_path if idx_path.exists() else single - if not target.exists(): + lib = _load() + if idx_path.exists(): + if not hasattr(lib, "sk_qwen_load_safetensors_index"): + raise RuntimeError("libsk.dylib has no sk_qwen_load_safetensors_index symbol; rebuild dylib") + rc = lib.sk_qwen_load_safetensors_index(m._h, str(idx_path).encode()) + if rc: + raise RuntimeError(f"sk_qwen_load_safetensors_index failed: {rc}") + elif single.exists(): + rc = lib.sk_qwen_load_safetensors(m._h, str(single).encode()) + if rc: + raise RuntimeError(f"sk_qwen_load_safetensors failed: {rc}") + else: raise FileNotFoundError(f"no safetensors in {snap}") - rc = _load().sk_qwen_load_safetensors(m._h, str(target).encode()) - if rc: - raise RuntimeError(f"sk_qwen_load_safetensors failed: {rc}") # RoPE tables m.bake_and_set_rope() @@ -94,7 +103,7 @@ def from_pretrained(**kwargs): from SuperKittens.api import register -for _spec in ("qwen3-0.6b",): +for _spec in ("qwen3-0.6b", "qwen3-8b"): try: register(_spec, _QwenFactory, variant=_spec) except ValueError: From afd8b866a651e5ab96fe92804974b5b348ed946c Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 21:28:45 -0400 Subject: [PATCH 2/3] qwen3-8b: default to Q8_0 GGUF; tolerate missing snapshot dir The prior 8B registration defaulted to bf16 safetensors, which sends 16 GB through the python-driven copy_into loop (10-15+ min per load on lexie/derek). Switch the 8B factory to default quant="q8_0" so it mmaps the GGUF via the existing sk_qwen_load_gguf path (matches the 0.6B fast path). - _VARIANT_DIMS hardcodes Qwen3-0.6B/8B dims so Config can be built without an on-disk config.json (Q8_0-only deployments commonly skip the snapshot). - _resolve_tokenizer falls back to huggingface_hub.hf_hub_download for tokenizer.json when the snapshot dir lacks one. - tie_word_embeddings=0 baked into the 8B variant dims (untied LM head); the existing C++ sk_qwen_load_gguf already reads output.weight in this case (weights.c++:374). - 8B snapshot dir name now resolves to Qwen3-8B-GGUF to match the hf download layout used on lexie. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/models/qwen/__init__.py | 115 ++++++++++++++++++++------- 1 file changed, 87 insertions(+), 28 deletions(-) diff --git a/SuperKittens/models/qwen/__init__.py b/SuperKittens/models/qwen/__init__.py index 57e94aa..c53d059 100644 --- a/SuperKittens/models/qwen/__init__.py +++ b/SuperKittens/models/qwen/__init__.py @@ -6,6 +6,23 @@ from .qwen import Qwen, Config +# Hardcoded Qwen3 variant dims (HF config.json values). Used when no on-disk +# snapshot is available — e.g. when loading a Q8_0 GGUF directly without the +# fp16 safetensors snapshot present. +_VARIANT_DIMS = { + "qwen3-0.6b": dict( + n_layers=28, d_model=1024, n_heads=16, n_kv_heads=8, head_dim=128, + n_int=3072, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, + tie_word_embeddings=1, + ), + "qwen3-8b": dict( + n_layers=36, d_model=4096, n_heads=32, n_kv_heads=8, head_dim=128, + n_int=12288, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, + tie_word_embeddings=0, # Qwen3-8B has UNTIED LM head + ), +} + + def _build_cfg_from_snapshot(snap: Path, **overrides) -> Config: cfgj = json.loads((snap / "config.json").read_text()) cfg = Config( @@ -27,37 +44,84 @@ def _build_cfg_from_snapshot(snap: Path, **overrides) -> Config: return cfg +def _build_cfg_from_variant(variant: str, **overrides) -> Config: + dims = dict(_VARIANT_DIMS[variant.lower()]) + dims["seq_max"] = overrides.pop("seq_max", 128) + dims["cache_max"] = overrides.pop("cache_max", 512) + cfg = Config(**dims) + for k, v in overrides.items(): + setattr(cfg, k, v) + return cfg + + _VARIANT_TO_DIR = { "qwen3-0.6b": "Qwen3-0.6B", - "qwen3-8b": "Qwen3-8B", + "qwen3-8b": "Qwen3-8B-GGUF", } _VARIANT_TO_GGUF = { "qwen3-0.6b": "Qwen3-0.6B-Q8_0.gguf", "qwen3-8b": "Qwen3-8B-Q8_0.gguf", } +_VARIANT_TO_HF_REPO = { + "qwen3-0.6b": "Qwen/Qwen3-0.6B", + "qwen3-8b": "Qwen/Qwen3-8B", +} + + +def _resolve_tokenizer(snap: Path, variant: str): + """Find tokenizer.json: prefer snapshot, then HF hub cache download.""" + from SuperKittens.models.load.tokenizer import Tokenizer + cand = snap / "tokenizer.json" + if cand.exists(): + return Tokenizer.from_hf_json(str(cand), family="qwen3") + # GGUF-only layout: tokenizer.json may sit alongside the .gguf inside the + # same directory (huggingface-cli download fetches the whole repo). + for f in snap.glob("tokenizer.json"): + return Tokenizer.from_hf_json(str(f), family="qwen3") + # Last resort: pull tokenizer.json from HF hub. + repo = _VARIANT_TO_HF_REPO.get(variant.lower()) + if repo: + try: + from huggingface_hub import hf_hub_download + tok_path = hf_hub_download(repo_id=repo, filename="tokenizer.json") + return Tokenizer.from_hf_json(tok_path, family="qwen3") + except Exception as e: + print(f"[qwen] hf_hub_download tokenizer failed: {e}") + return None def _from_pretrained(variant: str = "qwen3-0.6b", quant: str | None = None, snapshot: str | None = None, gguf_path: str | None = None, **cfg_overrides) -> Qwen: - spec = variant - dir_name = _VARIANT_TO_DIR.get(spec.lower(), spec) + spec = variant.lower() + dir_name = _VARIANT_TO_DIR.get(spec, spec) sk_root = Path(__file__).resolve().parents[3] snap = Path(snapshot) if snapshot else (sk_root / "SuperKittens" / "model_weights" / dir_name) - if not snap.exists(): - raise FileNotFoundError(f"snapshot dir not found: {snap}") - cfg = _build_cfg_from_snapshot(snap, **cfg_overrides) + is_gguf = quant in ("q8_0", "Q8_0", "gguf") + + # Build Config: prefer on-disk config.json, else fall back to hardcoded + # variant dims (Q8_0 GGUF-only deployments often skip the snapshot dir). + if (snap / "config.json").exists(): + cfg = _build_cfg_from_snapshot(snap, **cfg_overrides) + elif spec in _VARIANT_DIMS: + cfg = _build_cfg_from_variant(spec, **cfg_overrides) + else: + raise FileNotFoundError(f"no config.json at {snap} and no hardcoded dims for {spec}") + m = Qwen(cfg) - if quant in ("q8_0", "Q8_0", "gguf"): - gpath = Path(gguf_path) if gguf_path else ( - sk_root / "SuperKittens" / "model_weights" / _VARIANT_TO_GGUF.get(spec.lower(), "Qwen3-0.6B-Q8_0.gguf")) + if is_gguf: + gpath = Path(gguf_path) if gguf_path else (snap / _VARIANT_TO_GGUF.get(spec, "")) if not gpath.exists(): - raise FileNotFoundError(f"gguf file not found: {gpath}") + # Fallback: search snapshot dir for any *.gguf. + ggs = list(snap.glob("*.gguf")) if snap.exists() else [] + if ggs: + gpath = ggs[0] + else: + raise FileNotFoundError(f"gguf file not found: {gpath}") m.load_gguf(str(gpath)) else: - # fp16 safetensors path from .qwen import _load idx_path = snap / "model.safetensors.index.json" single = snap / "model.safetensors" @@ -75,27 +139,17 @@ def _from_pretrained(variant: str = "qwen3-0.6b", quant: str | None = None, else: raise FileNotFoundError(f"no safetensors in {snap}") - # RoPE tables m.bake_and_set_rope() - # Tokenizer - try: - from SuperKittens.models.load.tokenizer import Tokenizer - json_path = snap / "tokenizer.json" - sp_path = snap / "tokenizer.model" - if json_path.exists(): - m.tokenizer = Tokenizer.from_hf_json(str(json_path), family="qwen3") - elif sp_path.exists(): - m.tokenizer = Tokenizer.from_sentencepiece(str(sp_path)) - else: - print(f"[qwen] no tokenizer.json or tokenizer.model in {snap}") - except Exception as e: - print(f"[qwen] tokenizer attach failed: {e}") + tok = _resolve_tokenizer(snap, spec) + if tok is not None: + m.tokenizer = tok + else: + print(f"[qwen] no tokenizer available for {spec}") return m -# Adapter so MODEL_REGISTRY's cls.from_pretrained(**defaults, **kwargs) routes here. class _QwenFactory: @staticmethod def from_pretrained(**kwargs): @@ -103,9 +157,14 @@ def from_pretrained(**kwargs): from SuperKittens.api import register -for _spec in ("qwen3-0.6b", "qwen3-8b"): +# 8B defaults to Q8_0 GGUF (16 GB bf16 safetensors load via python is unworkable). +_DEFAULTS = { + "qwen3-0.6b": {"variant": "qwen3-0.6b"}, + "qwen3-8b": {"variant": "qwen3-8b", "quant": "q8_0"}, +} +for _spec, _defs in _DEFAULTS.items(): try: - register(_spec, _QwenFactory, variant=_spec) + register(_spec, _QwenFactory, **_defs) except ValueError: pass From 12e8d4c579f7203acf5dff2aa1470680b3ff1b36 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 21:54:29 -0400 Subject: [PATCH 3/3] qwen3-8b: bind safetensors_index through QWEN_ABI; drop 0.6b refs; collapse 4 parallel dicts into one spec table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes addressing PR #13 review: 1. Add `load_safetensors_index` to QWEN_ABI (qwen.py) so the central binder sets argtypes/restype. Previous code called `lib.sk_qwen_load_safetensors_index` directly without registering the signature — ctypes default int coercion would corrupt the handle pointer on 64-bit builds. 2. Drop 0.6b mentions per review. Consolidate the four parallel per-variant dicts (_VARIANT_DIMS, _VARIANT_TO_DIR, _VARIANT_TO_GGUF, _VARIANT_TO_HF_REPO) into a single _QWEN3_VARIANTS spec table with one row per variant. Adding a new qwen3 size now means one entry, not five. This is the same anti-pattern as the dropped 0.6b lit-everywhere: family-specific knowledge funnels through ONE seam (QWEN_ABI for ctypes, _QWEN3_VARIANTS for per-variant metadata) instead of leaking across files. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/models/qwen/__init__.py | 121 ++++++++++----------------- SuperKittens/models/qwen/qwen.py | 5 +- 2 files changed, 47 insertions(+), 79 deletions(-) diff --git a/SuperKittens/models/qwen/__init__.py b/SuperKittens/models/qwen/__init__.py index c53d059..8d9d7e0 100644 --- a/SuperKittens/models/qwen/__init__.py +++ b/SuperKittens/models/qwen/__init__.py @@ -3,27 +3,25 @@ import json from pathlib import Path -from .qwen import Qwen, Config - - -# Hardcoded Qwen3 variant dims (HF config.json values). Used when no on-disk -# snapshot is available — e.g. when loading a Q8_0 GGUF directly without the -# fp16 safetensors snapshot present. -_VARIANT_DIMS = { - "qwen3-0.6b": dict( - n_layers=28, d_model=1024, n_heads=16, n_kv_heads=8, head_dim=128, - n_int=3072, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, - tie_word_embeddings=1, - ), - "qwen3-8b": dict( - n_layers=36, d_model=4096, n_heads=32, n_kv_heads=8, head_dim=128, - n_int=12288, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, - tie_word_embeddings=0, # Qwen3-8B has UNTIED LM head - ), +from .qwen import Qwen, Config, _load + + +_QWEN3_VARIANTS = { + "qwen3-8b": { + "hf_repo": "Qwen/Qwen3-8B", + "weight_dir": "Qwen3-8B-GGUF", + "gguf_name": "Qwen3-8B-Q8_0.gguf", + "default_quant": "q8_0", + "dims": dict( + n_layers=36, d_model=4096, n_heads=32, n_kv_heads=8, head_dim=128, + n_int=12288, vocab_size=151936, eps=1e-6, rope_freq_base=1_000_000.0, + tie_word_embeddings=0, + ), + }, } -def _build_cfg_from_snapshot(snap: Path, **overrides) -> Config: +def _cfg_from_snapshot(snap: Path, **overrides) -> Config: cfgj = json.loads((snap / "config.json").read_text()) cfg = Config( n_layers = cfgj["num_hidden_layers"], @@ -44,8 +42,8 @@ def _build_cfg_from_snapshot(snap: Path, **overrides) -> Config: return cfg -def _build_cfg_from_variant(variant: str, **overrides) -> Config: - dims = dict(_VARIANT_DIMS[variant.lower()]) +def _cfg_from_dims(variant: str, **overrides) -> Config: + dims = dict(_QWEN3_VARIANTS[variant]["dims"]) dims["seq_max"] = overrides.pop("seq_max", 128) dims["cache_max"] = overrides.pop("cache_max", 512) cfg = Config(**dims) @@ -54,78 +52,52 @@ def _build_cfg_from_variant(variant: str, **overrides) -> Config: return cfg -_VARIANT_TO_DIR = { - "qwen3-0.6b": "Qwen3-0.6B", - "qwen3-8b": "Qwen3-8B-GGUF", -} -_VARIANT_TO_GGUF = { - "qwen3-0.6b": "Qwen3-0.6B-Q8_0.gguf", - "qwen3-8b": "Qwen3-8B-Q8_0.gguf", -} -_VARIANT_TO_HF_REPO = { - "qwen3-0.6b": "Qwen/Qwen3-0.6B", - "qwen3-8b": "Qwen/Qwen3-8B", -} - - def _resolve_tokenizer(snap: Path, variant: str): - """Find tokenizer.json: prefer snapshot, then HF hub cache download.""" from SuperKittens.models.load.tokenizer import Tokenizer - cand = snap / "tokenizer.json" - if cand.exists(): - return Tokenizer.from_hf_json(str(cand), family="qwen3") - # GGUF-only layout: tokenizer.json may sit alongside the .gguf inside the - # same directory (huggingface-cli download fetches the whole repo). - for f in snap.glob("tokenizer.json"): - return Tokenizer.from_hf_json(str(f), family="qwen3") - # Last resort: pull tokenizer.json from HF hub. - repo = _VARIANT_TO_HF_REPO.get(variant.lower()) - if repo: - try: - from huggingface_hub import hf_hub_download - tok_path = hf_hub_download(repo_id=repo, filename="tokenizer.json") - return Tokenizer.from_hf_json(tok_path, family="qwen3") - except Exception as e: - print(f"[qwen] hf_hub_download tokenizer failed: {e}") + for cand in (snap / "tokenizer.json", *snap.glob("tokenizer.json")): + if cand.exists(): + return Tokenizer.from_hf_json(str(cand), family="qwen3") + repo = _QWEN3_VARIANTS[variant]["hf_repo"] + try: + from huggingface_hub import hf_hub_download + tok_path = hf_hub_download(repo_id=repo, filename="tokenizer.json") + return Tokenizer.from_hf_json(tok_path, family="qwen3") + except Exception as e: + print(f"[qwen] hf_hub_download tokenizer failed: {e}") return None -def _from_pretrained(variant: str = "qwen3-0.6b", quant: str | None = None, +def _from_pretrained(variant: str = "qwen3-8b", quant: str | None = None, snapshot: str | None = None, gguf_path: str | None = None, **cfg_overrides) -> Qwen: spec = variant.lower() - dir_name = _VARIANT_TO_DIR.get(spec, spec) - sk_root = Path(__file__).resolve().parents[3] - snap = Path(snapshot) if snapshot else (sk_root / "SuperKittens" / "model_weights" / dir_name) + if spec not in _QWEN3_VARIANTS: + raise ValueError(f"unknown qwen3 variant {spec!r}; known: {list(_QWEN3_VARIANTS)}") + meta = _QWEN3_VARIANTS[spec] - is_gguf = quant in ("q8_0", "Q8_0", "gguf") + sk_root = Path(__file__).resolve().parents[3] + snap = Path(snapshot) if snapshot else (sk_root / "SuperKittens" / "model_weights" / meta["weight_dir"]) - # Build Config: prefer on-disk config.json, else fall back to hardcoded - # variant dims (Q8_0 GGUF-only deployments often skip the snapshot dir). if (snap / "config.json").exists(): - cfg = _build_cfg_from_snapshot(snap, **cfg_overrides) - elif spec in _VARIANT_DIMS: - cfg = _build_cfg_from_variant(spec, **cfg_overrides) + cfg = _cfg_from_snapshot(snap, **cfg_overrides) else: - raise FileNotFoundError(f"no config.json at {snap} and no hardcoded dims for {spec}") + cfg = _cfg_from_dims(spec, **cfg_overrides) m = Qwen(cfg) + quant = quant or meta["default_quant"] - if is_gguf: - gpath = Path(gguf_path) if gguf_path else (snap / _VARIANT_TO_GGUF.get(spec, "")) + if quant in ("q8_0", "Q8_0", "gguf"): + gpath = Path(gguf_path) if gguf_path else (snap / meta["gguf_name"]) if not gpath.exists(): - # Fallback: search snapshot dir for any *.gguf. ggs = list(snap.glob("*.gguf")) if snap.exists() else [] - if ggs: - gpath = ggs[0] - else: + if not ggs: raise FileNotFoundError(f"gguf file not found: {gpath}") + gpath = ggs[0] m.load_gguf(str(gpath)) else: - from .qwen import _load + lib = _load() # ABI bound via QWEN_ABI; argtypes/restype set centrally idx_path = snap / "model.safetensors.index.json" single = snap / "model.safetensors" - lib = _load() if idx_path.exists(): if not hasattr(lib, "sk_qwen_load_safetensors_index"): raise RuntimeError("libsk.dylib has no sk_qwen_load_safetensors_index symbol; rebuild dylib") @@ -157,14 +129,9 @@ def from_pretrained(**kwargs): from SuperKittens.api import register -# 8B defaults to Q8_0 GGUF (16 GB bf16 safetensors load via python is unworkable). -_DEFAULTS = { - "qwen3-0.6b": {"variant": "qwen3-0.6b"}, - "qwen3-8b": {"variant": "qwen3-8b", "quant": "q8_0"}, -} -for _spec, _defs in _DEFAULTS.items(): +for _spec in _QWEN3_VARIANTS: try: - register(_spec, _QwenFactory, **_defs) + register(_spec, _QwenFactory, variant=_spec) except ValueError: pass diff --git a/SuperKittens/models/qwen/qwen.py b/SuperKittens/models/qwen/qwen.py index 8d4ea3a..84294cb 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -67,8 +67,9 @@ class _Weights(ctypes.Structure): ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)], ctypes.c_int), "reset": ([ctypes.c_void_p], None), "destroy": ([ctypes.c_void_p], None), - "load_safetensors": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), - "load_gguf": optional([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "load_safetensors": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "load_safetensors_index": optional([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "load_gguf": optional([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), "set_rope_tables": optional([ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), "get_last_logits": optional([ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), }