Hugging Face's logo

androidli
/
laya-multilingual-onnx-int4

ONNX
multilingual
Chinese
English
onnxruntime
int4
4-bit precision
weight-only-quantization
MatMulNBits
GatherBlockQuantized
system-one
decision-model
cpu-inference
Model card Files
xet
 Community
laya-multilingual-onnx-int4 / scripts /emb_quant.py
androidli's picture
androidli
laya multilingual ONNX int4 (MatMulNBits + GatherBlockQuantized), tokenizer/config + scripts
ad024da verified
Raw History Blame Contribute Delete
8.94 kB
"""Quantize the vocabulary embedding of the laya ONNX graph, keeping the int4 layer weights intact.
Motivation: MatMulNBits only rewrites MatMul weights, so `tok_embeddings` [256000, 768] stays fp32
= 786MB = 61% of the model. ORT's own dynamic quantization does cover Gather, but with a SINGLE
per-tensor scale for the whole 256k x 768 table (that is part of why the int8 arm lost 8pp).
This does graph surgery instead: replace the fp32 embedding initializer with a quantized one plus a
DequantizeLinear feeding the Gather, with a scale PER ROW (每行一个 scale = 每个 token 向量一个).
python emb_quant.py --src laya-zh-int4-b32.onnx --dst laya-zh-int4-b32-emb8r.onnx --mode int8-row
modes: int8-row (per-row int8) | fp8-row (E4M3 per-row, needs opset 21 + ORT fp8 kernel)
u4-row (per-row uint4) | u4-block32 (4-bit blocked along K, ONNX blocked-quant layout)
The embedding is a pure lookup table, so quantization error here = a small perturbation of the
token vector; per-row scales keep that perturbation proportional to the row's own magnitude.
"""
import argparse, os, time
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
EMB = "encoder.embeddings.tok_embeddings.weight"
DQ_OUT = EMB + "_dqout" # the DequantizeLinear result the Gather must consume
def get_emb(model):
for t in model.graph.initializer:
if t.name == EMB:
return numpy_helper.to_array(t)
raise SystemExit(f"{EMB} not found")
def replace_emb(model, new_inits):
"""Insert DequantizeLinear(<q>, <scale>, [zp]) -> replaces the fp32 initializer as Gather input."""
# drop the old initializer
old = [i for i in model.graph.initializer if i.name == EMB]
for i in old:
model.graph.initializer.remove(i)
model.graph.initializer.extend(new_inits)
# find the Gather that consumes EMB
for n in model.graph.node:
if n.op_type == "Gather" and EMB in n.input:
idx = list(n.input).index(EMB)
n.input[idx] = DQ_OUT
return n
raise SystemExit("no Gather consumes the embedding")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", required=True)
ap.add_argument("--dst", required=True)
ap.add_argument("--mode", default="int8-row")
ap.add_argument("--block", type=int, default=32)
ap.add_argument("--axis", type=int, default=1,
help="1 = per-column scale [D] (ORT default per-axis), 0 = per-row scale [V]")
a = ap.parse_args()
src, dst = os.path.expanduser(a.src), os.path.expanduser(a.dst)
t0 = time.perf_counter()
model = onnx.load(src, load_external_data=True)
if a.mode.startswith("u4"):
del model.graph.value_info[:] # packed/blocked layouts need ORT to re-infer shapes
w = get_emb(model).astype(np.float32)
N, K = w.shape
print(f"embedding {w.shape} {w.nbytes/1e6:.0f}MB mode={a.mode}", flush=True)
new_inits, dq_inputs, attrs, extra = [], [], {}, {}
if a.mode == "int8-row":
if a.axis == 1: # one scale per hidden dim (ORT's per-axis default)
sc = (np.abs(w).max(axis=0) / 127.0).astype(np.float32) # [D]
q = np.clip(np.round(w / sc[None, :]), -127, 127).astype(np.int8)
else: # one scale per token row
sc = (np.abs(w).max(axis=1) / 127.0).astype(np.float32) # [V]
q = np.clip(np.round(w / sc[:, None]), -127, 127).astype(np.int8)
attrs["axis"] = 0
dq_inputs = [EMB + "_q8", EMB + "_scale"]
new_inits = [numpy_helper.from_array(q, EMB + "_q8"),
numpy_helper.from_array(sc, EMB + "_scale")]
elif a.mode == "fp8-row":
import ml_dtypes
if a.axis == 1:
sc = (np.abs(w).max(axis=0) / 448.0).astype(np.float32) # E4M3 max 448
q = (w / sc[None, :]).astype(ml_dtypes.float8_e4m3fn)
else:
sc = (np.abs(w).max(axis=1) / 448.0).astype(np.float32)
q = (w / sc[:, None]).astype(ml_dtypes.float8_e4m3fn)
attrs["axis"] = 0
dq_inputs = [EMB + "_f8", EMB + "_scale"]
new_inits = [numpy_helper.from_array(q, EMB + "_f8"),
numpy_helper.from_array(sc, EMB + "_scale")]
extra["opset"] = 21 # float8 types land in 19+, DQ/blocked in 21
elif a.mode == "u4-row":
sc = np.abs(w).max(axis=1) / 7.0
q = np.clip(np.round(w / sc[:, None]) + 8, 0, 15).astype(np.uint8) # zp=8 -> symmetric
packed = (q[:, 0::2] | (q[:, 1::2] << 4)).astype(np.uint8) # low nibble = first elem
zp = np.full((N, 1), 8, np.uint8)
dq_inputs = [EMB + "_q4", EMB + "_scale", EMB + "_zp"]
new_inits = [numpy_helper.from_array(packed, EMB + "_q4"),
numpy_helper.from_array(sc.reshape(N, 1).astype(np.float32), EMB + "_scale"),
numpy_helper.from_array(zp.reshape(N, 1), EMB + "_zp")]
extra["opset"] = 21
attrs["block_size"] = K # one block == the whole row
elif a.mode == "u4-gbq":
# 4-bit vocabulary via com.microsoft::GatherBlockQuantized.
# Layout (verified on this box with a synthetic probe): data uint8 [V, D/2] with the LOW
# nibble = first element; scales fp32 [V, D/block]; zero_points are OPTIONAL for 4-bit and
# when omitted the op uses 8 (symmetric) -- exactly what we want, and it keeps the table
# packed at inference (no fp32 materialisation, unlike DequantizeLinear).
B = a.block
assert B >= 16 and (B & (B - 1)) == 0 and K % B == 0, "block_size must be a power of 2 >= 16 dividing D"
wb = w.reshape(N, K // B, B)
sc = np.abs(wb).max(axis=2) / 7.0
q = np.clip(np.round(wb / sc[:, :, None]) + 8, 0, 15).astype(np.uint8)
packed = (q[:, :, 0::2] | (q[:, :, 1::2] << 4)).reshape(N, K // 2)
new_inits = [numpy_helper.from_array(packed.astype(np.uint8), EMB + "_q4packed"),
numpy_helper.from_array(sc.astype(np.float32), EMB + "_scales")]
gbq_attrs = dict(bits=4, block_size=B, gather_axis=0, quantize_axis=1)
elif a.mode == "u4-block32":
B = a.block
assert K % B == 0 and B % 2 == 0
wb = w.reshape(N, K // B, B)
sc = np.abs(wb).max(axis=2) / 7.0
q = np.clip(np.round(wb / sc[:, :, None]) + 8, 0, 15).astype(np.uint8)
packed = (q[:, :, 0::2] | (q[:, :, 1::2] << 4)).reshape(N, K // 2)
zp = np.full((N, K // B), 8, np.uint8)
dq_inputs = [EMB + "_q4", EMB + "_scale", EMB + "_zp"]
new_inits = [numpy_helper.from_array(packed, EMB + "_q4"),
numpy_helper.from_array(sc.astype(np.float32), EMB + "_scale"),
numpy_helper.from_array(zp, EMB + "_zp")]
extra["opset"] = 21
attrs["block_size"] = B
else:
raise SystemExit("unknown mode")
if a.mode == "u4-gbq":
gathered = [n for n in model.graph.node if n.op_type == "Gather" and EMB in n.input][0]
idx_in = list(gathered.input)[1 - list(gathered.input).index(EMB)]
node = helper.make_node("GatherBlockQuantized", [EMB + "_q4packed", idx_in, EMB + "_scales"],
list(gathered.output), name="emb_gbq", domain="com.microsoft", **gbq_attrs)
model.graph.node.remove(gathered)
model.graph.node.append(node)
for i in [t for t in model.graph.initializer if t.name == EMB]:
model.graph.initializer.remove(i)
model.graph.initializer.extend(new_inits)
doms = {op.domain for op in model.opset_import}
if "com.microsoft" not in doms:
model.opset_import.append(helper.make_opsetid("com.microsoft", 1))
g = node
else:
dq = helper.make_node("DequantizeLinear", dq_inputs, [DQ_OUT], name="emb_dq", **attrs)
g = replace_emb(model, new_inits)
model.graph.node.insert(list(model.graph.node).index(g), dq)
if extra.get("opset"):
for op in model.opset_import:
if op.domain in ("", "ai.onnx"):
op.version = max(op.version, extra["opset"])
for stale in (dst, dst + ".data"): # save_model refuses to overwrite an existing .data
if os.path.exists(stale):
os.remove(stale)
onnx.save_model(model, dst, save_as_external_data=True, all_tensors_to_one_file=True,
location=os.path.basename(dst) + ".data", size_threshold=1024)
got = os.path.getsize(dst) + (os.path.getsize(dst + ".data") if os.path.exists(dst + ".data") else 0)
print(f"done in {time.perf_counter()-t0:.1f}s -> {got/1e6:.0f}MB (embedding part: "
f"{sum(len(i.raw_data) for i in new_inits)/1e6:.0f}MB, was {w.nbytes/1e6:.0f}MB)")
print("EMB_QUANT_DONE")
if __name__ == "__main__":
main()