huggingface--transformers/tests/heterogeneity/test_configuration_utils.py

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
import tempfile
import unittest

from parameterized import parameterized

from transformers import LlamaConfig
from transformers.integrations.heterogeneity import AmbiguousGlobalPerLayerAttributeError
from transformers.utils import logging as transformers_logging


# ──────────────────────────────────────────────────────────────────────
# Tiny config factories
# ──────────────────────────────────────────────────────────────────────


def _tiny_llama_config(per_layer_config=None, **overrides):
    defaults = {
        "hidden_size": 64,
        "intermediate_size": 128,
        "num_hidden_layers": 4,
        "num_attention_heads": 4,
        "num_key_value_heads": 4,
        "head_dim": 16,
        "vocab_size": 32,
        "max_position_embeddings": 64,
        **overrides,
    }
    return LlamaConfig(per_layer_config=per_layer_config, **defaults)


# ──────────────────────────────────────────────────────────────────────
# Tests: Config
# ──────────────────────────────────────────────────────────────────────


class TestHeterogeneousConfig(unittest.TestCase):
    def test_per_layer_config_skip_normalization(self):
        config = _tiny_llama_config(per_layer_config={1: {"skip": ["mlp", "attention"]}, 2: {"skip": []}})

        self.assertEqual(config.to_dict()["per_layer_config"], {"1": {"skip": ["attention", "mlp"]}})
        self.assertEqual(config.per_layer_config[0].skip, [])
        self.assertEqual(config.per_layer_config[1].skip, ["attention", "mlp"])
        self.assertEqual(config.per_layer_config[2].skip, [])

        set_config = _tiny_llama_config(per_layer_config={1: {"skip": {"attention"}}})
        self.assertEqual(set_config.to_dict()["per_layer_config"], {"1": {"skip": ["attention"]}})
        self.assertEqual(set_config.per_layer_config[1].skip, ["attention"])

        no_skip_config = _tiny_llama_config(per_layer_config={1: {"intermediate_size": 96}})
        self.assertNotIn("skip", no_skip_config.to_dict()["per_layer_config"]["1"])
        self.assertEqual(no_skip_config.per_layer_config[1].skip, [])

    @parameterized.expand(
        [
            ("string", "attention"),
            ("non_string_item", ["attention", 1]),
        ]
    )
    def test_per_layer_config_invalid_skip_raises(self, _name, invalid_skip):
        with self.assertRaises(TypeError):
            _tiny_llama_config(per_layer_config={0: {"skip": invalid_skip}})

    def test_per_layer_config_string_indices_are_normalized(self):
        config = _tiny_llama_config(per_layer_config={"01": {"num_key_value_heads": 2}})

        self.assertEqual(config.to_dict()["per_layer_config"], {"1": {"num_key_value_heads": 2}})
        self.assertEqual(config.per_layer_config[1].num_key_value_heads, 2)

    def test_per_layer_config_and_fallback(self):
        """Per-layer values should override, and non-overridden layers should fall back to global."""
        config = _tiny_llama_config(per_layer_config={1: {"num_key_value_heads": 2}, 3: {"num_key_value_heads": 1}})
        self.assertTrue(config.is_heterogeneous)
        self.assertEqual(config.per_layer_attributes, {"num_key_value_heads"})
        # Per-layer configs
        self.assertEqual(config.per_layer_config[1].num_key_value_heads, 2)
        self.assertEqual(config.per_layer_config[3].num_key_value_heads, 1)
        # Fallback to original global value
        self.assertEqual(config.per_layer_config[0].num_key_value_heads, 4)
        # Other attributes are unaffected
        self.assertEqual(config.per_layer_config[0].hidden_size, 64)

    def test_per_layer_config_reassignment_uses_existing_global_fallback(self):
        config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 2}})

        config.per_layer_config = {1: {"num_key_value_heads": 1}}

        self.assertEqual(config.to_dict()["per_layer_config"], {"1": {"num_key_value_heads": 1}})
        self.assertEqual(config.per_layer_config[0].num_key_value_heads, 4)
        self.assertEqual(config.per_layer_config[1].num_key_value_heads, 1)

    def test_per_layer_values_matching_global_are_removed_from_sparse_config(self):
        config = _tiny_llama_config(
            per_layer_config={
                0: {"num_key_value_heads": 4},
                1: {"num_key_value_heads": 2},
                2: {"num_key_value_heads": 4},
            }
        )

        self.assertEqual(config.to_dict()["per_layer_config"], {"1": {"num_key_value_heads": 2}})
        self.assertEqual(config.per_layer_attributes, {"num_key_value_heads"})
        self.assertEqual(config.per_layer_config[0].num_key_value_heads, 4)
        self.assertEqual(config.per_layer_config[1].num_key_value_heads, 2)
        self.assertEqual(config.per_layer_config[2].num_key_value_heads, 4)

    def test_uniform_per_layer_values_do_not_overwrite_global(self):
        per_layer = {layer_idx: {"num_key_value_heads": 2} for layer_idx in range(4)}
        config = _tiny_llama_config(per_layer_config=per_layer)

        self.assertEqual(object.__getattribute__(config, "num_key_value_heads"), 4)
        self.assertEqual(
            config.to_dict()["per_layer_config"],
            {str(i): {"num_key_value_heads": 2} for i in range(4)},
        )
        for layer_idx in range(4):
            self.assertEqual(config.per_layer_config[layer_idx].num_key_value_heads, 2)

    def test_explicit_serialization_restores_pruned_global_values(self):
        per_layer = {layer_idx: {"num_key_value_heads": 4} for layer_idx in range(4)}
        sparse_config = _tiny_llama_config(per_layer_config=per_layer)
        explicit_config = _tiny_llama_config(
            per_layer_config=per_layer,
            serialize_explicit_per_layer_config=True,
        )

        self.assertEqual(sparse_config.to_dict()["per_layer_config"], {})
        self.assertEqual(sparse_config.per_layer_attributes, set())
        self.assertEqual(
            explicit_config.to_dict()["per_layer_config"],
            {str(i): {"num_key_value_heads": 4} for i in range(4)},
        )

    def test_per_layer_config_reflects_current_global_config_state(self):
        config = _tiny_llama_config(per_layer_config={0: {"intermediate_size": 64}})

        # PreTrainedModel.__init__ updates this after config construction.
        config._attn_implementation_internal = "sdpa"
        config.hidden_size = 96
        config.intermediate_size = 192

        self.assertIs(type(config.per_layer_config[0]), type(config))
        self.assertFalse(config.per_layer_config[0].is_heterogeneous)
        self.assertIsNone(config.per_layer_config[0].per_layer_config)
        self.assertEqual(config.per_layer_config[0]._attn_implementation, "sdpa")
        self.assertEqual(config.per_layer_config[1]._attn_implementation, "sdpa")
        self.assertEqual(config.per_layer_config[0].hidden_size, 96)
        self.assertEqual(config.per_layer_config[1].hidden_size, 96)
        self.assertEqual(config.per_layer_config[0].intermediate_size, 64)
        self.assertEqual(config.per_layer_config[1].intermediate_size, 192)

        layer_dict = config.per_layer_config[0].to_dict()
        self.assertNotIn("per_layer_config", layer_dict)
        self.assertEqual(layer_dict["hidden_size"], 96)
        self.assertEqual(layer_dict["intermediate_size"], 64)

    def test_accessing_per_layer_attr_raises(self):
        config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}})
        with self.assertRaisesRegex(
            AmbiguousGlobalPerLayerAttributeError, "allow_global_per_layer_attribute_access.*global value incorrectly"
        ):
            _ = config.num_key_value_heads

    def test_ambiguous_global_attr_access_is_not_treated_as_missing(self):
        config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}})

        self.assertIn("num_key_value_heads", config.__dict__)
        with self.assertRaises(AmbiguousGlobalPerLayerAttributeError):
            getattr(config, "num_key_value_heads", "default")
        with self.assertRaises(AmbiguousGlobalPerLayerAttributeError):
            hasattr(config, "num_key_value_heads")

    def test_allow_global_per_layer_attribute_access(self):
        config = _tiny_llama_config(
            per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}},
            allow_global_per_layer_attribute_access=True,
        )
        logger = transformers_logging.get_logger("transformers.integrations.heterogeneity.configuration_utils")
        logger.warning_once.cache_clear()

        with self.assertLogs(logger=logger, level="WARNING") as logs:
            self.assertEqual(config.num_key_value_heads, 4)

        self.assertIn("Reading global config value for per-layer attribute `num_key_value_heads`", logs.output[0])
        self.assertEqual(config.per_layer_config[0].num_key_value_heads, 2)
        self.assertEqual(config.per_layer_config[1].num_key_value_heads, 1)
        self.assertEqual(config.per_layer_config[2].num_key_value_heads, 4)

    def test_flags_are_applied_from_pretrained_kwargs(self):
        """The flags are properties, so `from_dict`'s kwargs handling applies them like `per_layer_config`."""
        with tempfile.TemporaryDirectory() as tmpdir:
            _tiny_llama_config().save_pretrained(tmpdir)
            config = LlamaConfig.from_pretrained(
                tmpdir,
                per_layer_config={1: {"num_key_value_heads": 2}},
                allow_global_per_layer_attribute_access=True,
                serialize_explicit_per_layer_config=True,
            )

        self.assertEqual(config.num_key_value_heads, 4)
        self.assertEqual(
            config.to_dict()["per_layer_config"],
            {str(i): {"num_key_value_heads": 2 if i == 1 else 4} for i in range(4)},
        )
        # The property defaults must not leak into the serialization of configs that never set the flags.
        self.assertNotIn("allow_global_per_layer_attribute_access", _tiny_llama_config().to_dict())

    def test_non_per_layer_attributes_do_not_warn(self):
        config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}})
        logger = transformers_logging.get_logger("transformers.integrations.heterogeneity.configuration_utils")
        logger.warning_once.cache_clear()

        with self.assertNoLogs(logger=logger, level="WARNING"):
            self.assertEqual(config.hidden_size, 64)

    def test_iter_skips_per_layer_attributes_by_default(self):
        config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}})

        keys = list(config)

        self.assertNotIn("num_key_value_heads", keys)
        self.assertIn("hidden_size", keys)

    def test_iter_includes_per_layer_attributes_when_global_access_allowed(self):
        config = _tiny_llama_config(
            per_layer_config={0: {"num_key_value_heads": 2}, 1: {"num_key_value_heads": 1}},
            allow_global_per_layer_attribute_access=True,
        )

        self.assertIn("num_key_value_heads", list(config))

    def test_validation_missing_global_attr(self):
        # "fake_attr" in layer 0 but not in layer 1, and not global → should fail
        with self.assertRaises(ValueError):
            _tiny_llama_config(
                per_layer_config={
                    0: {"fake_attr": 42, "intermediate_size": 64},
                    1: {"intermediate_size": 96},
                }
            )

    @parameterized.expand(
        [
            ("negative", -1),
            ("too_large", 4),
        ]
    )
    def test_validation_layer_idx_out_of_range(self, _name, layer_idx):
        with self.assertRaises(ValueError):
            _tiny_llama_config(per_layer_config={layer_idx: {"num_key_value_heads": 2}})

    def test_save_pretrained_config_round_trip(self):
        """Config should survive save_pretrained → from_pretrained on disk."""
        per_layer = {i: {"intermediate_size": 64 + i} for i in range(0, 12, 2)}
        config = _tiny_llama_config(per_layer_config=per_layer, num_hidden_layers=12)

        # Keys are zero-padded so they sort numerically in JSON (0,1,...,10 not 0,1,10,2,...)
        d = config.to_dict()
        self.assertEqual(list(d["per_layer_config"].keys()), sorted(d["per_layer_config"].keys()))

        with tempfile.TemporaryDirectory() as tmpdir:
            config.save_pretrained(tmpdir)
            loaded = LlamaConfig.from_pretrained(tmpdir)

        self.assertTrue(loaded.is_heterogeneous)
        for i in range(config.num_hidden_layers):
            self.assertEqual(
                config.per_layer_config[i].intermediate_size,
                loaded.per_layer_config[i].intermediate_size,
            )

    @parameterized.expand(
        [
            (
                "global_sw_global_acs",
                {"sliding_window": 4096, "attention_chunk_size": 2048},
                {0: {"intermediate_size": 64}},
                True,
            ),
            ("global_sw_per_layer_acs", {"sliding_window": 4096}, {0: {"attention_chunk_size": 2048}}, True),
            (
                "per_layer_sw_per_layer_acs_same_layer",
                {},
                {0: {"sliding_window": 4096, "attention_chunk_size": 2048}},
                True,
            ),
            (
                "per_layer_sw_per_layer_acs_different_layers",
                {"sliding_window": None, "attention_chunk_size": None},
                {0: {"sliding_window": 4096}, 1: {"attention_chunk_size": 2048}},
                False,
            ),
            (
                "global_conflict_resolved_by_per_layer_override",
                {"sliding_window": 4096, "attention_chunk_size": 2048},
                {
                    0: {"sliding_window": None},
                    1: {"sliding_window": None},
                    2: {"attention_chunk_size": None},
                    3: {"attention_chunk_size": None},
                },
                False,
            ),
        ],
    )
    def test_validation_sliding_window_and_attention_chunk_size(
        self, _name, overrides, per_layer_config, should_raise
    ):
        ctx = self.assertRaises(ValueError) if should_raise else contextlib.nullcontext()
        with ctx:
            _tiny_llama_config(per_layer_config=per_layer_config, **overrides)

    def test_all_layers_overridden_no_global_default(self):
        """Custom attribute on every layer without a global default should be accessible per layer."""
        config = _tiny_llama_config(
            per_layer_config={
                0: {"custom_attr": 10},
                1: {"custom_attr": 20},
                2: {"custom_attr": 30},
                3: {"custom_attr": 40},
            },
        )
        self.assertTrue(config.is_heterogeneous)
        self.assertEqual(config.per_layer_config[0].custom_attr, 10)
        self.assertEqual(config.per_layer_config[1].custom_attr, 20)
        self.assertEqual(config.per_layer_config[2].custom_attr, 30)
        self.assertEqual(config.per_layer_config[3].custom_attr, 40)

    def test_per_layer_config_can_serialize_explicit_layer_overrides(self):
        sparse_config = _tiny_llama_config(per_layer_config={0: {"num_key_value_heads": 1}})
        explicit_config = _tiny_llama_config(
            per_layer_config={0: {"num_key_value_heads": 1}},
            serialize_explicit_per_layer_config=True,
        )
        explicit_config.num_key_value_heads = 8
        explicit_config_dict = explicit_config.to_dict()

        self.assertEqual(sparse_config.to_dict()["per_layer_config"], {"0": {"num_key_value_heads": 1}})
        self.assertEqual(
            explicit_config_dict["per_layer_config"],
            {
                "0": {"num_key_value_heads": 1},
                "1": {"num_key_value_heads": 8},
                "2": {"num_key_value_heads": 8},
                "3": {"num_key_value_heads": 8},
            },
        )