File size: 8,040 Bytes
b0bebef e7fce3b b0bebef 171b968 b0bebef 171b968 b0bebef d639bbb b0bebef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | # /// script
# requires-python = ">=3.11,<3.13"
# dependencies = [
# "coreai-core==1.0.0b2",
# "coreai-torch==0.4.1",
# "coreai-opt==0.2.1",
# "scikit-learn>=1.7.2",
# "torch",
# "transformers>=5.9,<6",
# "huggingface_hub",
# "safetensors",
# "numpy",
# "pillow",
# "ml-dtypes<=0.6.0",
# "sentencepiece",
# "protobuf",
# "datasets",
# "tokenizers",
# ]
# ///
"""North Star: export LiquidAI/LFM2.5-8B-A1B (MoE) to a CoreAI bundle on HF Jobs.
Uses OUR lfm2_moe export def (routed SwitchGLU/GatherMM = reads only top-4 experts,
short-conv as native conv1d). VETS bit-exact vs the HF reference on a truncated model
FIRST β aborts the (expensive) full export if the def is wrong β then exports the full
8B bounded to 4096 ctx and pushes the bundle to the Hub.
"""
import os, sys, shutil, subprocess, glob, json
ROOT = "/tmp/work"
os.makedirs(ROOT, exist_ok=True)
os.chdir(ROOT)
BASE = "LiquidAI/LFM2.5-8B-A1B"
OUT_REPO = "AbdallaHolmes/LFM2.5-8B-A1B-CoreAI-NorthStar"
CTX = 4096
def run(cmd, **kw):
print(f"\n$ {' '.join(cmd) if isinstance(cmd,list) else cmd}", flush=True)
return subprocess.run(cmd, check=True, **kw)
# ββ 1. clone Apple's coreai-models + drop in OUR patch ββββββββββββββββββββββββ
run(["git", "clone", "--depth", "1", "https://github.com/john-rocky/coreai-models"])
from huggingface_hub import snapshot_download
patch = snapshot_download("AbdallaHolmes/coreai-models-lfm-patch", repo_type="model")
pkg = f"{ROOT}/coreai-models/python/src/coreai_models"
shutil.copy(f"{patch}/lfm2.py", f"{pkg}/models/macos/lfm2.py")
shutil.copy(f"{patch}/lfm2_moe.py", f"{pkg}/models/macos/lfm2_moe.py")
shutil.copy(f"{patch}/registry.py", f"{pkg}/models/registry.py")
shutil.copy(f"{patch}/macos.py", f"{pkg}/export/macos.py")
shutil.copy(f"{patch}/pipeline.py", f"{pkg}/export/pipeline.py")
print("patch applied", flush=True)
# make the package importable + CLI available without pulling its old pins
sys.path.insert(0, f"{ROOT}/coreai-models/python/src")
# ββ 2. VET the MoE def bit-exact vs HF reference (truncated) ββββββββββββββββββ
import torch
from transformers import AutoConfig
from transformers.models.lfm2_moe.modeling_lfm2_moe import Lfm2MoeForCausalLM as HFRef
from safetensors import safe_open
from coreai_models.models.macos.lfm2_moe import Lfm2MoeForCausalLM
from coreai_models.models.macos.lfm2 import num_attention_layers, num_conv_layers
N = 6 # covers dense(0,1)+MoE + conv + attention layers
DT = torch.float32
cfg = AutoConfig.from_pretrained(BASE)
cfg.num_hidden_layers = N
cfg.layer_types = list(cfg.layer_types[:N])
repo = snapshot_download(BASE, allow_patterns=["*.safetensors", "config.json"])
tensors = {}
for f in glob.glob(os.path.join(repo, "*.safetensors")):
with safe_open(f, framework="pt", device="cpu") as sf:
for k in sf.keys():
if k.startswith("model.layers.") and int(k.split(".")[2]) >= N:
continue
tensors[k] = sf.get_tensor(k).to(DT)
hf = HFRef(cfg).to(DT)
hf_miss, hf_unexp = hf.load_state_dict(dict(tensors), strict=False)
hf.tie_weights(); hf.eval()
# DIAGNOSTIC: if HF's own reference did not load the experts (per-expert checkpoint vs
# fused gate_up_proj in this transformers version), the reference is garbage and the vet
# would false-negative. Report it, and if so, convert the checkpoint to fused for HF.
_hf_real_miss = [m for m in hf_miss if "rotary" not in m and "lm_head" not in m]
print(f"[VET] HF ref load: {len(_hf_real_miss)} missing, {len(hf_unexp)} unexpected", flush=True)
if _hf_real_miss:
print(f"[VET] HF missing sample: {_hf_real_miss[:6]}", flush=True)
# HF uses fused experts.gate_up_proj / experts.down_proj β build them from per-expert w1/w3/w2
fused = dict(tensors)
import re as _re
for i in range(N):
pfx = f"model.layers.{i}.feed_forward.experts"
if f"{pfx}.0.w1.weight" not in fused:
continue
ne = 0
while f"{pfx}.{ne}.w1.weight" in fused:
ne += 1
gate_up = torch.stack([
torch.cat([fused.pop(f"{pfx}.{e}.w1.weight"), fused.pop(f"{pfx}.{e}.w3.weight")], dim=0)
for e in range(ne)
]) # (E, 2*inter, hidden)
down = torch.stack([fused.pop(f"{pfx}.{e}.w2.weight") for e in range(ne)]) # (E, hidden, inter)
fused[f"model.layers.{i}.feed_forward.experts.gate_up_proj"] = gate_up
fused[f"model.layers.{i}.feed_forward.experts.down_proj"] = down
hf = HFRef(cfg).to(DT)
hf_miss2, _ = hf.load_state_dict(fused, strict=False)
hf.tie_weights(); hf.eval()
still = [m for m in hf_miss2 if "rotary" not in m and "lm_head" not in m]
print(f"[VET] HF ref after fused-convert: {len(still)} missing (want 0)", flush=True)
ours = Lfm2MoeForCausalLM(cfg)
sd = dict(tensors); ours._mutate_state_dict(sd)
miss, unexp = ours.load_state_dict(sd, strict=False)
real_miss = [m for m in miss if "lm_head" not in m]
assert not real_miss, f"ours missing keys: {real_miss[:8]}"
assert not unexp, f"ours unexpected keys: {unexp[:8]}"
ours = ours.to(DT).eval()
torch.manual_seed(0)
prompt = torch.randint(1, cfg.vocab_size, (1, 7), dtype=torch.int32)
step = torch.randint(1, cfg.vocab_size, (1, 1), dtype=torch.int32)
with torch.no_grad():
ho = hf(prompt.long(), use_cache=True)
hs = hf(step.long(), past_key_values=ho.past_key_values, use_cache=True)
n_attn = num_attention_layers(cfg); n_conv = num_conv_layers(cfg)
hd = getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads
kc = torch.zeros(n_attn, 1, cfg.num_key_value_heads, 64, hd, dtype=DT)
vc = torch.zeros_like(kc)
cd = getattr(cfg, "conv_dim", cfg.hidden_size)
cv = torch.zeros(n_conv, 1, cd, cfg.conv_L_cache - 1, dtype=DT)
op = ours(prompt, torch.arange(7, dtype=torch.int32).unsqueeze(0), kc, vc, cv)
os_ = ours(step, torch.arange(8, dtype=torch.int32).unsqueeze(0), kc, vc, cv)
d1 = (ho.logits - op).abs().max().item()
d2 = (hs.logits - os_).abs().max().item()
scale = ho.logits.abs().max().item()
tol = 3e-3 * max(scale, 1.0)
print(f"\n[VET] prefill Ξ={d1:.6f} decode Ξ={d2:.6f} scale={scale:.1f} tol={tol:.4f}", flush=True)
if not (d1 < tol and d2 < tol):
print("[VET] β MISMATCH β our MoE def diverges from HF. ABORTING export (credit saved).", flush=True)
sys.exit(1)
print("[VET] β
MoE def is BIT-EXACT vs HF reference. Proceeding to full export.", flush=True)
del hf, ours, tensors, sd
# ββ 3. full 8B export (routed MoE + conv1d, 4-bit, bounded 4096, mmap path) ββββ
env = dict(os.environ)
env["PYTHONPATH"] = f"{ROOT}/coreai-models/python/src:" + env.get("PYTHONPATH", "")
run([sys.executable, "-m", "coreai_models.llm.export", BASE,
"--experimental", "--compute-precision", "float16",
"--max-context-length", str(CTX),
"--output-dir", f"{ROOT}/out"], env=env)
# inject chat_template into tokenizer_config (FIX4) so the bundle chats correctly
bdir = glob.glob(f"{ROOT}/out/*")[0]
tc_path = os.path.join(bdir, "tokenizer", "tokenizer_config.json")
jinja = os.path.join(bdir, "tokenizer", "chat_template.jinja")
if os.path.exists(tc_path) and os.path.exists(jinja):
tc = json.load(open(tc_path))
if "chat_template" not in tc:
tc["chat_template"] = open(jinja).read()
json.dump(tc, open(tc_path, "w"), indent=2, ensure_ascii=False)
print("chat_template injected", flush=True)
# ββ 4. push bundle to the Hub βββββββββββββββββββββββββββββββββββββββββββββββββ
from huggingface_hub import HfApi
api = HfApi()
api.create_repo(OUT_REPO, repo_type="model", exist_ok=True)
api.upload_folder(folder_path=f"{ROOT}/out", repo_id=OUT_REPO, repo_type="model")
print(f"\nβ
DONE β bundle pushed to https://huggingface.co/{OUT_REPO}", flush=True)
print("Download to the Mac and compile+bench with Tools/CoreAIBench.", flush=True)
|