File size: 8,936 Bytes
ad024da | 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 | """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()
|