File size: 3,078 Bytes
78c877d | 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 | """支援量化 draft 的載入器（原 local_load 只支援 bf16）。"""
import json
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
from dflash_port.model_mlx import DFlash2DraftModel, DFlashDraftModel, DFlashConfig
from dflash_port import local_load as _ll
def load_draft_any(path):
path = Path(path)
cfg = json.loads((path / "config.json").read_text())
weights = {k: v for f in path.glob("*.safetensors") for k, v in mx.load(str(f)).items()}
# codebook 鍵名在兩種發佈裡不一致，統一成 .weight 結尾
for name in ("predecessor_codebook", "successor_codebook"):
bare = f"candidate_selector.{name}"
if bare in weights and f"{bare}.weight" not in weights:
weights[f"{bare}.weight"] = weights.pop(bare)
q = cfg.get("quantization") or {}
if not q: # 非量化 → 走原本的路徑
return _ll.load_draft_from_dir(path)
# 重用原載入器的 config 組裝（複製其邏輯的關鍵欄位）
d = cfg.get("dflash_config", {})
rope = cfg.get("rope_parameters") or cfg.get("rope_scaling")
layer_types = tuple(cfg.get("layer_types") or ["full_attention"] * cfg["num_hidden_layers"])
config = DFlashConfig(
hidden_size=cfg["hidden_size"], num_hidden_layers=cfg["num_hidden_layers"],
num_attention_heads=cfg["num_attention_heads"], num_key_value_heads=cfg["num_key_value_heads"],
head_dim=cfg["head_dim"], intermediate_size=cfg["intermediate_size"],
vocab_size=cfg["vocab_size"], rms_norm_eps=cfg["rms_norm_eps"],
rope_theta=cfg.get("rope_theta", (rope or {}).get("rope_theta", 10000.0)),
max_position_embeddings=cfg["max_position_embeddings"],
block_size=int(d.get("block_size", 16)), target_layer_ids=tuple(d["target_layer_ids"]),
num_target_layers=cfg["num_target_layers"], mask_token_id=d["mask_token_id"],
rope_scaling=rope, layer_types=layer_types, sliding_window=cfg.get("sliding_window"),
final_logit_softcapping=d.get("final_logit_softcapping", cfg.get("final_logit_softcapping")),
input_embedding_scale=float(d.get("input_embedding_scale", 1.0)),
output_multiplier=float(d.get("output_multiplier", 1.0)),
conv_kernel_size=int(d.get("conv_kernel_size", 0)), conv_group_size=int(d.get("conv_group_size", 0)),
selector_rank=int(d.get("selector_rank", 0)), selector_top_k=int(d.get("selector_top_k", 0)),
is_causal=cfg.get("is_causal"),
)
klass = DFlash2DraftModel if "DFlash2DraftModel" in (cfg.get("architectures") or []) else DFlashDraftModel
model = klass(config)
# 只量化「權重檔裡真的有 scales」的模組 —— 與 mlx-lm 的判斷方式一致
qset = {k[:-len(".scales")] for k in weights if k.endswith(".scales")}
nn.quantize(model, group_size=int(q["group_size"]), bits=int(q["bits"]),
class_predicate=lambda p, m: p in qset)
model.eval()
model.load_weights(list(weights.items()))
mx.eval(model.parameters())
return model
|