wehub-resource-sync
84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for _block_dequant_fp8 scale decoding.
|
|
|
|
MXFP8 checkpoints (e.g. MiniMax-M3) store their e8m0 block scales with
|
|
safetensors dtype U8. Those bytes are shared exponents and must decode as
|
|
2^(s - 127), the same as the F8_E8M0 branch. Treating them as linear
|
|
scales blows the weights up by orders of magnitude.
|
|
|
|
DeepSeek-style FP8 checkpoints also use weight_scale_inv keys but store
|
|
the scales as real floats (block 128). Those must keep multiplying
|
|
linearly, so the discriminator is the scale dtype, not the key name.
|
|
"""
|
|
|
|
import glob
|
|
import os
|
|
|
|
import mlx.core as mx
|
|
import pytest
|
|
|
|
from omlx.oq import _block_dequant_fp8, _LazyTensorIndex
|
|
|
|
M3_DIR = "/Volumes/Scratch/models/MiniMax-M3-MXFP8"
|
|
|
|
|
|
def test_u8_scale_decodes_as_e8m0_exponent():
|
|
mx.random.seed(0)
|
|
w = mx.random.normal((64, 128)).astype(mx.bfloat16)
|
|
qw, scales = mx.quantize(w, group_size=32, bits=8, mode="mxfp8")
|
|
ref = mx.dequantize(qw, scales, group_size=32, bits=8, mode="mxfp8")
|
|
assert scales.dtype == mx.uint8
|
|
|
|
# On-disk view of the same data: raw e4m3 bytes, one per element.
|
|
raw_fp8 = qw.view(mx.uint8)
|
|
assert raw_fp8.shape == (64, 128)
|
|
|
|
# Sanity check the target first: the explicit from_fp8 * 2^(s-127)
|
|
# formula must reproduce mx.dequantize exactly, otherwise ref is not
|
|
# a valid oracle for the function under test.
|
|
explicit = (
|
|
mx.from_fp8(raw_fp8, dtype=mx.bfloat16).reshape(64, 4, 32).astype(mx.float32)
|
|
* mx.power(mx.array(2.0), scales.astype(mx.float32) - 127.0)[:, :, None]
|
|
).reshape(64, 128)
|
|
assert mx.array_equal(explicit.astype(mx.bfloat16), ref).item()
|
|
|
|
got = _block_dequant_fp8(raw_fp8, scales, "F8_E4M3", "U8")
|
|
assert got.shape == ref.shape
|
|
assert mx.allclose(
|
|
got.astype(mx.float32), ref.astype(mx.float32), atol=1e-2, rtol=1e-2
|
|
).item(), (
|
|
f"mean|got|={mx.abs(got).mean().item():.4g} vs "
|
|
f"mean|ref|={mx.abs(ref).mean().item():.4g}"
|
|
)
|
|
|
|
|
|
def test_f32_scale_stays_linear():
|
|
# DeepSeek-style pair: e4m3 weight with a float block scale
|
|
# (block 128). The scale is a linear multiplier and must be applied
|
|
# as-is, untouched by the U8 exponent decoding.
|
|
mx.random.seed(1)
|
|
w = mx.random.normal((256, 128)).astype(mx.bfloat16)
|
|
qw, _ = mx.quantize(w, group_size=32, bits=8, mode="mxfp8")
|
|
raw_fp8 = qw.view(mx.uint8)
|
|
scale = mx.array([[0.5], [2.0]], dtype=mx.float32)
|
|
|
|
got = _block_dequant_fp8(raw_fp8, scale, "F8_E4M3", "F32")
|
|
|
|
wf = mx.from_fp8(raw_fp8, dtype=mx.bfloat16).astype(mx.float32)
|
|
expected = mx.concatenate([wf[:128] * 0.5, wf[128:] * 2.0], axis=0)
|
|
assert mx.allclose(got.astype(mx.float32), expected, atol=1e-2, rtol=1e-2).item()
|
|
|
|
|
|
@pytest.mark.skipif(not os.path.isdir(M3_DIR), reason="M3 not present")
|
|
def test_minimax_m3_k_proj_magnitude():
|
|
# Grounded check on a real MXFP8 checkpoint. Pre-fix this layer
|
|
# dequantized to mean|w| ~13410; the correct value is ~0.03.
|
|
shards = sorted(glob.glob(os.path.join(M3_DIR, "model-*.safetensors")))
|
|
idx = _LazyTensorIndex(shards)
|
|
key = "language_model.model.layers.3.self_attn.k_proj.weight"
|
|
weight = idx._dequant_one(key)
|
|
mean_abs = mx.abs(weight).mean().item()
|
|
max_abs = mx.abs(weight).max().item()
|
|
assert mean_abs < 1.0, f"mean|w|={mean_abs}"
|
|
assert max_abs < 2.0, f"max|w|={max_abs}"
|