tests/test_basic_merges.py

from typing import Dict, Optional

import pytest
from transformers import AutoConfig

from mergekit.config import (
    InputModelDefinition,
    InputSliceDefinition,
    MergeConfiguration,
    OutputSliceDefinition,
    ParameterSetting,
)
from mergekit.io import LazyTensorLoader
from tests.common import (
    make_gpt2size,
    make_picogranite,
    make_picollama,
    make_picoLlaVa,
    run_and_check_merge,
)


@pytest.fixture(scope="session")
def model_a(tmp_path_factory):
    return make_picollama(tmp_path_factory.mktemp("model_a"))


@pytest.fixture(scope="session")
def model_b(tmp_path_factory):
    return make_picollama(tmp_path_factory.mktemp("model_b"))


@pytest.fixture(scope="session")
def model_c(tmp_path_factory):
    return make_picollama(tmp_path_factory.mktemp("model_c"))


@pytest.fixture(scope="session")
def vlm_a(tmp_path_factory):
    return make_picoLlaVa(tmp_path_factory.mktemp("vlm_a"))


@pytest.fixture(scope="session")
def vlm_b(tmp_path_factory):
    return make_picoLlaVa(tmp_path_factory.mktemp("vlm_b"))


@pytest.fixture(scope="session")
def vlm_c(tmp_path_factory):
    return make_picoLlaVa(tmp_path_factory.mktemp("vlm_c"))


@pytest.fixture(scope="session")
def gpt2_like(tmp_path_factory):
    return make_gpt2size(tmp_path_factory.mktemp("gpt2_like"))


@pytest.fixture(scope="session")
def granite_a(tmp_path_factory):
    return make_picogranite(tmp_path_factory.mktemp("granite_a"))


@pytest.fixture(scope="session")
def granite_b(tmp_path_factory):
    return make_picogranite(tmp_path_factory.mktemp("granite_b"))


class TestGraniteMerges:
    def test_granite_copy(self, granite_a):
        config = MergeConfiguration(
            merge_method="passthrough",
            models=[InputModelDefinition(model=granite_a)],
            dtype="bfloat16",
        )
        run_and_check_merge(config)

    def test_granite_linear_merge(self, granite_a, granite_b):
        config = MergeConfiguration(
            merge_method="linear",
            models=[
                InputModelDefinition(model=granite_a, parameters={"weight": 0.6}),
                InputModelDefinition(model=granite_b, parameters={"weight": 0.4}),
            ],
            dtype="bfloat16",
        )
        run_and_check_merge(config)

    def test_granite_slerp(self, granite_a, granite_b):
        config = MergeConfiguration(
            merge_method="slerp",
            base_model=granite_a,
            models=[
                InputModelDefinition(model=granite_a),
                InputModelDefinition(model=granite_b),
            ],
            parameters={"t": 0.5},
            dtype="bfloat16",
        )
        run_and_check_merge(config)


