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

fusion_bench/modelpool/base_pool.py

@@ -1,13 +1,13 @@
 import logging
 from copy import deepcopy
-from typing import Dict, List, Optional, Union
+from typing import Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 from omegaconf import DictConfig
 from torch import nn
 from torch.utils.data import Dataset
 
-from fusion_bench.mixins import BaseYAMLSerializableModel, HydraConfigMixin
+from fusion_bench.mixins import BaseYAMLSerializable, HydraConfigMixin
 from fusion_bench.utils import instantiate, timeit_context
 
 __all__ = ["BaseModelPool"]
@@ -15,7 +15,10 @@ __all__ = ["BaseModelPool"]
 log = logging.getLogger(__name__)
 
 
-class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
+class BaseModelPool(
+    HydraConfigMixin,
+    BaseYAMLSerializable,
+):
     """
     A class for managing and interacting with a pool of models along with their associated datasets or other specifications. For example, a model pool may contain multiple models, each with its own training, validation, and testing datasets. As for the specifications, a vision model pool may contain image preprocessor, and a language model pool may contain a tokenizer.
 
@@ -31,7 +34,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
     _program = None
     _config_key = "modelpool"
     _models: Union[DictConfig, Dict[str, nn.Module]]
-    _config_mapping = BaseYAMLSerializableModel._config_mapping | {
+    _config_mapping = BaseYAMLSerializable._config_mapping | {
         "_models": "models",
         "_train_datasets": "train_datasets",
         "_val_datasets": "val_datasets",
@@ -56,7 +59,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         super().__init__(**kwargs)
 
     @property
-    def has_pretrained(self):
+    def has_pretrained(self) -> bool:
         """
         Check if the model pool contains a pretrained model.
 
@@ -125,7 +128,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         return len(self.model_names)
 
     @staticmethod
-    def is_special_model(model_name: str):
+    def is_special_model(model_name: str) -> bool:
         """
         Determine if a model is special based on its name.
 
@@ -152,6 +155,23 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
             model_config = deepcopy(model_config)
         return model_config
 
+    def get_model_path(self, model_name: str) -> str:
+        """
+        Get the path for the specified model.
+
+        Args:
+            model_name (str): The name of the model.
+
+        Returns:
+            str: The path for the specified model.
+        """
+        if isinstance(self._models[model_name], str):
+            return self._models[model_name]
+        else:
+            raise ValueError(
+                "Model path is not a string. Try to override this method in derived modelpool class."
+            )
+
     def load_model(
         self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
     ) -> nn.Module:
@@ -159,7 +179,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
         Load a model from the pool based on the provided configuration.
 
         Args:
-            model (Union[str, DictConfig]): The model name or configuration.
+            model_name_or_config (Union[str, DictConfig]): The model name or configuration.
 
         Returns:
             nn.Module: The instantiated model.
@@ -201,11 +221,11 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
             model = self.load_model(self.model_names[0], *args, **kwargs)
         return model
 
-    def models(self):
+    def models(self) -> Generator[nn.Module, None, None]:
         for model_name in self.model_names:
             yield self.load_model(model_name)
 
-    def named_models(self):
+    def named_models(self) -> Generator[Tuple[str, nn.Module], None, None]:
         for model_name in self.model_names:
             yield model_name, self.load_model(model_name)
 

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -57,6 +57,15 @@ class CausalLMPool(BaseModelPool):
                 )
         self.load_lazy = load_lazy
 
+    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")
+
     @override
     def load_model(
         self,

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,5 @@
 # flake8: noqa F401
+from fusion_bench.utils import LazyStateDict
+
 from . import separate_io, utils
 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,182 @@
+"""
+This module contains utilities for working with Hugging Face models.
+"""
+
+import inspect
+import os
+import shutil
+from typing import Optional, cast
+
+from omegaconf import OmegaConf
+from transformers.modeling_utils import PreTrainedModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.utils.pylogger import getRankZeroLogger
+
+log = getRankZeroLogger(__name__)
+
+__all__ = [
+    "save_pretrained_with_remote_code",
+    "generate_readme_head",
+    "generate_readme_body",
+    "generate_complete_readme",
+]
+
+
+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 generate_readme_head(
+    models: list[str] | BaseModelPool,
+    library_name: str = "transformers",
+    tags: list[str] = ["fusion-bench", "merge"],
+):
+    text = "---\nbase_model:\n"
+    for model_name in models:
+        text += f"- {model_name}\n"
+    if library_name:
+        text += f"library_name: {library_name}\n"
+    text += "tags:\n"
+    for tag in tags:
+        text += f"- {tag}\n"
+    text += "---\n"
+    return text
+
+
+def generate_readme_body(
+    algorithm: BaseAlgorithm,
+    models_or_modelpool: Optional[list[str] | BaseModelPool] = None,
+    models: list[str] = None,
+):
+    text = """\
+# Merge
+
+This is a merge of pre-trained language models created using [fusion-bench](https://github.com/tanganke/fusion_bench).
+
+"""
+
+    if models is not None:
+        text += """
+## Models Merged
+
+The following models were included in the merge:
+
+"""
+        for model_name in models:
+            text += f"- {model_name}\n"
+        text += "\n"
+
+        try:
+            text += f"""\
+    ## Configuration
+
+    The following YAML configuration was used to produce this model:
+
+    ```yaml
+    {OmegaConf.to_yaml(algorithm.config, resolve=True, sort_keys=True)}
+    ```
+    """
+        except Exception as e:
+            return (
+                text  # If the algorithm config cannot be converted to YAML, we skip it.
+            )
+
+    if isinstance(models_or_modelpool, BaseModelPool):
+        try:
+            text += f"""
+```yaml
+{OmegaConf.to_yaml(models_or_modelpool.config, resolve=True, sort_keys=True)}
+```
+"""
+        except Exception as e:
+            pass  # If the model pool config cannot be converted to YAML, we skip it.
+    return text
+
+
+def generate_complete_readme(
+    algorithm: BaseAlgorithm, modelpool: BaseModelPool, models: list[str]
+):
+    # Generate the complete README content
+    text = generate_readme_head(
+        [modelpool.get_model_path(m) for m in modelpool.model_names]
+    )
+    readme_body = generate_readme_body(
+        algorithm,
+        models_or_modelpool=modelpool,
+        models=[modelpool.get_model_path(m) for m in modelpool.model_names],
+    )
+    complete_readme = text + "\n" + readme_body
+    return complete_readme

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/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)