mini-deepseek-v4-flash / validation /tests_fuse2_mlx.py
Akahsizrr's picture
Publish compatibility source validation/tests_fuse2_mlx.py
f76fbf8 verified
Raw
History Blame Contribute Delete
1.04 kB
from __future__ import annotations
import mlx.core as mx
from microscope.fuse2_mlx import Model, ModelArgs
def main():
args = ModelArgs(
model_type="fuse2",
hidden_size=32,
num_hidden_layers=2,
intermediate_size=64,
num_attention_heads=4,
rms_norm_eps=1e-6,
vocab_size=128,
num_key_value_heads=2,
max_position_embeddings=64,
rope_theta=1000000.0,
head_dim=8,
tie_word_embeddings=True,
experts_per_layer={"0": [0, 1]},
expert_hidden_size=32,
expert_intermediate_size=16,
)
model = Model(args)
ids = mx.array([[1, 2, 3]])
uncached = model(ids)
cache = model.make_cache()
model(ids[:, :-1], cache)
cached = model(ids[:, -1:], cache)
mx.eval(uncached, cached)
error = float(mx.max(mx.abs(uncached[:, -1] - cached[:, -1])))
if error > 1e-4:
raise AssertionError(error)
print({"status": "validated", "max_logit_error": error})
if __name__ == "__main__":
main()