class TestBasicMerges:
    def test_gpt2_copy(self, gpt2_like):
        config = MergeConfiguration(
            merge_method="passthrough",
            models=[InputModelDefinition(model=gpt2_like)],
            dtype="bfloat16",
        )
        run_and_check_merge(config)

    def test_gpt2_stack(self, gpt2_like):
        config = MergeConfiguration(
            merge_method="passthrough",
            slices=[
                OutputSliceDefinition(
                    sources=[InputSliceDefinition(model=gpt2_like, layer_range=[0, 12])]
                )
            ]
            * 2,
            dtype="bfloat16",
        )

        def _check_config_layers(p: str):
            config = AutoConfig.from_pretrained(p)
            assert config.n_layer == 24

        run_and_check_merge(config, validate=_check_config_layers)

    def test_passthrough_scale(self, model_a):
        config = MergeConfiguration(
            merge_method="passthrough",
            models=[
                InputModelDefinition(
                    model=model_a,
                    parameters={
                        "scale": [
                            {"filter": "o_proj", "value": 0},
                            {"value": 1},
                        ]
                    },
                )
            ],
        )

        def _check_o_proj(p: str):
            loader = LazyTensorLoader.from_disk(p)
            saw_any = False
            for name in loader.index.tensor_paths:
                if "o_proj" in name:
                    param = loader.get_tensor(name)
                    assert (param == 0).all()
                    saw_any = True
                elif "lm_head" in name:
                    param = loader.get_tensor(name)
                    assert param.count_nonzero() > 0

            assert saw_any, "No o_proj parameters found"

        run_and_check_merge(config, validate=_check_o_proj)

    def test_linear_merge(self, model_a, model_b):
        config = self.two_model_config(model_a, model_b, merge_method="linear")
        run_and_check_merge(config)

    def test_slerp_merge(self, model_a, model_b):
        config = self.two_model_config(
            model_a, model_b, merge_method="slerp", base_model=model_a
        )
        config.parameters = {"t": 0.35}
        run_and_check_merge(config)

    def test_nuslerp_merges(self, model_a, model_b, model_c):
        for base_model in [None, model_c]:
            for row_wise in [False, True]:
                for flatten in [False, True]:
                    print(
                        f"Testing nuslerp with row_wise={row_wise}, flatten={flatten}, base_model={base_model}"
                    )
                    run_and_check_merge(
                        self.two_model_config(
                            model_a,
                            model_b,
                            merge_method="nuslerp",
                            base_model=base_model,
                            params={
                                "nuslerp_row_wise": row_wise,
                                "nuslerp_flatten": flatten,
                            },
                        )
                    )

        # test weights that sum to zero
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="nuslerp",
            base_model=model_c,
            params={"nuslerp_row_wise": False, "nuslerp_flatten": False},
        )
        config.models[0].parameters["weight"] = -0.5
        config.models[1].parameters["weight"] = 0.5
        run_and_check_merge(config)

    def test_task_arithmetic_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a, model_b, merge_method="task_arithmetic", base_model=model_c
        )
        run_and_check_merge(config)

    def test_breadcrumbs_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a, model_b, merge_method="breadcrumbs", base_model=model_c
        )
        run_and_check_merge(config)

    def test_ties_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="ties",
            base_model=model_c,
            params={"density": 0.3},
        )
        run_and_check_merge(config)

    def test_sce_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="sce",
            base_model=model_c,
            params={"select_topk": 0.5},
        )
        run_and_check_merge(config)

    def test_ram_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="ram",
            base_model=model_c,
        )
        run_and_check_merge(config)

    def test_ramplus_tl_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="ramplus_tl",
            base_model=model_c,
            params={"r": 0.05, "alpha": 1.0},
        )
        run_and_check_merge(config)

    def test_multislerp_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="multislerp",
            base_model=model_c,
        )
        run_and_check_merge(config)

    def test_dare_ties_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="dare_ties",
            base_model=model_c,
            params={"density": 0.66},
        )
        run_and_check_merge(config)

    def test_model_stock_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_b, model_c, merge_method="model_stock", base_model=model_a
        )
        run_and_check_merge(config)

    def test_model_stock_filterwise_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_b,
            model_c,
            merge_method="model_stock",
            base_model=model_a,
            params={"filter_wise": True},
        )
        run_and_check_merge(config)

    def test_arcee_fusion_merge(self, model_a, model_b):
        config = self.two_model_config(
            model_a, model_b, merge_method="arcee_fusion", base_model=model_a
        )
        run_and_check_merge(config)

    def test_nearswap_merge(self, model_a, model_b):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="nearswap",
            base_model=model_a,
            params={"t": 0.0001},
        )
        run_and_check_merge(config)

    def test_della_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="della",
            base_model=model_c,
            params={"density": 0.66, "epsilon": 0.05, "lambda": 0.5},
        )
        run_and_check_merge(config)

    def test_della_invalid_epsilon(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="della",
            base_model=model_c,
            params={"density": 0.66, "epsilon": 1.35, "lambda": 0.5},
        )
        with pytest.raises(ValueError):
            run_and_check_merge(config)

    def test_karcher_merge(self, model_a, model_b, model_c):
        config = self.two_model_config(
            model_a,
            model_b,
            merge_method="karcher",
            base_model=model_c,
            params={"max_iter": 5, "tol": 1e-5},
        )
        run_and_check_merge(config)

    def two_model_config(
        self,
        model_a,
        model_b,
        merge_method: str,
        base_model: Optional[str] = None,
        params: Optional[Dict[str, ParameterSetting]] = None,
    ):
        config = MergeConfiguration(
            merge_method=merge_method,
            base_model=base_model,
            models=[
                InputModelDefinition(
                    model=model_a,
                    parameters={"weight": 0.6},
                ),
                InputModelDefinition(
                    model=model_b,
                    parameters={"weight": 0.4},
                ),
            ],
            dtype="bfloat16",
            parameters=params,
        )

        return config

    def test_linear_VLM_merge(self, vlm_a, vlm_b):
        config = self.two_model_config(vlm_a, vlm_b, merge_method="linear")
        run_and_check_merge(config, auto_arch=True)

    def test_slerp_VLM_merge(self, vlm_a, vlm_b):
        config = self.two_model_config(
            vlm_a, vlm_b, merge_method="slerp", base_model=vlm_a
        )
        config.parameters = {"t": 0.35}
        run_and_check_merge(config, auto_arch=True)

    def test_nuslerp_VLM_merges(self, vlm_a, vlm_b, vlm_c):
        for base_model in [None, vlm_c]:
            for row_wise in [False, True]:
                for flatten in [False, True]:
                    print(
                        f"Testing nuslerp with row_wise={row_wise}, flatten={flatten}, base_model={base_model}"
                    )
                    run_and_check_merge(
                        self.two_model_config(
                            vlm_a,
                            vlm_b,
                            merge_method="nuslerp",
                            base_model=base_model,
                            params={
                                "nuslerp_row_wise": row_wise,
                                "nuslerp_flatten": flatten,
                            },
                        ),
                        auto_arch=True,
                    )

        # test weights that sum to zero
        config = self.two_model_config(
            vlm_a,
            vlm_b,
            merge_method="nuslerp",
            base_model=vlm_c,
            params={"nuslerp_row_wise": False, "nuslerp_flatten": False},
        )
        config.models[0].parameters["weight"] = -0.5
        config.models[1].parameters["weight"] = 0.5
        run_and_check_merge(config, auto_arch=True)

    def test_task_arithmetic_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_a, vlm_b, merge_method="task_arithmetic", base_model=vlm_c
        )
        run_and_check_merge(config, auto_arch=True)

    def test_breadcrumbs_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_a, vlm_b, merge_method="breadcrumbs", base_model=vlm_c
        )
        run_and_check_merge(config, auto_arch=True)

    def test_ties_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_a,
            vlm_b,
            merge_method="ties",
            base_model=vlm_c,
            params={"density": 0.3},
        )
        run_and_check_merge(config, auto_arch=True)

    def test_dare_ties_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_a,
            vlm_b,
            merge_method="dare_ties",
            base_model=vlm_c,
            params={"density": 0.66},
        )
        run_and_check_merge(config, auto_arch=True)

    def test_model_stock_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_b, vlm_c, merge_method="model_stock", base_model=vlm_a
        )
        run_and_check_merge(config, auto_arch=True)

    def test_model_stock_filterwise_VLM_merge(self, vlm_a, vlm_b, vlm_c):
        config = self.two_model_config(
            vlm_b,
            vlm_c,
            merge_method="model_stock",
            base_model=vlm_a,
            params={"filter_wise": True},
        )
        run_and_check_merge(config, auto_arch=True)