fusion-bench 0.2.15__py3-none-any.whl → 0.2.16__py3-none-any.whl

fusion_bench/method/randes/modelsoup.py

@@ -0,0 +1,126 @@
+import logging
+from copy import deepcopy
+
+import torch
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils.parameters import count_parameters
+from fusion_bench.utils.state_dict_arithmetic import (
+    state_dict_mul,
+)
+
+from .base_algorithm import SuperposedAlgorithmBase, compare_models
+
+log = logging.getLogger(__name__)
+
+
+class SuperposedModelSoupAlgorithm(
+    SuperposedAlgorithmBase,
+):
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(models=modelpool)
+
+        log.info(
+            f"Compressing models using superposed model soup.\n"
+            f"Models: {modelpool.model_names}"
+        )
+        models = {}
+
+        # load state dicts
+        state_dicts = self._load_state_dicts(modelpool)
+        with self.profile("load model"):
+            pretrained_model = modelpool.load_model("_pretrained_")
+        absorber_state_dict = self._compute_absorber(state_dicts, pretrained_model)
+        if absorber_state_dict is not None:
+            state_dicts["absorber"] = absorber_state_dict
+
+        with self.profile("compress and retrieve"):
+            retrieved_state_dicts, metadata = self._compress_and_retrieve(
+                deepcopy(state_dicts), mode="superposed_model_soup"
+            )
+
+        with self.profile("retrieve models"):
+            for model_idx, model_name in enumerate(modelpool.model_names):
+                if self.ms_mode == "average":
+                    coefficient = 1 / len(modelpool.model_names)
+                    retrieved_state_dict = state_dict_mul(
+                        retrieved_state_dicts[model_name], coefficient
+                    )
+                elif self.ms_mode == "original":
+                    retrieved_state_dict = retrieved_state_dicts[model_name]
+                else:
+                    raise ValueError(f"Unsupported ms_mode: {self.ms_mode}")
+                retrieved_model = modelpool.load_model(
+                    model_name
+                )  # TODO: avoid repeated loading
+                # FIXME: for 'all' mode
+                for k, v in retrieved_state_dict.items():
+                    if v.shape[0] == 1:
+                        retrieved_state_dict[k] = v.squeeze(0)
+                retrieved_model.load_state_dict(retrieved_state_dict)
+                models[model_name] = retrieved_model
+                if self.debug >= 1:
+                    with self.profile("metadata"):
+                        if torch.cuda.is_available():
+                            retrieved_state_dicts[model_name] = {
+                                k: v.cuda()
+                                for k, v in retrieved_state_dicts[model_name].items()
+                            }
+                            state_dicts[model_name] = {
+                                k: v.cuda() for k, v in state_dicts[model_name].items()
+                            }
+                            retrieved_state_dict = {
+                                k: v.cuda() for k, v in retrieved_state_dict.items()
+                            }
+
+                        target_layers = metadata["target_layers"]
+                        # focus on the superposition retrieval performance on the target layers
+                        metadata["superposed_model_retrieval_similarity"][
+                            model_name
+                        ] = compare_models(
+                            retrieved_state_dicts[model_name],
+                            state_dicts[model_name],
+                            target_layers,
+                        )
+                        metadata["superposed_model_svd_subspace_similarities"][
+                            model_name
+                        ] = self._compute_svd_subspace_similarities(
+                            state_dicts[model_name],
+                            retrieved_state_dicts[model_name],
+                            target_layers,
+                        )
+                        # overall retrieval performance
+                        metadata["model_retrieval_similarity"][model_name] = (
+                            compare_models(
+                                retrieved_state_dict, state_dicts[model_name]
+                            )
+                        )
+                        metadata["model_svd_subspace_similarities"][model_name] = (
+                            self._compute_svd_subspace_similarities(
+                                state_dicts[model_name], retrieved_state_dict
+                            )
+                        )
+                        # delete the cuda tensors
+                        del (
+                            retrieved_state_dicts[model_name],
+                            state_dicts[model_name],
+                            retrieved_state_dict,
+                        )
+        with self.profile("metadata"):
+            if self.debug >= 0:
+                (
+                    metadata["trainable_param_count_pretrained_model"],
+                    metadata["active_param_count_pretrained_model"],
+                ) = count_parameters(pretrained_model)
+                (
+                    metadata["trainable_param_count_retrieved_model"],
+                    metadata["active_param_count_retrieved_model"],
+                ) = count_parameters(models[modelpool.model_names[0]])
+                print(
+                    f"Total storage (Gbs) for retrieval and original: {metadata['total_gb_retrieved']} | {metadata['total_gb_original']}"
+                )
+        self.print_profile_summary()
+        return {"models": models, "metadata": metadata}

