"""Streaming HF -> MLX conversion + quantization for Inkling. The model is far too large (~1.9 TB bf16) to instantiate in RAM, so we convert tensor-by-tensor: read each source shard (mmap), remap the name, apply the layout transform, optionally affine-quantize the weight, and write output shards. Affine quantization has no cross-tensor dependency, so per-tensor streaming is exactly equivalent to ``nn.quantize(model)``. Name/layout transforms vs. the checkpoint: * ``*_sconv.weight`` [C,1,K] -> [C,K,1] (MLX conv1d layout) * ``mlp.w13_dn`` [2I,H] -> gate_proj/up_proj (split dense fused gate+up) * ``experts.w13_weight`` [E,2I,H] -> gate_proj/up_proj (split) * ``experts.w2_weight`` [E,H,I] -> down_proj (identity) * ``shared_experts.shared_w13`` [2,2I,H] -> gate_proj/up_proj (split) * ``model.mtp.*`` dropped (inference-irrelevant) * everything else: identity """ from __future__ import annotations import glob import json import os import re import shutil import mlx.core as mx _LAYER_RE = re.compile(r"model\.llm\.layers\.(\d+)\.") _N_SHARED = 2 # Inkling: 2 shared experts (router rows after the routed ones) def prune_moe_tensor(name: str, w: mx.array, keep, dense_mlp_idx: int) -> mx.array: """Subset a MoE tensor to the kept experts (REAP). ``keep`` maps sparse-layer index -> array of kept routed-expert indices. Applied BEFORE map_name/transform. Routed experts + the router (gate) are subset; shared experts are untouched.""" m = _LAYER_RE.search(name) if m is None: return w L = int(m.group(1)) if L < dense_mlp_idx: # dense layer — no routed experts return w kidx = mx.array(keep[L - dense_mlp_idx]) # [K] if name.endswith(("experts.w13_weight", "experts.w2_weight")) and "shared" not in name: return w[kidx] # [E, ...] -> [K, ...] if name.endswith("mlp.gate.weight"): # [n_routed + n_shared, hidden] n_routed = w.shape[0] - _N_SHARED # rows: [routed... , shared...] return mx.concatenate([w[kidx], w[n_routed:]], axis=0) # [K + n_shared, hidden] if name.endswith("mlp.gate.bias"): # [n_routed] correction bias return w[kidx] return w # shared_experts.*, gate.global_scale, etc. def map_name(name: str): """HF checkpoint tensor name -> list of (out_name, kind) for the MLX model.""" if name.startswith("model.mtp."): return [] # drop MTP head if name.endswith(("k_sconv.weight", "v_sconv.weight", "attn_sconv.weight", "mlp_sconv.weight")): return [(name, "sconv")] # dense MLP fused gate+up / down if name.endswith("mlp.w13_dn.weight"): base = name[: -len("w13_dn.weight")] return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")] if name.endswith("mlp.w2_md.weight"): return [(name[: -len("w2_md.weight")] + "down_proj.weight", "identity")] # routed experts fused if name.endswith("experts.w13_weight"): base = name[: -len("w13_weight")] return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")] if name.endswith("experts.w2_weight"): return [(name[: -len("w2_weight")] + "down_proj.weight", "identity")] # shared experts fused if name.endswith("shared_experts.shared_w13_weight"): base = name[: -len("shared_w13_weight")] return [(base + "gate_proj.weight", "w13_gate"), (base + "up_proj.weight", "w13_up")] if name.endswith("shared_experts.shared_w2_weight"): return [(name[: -len("shared_w2_weight")] + "down_proj.weight", "identity")] return [(name, "identity")] def transform(w: mx.array, kind: str) -> mx.array: if kind == "identity": return w if kind == "sconv": # [C, 1, K] -> [C, K, 1] return mx.swapaxes(w, 1, 2) if kind in ("w13_gate", "w13_up"): # The checkpoint stores gate/up INTERLEAVED row-wise: [g0, u0, g1, u1, ...] # (SGLang `deinterleave_w13`). De-interleave: gate = rows 0::2, up = rows 1::2. # A contiguous [:half]/[half:] split scrambles gate<->up in every MLP. n = w.shape[-2] // 2 g = w.reshape(*w.shape[:-2], n, 2, w.shape[-1]) return g[..., 0, :] if kind == "w13_gate" else g[..., 1, :] raise ValueError(kind) # ---- quantization target predicate (must be identical in convert and load) ---- # Quant "recipes" — which module leaves get affine-quantized. # uniform : everything (attention, MLP/experts, embed/unembed, audio, vision) # experts_only : ONLY the MLP/expert matmuls (+ audio/vision); attention and # embed/unembed stay bf16. Inkling attention dominates 4-bit error # (~58% per layer vs ~15% for experts), so this keeps a 4-bit-sized # build coherent while the ~927 B experts still fit in 512 GB. _RECIPES = { "uniform": {"wq_du", "wk_dv", "wv_dv", "wr_du", "wo_ud", "gate_proj", "up_proj", "down_proj", "embed", "unembed", "encoder"}, "experts_only": {"gate_proj", "up_proj", "down_proj", "encoder"}, } def is_quant_target(out_name: str, quant_axis_size: int, group_size: int, recipe: str = "uniform") -> bool: """Whether ``out_name`` (a converted param path) should be affine-quantized.""" if not out_name.endswith(".weight"): return False leaf = out_name[: -len(".weight")].rsplit(".", 1)[-1] leaves = _RECIPES[recipe] # vision projection layers (linear_0 .. linear_3) — quantized in both recipes is_vision_linear = leaf.startswith("linear_") and ".visual." in out_name if leaf not in leaves and not is_vision_linear: return False # router gate stays fp (leaf == "gate", excluded above); norms/sconv excluded by leaf # can only group-quantize when the input dim is a multiple of group_size return quant_axis_size % group_size == 0 # ------------------------------ streaming driver ------------------------------ _SHARD_CAP_BYTES = 5_000_000_000 # ~5 GB per output shard def _process_tensor(name, w, bits, group_size, out_dtype, recipe="uniform", keep=None, dmi=2): """Yield (out_name, array) pairs for one source tensor.""" if keep is not None: w = prune_moe_tensor(name, w, keep, dmi) # REAP: subset to kept experts for out_name, kind in map_name(name): wt = transform(w, kind) quantize = bits is not None and is_quant_target(out_name, wt.shape[-1], group_size, recipe) if quantize: qw, scales, biases = mx.quantize(wt, group_size=group_size, bits=bits) base = out_name[: -len(".weight")] yield out_name, qw yield base + ".scales", scales yield base + ".biases", biases else: # keep norms/router/sconv/rel-proj in fp32-safe dtype; matmul weights in out_dtype keep_hi = wt.dtype == mx.float32 and (".global_scale" in out_name or ".bias" in out_name or out_name.endswith(("_norm.weight", "norm.weight"))) yield out_name, wt.astype(mx.float32 if keep_hi else out_dtype) def convert_model(src: str, dst: str, bits=None, group_size: int = 64, out_dtype=mx.bfloat16, recipe: str = "uniform", keep_path=None): """Stream-convert an Inkling checkpoint from ``src`` to ``dst``. ``bits=None`` -> plain dtype cast (bf16). ``bits in {4,6,8}`` -> affine quant. ``recipe`` selects which modules are quantized (see ``_RECIPES``). ``keep_path`` (REAP): npz with ``keep`` [n_sparse_layers, K] + ``dense_mlp_idx`` -> prune each MoE layer to its K kept experts and set ``n_routed_experts=K`` in config. Processes one source shard at a time; never holds the whole model in RAM. """ os.makedirs(dst, exist_ok=True) keep = dmi = new_ne = None if keep_path is not None: import numpy as np kd = np.load(keep_path) keep = kd["keep"]; dmi = int(kd["dense_mlp_idx"]); new_ne = int(kd["K"]) index = json.load(open(os.path.join(src, "model.safetensors.index.json"))) weight_map = index["weight_map"] shard_to_names: dict[str, list[str]] = {} for n, s in weight_map.items(): shard_to_names.setdefault(s, []).append(n) out_index: dict[str, str] = {} buffer: dict[str, mx.array] = {} buffer_bytes = 0 out_shard_id = 0 total_out_shards_placeholder = "{:05d}" def flush(final=False): nonlocal buffer, buffer_bytes, out_shard_id if not buffer: return out_shard_id += 1 fname = f"model-{total_out_shards_placeholder.format(out_shard_id)}.safetensors" mx.save_safetensors(os.path.join(dst, fname), buffer, metadata={"format": "mlx"}) for k in buffer: out_index[k] = fname buffer = {} buffer_bytes = 0 for shard in sorted(shard_to_names): path = os.path.join(src, shard) tensors = mx.load(path) # mmap for name in shard_to_names[shard]: w = tensors[name] for out_name, arr in _process_tensor(name, w, bits, group_size, out_dtype, recipe, keep, dmi): mx.eval(arr) buffer[out_name] = arr buffer_bytes += arr.nbytes if buffer_bytes >= _SHARD_CAP_BYTES: flush() del tensors flush(final=True) # rename shards with correct total, build index.json _finalize_index(dst, out_index, out_shard_id) _write_config(src, dst, bits, group_size, recipe, new_ne) _copy_aux(src, dst) return dst def _finalize_index(dst, out_index, n_shards): # rewrite shard filenames to model-XXXXX-of-YYYYY.safetensors remap = {} for i in range(1, n_shards + 1): old = f"model-{i:05d}.safetensors" new = f"model-{i:05d}-of-{n_shards:05d}.safetensors" if old != new and os.path.exists(os.path.join(dst, old)): os.rename(os.path.join(dst, old), os.path.join(dst, new)) remap[old] = new weight_map = {k: remap[v] for k, v in out_index.items()} total = sum(os.path.getsize(os.path.join(dst, f)) for f in set(weight_map.values())) with open(os.path.join(dst, "model.safetensors.index.json"), "w") as f: json.dump({"metadata": {"total_size": total}, "weight_map": weight_map}, f, indent=2) def _write_config(src, dst, bits, group_size, recipe="uniform", new_ne=None): cfg = json.load(open(os.path.join(src, "config.json"))) if bits is not None: cfg["quantization"] = {"group_size": group_size, "bits": bits, "recipe": recipe} if new_ne is not None: # REAP: fewer routed experts cfg["text_config"]["n_routed_experts"] = new_ne cfg.setdefault("reap", {})["kept_experts"] = new_ne with open(os.path.join(dst, "config.json"), "w") as f: json.dump(cfg, f, indent=2) def _copy_aux(src, dst): for pat in ("tokenizer*", "special_tokens_map.json", "*.tiktoken", "tiktoken", "chat_template.jinja", "processor_config.json"): for p in glob.glob(os.path.join(src, pat)): base = os.path.basename(p) target = os.path.join(dst, base) if os.path.isdir(p): shutil.copytree(p, target, dirs_exist_ok=True) else: shutil.copy2(p, target)