124 lines
4.9 KiB
Python
124 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Round-trip tests for pre-quantized MLX DiT checkpoints.
|
|
|
|
Saving an already-quantized DiT and reloading it must reproduce the original
|
|
forward pass exactly (same packed weights, scales, and config), for both plain
|
|
fp32 and int8-quantized models. Runs on any MLX backend (Metal or CPU).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
pytest.importorskip("mlx.core", reason="MLX is required for checkpoint tests")
|
|
|
|
from fastvideo.mlx_runtime.checkpoint import ( # noqa: E402
|
|
MANIFEST_FILENAME,
|
|
load_mlx_dit_checkpoint,
|
|
save_mlx_dit_checkpoint,
|
|
)
|
|
from fastvideo.mlx_runtime.fastwan import ( # noqa: E402
|
|
MLXQuantizationSpec,
|
|
QuantizedMatrix,
|
|
)
|
|
from fastvideo.tests.mlx.tiny_wan import ( # noqa: E402
|
|
build_hf_config,
|
|
build_inputs,
|
|
build_tiny_wan_config,
|
|
build_torch_model,
|
|
mlx_dit_from_torch_model,
|
|
mlx_output,
|
|
mlx_rotary_embeddings,
|
|
)
|
|
|
|
|
|
def _forward(dit) -> np.ndarray:
|
|
"""Run the DiT model with standard inputs and return its output as a NumPy array."""
|
|
hidden_states, encoder_hidden_states, timestep = build_inputs()
|
|
return mlx_output(dit, hidden_states, encoder_hidden_states, timestep, mlx_rotary_embeddings(hidden_states))
|
|
|
|
|
|
def test_plain_checkpoint_round_trip(distributed_setup, tmp_path) -> None:
|
|
dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
|
|
before = _forward(dit)
|
|
|
|
save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
|
|
loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt")
|
|
|
|
# JSON canonicalizes tuples to lists (matching what the Diffusers
|
|
# config.json loader produces), so compare JSON-normalized configs.
|
|
assert json.loads(json.dumps(loaded.config)) == json.loads(json.dumps(dit.config))
|
|
assert len(loaded.blocks) == len(dit.blocks)
|
|
after = _forward(loaded)
|
|
np.testing.assert_array_equal(before, after)
|
|
|
|
|
|
def test_int8_checkpoint_round_trip_skips_requantization(distributed_setup, tmp_path) -> None:
|
|
spec = MLXQuantizationSpec.from_name("int8")
|
|
dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()),
|
|
quantization=spec)
|
|
before = _forward(dit)
|
|
|
|
save_mlx_dit_checkpoint(dit, tmp_path / "ckpt-int8")
|
|
loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt-int8")
|
|
|
|
# The loaded model must carry the same packed quantized weights -- no
|
|
# requantization pass -- and produce a bitwise-identical forward.
|
|
reloaded_q = loaded.weights["proj_out.weight"]
|
|
original_q = dit.weights["proj_out.weight"]
|
|
assert isinstance(reloaded_q, QuantizedMatrix)
|
|
assert reloaded_q.spec == spec
|
|
np.testing.assert_array_equal(np.array(reloaded_q.weight), np.array(original_q.weight))
|
|
np.testing.assert_array_equal(np.array(reloaded_q.scales), np.array(original_q.scales))
|
|
|
|
block_q = loaded.blocks[0].weights["to_q.weight"]
|
|
assert isinstance(block_q, QuantizedMatrix)
|
|
|
|
after = _forward(loaded)
|
|
np.testing.assert_array_equal(before, after)
|
|
|
|
|
|
def test_missing_checkpoint_raises_clear_error(tmp_path) -> None:
|
|
with pytest.raises(FileNotFoundError, match="Not an MLX DiT checkpoint"):
|
|
load_mlx_dit_checkpoint(tmp_path / "does-not-exist")
|
|
|
|
|
|
def test_existing_checkpoint_is_replaced(distributed_setup, tmp_path) -> None:
|
|
dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
|
|
checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
|
|
dit.config["checkpoint_revision"] = "second"
|
|
|
|
save_mlx_dit_checkpoint(dit, checkpoint_dir)
|
|
|
|
assert load_mlx_dit_checkpoint(checkpoint_dir).config["checkpoint_revision"] == "second"
|
|
|
|
|
|
def test_failed_checkpoint_overwrite_preserves_existing(distributed_setup, tmp_path) -> None:
|
|
dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
|
|
checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
|
|
weights_before = (checkpoint_dir / "mlx_dit.safetensors").read_bytes()
|
|
manifest_before = (checkpoint_dir / MANIFEST_FILENAME).read_bytes()
|
|
dit.config["not_json_serializable"] = object()
|
|
|
|
with pytest.raises(TypeError, match="not JSON serializable"):
|
|
save_mlx_dit_checkpoint(dit, checkpoint_dir)
|
|
|
|
assert (checkpoint_dir / "mlx_dit.safetensors").read_bytes() == weights_before
|
|
assert (checkpoint_dir / MANIFEST_FILENAME).read_bytes() == manifest_before
|
|
|
|
|
|
def test_future_format_version_is_rejected(distributed_setup, tmp_path) -> None:
|
|
dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()))
|
|
ckpt = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt")
|
|
|
|
manifest_path = ckpt / MANIFEST_FILENAME
|
|
manifest = json.loads(manifest_path.read_text())
|
|
manifest["format_version"] = 99
|
|
manifest_path.write_text(json.dumps(manifest))
|
|
|
|
with pytest.raises(ValueError, match="format_version=99"):
|
|
load_mlx_dit_checkpoint(ckpt)
|