i added native support for
@NVIDIAAI's Nemotron Puzzle 75B to mlx-lm. it now runs natively on an M2 Max 64GB:
โก๏ธ22 tok/s
๐พ45.5 GB peak memory usage
๐4-bit experts + 6-bit dense + BF16 head
i also fixed an annoying numerical bug in mlx-lm. outputs were subtly wrong, cosine similarity was 0.8832 vs NVIDIA's reference (identical inputs).
the culprit was one dtype cast in the Mamba layers happening in a different spot than NVIDIA's. once moved, the cosine similarity improved to a satisfying level (0.999...).
related PR:
weights: