88 lines
3.8 KiB
Python
88 lines
3.8 KiB
Python
"""直接把 per-expert safetensors 打包成「每专家一个连续 blob」(按层一个文件)。
|
||
|
||
字节布局由 prep/blob_layout.py 统一描述（v1 affine / v2 mxfp4），与 blob_loader 完全一致：
|
||
按 _split_meta.json 的 blob_format 选段表，每专家依次写各 proj 的 [weight, scales(, biases)]。
|
||
- weight: uint32；v1 的 scales/biases 为 bf16 按 uint16 原始 2 字节；v2 mxfp4 的 scales 为 uint8。
|
||
|
||
环境变量：EXPERT_DIR(源 per-expert) / BLOB_DIR(输出) / BITS / GROUP / LAYERS(逗号或 all)
|
||
"""
|
||
import json
|
||
import os
|
||
|
||
import mlx.core as mx
|
||
import numpy as np
|
||
|
||
from mlx_streaming.prep.blob_layout import layout_for, BLOB_V1_AFFINE, BLOB_V2_MXFP4
|
||
|
||
EXPERT_DIR = os.environ.get("EXPERT_DIR", "/tmp/qwen3_next_experts_8bit_g128")
|
||
OUT = os.environ.get("BLOB_DIR", "/tmp/cb_8bit_blob")
|
||
BITS = int(os.environ.get("BITS", "8"))
|
||
GROUP = int(os.environ.get("GROUP", "128"))
|
||
|
||
|
||
def _meta():
|
||
m = json.load(open(os.path.join(EXPERT_DIR, "_split_meta.json")))
|
||
d = m["dims"]
|
||
fmt = m.get("blob_format", BLOB_V1_AFFINE)
|
||
return (int(d["num_experts"]), int(d["hidden"]), int(d["moe_intermediate"]),
|
||
fmt, d.get("quant_mode", "affine"))
|
||
|
||
|
||
def _raw_bytes(arr) -> bytes:
|
||
"""uint32 → 4B/元素；uint8 → 1B/元素；其余(affine 的 bf16 scales/biases) → uint16 2B/元素。"""
|
||
if arr.dtype == mx.uint32:
|
||
return np.array(arr, copy=False).tobytes()
|
||
if arr.dtype == mx.uint8:
|
||
return np.array(arr, copy=False).tobytes()
|
||
return np.array(arr.view(mx.uint16), copy=False).tobytes()
|
||
|
||
|
||
def pack_layer(layer, num_experts, segs, stride, fmt, quant_mode) -> None:
|
||
out_path = os.path.join(OUT, f"layer{layer:02d}.blob")
|
||
with open(out_path, "wb") as f:
|
||
for e in range(num_experts):
|
||
w = mx.load(os.path.join(EXPERT_DIR, f"layer{layer:02d}_expert{e:03d}.safetensors"))
|
||
for proj, tensor, dt, shape, nb in segs:
|
||
b = _raw_bytes(w[f"{proj}.{tensor}"])
|
||
assert len(b) == nb, f"L{layer} e{e} {proj}.{tensor}: {len(b)}!={nb}"
|
||
f.write(b)
|
||
assert os.path.getsize(out_path) == stride * num_experts
|
||
index = {"format": fmt, "quant_mode": quant_mode, "layer": layer,
|
||
"num_experts": num_experts, "stride": stride,
|
||
"page_aligned": stride % 16384 == 0,
|
||
"segments": [{"proj": p, "tensor": t, "nbytes": n} for p, t, _, _, n in segs]}
|
||
with open(os.path.join(OUT, f"layer{layer:02d}.blob.index.json"), "w") as f:
|
||
json.dump(index, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def _resolve_layers(spec: str, num_experts: int) -> list:
|
||
spec = spec.strip()
|
||
if spec.lower() == "all":
|
||
layers = set()
|
||
for name in os.listdir(EXPERT_DIR):
|
||
if name.startswith("layer") and name.endswith("_expert000.safetensors"):
|
||
layers.add(int(name[5:7]))
|
||
return sorted(layers)
|
||
return [int(x) for x in spec.split(",") if x.strip()]
|
||
|
||
|
||
def main():
|
||
os.makedirs(OUT, exist_ok=True)
|
||
num_experts, hidden, inter, fmt, quant_mode = _meta()
|
||
segs, stride = layout_for(fmt, hidden, inter, BITS, GROUP)
|
||
layers = _resolve_layers(os.environ.get("LAYERS", "all"), num_experts)
|
||
for i, L in enumerate(layers):
|
||
pack_layer(L, num_experts, segs, stride, fmt, quant_mode)
|
||
print(f" layer {L} packed ({i+1}/{len(layers)})", flush=True)
|
||
summary = {"format": fmt, "quant_mode": quant_mode, "stride": stride,
|
||
"page_aligned": stride % 16384 == 0, "num_experts": num_experts,
|
||
"layers": layers, "bits": BITS, "group_size": GROUP}
|
||
with open(os.path.join(OUT, "blob_index.json"), "w") as f:
|
||
json.dump(summary, f, ensure_ascii=False, indent=2)
|
||
print(json.dumps({"out": OUT, "n_layers": len(layers), "stride_bytes": stride,
|
||
"page_aligned": stride % 16384 == 0, "format": fmt}, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|