fusion-bench 0.2.14__tar.gz → 0.2.15__tar.gz

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.14
+Version: 0.2.15
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -63,7 +63,7 @@ Dynamic: license-file
 
 </div>
 
-> [!TIP]  
+> [!TIP]
 > Documentation is available at [tanganke.github.io/fusion_bench/](https://tanganke.github.io/fusion_bench/).
 
 ## Overview
@@ -157,6 +157,9 @@ pip install -e ".[lm-eval-harness]"
 This will install the latest version of fusion-bench and the dependencies required for LM-Eval Harness.
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
+> [!TIP]
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+
 ## Introduction to Deep Model Fusion
 
 Deep model fusion is a technique that merges, ensemble, or fuse multiple deep neural networks to obtain a unified model.

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/README.md

@@ -12,7 +12,7 @@
 
 </div>
 
-> [!TIP]  
+> [!TIP]
 > Documentation is available at [tanganke.github.io/fusion_bench/](https://tanganke.github.io/fusion_bench/).
 
 ## Overview
@@ -106,6 +106,9 @@ pip install -e ".[lm-eval-harness]"
 This will install the latest version of fusion-bench and the dependencies required for LM-Eval Harness.
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
+> [!TIP]
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+
 ## Introduction to Deep Model Fusion
 
 Deep model fusion is a technique that merges, ensemble, or fuse multiple deep neural networks to obtain a unified model.

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -1,15 +1,19 @@
+"""
+Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/causal_lm
+"""
+
 import logging
 import os
 from copy import deepcopy
-from typing import Any, Optional, TypeAlias, Union, cast  # noqa: F401
+from typing import Any, Dict, Optional, TypeAlias, Union, cast  # noqa: F401
 
 import peft
 from omegaconf import DictConfig, flag_override
 from torch import nn
 from torch.nn.modules import Module
 from transformers import (
-    LlamaForCausalLM,
-    MistralForCausalLM,
+    AutoModelForCausalLM,
+    AutoTokenizer,
     PreTrainedModel,
     PreTrainedTokenizer,
 )
@@ -21,8 +25,6 @@ from fusion_bench.utils.dtype import parse_dtype
 
 log = logging.getLogger(__name__)
 
-CausalLM: TypeAlias = Union[LlamaForCausalLM, MistralForCausalLM, Any]
-
 
 class CausalLMPool(BaseModelPool):
     _config_mapping = BaseModelPool._config_mapping | {
@@ -56,17 +58,78 @@ class CausalLMPool(BaseModelPool):
         model_name_or_config: str | DictConfig,
         *args,
         **kwargs,
-    ) -> LlamaForCausalLM | MistralForCausalLM | nn.Module:
+    ) -> PreTrainedModel:
+        """
+        Example of YAML config:
+
+        ```yaml
+        models:
+          _pretrained_: path_to_pretrained_model # if a plain string, it will be passed to AutoModelForCausalLM.from_pretrained
+          model_a: path_to_model_a
+          model_b: path_to_model_b
+        ```
+
+        or equivalently,
+
+        ```yaml
+        models:
+          _pretrained_:
+            _target_: transformers.AutoModelForCausalLM # any callable that returns a model
+            pretrained_model_name_or_path: path_to_pretrained_model
+          model_a:
+            _target_: transformers.AutoModelForCausalLM
+            pretrained_model_name_or_path: path_to_model_a
+          model_b:
+            _target_: transformers.AutoModelForCausalLM
+            pretrained_model_name_or_path: path_to_model_b
+        ```
+        """
         model_kwargs = deepcopy(self._model_kwargs)
         model_kwargs.update(kwargs)
+
         if isinstance(model_name_or_config, str):
             log.info(f"Loading model: {model_name_or_config}", stacklevel=2)
-        return super().load_model(model_name_or_config, *args, **model_kwargs)
+            if model_name_or_config in self._models.keys():
+                model_config = self._models[model_name_or_config]
+                if isinstance(model_config, str):
+                    # model_config is a string
+                    model = AutoModelForCausalLM.from_pretrained(
+                        model_config,
+                        *args,
+                        **model_kwargs,
+                    )
+                    return model
+        elif isinstance(model_name_or_config, (DictConfig, Dict)):
+            model_config = model_name_or_config
+
+        model = instantiate(model_config, *args, **model_kwargs)
+        return model
 
     def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
+        """
+        Example of YAML config:
+
+        ```yaml
+        tokenizer: google/gemma-2-2b-it # if a plain string, it will be passed to AutoTokenizer.from_pretrained
+        ```
+
+        or equivalently,
+
+        ```yaml
+        tokenizer:
+          _target_: transformers.AutoTokenizer # any callable that returns a tokenizer
+          pretrained_model_name_or_path: google/gemma-2-2b-it
+        ```
+
+        Returns:
+            PreTrainedTokenizer: The tokenizer.
+        """
         assert self._tokenizer is not None, "Tokenizer is not defined in the config"
         log.info("Loading tokenizer.", stacklevel=2)
-        tokenizer = instantiate(self._tokenizer, *args, **kwargs)
+        if isinstance(self._tokenizer, str):
+            tokenizer = AutoTokenizer.from_pretrained(self._tokenizer, *args, **kwargs)
+        else:
+            tokenizer = instantiate(self._tokenizer, *args, **kwargs)
         return tokenizer
 
     @override
@@ -113,7 +176,7 @@ class CausalLMBackbonePool(CausalLMPool):
     def load_model(
         self, model_name_or_config: str | DictConfig, *args, **kwargs
     ) -> Module:
-        model: Union[MistralForCausalLM, LlamaForCausalLM, Any] = super().load_model(
+        model: AutoModelForCausalLM = super().load_model(
             model_name_or_config, *args, **kwargs
         )
         return model.model.layers
@@ -126,7 +189,7 @@ def load_peft_causal_lm(
     is_trainable: bool = True,
     merge_and_unload: bool = False,
 ):
-    base_model = LlamaForCausalLM.from_pretrained(
+    base_model = AutoModelForCausalLM.from_pretrained(
         base_model_path, torch_dtype=torch_dtype
     )
     model = peft.PeftModel.from_pretrained(

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/fusion_bench.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.14
+Version: 0.2.15
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -63,7 +63,7 @@ Dynamic: license-file
 
 </div>
 
-> [!TIP]  
+> [!TIP]
 > Documentation is available at [tanganke.github.io/fusion_bench/](https://tanganke.github.io/fusion_bench/).
 
 ## Overview
@@ -157,6 +157,9 @@ pip install -e ".[lm-eval-harness]"
 This will install the latest version of fusion-bench and the dependencies required for LM-Eval Harness.
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
+> [!TIP]
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+
 ## Introduction to Deep Model Fusion
 
 Deep model fusion is a technique that merges, ensemble, or fuse multiple deep neural networks to obtain a unified model.

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/fusion_bench.egg-info/SOURCES.txt

@@ -344,6 +344,14 @@ fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/llama_ultrachat.yaml
 fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/qwen2_math_1.5B_and_R1.yaml
 fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/simle_mixtral_exp_v4.yaml
 fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/single_llama_model.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.1-8B-Instruct.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.1-8B.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.2-3B-Instruct.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.2-3B.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-2b-it.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-2b.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-9b-it.yaml
+fusion_bench/../fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-9b.yaml
 fusion_bench/../fusion_bench_config/modelpool/OpenCLIPVisionModelPool/README.md
 fusion_bench/../fusion_bench_config/modelpool/OpenCLIPVisionModelPool/ViT-B-16_TA8.yaml
 fusion_bench/../fusion_bench_config/modelpool/OpenCLIPVisionModelPool/ViT-B-32_TA8.yaml

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.1-8B-Instruct.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: meta-llama/Llama-3.1-8B-Instruct
+  instruction: MergeBench/Llama-3.1-8B-Instruct_instruction
+  math: MergeBench/Llama-3.1-8B-Instruct_math
+  coding: MergeBench/Llama-3.1-8B-Instruct_coding
+  multilingual: MergeBench/Llama-3.1-8B-Instruct_multilingual
+  safety: MergeBench/Llama-3.1-8B-Instruct_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: meta-llama/Llama-3.1-8B-Instruct

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.1-8B.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: meta-llama/Llama-3.1-8B
+  instruction: MergeBench/Llama-3.1-8B_instruction
+  math: MergeBench/Llama-3.1-8B_math
+  coding: MergeBench/Llama-3.1-8B_coding
+  multilingual: MergeBench/Llama-3.1-8B_multilingual
+  safety: MergeBench/Llama-3.1-8B_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: meta-llama/Llama-3.1-8B

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.2-3B-Instruct.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: meta-llama/Llama-3.2-3B-Instruct
+  instruction: MergeBench/Llama-3.2-3B-Instruct_instruction
+  math: MergeBench/Llama-3.2-3B-Instruct_math
+  coding: MergeBench/Llama-3.2-3B-Instruct_coding
+  multilingual: MergeBench/Llama-3.2-3B-Instruct_multilingual
+  safety: MergeBench/Llama-3.2-3B-Instruct_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: meta-llama/Llama-3.2-3B-Instruct

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/Llama-3.2-3B.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: meta-llama/Llama-3.2-3B
+  instruction: MergeBench/Llama-3.2-3B_instruction
+  math: MergeBench/Llama-3.2-3B_math
+  coding: MergeBench/Llama-3.2-3B_coding
+  multilingual: MergeBench/Llama-3.2-3B_multilingual
+  safety: MergeBench/Llama-3.2-3B_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: meta-llama/Llama-3.2-3B

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-2b-it.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: google/gemma-2-2b-it
+  instruction: MergeBench/gemma-2-2b-it_instruction
+  math: MergeBench/gemma-2-2b-it_math
+  coding: MergeBench/gemma-2-2b-it_coding
+  multilingual: MergeBench/gemma-2-2b-it_multilingual
+  safety: MergeBench/gemma-2-2b-it_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: google/gemma-2-2b-it

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-2b.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: google/gemma-2-2b
+  instruction: MergeBench/gemma-2-2b_instruction
+  math: MergeBench/gemma-2-2b_math
+  coding: MergeBench/gemma-2-2b_coding
+  multilingual: MergeBench/gemma-2-2b_multilingual
+  safety: MergeBench/gemma-2-2b_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: google/gemma-2-2b

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-9b-it.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: google/gemma-2-9b-it
+  instruction: MergeBench/gemma-2-9b-it_instruction
+  math: MergeBench/gemma-2-9b-it_math
+  coding: MergeBench/gemma-2-9b-it_coding
+  multilingual: MergeBench/gemma-2-9b-it_multilingual
+  safety: MergeBench/gemma-2-9b-it_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: google/gemma-2-9b-it

fusion_bench-0.2.15/fusion_bench_config/modelpool/CausalLMPool/mergebench/gemma-2-9b.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+models:
+  _pretrained_: google/gemma-2-9b
+  instruction: MergeBench/gemma-2-9b_instruction
+  math: MergeBench/gemma-2-9b_math
+  coding: MergeBench/gemma-2-9b_coding
+  multilingual: MergeBench/gemma-2-9b_multilingual
+  safety: MergeBench/gemma-2-9b_safety
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: google/gemma-2-9b

{fusion_bench-0.2.14 → fusion_bench-0.2.15}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "fusion_bench"
 description = "A Comprehensive Benchmark of Deep Model Fusion"
 readme = "README.md"
-version = "0.2.14"
+version = "0.2.15"
 authors = [{ name = "Anke Tang", email = "tang.anke@foxmail.com" }]
 license = { file = 'LICENSE' }
 requires-python = ">= 3.10"