93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
"""把 compute buffer 重打包成"每专家一个连续 blob"格式（按层一个文件）。
|
||
|
||
blob 内顺序：gate[w,s,b] + up[w,s,b] + down[w,s,b]，正好 864KB（= 54×16KB 页对齐）。
|
||
读一个专家 = 1 次 pread(stride, e*stride)，而不是当前的 9 次散读。
|
||
|
||
环境变量：COMPUTE_BUFFER_DIR(源)、BLOB_DIR(输出)、LAYERS(逗号分隔，默认 15,25)
|
||
"""
|
||
import json
|
||
import os
|
||
|
||
import numpy as np
|
||
|
||
HIDDEN = 2048
|
||
INTER = 512
|
||
GROUP = 128
|
||
BITS = 2
|
||
NUM_EXPERTS = 512
|
||
SRC = os.environ.get("COMPUTE_BUFFER_DIR", "/tmp/cb_2bit_g128")
|
||
OUT = os.environ.get("BLOB_DIR", "/tmp/cb_2bit_blob")
|
||
PROJS = (("gate_proj", INTER, HIDDEN), ("up_proj", INTER, HIDDEN), ("down_proj", HIDDEN, INTER))
|
||
|
||
|
||
def _layout():
|
||
"""返回 [(proj, tensor, nbytes_per_expert), ...] 和单专家总字节。"""
|
||
segs = []
|
||
for proj, out_dim, in_dim in PROJS:
|
||
words = in_dim * BITS // 32
|
||
groups = in_dim // GROUP
|
||
segs.append((proj, "weight", out_dim * words * 4))
|
||
segs.append((proj, "scales", out_dim * groups * 2))
|
||
segs.append((proj, "biases", out_dim * groups * 2))
|
||
return segs, sum(s[2] for s in segs)
|
||
|
||
|
||
def repack_layer(layer: int) -> dict:
|
||
os.makedirs(OUT, exist_ok=True)
|
||
segs, stride = _layout()
|
||
# 源 .bin 以 uint8 原始字节映射，按专家字节偏移切片
|
||
raw = {}
|
||
for proj, _, _ in PROJS:
|
||
base = os.path.join(SRC, f"layer{layer:02d}.{proj}")
|
||
for tensor in ("weight", "scales", "biases"):
|
||
raw[(proj, tensor)] = np.memmap(f"{base}.{tensor}.bin", dtype=np.uint8, mode="r")
|
||
out_path = os.path.join(OUT, f"layer{layer:02d}.blob")
|
||
with open(out_path, "wb") as f:
|
||
for e in range(NUM_EXPERTS):
|
||
for proj, tensor, nb in segs:
|
||
buf = raw[(proj, tensor)]
|
||
f.write(buf[e * nb:(e + 1) * nb].tobytes())
|
||
assert os.path.getsize(out_path) == stride * NUM_EXPERTS
|
||
index = {
|
||
"format": "expert_blob_v1",
|
||
"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)
|
||
return index
|
||
|
||
|
||
def _resolve_layers(spec: str) -> list[int]:
|
||
"""LAYERS=all → 读源 compute buffer 现有的全部层；否则按逗号解析。"""
|
||
spec = spec.strip()
|
||
if spec.lower() == "all":
|
||
layers = []
|
||
for name in os.listdir(SRC):
|
||
if name.endswith(".gate_proj.weight.bin") and name.startswith("layer"):
|
||
layers.append(int(name[len("layer"):len("layer") + 2]))
|
||
return sorted(set(layers))
|
||
return [int(x) for x in spec.split(",") if x.strip()]
|
||
|
||
|
||
def main():
|
||
layers = _resolve_layers(os.environ.get("LAYERS", "15,25"))
|
||
_, stride = _layout()
|
||
for L in layers:
|
||
repack_layer(L)
|
||
# 汇总 index，供 loader 校验/发现
|
||
summary = {"format": "expert_blob_v1", "stride": stride,
|
||
"page_aligned": stride % 16384 == 0, "num_experts": NUM_EXPERTS,
|
||
"layers": layers}
|
||
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}, ensure_ascii=False))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|