File size: 2,026 Bytes
3da88e3 | 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 | 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)
|