108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for reusable and Anima-specific quantization profiles."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from simple_syrup.domain.anima_quantization import (
|
|
FP8_E4M3_PROFILE,
|
|
FP8_E5M2_PROFILE,
|
|
MXFP8_PROFILE,
|
|
NVFP4_MIXED_PROFILE,
|
|
ORIGINAL_PROFILE,
|
|
AnimaQuantizationRecipe,
|
|
)
|
|
from simple_syrup.domain.model_quantization import (
|
|
QuantizationFormat,
|
|
QuantizationProfile,
|
|
TensorDescriptor,
|
|
)
|
|
|
|
|
|
def test_anima_profiles_parse_labels_and_stable_ids() -> None:
|
|
"""Workflow labels and persisted identifiers resolve consistently."""
|
|
|
|
recipe = AnimaQuantizationRecipe()
|
|
|
|
assert recipe.profile_from_selection("Original") == ORIGINAL_PROFILE
|
|
assert recipe.profile_from_selection("nvfp4-mixed") == NVFP4_MIXED_PROFILE
|
|
assert QuantizationFormat.NVFP4.label == "NVFP4"
|
|
with pytest.raises(ValueError, match="quantization profile must be one of"):
|
|
recipe.profile_from_selection("unknown")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"net.blocks.0.attn.q_proj.weight",
|
|
"net.blocks.1.attn.q_proj.weight",
|
|
"net.blocks.27.attn.q_proj.weight",
|
|
"net.blocks.14.adaln_modulation.1.weight",
|
|
"net.final_layer.linear.weight",
|
|
"net.llm_adapter.proj.weight",
|
|
"net.t_embedder.1.weight",
|
|
"net.x_embedder.proj.weight",
|
|
"some_other_model.blocks.14.attn.q_proj.weight",
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"profile",
|
|
[FP8_E4M3_PROFILE, MXFP8_PROFILE, NVFP4_MIXED_PROFILE],
|
|
)
|
|
def test_anima_profiles_preserve_quality_sensitive_weights(
|
|
name: str,
|
|
profile: QuantizationProfile,
|
|
) -> None:
|
|
"""Every profile keeps known sensitive and out-of-envelope weights."""
|
|
|
|
recipe = AnimaQuantizationRecipe()
|
|
selected = recipe.profile_from_selection(profile.profile_id)
|
|
|
|
assert recipe.policy_for(TensorDescriptor(name, (16, 16), "BF16"), selected) is None
|
|
|
|
|
|
def test_anima_mixed_profile_assigns_projection_specific_formats() -> None:
|
|
"""Recommended mixed precision follows the intended attention/MLP split."""
|
|
|
|
recipe = AnimaQuantizationRecipe()
|
|
|
|
def assigned(name: str) -> QuantizationFormat | None:
|
|
return recipe.policy_for(
|
|
TensorDescriptor(f"net.blocks.14.{name}.weight", (16, 16), "BF16"),
|
|
NVFP4_MIXED_PROFILE,
|
|
)
|
|
|
|
assert assigned("attn.q_proj") is QuantizationFormat.NVFP4
|
|
assert assigned("attn.k_proj") is QuantizationFormat.NVFP4
|
|
assert assigned("attn.output_proj") is QuantizationFormat.NVFP4
|
|
assert assigned("attn.v_proj") is QuantizationFormat.FP8_E4M3
|
|
assert assigned("mlp.fc1") is QuantizationFormat.FP8_E4M3
|
|
assert assigned("unmatched") is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("profile", "expected"),
|
|
[
|
|
(FP8_E4M3_PROFILE, QuantizationFormat.FP8_E4M3),
|
|
(FP8_E5M2_PROFILE, QuantizationFormat.FP8_E5M2),
|
|
(MXFP8_PROFILE, QuantizationFormat.MXFP8),
|
|
],
|
|
)
|
|
def test_uniform_profiles_quantize_only_eligible_middle_block_matrices(
|
|
profile: QuantizationProfile,
|
|
expected: QuantizationFormat,
|
|
) -> None:
|
|
"""Uniform profiles share the safety envelope while selecting their format."""
|
|
|
|
descriptor = TensorDescriptor(
|
|
"diffusion_model.blocks.14.self_attn.q_proj.weight",
|
|
(16, 16),
|
|
"BF16",
|
|
)
|
|
|
|
assert AnimaQuantizationRecipe().policy_for(descriptor, profile) is expected
|