70 lines
2.7 KiB
Python
70 lines
2.7 KiB
Python
"""路线 B 离线工具：把模型里堆叠的 switch_mlp 专家权重按专家拆成 per-expert 小文件。
|
||
|
||
拆分后每个文件 layer{L:02d}_expert{E:03d}.safetensors 含扁平 dict：
|
||
gate_proj.weight / gate_proj.scales / gate_proj.biases
|
||
up_proj.weight / up_proj.scales / up_proj.biases
|
||
down_proj.weight / down_proj.scales / down_proj.biases
|
||
（非量化模型则只有 .weight，可能还有 .bias）
|
||
|
||
一次拆分、逐专家物化，全程低内存。
|
||
"""
|
||
import os
|
||
import sys
|
||
import json
|
||
|
||
import mlx.core as mx
|
||
|
||
PROJ_NAMES = ["gate_proj", "up_proj", "down_proj"]
|
||
|
||
|
||
def split_switch_glu(switch_glu, out_dir: str, layer: int) -> int:
|
||
"""把一个 SwitchGLU 的三组 SwitchLinear 沿专家维拆成 per-expert 文件。返回专家数。"""
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
E = switch_glu.gate_proj.num_experts
|
||
for e in range(E):
|
||
d = {}
|
||
for proj_name in PROJ_NAMES:
|
||
proj = getattr(switch_glu, proj_name)
|
||
for pname, p in proj.parameters().items():
|
||
if isinstance(p, mx.array) and p.ndim >= 1 and p.shape[0] == E:
|
||
d[f"{proj_name}.{pname}"] = p[e]
|
||
mx.eval(d) # 只物化这一个专家
|
||
path = os.path.join(out_dir, f"layer{layer:02d}_expert{e:03d}.safetensors")
|
||
mx.save_safetensors(path, d)
|
||
return E
|
||
|
||
|
||
def split_model(model_path: str, out_dir: str) -> dict:
|
||
"""加载模型（lazy）并把所有 MoE 层的专家拆到 out_dir。返回 {dims/统计}。"""
|
||
from mlx_lm import load
|
||
model, _ = load(model_path, lazy=True)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
moe_layers = []
|
||
dims = None
|
||
for l, layer in enumerate(model.layers):
|
||
mlp = getattr(layer, "mlp", None)
|
||
if mlp is not None and hasattr(mlp, "switch_mlp") and hasattr(mlp, "gate"):
|
||
sm = mlp.switch_mlp
|
||
split_switch_glu(sm, out_dir, l)
|
||
moe_layers.append(l)
|
||
if dims is None:
|
||
gp = sm.gate_proj
|
||
dims = {
|
||
"hidden": gp.input_dims,
|
||
"moe_intermediate": gp.output_dims,
|
||
"num_experts": gp.num_experts,
|
||
"group_size": getattr(gp, "group_size", None),
|
||
"bits": getattr(gp, "bits", None),
|
||
}
|
||
meta = {"out_dir": out_dir, "moe_layers": moe_layers, "dims": dims}
|
||
with open(os.path.join(out_dir, "_split_meta.json"), "w") as f:
|
||
json.dump(meta, f, ensure_ascii=False, indent=2)
|
||
return meta
|
||
|
||
|
||
if __name__ == "__main__":
|
||
mp = sys.argv[1]
|
||
od = sys.argv[2] if len(sys.argv) > 2 else "/tmp/mlx_qwen3_experts"
|
||
m = split_model(mp, od)
|
||
print(json.dumps(m, ensure_ascii=False, indent=2))
|