fusion_bench/method/randes/task_arithmetic.py

@@ -0,0 +1,318 @@
+import logging
+import os
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Optional
+
+import torch
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils.parameters import count_parameters
+from fusion_bench.utils.state_dict_arithmetic import (
+    state_dict_add,
+    state_dict_mul,
+    state_dict_sub,
+)
+
+from .base_algorithm import SuperposedAlgorithmBase, compare_models
+
+log = logging.getLogger(__name__)
+
+
+class SuperposedTaskArithmeticAlgorithm(
+    SuperposedAlgorithmBase,
+):
+    _config_mapping = SuperposedAlgorithmBase._config_mapping | {
+        "scaling_factor": "scaling_factor",
+        "model_path": "model_path",
+    }
+
+    def __init__(
+        self,
+        scaling_factor: float,
+        model_path: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.scaling_factor = scaling_factor
+        self.model_path = model_path
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(models=modelpool)
+
+        log.info("Compressing models using superposed task arithmetic.")
+        task_vector = None
+        with self.profile("load model"):
+            pretrained_model = modelpool.load_model("_pretrained_")
+
+        # Calculate the task vector superposition
+        task_vectors = {}
+        models = {}
+        for model_name in modelpool.model_names:
+            with self.profile("load model"):
+                model = modelpool.load_model(model_name)
+            for layer_name, layer in model.state_dict(keep_vars=True).items():
+                if self.verbose >= 1:
+                    log.info(f"{layer_name} | {layer.shape}")
+            task_vector = state_dict_sub(
+                model.state_dict(keep_vars=True),
+                pretrained_model.state_dict(keep_vars=True),
+            )
+            task_vectors[model_name] = task_vector
+
+        with self.profile("compress and retrieve"):
+            retrieved_task_vectors, metadata = self._compress_and_retrieve(
+                deepcopy(task_vectors), mode="superposed_task_arithmetic"
+            )
+        with self.profile("retrieve models"):
+            for model_name in modelpool.model_names:
+                retrieved_task_vector = state_dict_mul(
+                    retrieved_task_vectors[model_name], self.scaling_factor
+                )
+                retrieved_state_dict = state_dict_add(
+                    pretrained_model.state_dict(keep_vars=True), retrieved_task_vector
+                )
+                retrieved_model = deepcopy(pretrained_model)
+                # FIXME: for 'all' mode
+                for k, v in retrieved_state_dict.items():
+                    if v.shape[0] == 1:
+                        retrieved_state_dict[k] = v.squeeze(0)
+                retrieved_model.load_state_dict(retrieved_state_dict)
+                models[model_name] = retrieved_model
+
+                if self.debug >= 1:
+                    with self.profile("metadata"):
+                        model = modelpool.load_model(model_name)
+                        if torch.cuda.is_available():
+                            retrieved_state_dict = {
+                                k: v.cuda() for k, v in retrieved_state_dict.items()
+                            }
+                            retrieved_task_vectors[model_name] = {
+                                k: v.cuda()
+                                for k, v in retrieved_task_vectors[model_name].items()
+                            }
+                            task_vectors[model_name] = {
+                                k: v.cuda() for k, v in task_vectors[model_name].items()
+                            }
+                            model_state_dict = {
+                                k: v.cuda()
+                                for k, v in model.state_dict(keep_vars=True).items()
+                            }
+                        # target_layers = metadata['target_layers']
+                        metadata["task_vector_retrieval_similarity"][model_name] = (
+                            compare_models(
+                                retrieved_task_vectors[model_name],
+                                task_vectors[model_name],
+                            )
+                        )
+                        metadata["task_vector_svd_subspace_similarities"][
+                            model_name
+                        ] = self._compute_svd_subspace_similarities(
+                            task_vectors[model_name], retrieved_task_vectors[model_name]
+                        )
+                        # overall retrieval performance
+                        metadata["model_retrieval_similarity"][model_name] = (
+                            compare_models(retrieved_state_dict, model_state_dict)
+                        )
+                        metadata["model_svd_subspace_similarities"][model_name] = (
+                            self._compute_svd_subspace_similarities(
+                                model_state_dict, retrieved_state_dict
+                            )
+                        )
+                        # delete the cuda tensors
+                        del (
+                            retrieved_state_dict,
+                            retrieved_task_vectors[model_name],
+                            task_vectors[model_name],
+                            model_state_dict,
+                        )
+
+        with self.profile("metadata"):
+            if self.debug >= 0:
+                (
+                    metadata["trainable_param_count_pretrained_model"],
+                    metadata["active_param_count_pretrained_model"],
+                ) = count_parameters(pretrained_model)
+                (
+                    metadata["trainable_param_count_retrieved_model"],
+                    metadata["active_param_count_retrieved_model"],
+                ) = count_parameters(models[modelpool.model_names[0]])
+                metadata["nonzero_parameter_count"] += metadata[
+                    "active_param_count_pretrained_model"
+                ]
+                metadata["total_gb_retrieved"] += metadata["total_gb_original"]
+                print(
+                    f"Total storage (Gbs) for retrieval and original: {metadata['total_gb_retrieved']} | {metadata['total_gb_original']}"
+                )
+
+        if self.model_path is not None:
+            os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
+            torch.save(models, self.model_path)
+
+        self.print_profile_summary()
+        return {"models": models, "metadata": metadata}
+
+
+class SuperposedTaskArithmeticLoRAAlgorithm(
+    SuperposedAlgorithmBase,
+):
+    _config_mapping = SuperposedAlgorithmBase._config_mapping | {
+        "scaling_factor": "scaling_factor",
+        "model_path": "model_path",
+    }
+
+    def __init__(
+        self,
+        scaling_factor: float,
+        model_path: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.scaling_factor = scaling_factor
+        self.model_path = model_path
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(models=modelpool)
+
+        log.info("Compressing models using superposed task arithmetic.")
+        task_vector = None
+        with self.profile("load model"):
+            pretrained_model = modelpool.load_model("_pretrained_")
+
+        # Calculate the task vector superposition
+        loras = {}
+        models = {}
+        for model_name in modelpool.model_names:
+            with self.profile("load model"):
+                model = modelpool.load_model(model_name)
+            for layer_name, layer in model.items():
+                if self.verbose >= 1:
+                    log.info(f"{layer_name} | {layer.shape}")
+            # task_vector = state_dict_sub(
+            #     model.state_dict(keep_vars=True),
+            #     pretrained_model.state_dict(keep_vars=True),
+            # )
+            loras[model_name] = model
+
+        with self.profile("compress and retrieve"):
+            retrieved_loras, metadata = self._compress_and_retrieve(
+                deepcopy(loras), mode="superposed_task_arithmetic"
+            )
+        with self.profile("retrieve models"):
+            for model_name in modelpool.model_names:
+                retrieved_lora = retrieved_loras[model_name]
+                # retrieved_lora = state_dict_mul(retrieved_loras[model_name], self.config.scaling_factor)
+                # retrieved_state_dict = state_dict_add(pretrained_model.state_dict(keep_vars=True), retrieved_lora)
+                retrieved_model = deepcopy(pretrained_model)
+                sd = retrieved_model.state_dict(keep_vars=True)
+                # for layer_name, layer in sd.items():
+                #     print(layer_name)
+                # manually merge the lora back
+                lora_weights = {}
+                lora_weights_ready_to_merge = OrderedDict()
+                for layer_name, layer in retrieved_lora.items():
+                    parts = layer_name.split(".")
+                    # print(parts)
+                    base_name = ".".join(parts[2:-2] + [parts[-1]])
+                    if base_name not in lora_weights:
+                        lora_weights[base_name] = []
+                    lora_weights[base_name].append(layer)
+                for base_name, layers in lora_weights.items():
+                    lora_weight = layers[-1] @ layers[0]
+                    # sd[base_name] += lora_weight
+                    lora_weights_ready_to_merge[base_name] = lora_weight
+
+                retrieved_lora_ready = state_dict_mul(
+                    lora_weights_ready_to_merge, self.config.scaling_factor
+                )
+                for layer_name, layer in retrieved_lora_ready.items():
+                    sd[layer_name] += layer
+                retrieved_model.load_state_dict(sd)
+                models[model_name] = retrieved_model
+
+                # # FIXME: for 'all' mode
+                # for k, v in retrieved_state_dict.items():
+                #     if v.shape[0] == 1:
+                #         retrieved_state_dict[k] = v.squeeze(0)
+                # retrieved_model.load_state_dict(sd)
+                # models[model_name] = retrieved_model
+
+                if self.debug >= 1:
+                    with self.profile("metadata"):
+                        model = modelpool.load_model(model_name)
+                        if torch.cuda.is_available():
+                            retrieved_state_dict = {
+                                k: v.cuda() for k, v in retrieved_state_dict.items()
+                            }
+                            retrieved_loras[model_name] = {
+                                k: v.cuda()
+                                for k, v in retrieved_loras[model_name].items()
+                            }
+                            loras[model_name] = {
+                                k: v.cuda() for k, v in loras[model_name].items()
+                            }
+                            model_state_dict = {
+                                k: v.cuda()
+                                for k, v in model.state_dict(keep_vars=True).items()
+                            }
+                        # focus on the superposition retrieval performance on the target layers
+                        target_layers = metadata["target_layers"]
+                        metadata["lora_retrieval_similarity"][model_name] = (
+                            compare_models(
+                                retrieved_loras[model_name],
+                                loras[model_name],
+                                target_layers,
+                            )
+                        )
+                        metadata["lora_svd_subspace_similarities"][model_name] = (
+                            self._compute_svd_subspace_similarities(
+                                loras[model_name],
+                                retrieved_loras[model_name],
+                                target_layers,
+                            )
+                        )
+                        # overall retrieval performance
+                        metadata["model_retrieval_similarity"][model_name] = (
+                            compare_models(retrieved_state_dict, model_state_dict)
+                        )
+                        metadata["model_svd_subspace_similarities"][model_name] = (
+                            self._compute_svd_subspace_similarities(
+                                model_state_dict, retrieved_state_dict
+                            )
+                        )
+                        # delete the cuda tensors
+                        del (
+                            retrieved_state_dict,
+                            retrieved_loras[model_name],
+                            loras[model_name],
+                            model_state_dict,
+                        )
+
+        with self.profile("metadata"):
+            if self.debug >= 0:
+                (
+                    metadata["trainable_param_count_pretrained_model"],
+                    metadata["active_param_count_pretrained_model"],
+                ) = count_parameters(pretrained_model)
+                (
+                    metadata["trainable_param_count_retrieved_model"],
+                    metadata["active_param_count_retrieved_model"],
+                ) = count_parameters(models[modelpool.model_names[0]])
+                metadata["nonzero_parameter_count"] += metadata[
+                    "active_param_count_pretrained_model"
+                ]
+                metadata["total_gb_retrieved"] += metadata["total_gb_original"]
+                print(
+                    f"Total storage (Gbs) for retrieval and original: {metadata['total_gb_retrieved']} | {metadata['total_gb_original']}"
+                )
+
+        if self.model_path is not None:
+            os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
+            torch.save(models, self.model_path)
+
+        self.print_profile_summary()
+        return {"models": models, "metadata": metadata}

fusion_bench/method/sparselo/sparselo.py

@@ -32,6 +32,7 @@ from fusion_bench.models.modeling_losparse_llama.losparse_linear import LoSparse
 from fusion_bench.models.modeling_losparse_llama.utils import convert_to_losparse_llama
 from fusion_bench.utils import cache_to_disk, print_parameters, timeit_context
 from fusion_bench.utils.devices import get_device
+from fusion_bench.utils.dtype import get_dtype
 
 log = logging.getLogger(__name__)
 
@@ -141,6 +142,7 @@ class SparseLoForLlama(BaseAlgorithm, SimpleProfilerMixin):
 
     @override
     def run(self, modelpool: CausalLMPool):
+        self.modelpool = modelpool
         if self.seed is not None:
             L.seed_everything(self.seed)
 
@@ -691,12 +693,16 @@ class IterativeSparseLoForLlama(SparseLoForLlama):
         "num_iterations": "num_iterations",
     }
 
-    def __init__(self, num_iterations: int, **kwargs):
+    def __init__(
+        self, num_iterations: int, use_reference_model: bool = False, **kwargs
+    ):
         super().__init__(**kwargs)
         self.num_iterations = num_iterations
+        self.use_reference_model = use_reference_model
 
     @override
     def run(self, modelpool):
+        self.modelpool = modelpool
         if self.seed is not None:
             L.seed_everything(self.seed)
 
@@ -802,13 +808,25 @@ class IterativeSparseLoForLlama(SparseLoForLlama):
     @torch.no_grad()
     def iterative_magnitude_prune_(self, model):
         layers: nn.ModuleList = model.model.layers
+        if self.use_reference_model:
+            reference_model = self.modelpool.load_model(
+                "reference_model", torch_dtype="float16"
+            )
+            reference_layers: nn.ModuleList = reference_model.model.layers
         for layer_idx, layer in tqdm(
             enumerate(layers), "Pruning Layers", total=len(layers), dynamic_ncols=True
         ):
             for name, linear in layer.named_modules():
                 if isinstance(linear, LoSparseLinear):
                     log.info(f"Magnitude Pruning {name}")
-                    W = linear.weight.data.clone()
+                    W = (
+                        linear.weight.data.clone()
+                        if not self.use_reference_model
+                        else reference_layers[layer_idx]
+                        .get_submodule(name)
+                        .weight.data.clone()
+                        .to(linear.weight.data.device)
+                    )
                     if self.prune_type == PruningType.UNSTRUCTURED:
                         unstructured_magnitude_prune_(
                             linear.weight.data,

fusion_bench/method/tall_mask/__init__.py

@@ -0,0 +1 @@
+from .task_arithmetic import TallMaskTaskArithmeticAlgorithm

fusion_bench/method/tall_mask/task_arithmetic.py

@@ -0,0 +1,133 @@
+"""
+Modified from https://github.com/Zhou-Hangyu/randes/tree/main/benchmark/fusion_bench
+"""
+
+import logging
+from collections import OrderedDict
+from copy import deepcopy
+
+import torch
+
+from fusion_bench import BaseAlgorithm
+from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils.state_dict_arithmetic import (
+    state_dict_add,
+    state_dict_binary_mask,
+    state_dict_diff_abs,
+    state_dict_hadmard_product,
+    state_dict_mul,
+    state_dict_sub,
+    state_dict_sum,
+)
+
+log = logging.getLogger(__name__)
+
+
+def generate_task_masks(
+    multi_task_vector: OrderedDict,
+    ft_task_vector: OrderedDict,
+    pretrained_task_vector: OrderedDict,
+    tall_mask_lambda: float = 1.0,
+) -> OrderedDict:
+    """Adopted from https://github.com/nik-dim/tall_masks/tree/master.
+    Generate task-specific TALL masks
+    TALL masks are generated as: mask_t = |theta_0 - theta_t| > |theta_mt - theta_t| * lambda
+
+    Args:
+        multi_task_vector: multi-task vector
+        ft_task_vector: individual theta_t (fine-tuned weights)
+        pretrained_task_vector: theta_0 (pre-trained weight)
+        tall_mask_lambda: hyper-parameter lambda for generating TALL masks
+    Returns:
+        final_mask: generated TALL masks with the given lambda
+    """
+
+    print(f"Generating TALL masks.")
+
+    # generate masks by comparing the l1 distance between |theta_0 - theta_t| and |theta_mt - theta_t|
+    diff_pt_ft = state_dict_diff_abs(pretrained_task_vector, ft_task_vector)
+    diff_multi_ft = state_dict_diff_abs(multi_task_vector, ft_task_vector)
+    # compare the l1 distance, scaled with hyper-parameter lambda
+    final_mask = state_dict_binary_mask(
+        diff_pt_ft,
+        state_dict_mul(diff_multi_ft, tall_mask_lambda),
+    )
+    for key, value in final_mask.items():
+        final_mask[key] = value.float()
+    return final_mask
+
+
+class TallMaskTaskArithmeticAlgorithm(
+    BaseAlgorithm,
+    SimpleProfilerMixin,
+):
+    _config_mapping = BaseAlgorithm._config_mapping | {
+        "tall_mask_lambda": "tall_mask_lambda",
+        "debug": "debug",
+        "verbose": "verbose",
+    }
+
+    def __init__(
+        self,
+        tall_mask_lambda: float,
+        debug: int = 0,
+        verbose: int = 0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tall_mask_lambda = tall_mask_lambda
+        self.debug = debug
+        self.verbose = verbose
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(models=modelpool)
+
+        log.info("Compressing models using tall mask task arithmetic.")
+        task_vector = None
+        with self.profile("load model"):
+            pretrained_model = modelpool.load_model("_pretrained_")
+
+        task_vectors = {}
+        models = {}
+        for model_name in modelpool.model_names:
+            with self.profile("load model"):
+                model = modelpool.load_model(model_name)
+            for layer_name, layer in model.state_dict(keep_vars=True).items():
+                if self.verbose >= 1:
+                    log.info(f"{layer_name} | {layer.shape}")
+            task_vector = state_dict_sub(
+                model.state_dict(keep_vars=True),
+                pretrained_model.state_dict(keep_vars=True),
+            )
+            task_vectors[model_name] = task_vector
+
+        multi_task_vector = state_dict_sum(list(task_vectors.values()))
+
+        tall_masks = {model: {} for model in modelpool.model_names}
+
+        for model_name in modelpool.model_names:
+            tall_mask = generate_task_masks(
+                multi_task_vector,
+                task_vectors[model_name],
+                pretrained_model.state_dict(keep_vars=True),
+                tall_mask_lambda=self.tall_mask_lambda,
+            )
+            tall_masks[model_name] = tall_mask
+
+        with self.profile("compress and retrieve"):
+            for model_name in modelpool.model_names:
+                retrieved_task_vector = state_dict_hadmard_product(
+                    tall_masks[model_name], multi_task_vector
+                )
+                retrieved_state_dict = state_dict_add(
+                    pretrained_model.state_dict(keep_vars=True), retrieved_task_vector
+                )
+                retrieved_model = deepcopy(pretrained_model)
+                retrieved_model.load_state_dict(retrieved_state_dict)
+                models[model_name] = retrieved_model
+
+        self.print_profile_summary()
+        return {"models": models, "metadata": None}

fusion_bench/modelpool/lazy_state_dict_pool.py

@@ -0,0 +1,15 @@
+from fusion_bench import BaseModelPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.lazy_state_dict import LazyStateDict
+
+
+class LazyStateDictPool(BaseModelPool):
+    def load_model(self, model_name_or_config: str, *args, **kwargs) -> LazyStateDict:
+        if model_name_or_config in self._models:
+            checkpoint_config = self._models[model_name_or_config]
+        else:
+            checkpoint_config = model_name_or_config
+        if isinstance(checkpoint_config, str):
+            return LazyStateDict(checkpoint_config, *args, **kwargs)
+        else:
+            return instantiate(checkpoint_config, *args, **kwargs)

fusion_bench/models/modeling_deepseek_v2/__init__.py

@@ -0,0 +1,15 @@
+"""
+This is a direct copy of the DeepSeek-V2-Lite model from HuggingFace https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite/tree/main
+"""
+
+from .configuration_deepseek import DeepseekV2Config
+from .modeling_deepseek import (
+    DeepseekV2ForCausalLM,
+    DeepseekV2ForSequenceClassification,
+    DeepseekV2MLP,
+    DeepseekV2Model,
+    DeepseekV2MoE,
+    DeepseekV2DecoderLayer,
+)
+from .modeling_deepseek import MoEGate as DeepseekV2MoEGate
+from .tokenization_deepseek_fast import DeepseekTokenizerFast