| """Convert attention/MLP matrices one block at a time, retaining native precision islands.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import logging |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from .config import DiTConfig |
| from .runtime import configure_memory |
| from .weights import Linear, WeightStore |
|
|
| log = logging.getLogger(__name__) |
|
|
|
|
| def convert(model_path, destination, *, bits=8, group_size=64): |
| configure_memory() |
| model_path, destination = Path(model_path), Path(destination) |
| cfg = DiTConfig.load(model_path / "transformer/config.json") |
| source = WeightStore(model_path / "transformer") |
| if destination.exists() and any(destination.iterdir()): |
| raise FileExistsError(f"Conversion destination must be empty: {destination}") |
| destination.mkdir(parents=True, exist_ok=True) |
| prefixes = [f"blocks.{i}" for i in range(cfg.num_layers)] |
| prefixes += [ |
| f"token_refiner.blocks.{i}" for i in range(cfg.token_refiner_num_layers) |
| ] |
| metadata = { |
| "format": "h3-mlx-core-v1", |
| "bits": bits, |
| "group_size": group_size, |
| "qkv_layout": "qkv-head-dim", |
| "source_fingerprint": source.fingerprint(), |
| "source_config_sha256": hashlib.sha256( |
| (model_path / "transformer/config.json").read_bytes() |
| ).hexdigest(), |
| "precision_policy": "Only BF16 attention/MLP matrices quantized; AdaLN and FP32 islands remain in the original checkpoint", |
| } |
| files = [] |
| for i, prefix in enumerate(prefixes): |
| tensors = {} |
| for name in ( |
| "norm1.weight", |
| "norm2.weight", |
| "attn.q_norm.weight", |
| "attn.k_norm.weight", |
| ): |
| tensors[f"{prefix}.{name}"] = source.get(f"{prefix}.{name}") |
| for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"): |
| grouped = ( |
| (cfg.num_attention_heads, cfg.attention_head_dim) |
| if name == "attn.qkv_proj" |
| else None |
| ) |
| layer = Linear.load( |
| source, |
| f"{prefix}.{name}", |
| bits=bits, |
| group_size=group_size, |
| grouped_heads=grouped, |
| ) |
| for suffix, value in ( |
| ("weight", layer.weight), |
| ("scales", layer.scales), |
| ("biases", layer.qbiases), |
| ("bias", layer.bias), |
| ): |
| if value is not None: |
| tensors[f"{prefix}.{name}.{suffix}"] = value |
| del layer |
| mx.eval(tensors) |
| file = destination / f"core-{i:03d}.safetensors" |
| tmp = file.with_suffix(".tmp.safetensors") |
| mx.save_safetensors(str(tmp), tensors) |
| tmp.replace(file) |
| digest = hashlib.sha256() |
| with file.open("rb") as handle: |
| for data in iter(lambda: handle.read(8 * 1024**2), b""): |
| digest.update(data) |
| files.append( |
| { |
| "file": file.name, |
| "bytes": file.stat().st_size, |
| "sha256": digest.hexdigest(), |
| } |
| ) |
| del tensors |
| mx.clear_cache() |
| log.info("Converted block %d/%d", i + 1, len(prefixes)) |
| metadata["files"] = files |
| (destination / "quantization.json").write_text( |
| json.dumps(metadata, indent=2) + "\n" |
| ) |
| return metadata |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--output", required=True) |
| p.add_argument("--bits", type=int, choices=(4, 8), default=8) |
| p.add_argument("--group-size", type=int, choices=(32, 64, 128), default=64) |
| a = p.parse_args() |
| logging.basicConfig(level=logging.INFO) |
| convert(a.model, a.output, bits=a.bits, group_size=a.group_size) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
