| import json |
|
|
| import mlx.core as mx |
| import numpy as np |
| import pytest |
| import torch |
| from safetensors.torch import save_file |
|
|
| from h3_mlx.convert import convert |
| from h3_mlx.weights import Linear, WeightStore |
|
|
|
|
| def test_quantized_shards_roundtrip_and_preserve_qkv_layout(tmp_path): |
| model = tmp_path / "original" |
| transformer = model / "transformer" |
| transformer.mkdir(parents=True) |
| config = { |
| "hidden_size": 128, |
| "num_layers": 1, |
| "token_refiner_num_layers": 0, |
| "num_attention_heads": 2, |
| "attention_head_dim": 128, |
| "ffn_hidden_size": 128, |
| "audio_enabled": False, |
| } |
| (transformer / "config.json").write_text(json.dumps(config)) |
| torch.manual_seed(77) |
| tensors = {} |
| for name in ( |
| "norm1.weight", |
| "norm2.weight", |
| "attn.q_norm.weight", |
| "attn.k_norm.weight", |
| ): |
| tensors["blocks.0." + name] = torch.randn(128).bfloat16() |
| for name, shape in ( |
| ("attn.qkv_proj", (768, 128)), |
| ("attn.out_proj", (128, 256)), |
| ("mlp.fc1", (256, 128)), |
| ("mlp.fc2", (128, 128)), |
| ): |
| tensors[f"blocks.0.{name}.weight"] = torch.randn(*shape).bfloat16() |
| save_file(tensors, str(transformer / "model.safetensors")) |
| output = tmp_path / "converted" |
| meta = convert(model, output, bits=8, group_size=64) |
| assert meta["format"] == "h3-mlx-core-v1" and len(meta["files"]) == 1 |
| original = WeightStore(transformer) |
| converted = WeightStore(output) |
| a = Linear.load(original, "blocks.0.attn.qkv_proj", bits=8, grouped_heads=(2, 128)) |
| b = Linear.load(converted, "blocks.0.attn.qkv_proj", grouped_heads=(2, 128)) |
| for x, y in ((a.weight, b.weight), (a.scales, b.scales), (a.qbiases, b.qbiases)): |
| np.testing.assert_array_equal( |
| np.asarray(x.view(mx.uint16) if x.dtype == mx.bfloat16 else x), |
| np.asarray(y.view(mx.uint16) if y.dtype == mx.bfloat16 else y), |
| ) |
| with pytest.raises(FileExistsError): |
| convert(model, output) |
|
|
