49 lines
2.6 KiB
Python
49 lines
2.6 KiB
Python
|
"""Real inputs for the output head: the last layer's output, pre-final-norm."""
|
||
|
|
import importlib.util, json, sys, types
|
||
|
|
from pathlib import Path
|
||
|
|
import numpy as np, torch
|
||
|
|
|
||
|
|
ROOT = Path("/data/zhoutaichang/embedding_infer/vllm-omni")
|
||
|
|
HERE = Path("/data/zhoutaichang/embedding_infer/analysis/experiments/spark_edge")
|
||
|
|
def _load(n, p):
|
||
|
|
sp = importlib.util.spec_from_file_location(n, p); m = importlib.util.module_from_spec(sp)
|
||
|
|
sys.modules[n] = m; sp.loader.exec_module(m); return m
|
||
|
|
for pkg in ("vllm_omni", "vllm_omni.edge"):
|
||
|
|
mod = types.ModuleType(pkg); mod.__path__ = []; sys.modules[pkg] = mod
|
||
|
|
_load("vllm_omni.edge.qnn_export", ROOT / "vllm_omni/edge/qnn_export.py")
|
||
|
|
se = _load("vllm_omni.edge.spark_export", ROOT / "vllm_omni/edge/spark_export.py")
|
||
|
|
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
MODEL = "XHToken/Spark-X2.5-1.7B"
|
||
|
|
SNAP = ("/data/zhoutaichang/.cache/huggingface/hub/models--XHToken--Spark-X2.5-1.7B/"
|
||
|
|
"snapshots/448e61eb392c00f2c403185c5b56d5e0665bfaab")
|
||
|
|
TEXTS = ["The quick brown fox jumps over the lazy dog. " * 60,
|
||
|
|
"Engineers argued about memory bandwidth on the train home. " * 60,
|
||
|
|
"def forward(self, x):\n return self.down(self.act(self.gate(x)))\n" * 60,
|
||
|
|
"移动端推理的瓶颈通常是内存带宽而不是算力。" * 80]
|
||
|
|
|
||
|
|
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||
|
|
m = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True,
|
||
|
|
dtype=torch.float32).eval()
|
||
|
|
hf = json.loads(Path(f"{SNAP}/config.json").read_text())
|
||
|
|
cfg = se.config_from_spark(hf, slice(0, 0), include_lm_head=True)
|
||
|
|
step = se.SparkDecodeStep(cfg); se.load_spark_weights(step, SNAP); step.eval()
|
||
|
|
|
||
|
|
grabbed = {}
|
||
|
|
h = m.model.layers[-1].register_forward_hook(
|
||
|
|
lambda mod, inp, out: grabbed.__setitem__("h", out[0] if isinstance(out, tuple) else out))
|
||
|
|
xs, outs = [], []
|
||
|
|
for t in TEXTS:
|
||
|
|
ids = tok(t, add_special_tokens=False)["input_ids"][:256]
|
||
|
|
with torch.no_grad():
|
||
|
|
m(torch.tensor([ids]))
|
||
|
|
x = grabbed["h"][:, -1:, :].float()
|
||
|
|
logits = step(x, *[torch.zeros(1)] * 0) if False else step(
|
||
|
|
x, torch.zeros(1,1,1,4), torch.zeros(1,1,1,4),
|
||
|
|
torch.zeros(1,1,1,4), torch.zeros(1,1,1,4),
|
||
|
|
torch.zeros(1,1,1,1), torch.zeros(1,1,1,1))[0]
|
||
|
|
xs.append(x.numpy().astype(np.float32)); outs.append(logits.numpy().astype(np.float32))
|
||
|
|
h.remove()
|
||
|
|
np.savez(HERE / "real_acts_lmhead.npz", x=np.stack(xs), hidden_out=np.stack(outs))
|
||
|
|
print("captured", len(TEXTS), "| x absmax", np.abs(np.stack(xs)).max(),
|
||
|
|
"| logits absmax", np.abs(np.stack(outs)).max())
|