fusion-bench 0.2.21__py3-none-any.whl → 0.2.23__py3-none-any.whl

fusion_bench/method/slerp/slerp.py

@@ -1,16 +1,24 @@
 import logging
-from typing import Any, Dict
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch
 from torch import nn
+from tqdm import tqdm
 from typing_extensions import override
 
+from fusion_bench import LazyStateDict, create_default_model_card, timeit_context
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
+from fusion_bench.modelpool import BaseModelPool, CausalLMPool
 from fusion_bench.utils.type import StateDictType
 
 from .slerp_utils import slerp
 
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
 log = logging.getLogger(__name__)
 
 
@@ -21,6 +29,7 @@ def slerp_on_state_dicts(
     *,
     DOT_THRESHOLD: float = 0.9995,
     epsilon: float = 1e-8,
+    show_pbar: bool = False,
 ) -> StateDictType:
     """
     Perform spherical linear interpolation (slerp) on the state dictionaries of two models.
@@ -36,7 +45,8 @@ def slerp_on_state_dicts(
         dict: The interpolated state dictionary.
     """
     state_dict = {}
-    for key in secondary_state_dict:
+    pbar = secondary_state_dict if not show_pbar else tqdm(secondary_state_dict)
+    for key in pbar:
         v0 = primary_state_dict[key]
         v1 = secondary_state_dict[key]
         if v0.shape != v1.shape:
