fusion-bench 0.2.20__py3-none-any.whl → 0.2.22__py3-none-any.whl

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -8,7 +8,7 @@ from copy import deepcopy
 from typing import Any, Dict, Optional, TypeAlias, Union, cast  # noqa: F401
 
 import peft
-from omegaconf import DictConfig, flag_override
+from omegaconf import DictConfig, OmegaConf, flag_override
 from torch import nn
 from torch.nn.modules import Module
 from transformers import (
@@ -19,43 +19,51 @@ from transformers import (
 )
 from typing_extensions import override
 
-from fusion_bench.modelpool import BaseModelPool
-from fusion_bench.utils import instantiate
-from fusion_bench.utils.dtype import parse_dtype
+from fusion_bench import (
+    BaseModelPool,
+    auto_register_config,
+    import_object,
+    instantiate,
+    parse_dtype,
+)
 from fusion_bench.utils.lazy_state_dict import LazyStateDict
-from fusion_bench.utils.packages import import_object
 
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class CausalLMPool(BaseModelPool):
-    _config_mapping = BaseModelPool._config_mapping | {
-        "_tokenizer": "tokenizer",
-        "_model_kwargs": "model_kwargs",
-        "load_lazy": "load_lazy",
-    }
-
     def __init__(
         self,
         models,
         *,
-        tokenizer: Optional[DictConfig],
+        tokenizer: Optional[DictConfig | str],
         model_kwargs: Optional[DictConfig] = None,
-        load_lazy: bool = False,
+        enable_lazy_loading: bool = False,
         **kwargs,
     ):
         super().__init__(models, **kwargs)
-        # process `model_kwargs`
-        self._tokenizer = tokenizer
-        self._model_kwargs = model_kwargs
-        if self._model_kwargs is None:
-            self._model_kwargs = DictConfig({})
-        with flag_override(self._model_kwargs, "allow_objects", True):
-            if hasattr(self._model_kwargs, "torch_dtype"):
-                self._model_kwargs.torch_dtype = parse_dtype(
-                    self._model_kwargs.torch_dtype
-                )
-        self.load_lazy = load_lazy
+        if model_kwargs is None:
+            self.model_kwargs = DictConfig({})
+
+    def get_model_path(self, model_name: str):
+        model_name_or_config = self._models[model_name]
+        if isinstance(model_name_or_config, str):
+            return model_name_or_config
+        elif isinstance(model_name_or_config, (DictConfig, dict)):
+            return model_name_or_config.get("pretrained_model_name_or_path")
+        else:
+            raise RuntimeError("Invalid model configuration")
+
+    def get_model_kwargs(self):
+        model_kwargs = (
+            OmegaConf.to_container(self.model_kwargs, resolve=True)
+            if isinstance(self.model_kwargs, DictConfig)
+            else self.model_kwargs
+        )
+        if "torch_dtype" in model_kwargs:
+            model_kwargs["torch_dtype"] = parse_dtype(model_kwargs["torch_dtype"])
+        return model_kwargs
 
     @override
     def load_model(
@@ -89,7 +97,7 @@ class CausalLMPool(BaseModelPool):
             pretrained_model_name_or_path: path_to_model_b
         ```
         """
-        model_kwargs = deepcopy(self._model_kwargs)
+        model_kwargs = self.get_model_kwargs()
         model_kwargs.update(kwargs)
 
         if isinstance(model_name_or_config, str):
@@ -99,7 +107,7 @@ class CausalLMPool(BaseModelPool):
                 model_config = self._models[model_name_or_config]
                 if isinstance(model_config, str):
                     # model_config is a string
-                    if not self.load_lazy:
+                    if not self.enable_lazy_loading:
                         model = AutoModelForCausalLM.from_pretrained(
                             model_config,
                             *args,
@@ -117,7 +125,7 @@ class CausalLMPool(BaseModelPool):
         elif isinstance(model_name_or_config, (DictConfig, Dict)):
             model_config = model_name_or_config
 
-        if not self.load_lazy:
+        if not self.enable_lazy_loading:
             model = instantiate(model_config, *args, **model_kwargs)
         else:
             meta_module_class = model_config.pop("_target_")
@@ -149,12 +157,12 @@ class CausalLMPool(BaseModelPool):
         Returns:
             PreTrainedTokenizer: The tokenizer.
         """
-        assert self._tokenizer is not None, "Tokenizer is not defined in the config"
+        assert self.tokenizer is not None, "Tokenizer is not defined in the config"
         log.info("Loading tokenizer.", stacklevel=2)
-        if isinstance(self._tokenizer, str):
-            tokenizer = AutoTokenizer.from_pretrained(self._tokenizer, *args, **kwargs)
+        if isinstance(self.tokenizer, str):
+            tokenizer = AutoTokenizer.from_pretrained(self.tokenizer, *args, **kwargs)
         else:
-            tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+            tokenizer = instantiate(self.tokenizer, *args, **kwargs)
         return tokenizer
 
     @override
@@ -204,12 +212,12 @@ class CausalLMBackbonePool(CausalLMPool):
     def load_model(
         self, model_name_or_config: str | DictConfig, *args, **kwargs
     ) -> Module:
-        if self.load_lazy:
+        if self.enable_lazy_loading:
             log.warning(
                 "CausalLMBackbonePool does not support lazy loading. "
                 "Falling back to normal loading."
             )
-            self.load_lazy = False
+            self.enable_lazy_loading = False
         model: AutoModelForCausalLM = super().load_model(
             model_name_or_config, *args, **kwargs
         )

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -11,9 +11,7 @@ from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
 from typing_extensions import override
 
 from fusion_bench.utils import instantiate, timeit_context
-from fusion_bench.utils.modelscope import (
-    resolve_repo_path,
-)
+from fusion_bench.utils.modelscope import resolve_repo_path
 
 from ..base_pool import BaseModelPool
 

fusion_bench/modelpool/seq_classification_lm/__init__.py

@@ -1,2 +1,2 @@
 from .reward_model import create_reward_model_from_pretrained
-from .seq_classification_lm import SeqenceClassificationModelPool
+from .seq_classification_lm import SequenceClassificationModelPool

fusion_bench/modelpool/seq_classification_lm/seq_classification_lm.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-class SeqenceClassificationModelPool(BaseModelPool):
+class SequenceClassificationModelPool(BaseModelPool):
 
     def __init__(
         self,

fusion_bench/models/__init__.py

@@ -1,4 +1,10 @@
 # flake8: noqa F401
+from fusion_bench.utils import LazyStateDict
+
 from . import separate_io, utils
+from .hf_utils import (
+    create_default_model_card,
+    load_model_card_template,
+    save_pretrained_with_remote_code,
+)
 from .parameter_dict import ParameterDictModel
-from fusion_bench.utils import LazyStateDict

fusion_bench/models/expert_sparsity/mixtral/__init__.py

@@ -10,6 +10,6 @@ Reference:
 """
 
 from .wrapper import (
-    PrunableMixtralSparseMoeBlockWrapper,
     DynamicSkippingMixtralSparseMoeBlockWrapper,
+    PrunableMixtralSparseMoeBlockWrapper,
 )

fusion_bench/models/hf_utils.py

@@ -0,0 +1,160 @@
+"""
+This module contains utilities for working with Hugging Face models.
+"""
+
+import inspect
+import os
+import shutil
+from typing import List, Optional, cast
+
+from omegaconf import DictConfig, OmegaConf
+from transformers.modeling_utils import PreTrainedModel
+
+from fusion_bench.utils.pylogger import get_rankzero_logger
+
+log = get_rankzero_logger(__name__)
+
+__all__ = [
+    "load_model_card_template",
+    "save_pretrained_with_remote_code",
+    "create_default_model_card",
+]
+
+MODEL_CARD_TEMPLATE_DIRS = [
+    os.path.join(os.path.dirname(__file__), "model_card_templates")
+]
+
+
+def load_model_card_template(basename: str) -> str:
+    """
+    Load a model card template from file.
+
+    Searches for a template file by name, first checking if the name is a direct file path,
+    then searching through predefined template directories.
+
+    Args:
+        name (str): The name of the template file or a direct file path to the template.
+
+    Returns:
+        str: The contents of the template file as a string.
+
+    Raises:
+        FileNotFoundError: If the template file is not found in any of the search locations.
+    """
+    if os.path.exists(basename):
+        return open(basename).read()
+
+    for template_dir in MODEL_CARD_TEMPLATE_DIRS:
+        template_path = os.path.join(template_dir, basename)
+        if os.path.exists(template_path):
+            return open(template_path).read()
+
+    raise FileNotFoundError(f"Model card template '{basename}' not found.")
+
+
+def try_to_yaml(config):
+    if config is None:
+        return None
+
+    try:
+        return OmegaConf.to_yaml(config, resolve=True, sort_keys=True)
+    except Exception as e:
+        log.error(f"Failed to convert config to YAML: {e}. Return `None`.")
+        return None
+
+
+def save_pretrained_with_remote_code(
+    model: PreTrainedModel,
+    auto_map: dict[str, object],
+    save_directory,
+    **kwargs,
+):
+    """
+    Saves a model with custom code to a directory.
+
+    This function facilitates saving a Hugging Face `PreTrainedModel` along with its
+    associated custom code. It inspects the objects provided in the `auto_map`,
+    copies their source files to the `save_directory`, and generates an `__init__.py`
+    to make them importable. It also updates the model's configuration with an
+    `auto_map` attribute, which allows `AutoModel.from_pretrained` to correctly
+    instantiate the custom model classes when `trust_remote_code=True`.
+
+    Args:
+        model (PreTrainedModel): The model instance to be saved.
+        auto_map (dict[str, object]): A dictionary mapping auto class names
+            (e.g., "AutoModelForCausalLM") to the corresponding custom class objects.
+        save_directory (str or os.PathLike): The directory where the model and
+            custom code files will be saved.
+        **kwargs: Additional keyword arguments to be passed to the
+            `model.save_pretrained` method.
+
+    Example:
+        ```python
+        # Assuming `model` is an instance of `SmileQwen2ForCausalLM`
+        # and `SmileQwen2Config`, `SmileQwen2Model`, `SmileQwen2ForCausalLM`
+        # are custom classes defined in your project.
+
+        save_pretrained_with_remote_code(
+            model,
+            auto_map={
+                "AutoConfig": SmileQwen2Config,
+                "AutoModel": SmileQwen2Model,
+                "AutoModelForCausalLM": SmileQwen2ForCausalLM,
+            },
+            save_directory="./my-custom-model",
+        )
+
+        # The model can then be loaded with `trust_remote_code=True`:
+        # from transformers import AutoModelForCausalLM
+        # loaded_model = AutoModelForCausalLM.from_pretrained(
+        #     "./my-custom-model", trust_remote_code=True
+        # )
+        ```
+    """
+    auto_map_files = {}
+    auto_map_strs = {}
+    for key, obj in auto_map.items():
+        auto_map_files[key] = inspect.getfile(obj)
+
+    for key, obj in auto_map.items():
+        auto_map_strs[key] = (
+            f"{(inspect.getmodule(obj).__name__).split('.')[-1]}.{obj.__name__}"
+        )
+
+    model.config.auto_map = auto_map_strs
+
+    # save model to `save_directory`
+    model.save_pretrained(save_directory=save_directory, **kwargs)
+
+    # copy source files to `save_directory`
+    for key, file_path in auto_map_files.items():
+        shutil.copy(
+            src=file_path, dst=os.path.join(save_directory, os.path.basename(file_path))
+        )
+    # construct `__init__.py`
+    init_file = os.path.join(save_directory, "__init__.py")
+    with open(init_file, "w") as f:
+        for key, file_name in auto_map_files.items():
+            base_name = os.path.basename(file_name).split(".")[0]
+            f.write(f"from .{base_name} import {auto_map[key].__name__}\n")
+
+
+def create_default_model_card(
+    models: list[str],
+    description=None,
+    algorithm_config: DictConfig = None,
+    modelpool_config: DictConfig = None,
+):
+    from jinja2 import Template
+
+    template: Template = Template(load_model_card_template("default.md"))
+    card = template.render(
+        models=models,
+        library_name="transformers",
+        tags=["fusion-bench", "merge"],
+        title="Deep Model Fusion",
+        description=description,
+        algorithm_config_str=try_to_yaml(algorithm_config),
+        modelpool_config_str=try_to_yaml(modelpool_config),
+    )
+    return card

fusion_bench/models/linearized/linearized_model_utils.py

@@ -1,7 +1,7 @@
 import logging
 from collections import OrderedDict
 from copy import deepcopy
-from typing import Optional
+from typing import Any, Dict, Optional, Tuple
 
 import torch.nn as nn
 from torch.func import functional_call, jvp
@@ -9,7 +9,7 @@ from torch.func import functional_call, jvp
 log = logging.getLogger(__name__)
 
 
-def dict_params_to_tuple(dict_params: dict):
+def dict_params_to_tuple(dict_params: dict) -> Tuple:
     return tuple(v for k, v in dict_params.items())
 
 
@@ -33,7 +33,7 @@ class LinearizedModelWraper(nn.Module):
         for p in self.params0_values:
             p.requires_grad_(False)
 
-    def tuple_params_to_dict(self, tuple_params):
+    def tuple_params_to_dict(self, tuple_params) -> Dict[str, Any]:
         """
         Converts a tuple of parameters to a dictionary with keys corresponding to the parameter names.
 
@@ -50,7 +50,7 @@ class LinearizedModelWraper(nn.Module):
             state_dict[k] = p
         return state_dict
 
-    def forward(self, *args, **kwargs):
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
         """
         Computes the linearized model output using a first-order Taylor decomposition.
 

fusion_bench/models/linearized/vision_model.py

@@ -70,7 +70,7 @@ def load_lora_vision_model_hf(
     peft_name: str,
     merge_and_unload: bool = False,
     return_vison_model=True,
-):
+) -> PeftModel:
     """
     Load a LoRA (Low-Rank Adaptation) vision model from Hugging Face.
 

fusion_bench/models/model_card_templates/default.md

@@ -0,0 +1,46 @@
+---
+base_model:
+{%- for model in models %}
+- {{ model }}
+{%- endfor %}
+library_name: {{ library_name }}
+tags:
+{%- for tag in tags %}
+- {{ tag }}
+{%- endfor %}
+---
+# {{ title }}
+
+{% if description is not none %}{{ description }}{% endif %}
+
+## Models Merged
+
+This is a merged model created using [fusion-bench](https://github.com/tanganke/fusion_bench).
+
+The following models were included in the merge:
+{% for model in models %}
+- {{ model }}
+{%- endfor %}
+
+{% if algorithm_config_str is not none or modelpool_config_str is not none %}
+## Configuration
+
+The following YAML configuration was used to produce this model:
+
+{% if algorithm_config_str is not none -%}
+### Algorithm Configuration
+
+```yaml
+{{ algorithm_config_str -}}
+```
+{%- endif %}
+
+{% if modelpool_config_str is not none -%}
+### Model Pool Configuration
+
+```yaml
+{{ modelpool_config_str -}}
+```
+{%- endif %}
+
+{% endif %}

fusion_bench/models/modeling_deepseek_v2/__init__.py

@@ -4,12 +4,12 @@ This is a direct copy of the DeepSeek-V2-Lite model from HuggingFace https://hug
 
 from .configuration_deepseek import DeepseekV2Config
 from .modeling_deepseek import (
+    DeepseekV2DecoderLayer,
     DeepseekV2ForCausalLM,
     DeepseekV2ForSequenceClassification,
     DeepseekV2MLP,
     DeepseekV2Model,
     DeepseekV2MoE,
-    DeepseekV2DecoderLayer,
 )
 from .modeling_deepseek import MoEGate as DeepseekV2MoEGate
 from .tokenization_deepseek_fast import DeepseekTokenizerFast

fusion_bench/models/modeling_deepseek_v2/modeling_deepseek.py

@@ -17,17 +17,18 @@
 # 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.
-""" PyTorch DeepSeek model."""
+"""PyTorch DeepSeek model."""
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
 
+import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_attn_mask_utils import (
@@ -54,9 +55,8 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from transformers.utils.import_utils import is_torch_fx_available
+
 from .configuration_deepseek import DeepseekV2Config
-import torch.distributed as dist
-import numpy as np
 
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func

fusion_bench/models/modeling_deepseek_v2/tokenization_deepseek_fast.py

@@ -1,6 +1,5 @@
 from typing import List, Optional, Union
 
-
 from transformers.models.llama import LlamaTokenizerFast
 
 

fusion_bench/models/modeling_smile_gemma2/__init__.py

@@ -0,0 +1,9 @@
+from . import register
+from .configuration_smile_gemma2 import SmileGemma2Config
+from .modeling_smile_gemma2 import (
+    SmileGemma2ForCausalLM,
+    SmileGemma2ForSequenceClassification,
+    SmileGemma2ForTokenClassification,
+    SmileGemma2Model,
+    SmileGemma2PreTrainedModel,
+)

fusion_bench/models/modeling_smile_gemma2/configuration_smile_gemma2.py

@@ -0,0 +1,20 @@
+from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
+
+
+class SmileGemma2Config(Gemma2Config):
+    model_type = "smile_gemma2"
+
+    def __init__(
+        self,
+        num_experts_per_tok: int = 1,
+        rank_of_router: int = None,
+        rank_of_expert: int = None,
+        num_local_experts: int = None,
+        **kwargs,
+    ):
+        self.num_experts_per_tok = num_experts_per_tok
+        self.rank_of_router = rank_of_router
+        self.rank_of_expert = rank_of_expert
+        self.num_local_experts = num_local_experts
+
+        super().__init__(**kwargs)