File size: 2,195 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 59 60 61 62 63 64 | import mlx.core as mx
import numpy as np
from h3_mlx.config import DiTConfig
from h3_mlx.dit import DiTBlock
from h3_mlx.weights import WeightStore
def test_packed_block_cache_and_compiled_bf16_boundaries(tmp_path):
mx.random.seed(729)
c = DiTConfig(
hidden_size=32, num_attention_heads=2, attention_head_dim=16, ffn_hidden_size=64
)
tensors = {}
for name, shape in {
"norm1": (32,),
"norm2": (32,),
"attn.q_norm": (16,),
"attn.k_norm": (16,),
"attn.qkv_proj": (96, 32),
"attn.out_proj": (32, 32),
"mlp.fc1": (128, 32),
"mlp.fc2": (32, 64),
}.items():
tensors[f"blocks.0.{name}.weight"] = (mx.random.normal(shape) * 0.02).astype(
mx.bfloat16
)
if len(shape) == 2:
tensors[f"blocks.0.{name}.bias"] = (
mx.random.normal((shape[0],)) * 0.01
).astype(mx.bfloat16)
original = tmp_path / "original"
original.mkdir()
mx.save_safetensors(str(original / "model.safetensors"), tensors)
store = WeightStore(original)
native = DiTBlock(store, "blocks.0", c)
packed = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed")
restored = DiTBlock(store, "blocks.0", c, packed_cache=tmp_path / "packed")
x = mx.random.normal((17, 32)).astype(mx.bfloat16)
params = tuple(
(mx.random.normal((3, 32)) * 0.1).astype(mx.bfloat16) for _ in range(6)
)
ids = mx.arange(17) % 3
expected = native(x, params=params, indices=ids, mlp_chunk=0)
state = tuple(
p.weight.state()
for p in (restored.qkv, restored.out, restored.fc1, restored.fc2)
)
def forward(x, params, ids, state):
return restored(
x, params=params, indices=ids, mlp_chunk=0, projection_weights=state
)
compiled = mx.compile(forward)
for actual in (
packed(x, params=params, indices=ids, mlp_chunk=0),
restored(x, params=params, indices=ids, mlp_chunk=0),
compiled(x, params, ids, state),
):
np.testing.assert_array_equal(
np.asarray(actual.view(mx.uint16)), np.asarray(expected.view(mx.uint16))
)
|