@@ -49,18 +59,19 @@ def slerp_on_state_dicts(
     return state_dict
 
 
+@auto_register_config
 class SlerpMergeAlgorithm(BaseAlgorithm):
     """
     General purpose implementation of Slerp (Spherical Linear Interpolation) for PyTorch models.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "t": "t",
-        "DOT_THRESHOLD": "DOT_THRESHOLD",
-        "epsilon": "epsilon",
-    }
-
-    def __init__(self, t: float, DOT_THRESHOLD: float = 0.9995, epsilon: float = 1e-8):
+    def __init__(
+        self,
+        t: float,
+        DOT_THRESHOLD: float = 0.9995,
+        epsilon: float = 1e-8,
+        **kwargs,
+    ):
         """
         Initialize the SlerpMergeAlgorithm.
 
@@ -69,10 +80,7 @@ class SlerpMergeAlgorithm(BaseAlgorithm):
             DOT_THRESHOLD (float, optional): The threshold for the dot product of the two vectors. Defaults to 0.9995.
             epsilon (float, optional): The epsilon value for numerical stability. Defaults to 1e-8.
         """
-        self.t = t
-        self.DOT_THRESHOLD = DOT_THRESHOLD
-        self.epsilon = epsilon
-        super().__init__()
+        super().__init__(**kwargs)
 
     @override
     def run(self, modelpool: BaseModelPool) -> nn.Module:
@@ -102,3 +110,91 @@ class SlerpMergeAlgorithm(BaseAlgorithm):
 
         primary_model.load_state_dict(state_dict)
         return primary_model
+
+
+@auto_register_config
+class SlerpForCausalLM(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    """
+    Slerp (Spherical Linear Interpolation) for Causal Language Models.
+    """
+
+    def __init__(
+        self,
+        t: float,
+        DOT_THRESHOLD: float = 0.9995,
+        epsilon: float = 1e-8,
+        model_save_path: Optional[str] = None,
+        show_pbar: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize the SlerpForCausalLM algorithm.
+
+        Args:
+            t (float): The interpolation parameter. Must be in the range [0, 1].
+                      t=0 returns the first model, t=1 returns the second model,
+                      t=0.5 provides balanced interpolation.
+            DOT_THRESHOLD (float, optional): The threshold for the dot product of normalized vectors.
+                                           When the absolute dot product exceeds this threshold,
+                                           vectors are considered nearly collinear and linear
+                                           interpolation (LERP) is used instead of SLERP for
+                                           numerical stability. Defaults to 0.9995.
+            epsilon (float, optional): Small value used for numerical stability to avoid
+                                     division by zero during vector normalization.
+                                     Defaults to 1e-8.
+            model_save_path (Optional[str], optional): Path where the merged model should be saved.
+                                                     If None, the model is not saved to disk.
+                                                     Defaults to None.
+            show_pbar (bool, optional): Whether to display a progress bar during the interpolation
+                                      process. Useful for debugging or monitoring progress with
+                                      large models. Defaults to False.
+            **kwargs: Additional keyword arguments passed to the parent BaseAlgorithm class.
+        """
+        super().__init__(**kwargs)
+
+    @override
+    def run(self, modelpool: CausalLMPool):
+        assert len(modelpool.all_model_names) == 2, "Slerp expect exactly 2 models"
+        primary_model = modelpool.load_model(modelpool.all_model_names[0])
+        secondary_model = modelpool.load_model(modelpool.all_model_names[1])
+
+        with torch.no_grad():
+            primary_state_dict = primary_model.state_dict()
+            secondary_state_dict = secondary_model.state_dict()
+            state_dict = slerp_on_state_dicts(
+                self.t,
+                primary_state_dict,
+                secondary_state_dict,
+                DOT_THRESHOLD=self.DOT_THRESHOLD,
+                epsilon=self.epsilon,
+            )
+
+        if isinstance(primary_model, nn.Module):
+            model = primary_model
+            model.load_state_dict(state_dict)
+        elif isinstance(primary_model, LazyStateDict):
+            model: "PreTrainedModel" = deepcopy(primary_model.meta_module)
+            model.to(device=primary_model._device)
+            model.load_state_dict(state_dict)
+        else:
+            raise TypeError(
+                f"Unsupported model type: {type(primary_model)}. "
+                "Expected nn.Module or LazyStateDict."
+            )
+        if self.model_save_path is not None:
+            with timeit_context(f"Saving the model to {self.model_save_path}"):
+                tokenizer = modelpool.load_tokenizer()
+                tokenizer.save_pretrained(self.model_save_path)
+                model.save_pretrained(self.model_save_path)
+                model_card_str = create_default_model_card(
+                    models=[modelpool.get_model_path(m) for m in modelpool.model_names],
+                    description="Merged model using Slerp.",
+                    algorithm_config=self.config,
+                    modelpool_config=modelpool.config,
+                )
+                with open(os.path.join(self.model_save_path, "README.md"), "w") as f:
+                    f.write(model_card_str)
+        return model

fusion_bench/method/smile_upscaling/causal_lm_upscaling.py

@@ -0,0 +1,371 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Type, Union
+
+import torch
+from accelerate import init_empty_weights
+from tqdm.auto import tqdm
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    LlamaForCausalLM,
+    MistralForCausalLM,
+    PretrainedConfig,
+    PreTrainedModel,
+    Qwen2ForCausalLM,
+)
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+from transformers.models.mistral.modeling_mistral import MistralDecoderLayer
+from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.compat.modelpool import to_modelpool
+from fusion_bench.constants import RuntimeConstants
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
+from fusion_bench.modelpool import CausalLMPool
+from fusion_bench.models.hf_utils import (
+    create_default_model_card,
+    save_pretrained_with_remote_code,
+)
+from fusion_bench.models.modeling_smile_llama import (
+    SmileLlamaConfig,
+    SmileLlamaForCausalLM,
+    SmileLlamaModel,
+)
+from fusion_bench.models.modeling_smile_llama.modeling_smile_llama import (
+    SmileLlamaDecoderLayer,
+)
+from fusion_bench.models.modeling_smile_mistral import (
+    SmileMistralConfig,
+    SmileMistralForCausalLM,
+    SmileMistralModel,
+)
+from fusion_bench.models.modeling_smile_mistral.modeling_smile_mistral import (
+    SmileMistralDecoderLayer,
+)
+
+# Import all SMILE configurations and models
+from fusion_bench.models.modeling_smile_qwen2 import (
+    SmileQwen2Config,
+    SmileQwen2ForCausalLM,
+    SmileQwen2Model,
+)
+from fusion_bench.models.modeling_smile_qwen2.modeling_smile_qwen2 import (
+    SmileQwen2DecoderLayer,
+)
+from fusion_bench.models.smile_moe.linear_from_hf_config import (
+    ExpertNotTrainedError,
+    upscale_to_smile_linear,
+)
+from fusion_bench.utils.dtype import parse_dtype
+from fusion_bench.utils.parameters import print_parameters
+
+log = logging.getLogger(__name__)
+
+# Model type mappings
+MODEL_TYPE_MAPPINGS = {
+    "qwen2": {
+        "base_model_cls": Qwen2ForCausalLM,
+        "base_decoder_layer_cls": Qwen2DecoderLayer,
+        "smile_config_cls": SmileQwen2Config,
+        "smile_model_cls": SmileQwen2ForCausalLM,
+        "smile_base_model_cls": SmileQwen2Model,
+        "smile_decoder_layer_cls": SmileQwen2DecoderLayer,
+        "description": "Qwen2",
+    },
+    "llama": {
+        "base_model_cls": LlamaForCausalLM,
+        "base_decoder_layer_cls": LlamaDecoderLayer,
+        "smile_config_cls": SmileLlamaConfig,
+        "smile_model_cls": SmileLlamaForCausalLM,
+        "smile_base_model_cls": SmileLlamaModel,
+        "smile_decoder_layer_cls": SmileLlamaDecoderLayer,
+        "description": "Llama",
+    },
+    "mistral": {
+        "base_model_cls": MistralForCausalLM,
+        "base_decoder_layer_cls": MistralDecoderLayer,
+        "smile_config_cls": SmileMistralConfig,
+        "smile_model_cls": SmileMistralForCausalLM,
+        "smile_base_model_cls": SmileMistralModel,
+        "smile_decoder_layer_cls": SmileMistralDecoderLayer,
+        "description": "Mistral",
+    },
+}
+
+
+def detect_model_type(
+    model_or_config: Union[PreTrainedModel, PretrainedConfig, str],
+) -> str:
+    """
+    Detect the model type from a model, config, or model name/path.
+
+    Args:
+        model_or_config: Model, config, or model name/path to detect type from
+
+    Returns:
+        str: The detected model type ("qwen2", "llama", "mistral")
+
+    Raises:
+        ValueError: If model type cannot be detected or is not supported
+    """
+    if isinstance(model_or_config, str):
+        # Load config from path/name
+        config = AutoConfig.from_pretrained(model_or_config)
+    elif isinstance(model_or_config, PreTrainedModel):
+        config = model_or_config.config
+    elif isinstance(model_or_config, PretrainedConfig):
+        config = model_or_config
+    else:
+        raise ValueError(
+            f"Unsupported type for model type detection: {type(model_or_config)}"
+        )
+
+    model_type = getattr(config, "model_type", "").lower()
+
+    # Handle various model type variations
+    if model_type in MODEL_TYPE_MAPPINGS:
+        return model_type
+    else:
+        raise ValueError(
+            f"Unsupported model type: {model_type}. Supported types: {list(MODEL_TYPE_MAPPINGS.keys())}"
+        )
+
+
+@auto_register_config
+class SmileCausalLMUpscalingAlgorithm(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    R"""
+    SmileCausalLMUpscalingAlgorithm is a generic model fusion algorithm designed to upscale
+    a pretrained CausalLM model using a set of fine-tuned expert models. The algorithm
+    supports Qwen2, Llama, and Mistral model architectures and leverages Singular Value
+    Decomposition (SVD) to merge the weights of the pretrained model and the expert models
+    into a new upscaled model.
+
+    The algorithm automatically detects the model type and uses the appropriate SMILE
+    configuration and model classes.
+
+    Methods:
+        run(modelpool: BaseModelPool) -> Union[SmileQwen2ForCausalLM, SmileLlamaForCausalLM, SmileMistralForCausalLM]:
+            Executes the upscaling process and returns the upscaled model.
+
+        merge(pretrained_model: PreTrainedModel, finetuned_models: List[PreTrainedModel]) -> PreTrainedModel:
+            Merges the pretrained model with the fine-tuned models to create an upscaled model.
+    """
+
+    modelpool: CausalLMPool
+
+    def __init__(
+        self,
+        device,
+        accelerator,
+        model_save_path,
+        model_dtype,
+        num_experts_per_tok,
+        rank_of_router,
+        rank_of_expert,
+        save_with_remote_code: bool = True,
+        model_type: str = None,  # Optional: explicitly specify model type
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_mappings = None  # Will be set during run()
+
+        if not torch.cuda.is_available():
+            if "cuda" in self.device:
+                self.device = "cpu"
+            if "cuda" in self.accelerator:
+                self.accelerator = "cpu"
+
+    @torch.no_grad()
+    def run(self, modelpool) -> PreTrainedModel:
+        """
+        Executes the upscaling process.
+
+        Args:
+            modelpool (ModelPool): The pool of models to be used for upscaling.
+
+        Returns:
+            PreTrainedModel: The upscaled model (specific type depends on detected model architecture).
+        """
+        self.modelpool = modelpool = to_modelpool(modelpool)
+        config = self.config
+
+        # Auto-detect model type if not specified
+        if self.model_type is None:
+            self.model_type = detect_model_type(
+                modelpool.get_model_path("_pretrained_")
+            )
+            log.info(f"Auto-detected model type: {self.model_type}")
+
+        # Get the appropriate model mappings
+        if self.model_type not in MODEL_TYPE_MAPPINGS:
+            raise ValueError(
+                f"Unsupported model type: {self.model_type}. Supported: {list(MODEL_TYPE_MAPPINGS.keys())}"
+            )
+
+        self.model_mappings = MODEL_TYPE_MAPPINGS[self.model_type]
+        log.info(f"Using {self.model_mappings['description']} model architecture")
+
+        with self.profile("load pretrained model"):
+            pretrained_model = modelpool.load_pretrained_model()
+
+        with self.profile("load fine-tuned model"):
+            finetuned_models = [
+                m for m in tqdm(modelpool.models(), total=len(modelpool.model_names))
+            ]
+
+        if self.device == "cuda" and torch.cuda.is_available():
+            pretrained_model = pretrained_model.cuda()
+            print("parameter count of pretrained model:")
+            print_parameters(pretrained_model)
+            finetuned_models = [m.cuda() for m in finetuned_models]
+
+        with self.profile("merge model"):
+            model = self.merge(pretrained_model, finetuned_models)
+
+        self.print_profile_summary()
+        print("parameter count of upscaled MoE model:")
+        print_parameters(model)
+        print(model)
+
+        if self.model_dtype is not None:
+            model.to(dtype=parse_dtype(self.model_dtype))
+
+        if self.model_save_path is not None:
+            if os.path.dirname(self.model_save_path):
+                os.makedirs(os.path.dirname(self.model_save_path), exist_ok=True)
+            log.info(f"Saving model to {self.model_save_path}")
+            tokenizer = self.modelpool.load_tokenizer()
+            tokenizer.save_pretrained(self.model_save_path)
+            if not self.save_with_remote_code:
+                model.save_pretrained(self.model_save_path)
+            else:
+                # Use the appropriate auto_map for the detected model type
+                auto_map = {
+                    "AutoConfig": self.model_mappings["smile_config_cls"],
+                    "AutoModel": self.model_mappings["smile_base_model_cls"],
+                    "AutoModelForCausalLM": self.model_mappings["smile_model_cls"],
+                }
+                save_pretrained_with_remote_code(
+                    model,
+                    auto_map=auto_map,
+                    save_directory=self.model_save_path,
+                )
+
+            # save readme
+            model_card_str = create_default_model_card(
+                models=[modelpool.get_model_path(m) for m in modelpool.all_model_names],
+                description=f"Merged {self.model_mappings['description']} model using SMILE Upscaling",
+                algorithm_config=self.config,
+                modelpool_config=modelpool.config,
+            )
+            with open(os.path.join(self.model_save_path, "README.md"), "w") as f:
+                f.write(model_card_str)
+
+        return model
+
+    def merge(
+        self,
+        pretrained_model: PreTrainedModel,
+        finetuned_models: List[PreTrainedModel],
+    ) -> PreTrainedModel:
+        """
+        Merges the pretrained model with the fine-tuned models to create an upscaled model.
+
+        Args:
+            pretrained_model (PreTrainedModel): The pretrained model.
+            finetuned_models (List[PreTrainedModel]): A list of fine-tuned models.
+
+        Returns:
+            PreTrainedModel: The upscaled model (specific type depends on model architecture).
+        """
+        with init_empty_weights():
+            pretrained_model_config = self.modelpool.get_model_config("_pretrained_")
+            if isinstance(pretrained_model_config, str):
+                pretrained_path = pretrained_model_config
+            else:
+                pretrained_path = pretrained_model_config.get(
+                    "path", pretrained_model_config["pretrained_model_name_or_path"]
+                )
+            base_config = AutoConfig.from_pretrained(pretrained_path)
+
+            # Create the appropriate SMILE config for the detected model type
+            SmileConfigClass = self.model_mappings["smile_config_cls"]
+            model_config = SmileConfigClass(
+                num_experts_per_tok=self.num_experts_per_tok,
+                rank_of_router=self.rank_of_router,
+                rank_of_expert=self.rank_of_expert,
+                num_local_experts=len(finetuned_models),
+                **base_config.to_dict(),
+            )
+
+            # Create the appropriate SMILE model for the detected model type
+            SmileModelClass = self.model_mappings["smile_model_cls"]
+            model = SmileModelClass(model_config)
+
+        model.to(dtype=pretrained_model.dtype).to_empty(device="cpu")
+
+        # copy pretrained model weights
+        state_dict = model.state_dict()
+        pretrained_state_dict = pretrained_model.state_dict()
+        for key in list(pretrained_state_dict.keys()):
+            if key not in state_dict:
+                pretrained_state_dict.pop(key)
+        model.load_state_dict(pretrained_state_dict, strict=False)
+
+        # upscale model
+        BaseDecoderLayerClass = self.model_mappings["base_decoder_layer_cls"]
+        SmileDecoderLayerClass = self.model_mappings["smile_decoder_layer_cls"]
+
+        for layer_idx in tqdm(
+            range(len(pretrained_model.model.layers)),
+            "Upscaling Modules (layer)",
+            dynamic_ncols=True,
+        ):
+            if RuntimeConstants.debug and layer_idx > 0:
+                log.info(
+                    "Debug mode enabled: processing only the first layer, skipping remaining layers"
+                )
+                break
+
+            pretrained_layer = pretrained_model.model.layers[layer_idx]
+            finetuned_layers = [m.model.layers[layer_idx] for m in finetuned_models]
+
+            target_layer = model.model.layers[layer_idx]
+
+            for n in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+                try:
+                    upscale_to_smile_linear(
+                        base=getattr(pretrained_layer.self_attn, n),
+                        experts=[getattr(m.self_attn, n) for m in finetuned_layers],
+                        target=getattr(target_layer.self_attn, n),
+                        accelerator=self.accelerator,
+                    )
+                except ExpertNotTrainedError:
+                    setattr(
+                        target_layer.self_attn,
+                        n,
+                        getattr(pretrained_layer.self_attn, n),
+                    )
+
+            for n in ["gate_proj", "up_proj", "down_proj"]:
+                try:
+                    upscale_to_smile_linear(
+                        base=getattr(pretrained_layer.mlp, n),
+                        experts=[getattr(m.mlp, n) for m in finetuned_layers],
+                        target=getattr(target_layer.mlp, n),
+                        accelerator=self.accelerator,
+                    )
+                except ExpertNotTrainedError:
+                    setattr(
+                        target_layer.mlp,
+                        n,
+                        getattr(pretrained_layer.mlp, n),
+                    )
+
+        return model

fusion_bench/method/smile_upscaling/projected_energy.py

@@ -3,12 +3,11 @@ from typing import Literal
 
 import pandas as pd
 import torch
+from tqdm import tqdm
 
 from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
 from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 
-from tqdm import tqdm
-
 
 class ProjectedEnergyAnalysis(
     SimpleProfilerMixin,

fusion_bench/method/smile_upscaling/smile_mistral_upscaling.py

@@ -20,6 +20,7 @@ from transformers.models.mistral.modeling_mistral import MistralDecoderLayer
 from fusion_bench.compat.modelpool import to_modelpool
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.method.simple_average import simple_average
+from fusion_bench.mixins import auto_register_config
 from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.models.modeling_smile_mistral import (
@@ -40,7 +41,10 @@ from fusion_bench.utils.parameters import print_parameters
 log = logging.getLogger(__name__)
 
 
-class SmileMistralUpscalingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
+class SmileMistralUpscalingAlgorithm(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
     R"""
     SmileMistralUpscalingAlgorithm is a model fusion algorithm designed to upscale
     a pretrained Mistral model using a set of fine-tuned expert models. The algorithm