File size: 6,555 Bytes
e5030dc | 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 126 127 128 129 130 131 132 | """Stream the official checkpoint into MLX-VLM-compatible 4/8-bit shards."""
import argparse
import copy
import json
import shutil
import struct
import time
from pathlib import Path
import mlx.core as mx
import mlx.nn as nn
import numpy as np
from mlx.utils import tree_flatten
from mlx_vlm.models.muse_glimmer import Model, ModelConfig
from mlx_vlm.utils import skip_multimodal_module
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--revision", required=True)
parser.add_argument("--wait-for-download", action="store_true")
args = parser.parse_args()
config = json.loads((args.source / "config.json").read_text())
model = Model(ModelConfig.from_dict(config))
shapes = {k: v.shape for k, v in tree_flatten(model.parameters())}
quantized_paths = {
path for path, module in model.named_modules()
if isinstance(module, (nn.Linear, nn.Embedding))
and not skip_multimodal_module(path)
and module.weight.shape[-1] % 64 == 0
}
source_index = json.loads((args.source / "model.safetensors.index.json").read_text())
mapped = model.sanitize({key: None for key in source_index["weight_map"]})
if set(mapped) != set(shapes):
raise ValueError(f"Schema mismatch: missing={set(shapes)-set(mapped)}, extra={set(mapped)-set(shapes)}")
del model
mx.clear_cache()
mx.set_cache_limit(256 * 1024**2)
states = {}
for bits in (4, 8):
folder = args.output / f"meta-encoder-mlx-{bits}bit"
folder.mkdir(parents=True, exist_ok=True)
if list(folder.glob("*.safetensors")):
raise FileExistsError(f"Refusing to overwrite existing weights in {folder}")
states[bits] = dict(folder=folder, pending={}, bytes=0, total=0, index={}, shards=[])
def flush(state):
if not state["pending"]:
return
name = f"part-{len(state['shards'])+1:05d}.safetensors"
mx.save_safetensors(str(state["folder"] / name), state["pending"], metadata={"format": "mlx"})
state["index"].update({key: name for key in state["pending"]})
state["shards"].append(name)
state["pending"].clear()
state["bytes"] = 0
count = 0
for shard_name in sorted(set(source_index["weight_map"].values()), reverse=True):
if args.wait_for_download:
while not (args.source / shard_name).exists():
print(f"Waiting for download: {shard_name}", flush=True)
time.sleep(30)
with (args.source / shard_name).open("rb") as source:
header_size = struct.unpack("<Q", source.read(8))[0]
header = json.loads(source.read(header_size))
for original, info in header.items():
if original == "__metadata__":
continue
key = next(iter(Model.sanitize(None, {original: None})))
if tuple(info["shape"]) != tuple(shapes[key]):
raise ValueError(f"Shape mismatch for {key}")
if info["dtype"] != "BF16":
raise ValueError(f"Unexpected dtype {info['dtype']} for {key}")
start, end = info["data_offsets"]
source.seek(8 + header_size + start)
raw = source.read(end - start)
if len(raw) != end - start:
raise EOFError(original)
weight = mx.array(np.frombuffer(raw, dtype=np.uint16).reshape(info["shape"])).view(mx.bfloat16)
mx.eval(weight)
del raw
for bits, state in states.items():
if key.removesuffix(".weight") in quantized_paths and key.endswith(".weight"):
packed, scales, biases = mx.quantize(weight, group_size=64, bits=bits, mode="affine")
arrays = {key: packed, key.removesuffix(".weight") + ".scales": scales,
key.removesuffix(".weight") + ".biases": biases}
else:
arrays = {key: weight}
mx.eval(list(arrays.values()))
size = sum(value.nbytes for value in arrays.values())
if state["bytes"] + size > 2 * 1024**3:
flush(state)
state["pending"].update(arrays)
state["bytes"] += size
state["total"] += size
del arrays, weight
count += 1
if count % 100 == 0:
print(f"Converted {count}/{len(shapes)} tensors", flush=True)
for bits, state in states.items():
flush(state)
names = {}
for number, old in enumerate(state["shards"], 1):
new = f"model-{number:05d}-of-{len(state['shards']):05d}.safetensors"
(state["folder"] / old).rename(state["folder"] / new)
names[old] = new
index = {"metadata": {"total_size": state["total"]},
"weight_map": {k: names[v] for k, v in sorted(state["index"].items())}}
(state["folder"] / "model.safetensors.index.json").write_text(json.dumps(index, indent=2) + "\n")
for name in ("tokenizer.json", "tokenizer_config.json", "processor_config.json",
"generation_config.json", "chat_template.jinja", "LICENSE", "NOTICE", "USAGE_POLICY.md"):
if (args.source / name).exists():
shutil.copy2(args.source / name, state["folder"] / name)
out_config = copy.deepcopy(config)
out_config["quantization"] = {"group_size": 64, "bits": bits, "mode": "affine"}
out_config["quantization_config"] = copy.deepcopy(out_config["quantization"])
out_config["vision_config"]["skip_vision"] = True
(state["folder"] / "config.json").write_text(json.dumps(out_config, indent=2) + "\n")
provenance = {"source": "facebook/meta-encoder", "revision": args.revision,
"bits": bits, "group_size": 64, "mode": "affine",
"vision": "unchanged BF16", "source_tensors": count,
"quantized_modules": len(quantized_paths), "weight_bytes": state["total"],
"mlx": mx.__version__}
(state["folder"] / "conversion.json").write_text(json.dumps(provenance, indent=2) + "\n")
print(json.dumps(provenance), flush=True)
if __name__ == "__main__":
main()
|