File size: 4,562 Bytes
a066584 | 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 | # Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Weight remapping from the official Mk1 safetensors export to MLX."""
from __future__ import annotations
from collections.abc import Iterable
import mlx.core as mx
# PyTorch → MLX
# *.experts.down_proj *.experts.down_proj.weight
# *.experts.gate_up_proj *.experts.gate_up_proj.weight
# *.in_proj_weight [3D, D] query/key/value_proj.weight
# *.in_proj_bias [3D] query/key/value_proj.bias
# token_ff.0 / token_ff.2 token_ff.layers.0 / token_ff.layers.2
# memory_ff.0 / memory_ff.2 memory_ff.layers.0 / memory_ff.layers.2
_ATTENTION_MODULES = (
"local_attention",
"token_memory_attention",
"memory_token_attention",
)
_SKIP_SUBSTRINGS = (
"rotary_emb",
"lm_head.weight",
)
_CLIP_MARKERS = ("input_max", "input_min", "output_max", "output_min")
def should_keep_source_key(key: str) -> bool:
if any(marker in key for marker in _SKIP_SUBSTRINGS):
return False
if key.startswith("model.encoder.language_model.") and not key.endswith(
".layer_scalar"
):
return False
if key.startswith("model.encoder.vision_tower.") or key.startswith(
"model.encoder.embed_vision."
):
if any(marker in key for marker in _CLIP_MARKERS):
return False
return True
def _split_qkv(prefix: str, value: mx.array) -> list[tuple[str, mx.array]]:
if value.ndim == 1:
width = value.shape[0]
if width % 3:
raise ValueError(f"Cannot split QKV bias for {prefix}: shape {value.shape}")
head = width // 3
pieces = (value[:head], value[head : 2 * head], value[2 * head :])
names = ("query_proj.bias", "key_proj.bias", "value_proj.bias")
elif value.ndim == 2:
width = value.shape[0]
if width % 3:
raise ValueError(
f"Cannot split QKV weight for {prefix}: shape {value.shape}"
)
head = width // 3
pieces = (value[:head], value[head : 2 * head], value[2 * head :])
names = ("query_proj.weight", "key_proj.weight", "value_proj.weight")
if pieces[0].shape[0] != pieces[0].shape[1]:
raise ValueError(
f"Split QKV weight for {prefix} is not square: {pieces[0].shape}"
)
else:
raise ValueError(f"Unexpected QKV tensor rank for {prefix}: {value.shape}")
return [(f"{prefix}.{name}", piece) for name, piece in zip(names, pieces)]
def remap_weight(key: str, value: mx.array) -> list[tuple[str, mx.array]]:
"""Map one official Mk1 tensor onto one or more MLX parameter names."""
if not should_keep_source_key(key):
return []
if key.endswith(".experts.down_proj"):
return [(key + ".weight", value)]
if key.endswith(".experts.gate_up_proj"):
return [(key + ".weight", value)]
for module in _ATTENTION_MODULES:
in_proj_weight = f".{module}.in_proj_weight"
in_proj_bias = f".{module}.in_proj_bias"
if key.endswith(in_proj_weight):
prefix = key[: -len(".in_proj_weight")]
return _split_qkv(prefix, value)
if key.endswith(in_proj_bias):
prefix = key[: -len(".in_proj_bias")]
return _split_qkv(prefix, value)
if ".token_ff.0." in key or key.endswith(".token_ff.0.weight") or key.endswith(
".token_ff.0.bias"
):
return [(key.replace(".token_ff.0.", ".token_ff.layers.0."), value)]
if ".token_ff.2." in key or key.endswith(".token_ff.2.weight") or key.endswith(
".token_ff.2.bias"
):
return [(key.replace(".token_ff.2.", ".token_ff.layers.2."), value)]
if ".memory_ff.0." in key or key.endswith(".memory_ff.0.weight") or key.endswith(
".memory_ff.0.bias"
):
return [(key.replace(".memory_ff.0.", ".memory_ff.layers.0."), value)]
if ".memory_ff.2." in key or key.endswith(".memory_ff.2.weight") or key.endswith(
".memory_ff.2.bias"
):
return [(key.replace(".memory_ff.2.", ".memory_ff.layers.2."), value)]
return [(key, value)]
def remap_state_dict(
source: Iterable[tuple[str, mx.array]],
) -> dict[str, mx.array]:
remapped: dict[str, mx.array] = {}
for key, value in source:
for new_key, new_value in remap_weight(key, value):
if new_key in remapped:
raise ValueError(f"Duplicate remapped key: {new_key}")
remapped[new_key] = new_value
return remapped
|