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

fusion_bench/method/__init__.py

@@ -111,6 +111,12 @@ _import_structure = {
         "SparseLoForLlama",
         "PCPSparseLoForLlama",
     ],
+    # MoE expert pruning
+    "expert_sparsity": [
+        "DynamicSkippingPruningForMixtral",
+        "LayerWisePruningForMixtral",
+        "ProgressivePruningForMixtral",
+    ],
 }
 
 
@@ -142,6 +148,11 @@ if TYPE_CHECKING:
         SimpleEnsembleAlgorithm,
         WeightedEnsembleAlgorithm,
     )
+    from .expert_sparsity import (
+        DynamicSkippingPruningForMixtral,
+        LayerWisePruningForMixtral,
+        ProgressivePruningForMixtral,
+    )
     from .fisher_merging import FisherMergingForCLIPVisionModel
     from .fw_merging import FrankWolfeHardAlgorithm, FrankWolfeSoftAlgorithm
     from .gossip import (

fusion_bench/method/adamerging/flan_t5_layer_wise_adamerging.py

@@ -29,7 +29,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .min_norm_solvers import MinNormSolver

fusion_bench/method/adamerging/gpt2_layer_wise_adamerging.py

@@ -29,7 +29,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .min_norm_solvers import MinNormSolver

fusion_bench/method/base_algorithm.py

@@ -19,6 +19,7 @@ class BaseAlgorithm(BaseYAMLSerializableModel):
     """
 
     _program = None
+    _config_key = "method"
 
     @abstractmethod
     def run(self, modelpool: BaseModelPool):

fusion_bench/method/dawe/dawe_for_clip.py

@@ -23,7 +23,7 @@ from fusion_bench.mixins import CLIPClassificationMixin
 from fusion_bench.modelpool import CLIPVisionModelPool
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .warppers.dawe_model import DataAdaptiveWeightEnsemblingCLIPVisionModel
 

fusion_bench/method/depth_upscaling/depth_upscaling_for_llama.py

@@ -1,9 +1,10 @@
 import os
 from typing import Optional
 
+from transformers import PreTrainedModel
 from typing_extensions import override
 
-from fusion_bench.modelpool.causal_lm.causal_lm import CausalLM, CausalLMPool
+from fusion_bench.modelpool.causal_lm.causal_lm import CausalLMPool
 from fusion_bench.utils import timeit_context
 
 from .depth_upscaling import DepthUpscalingAlgorithm
@@ -46,7 +47,7 @@ class DepthUpscalingForLlama(DepthUpscalingAlgorithm):
         if self.model_save_path is not None:
             tokenizer = modelpool.load_tokenizer()
 
-        model: CausalLM = modelpool.load_pretrained_or_first_model()
+        model: PreTrainedModel = modelpool.load_pretrained_or_first_model()
         model.model.layers = super().run(model.model.layers)
         model.config.num_hidden_layers = len(model.model.layers)
 

fusion_bench/method/expert_sparsity/__init__.py

@@ -0,0 +1,10 @@
+"""
+Original repo: https://github.com/Lucky-Lance/Expert_Sparsity
+
+Reference:
+    Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models.
+    ACL 2024.
+    http://arxiv.org/abs/2402.14800
+"""
+
+from .mixtral import *

fusion_bench/method/expert_sparsity/mixtral/__init__.py

@@ -0,0 +1,23 @@
+R"""
+```bash
+fusion_bench \
+    modelpool=CausalLMPool/mixtral-8x7b \
+    ...
+```
+
+if use flash attention 2, pass the following to the command line:
+
+```bash
++modelpool.models._pretrained_.attn_implementation=flash_attention_2
+```
+"""
+
+from .dynamic_skipping import DynamicSkippingPruningForMixtral
+from .layer_wise_pruning import LayerWisePruningForMixtral
+from .progressive_pruning import ProgressivePruningForMixtral
+
+__all__ = [
+    "DynamicSkippingPruningForMixtral",
+    "LayerWisePruningForMixtral",
+    "ProgressivePruningForMixtral",
+]

fusion_bench/method/expert_sparsity/mixtral/dynamic_skipping.py

@@ -0,0 +1,175 @@
+R"""
+Example:
+
+```bash
+fusion_bench \
+    fabric.loggers.name="mixtral_8x7b_expert_pruning/dynamic_skipping" \
+    method=expert_sparsity/mixtral \
+    method._target_=fusion_bench.method.DynamicSkippingPruningForMixtral \
+    modelpool=CausalLMPool/mixtral-8x7b
+```
+"""
+
+import logging
+import os
+
+import lightning as L
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import MixtralForCausalLM
+from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM
+
+import fusion_bench as fb
+from fusion_bench.method.expert_sparsity.utils.calibration_data import (
+    build_calib_loader,
+)
+from fusion_bench.models.expert_sparsity.mixtral.wrapper import (
+    PrunableMixtralSparseMoeBlockWrapper,
+)
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+
+
+def dynamic_skipping(
+    model: MixtralForCausalLM,
+    calib_loader: DataLoader,
+    batch_size: int,
+):
+    assert isinstance(
+        model, MixtralForCausalLM
+    ), "Currently only `Mixtral` is supported"
+
+    for l, layer in enumerate(model.model.layers):
+        layer.block_sparse_moe = PrunableMixtralSparseMoeBlockWrapper(
+            layer.block_sparse_moe
+        )
+        layer.block_sparse_moe.cache_logits = True
+        layer.block_sparse_moe.cache_X = True
+        layer.block_sparse_moe.cache_Z = True
+
+    with torch.inference_mode():
+        for i, batch in enumerate(
+            tqdm(calib_loader, desc="Model forwarding on sample set...")
+        ):
+            model_inputs = model.prepare_inputs_for_generation(**batch)
+            outputs = model(**model_inputs)
+            assert outputs is not None
+
+    res_median = {}
+    res_mean = {}
+
+    for layer_idx in range(len(model.model.layers)):
+        b = model.model.layers[layer_idx].block_sparse_moe
+        b.cache_space.prepare_for_loader()
+        dataloader = torch.utils.data.DataLoader(
+            b.cache_space,
+            batch_size=batch_size,
+            shuffle=True,
+        )
+        logger.info(len(dataloader))
+
+        ana_list = []
+        for i, (router_logits, X, Z) in enumerate(dataloader):
+            routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float).view(
+                -1, b.model.num_experts
+            )
+            for j in range(len(routing_weights)):
+                sorted_weights, sort_indices = torch.sort(
+                    routing_weights[j], descending=True
+                )
+                ana_list.append(float(sorted_weights[1] / sorted_weights[0]))
+
+        median = np.median(ana_list)
+        mean = np.mean(ana_list)
+        logger.info(f"layer {layer_idx} | mean: {mean}, median: {median}")
+        res_median[str(layer_idx)] = median
+        res_mean[str(layer_idx)] = mean
+
+    for l, layer in enumerate(model.model.layers):
+        layer.block_sparse_moe = layer.block_sparse_moe.model
+
+    model.config.betas = res_median
+    return model, (res_median, res_mean)
+
+
+class DynamicSkippingPruningForMixtral(
+    fb.BaseAlgorithm,
+    fb.mixins.LightningFabricMixin,
+    fb.mixins.SimpleProfilerMixin,
+):
+    modelpool: fb.modelpool.CausalLMPool
+
+    def __init__(
+        self,
+        calib_set: str,
+        max_block_size: int,
+        n_blocks_for_stat: int,
+        batch_size: int,
+        num_workers: int,
+        num_preserved_experts: int,
+        seed: int = 42,
+        model_save_path: str = R"{log_dir}/pruned_model",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_save_path = model_save_path
+        self.calib_set = calib_set
+        self.max_block_size = max_block_size
+        self.n_blocks_for_stat = n_blocks_for_stat
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.seed = seed
+        self.num_preserved_experts = num_preserved_experts
+
+    def run(self, modelpool: fb.modelpool.CausalLMPool):
+        """
+        Args:
+            modelpool (fb.modelpool.CausalLMPool): The model pool to run the algorithm on.
+                Example Config: config/modelpool/CausalLMPool/mixtral-8x7b.yaml
+        """
+        self.modelpool = modelpool
+        # set random seed
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        # parse model_save_path
+        self.model_save_path = self.model_save_path.format(log_dir=self.log_dir)
+
+        with self.profile("load model"):
+            model = modelpool.load_pretrained_or_first_model()
+            tokenizer = modelpool.load_tokenizer()
+
+        # Load the calibration data
+        with self.profile("load calibration data"):
+            calib_loader = build_calib_loader(
+                self.calib_set,
+                tokenizer=tokenizer,
+                max_block_size=self.max_block_size,
+                n_blocks_for_stat=self.n_blocks_for_stat,
+                batch_size=self.batch_size,
+                num_workers=self.num_workers,
+                seed=self.seed,
+            )
+
+        with self.profile("prune model"):
+            model, info = dynamic_skipping(
+                model,
+                calib_loader,
+                batch_size=self.batch_size,
+            )
+
+        if self.model_save_path is not None:
+            with self.profile("save model"):
+                modelpool.save_model(
+                    model,
+                    path=self.model_save_path,
+                    tokenizer=tokenizer,
+                )
+                torch.save(info, os.path.join(self.log_dir, "pruning_info.pt"))
+
+        self.print_profile_summary()
+        return model

fusion_bench/method/expert_sparsity/mixtral/layer_wise_pruning.py

@@ -0,0 +1,159 @@
+R"""
+Example:
+
+```bash
+fusion_bench \
+    fabric.loggers.name="mixtral_8x7b_expert_pruning/layer_wise_pruning" \
+    method=expert_sparsity/mixtral \
+    method._target_=fusion_bench.method.LayerWisePruningForMixtral \
+    modelpool=CausalLMPool/mixtral-8x7b
+```
+"""
+
+import logging
+import os
+from typing import cast
+
+import lightning as L
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import MixtralForCausalLM
+from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
+
+import fusion_bench as fb
+from fusion_bench.method.expert_sparsity.utils.calibration_data import (
+    build_calib_loader,
+)
+from fusion_bench.models.expert_sparsity.mixtral import (
+    PrunableMixtralSparseMoeBlockWrapper,
+)
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+
+
+def layerwise_pruning(
+    model: MixtralForCausalLM,
+    calib_loader: DataLoader,
+    r: int,
+):
+    assert isinstance(
+        model, MixtralForCausalLM
+    ), "Currently only `Mixtral` is supported"
+
+    for l, layer in enumerate(model.model.layers):
+        layer = cast(MixtralDecoderLayer, layer)
+        layer.block_sparse_moe = PrunableMixtralSparseMoeBlockWrapper(
+            layer.block_sparse_moe, r=r
+        )
+        layer.block_sparse_moe.cache_X = True
+        layer.block_sparse_moe.cache_Z = True
+
+    with torch.inference_mode():
+        for i, batch in enumerate(
+            tqdm(calib_loader, desc="Model forwarding on sample set...")
+        ):
+            model_inputs = model.prepare_inputs_for_generation(**batch)
+            outputs = model(**model_inputs)
+            assert outputs is not None
+
+    global_loss_history = dict()
+    for l, layer in tqdm(
+        list(enumerate(model.model.layers)), desc="Enumerating loss on sample set..."
+    ):
+        layer = cast(MixtralDecoderLayer, layer)
+        b: PrunableMixtralSparseMoeBlockWrapper = layer.block_sparse_moe
+        if not hasattr(b, "cache_space"):
+            continue
+        loss_history = b.enumerate()
+        global_loss_history[l] = loss_history
+        b.prune()
+
+    logger.info("Merging & saving...")
+    for l, layer in enumerate(model.model.layers):
+        layer.block_sparse_moe = layer.block_sparse_moe.model
+
+    model.num_experts = r
+    model.config.num_local_experts = r
+
+    return model, (global_loss_history,)
+
+
+class LayerWisePruningForMixtral(
+    fb.BaseAlgorithm,
+    fb.mixins.LightningFabricMixin,
+    fb.mixins.SimpleProfilerMixin,
+):
+    modelpool: fb.modelpool.CausalLMPool
+
+    def __init__(
+        self,
+        calib_set: str,
+        max_block_size: int,
+        n_blocks_for_stat: int,
+        batch_size: int,
+        num_workers: int,
+        num_preserved_experts: int,
+        seed: int = 42,
+        model_save_path: str = R"{log_dir}/pruned_model",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_save_path = model_save_path
+        self.calib_set = calib_set
+        self.max_block_size = max_block_size
+        self.n_blocks_for_stat = n_blocks_for_stat
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.seed = seed
+        self.num_preserved_experts = num_preserved_experts
+
+    def run(self, modelpool: fb.modelpool.CausalLMPool):
+        """
+        Args:
+            modelpool (fb.modelpool.CausalLMPool): The model pool to run the algorithm on.
+                Example Config: config/modelpool/CausalLMPool/mixtral-8x7b.yaml
+        """
+        self.modelpool = modelpool
+        # set random seed
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        # parse model_save_path
+        self.model_save_path = self.model_save_path.format(log_dir=self.log_dir)
+
+        with self.profile("load model"):
+            model = modelpool.load_pretrained_or_first_model()
+            tokenizer = modelpool.load_tokenizer()
+
+        # Load the calibration data
+        with self.profile("load calibration data"):
+            calib_loader = build_calib_loader(
+                self.calib_set,
+                tokenizer=tokenizer,
+                max_block_size=self.max_block_size,
+                n_blocks_for_stat=self.n_blocks_for_stat,
+                batch_size=self.batch_size,
+                num_workers=self.num_workers,
+                seed=self.seed,
+            )
+
+        with self.profile("prune model"):
+            model, info = layerwise_pruning(
+                model,
+                calib_loader,
+                r=self.num_preserved_experts,
+            )
+
+        if self.model_save_path is not None:
+            with self.profile("save model"):
+                modelpool.save_model(
+                    model,
+                    path=self.model_save_path,
+                    tokenizer=tokenizer,
+                )
+                torch.save(info, os.path.join(self.log_dir, "pruning_info.pt"))
+
+        self.print_profile_summary()
+        return model

fusion_bench/method/expert_sparsity/mixtral/progressive_pruning.py

@@ -0,0 +1,173 @@
+R"""
+Example:
+
+```bash
+fusion_bench \
+    fabric.loggers.name="mixtral_8x7b_expert_pruning/progressive_pruning" \
+    method=expert_sparsity/mixtral \
+    method._target_=fusion_bench.method.ProgressivePruningForMixtral \
+    modelpool=CausalLMPool/mixtral-8x7b
+```
+"""
+
+import logging
+import os
+
+import lightning as L
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import MixtralForCausalLM
+
+import fusion_bench as fb
+from fusion_bench.method.expert_sparsity.utils.calibration_data import (
+    build_calib_loader,
+)
+from fusion_bench.models.expert_sparsity.mixtral import (
+    PrunableMixtralSparseMoeBlockWrapper,
+)
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+logger = logging.getLogger(__name__)
+
+
+def progressive_pruning(
+    model: MixtralForCausalLM,
+    calib_loader: DataLoader,
+    r: int,
+):
+    assert isinstance(
+        model, MixtralForCausalLM
+    ), "Currently only `Mixtral` is supported"
+
+    for l, layer in enumerate(model.model.layers):
+        layer.block_sparse_moe = PrunableMixtralSparseMoeBlockWrapper(
+            layer.block_sparse_moe, r=r
+        )
+        layer.block_sparse_moe.cache_Z = True
+
+    with torch.inference_mode():
+        for i, batch in enumerate(
+            tqdm(calib_loader, desc="Computing Z activations on sample set...")
+        ):
+            model_inputs = model.prepare_inputs_for_generation(**batch)
+            outputs = model(**model_inputs)
+            assert outputs is not None
+
+    del model_inputs
+    del outputs
+    torch.cuda.empty_cache()
+
+    for l, layer in enumerate(model.model.layers):
+        layer.block_sparse_moe.cache_Z = False
+
+    # Drop
+    global_loss_history = dict()
+
+    for l, layer in tqdm(
+        list(enumerate(model.model.layers)), desc="Dropping layers..."
+    ):
+        b = layer.block_sparse_moe
+
+        b.cache_X = True
+        with torch.inference_mode():
+            for i, batch in enumerate(calib_loader):
+                model_inputs = model.prepare_inputs_for_generation(**batch)
+                outputs = model(**model_inputs)
+                assert outputs is not None
+
+        del model_inputs
+        del outputs
+        torch.cuda.empty_cache()
+        b.cache_X = False
+
+        loss_history = b.enumerate()
+        global_loss_history[l] = loss_history
+
+        b.prune()
+        layer.block_sparse_moe = b.model
+
+    # Prune & save
+    model.num_experts = r
+    model.config.num_local_experts = r
+
+    return model, (global_loss_history,)
+
+
+class ProgressivePruningForMixtral(
+    fb.BaseAlgorithm,
+    fb.mixins.LightningFabricMixin,
+    fb.mixins.SimpleProfilerMixin,
+):
+    modelpool: fb.modelpool.CausalLMPool
+
+    def __init__(
+        self,
+        calib_set: str,
+        max_block_size: int,
+        n_blocks_for_stat: int,
+        batch_size: int,
+        num_workers: int,
+        num_preserved_experts: int,
+        seed: int = 42,
+        model_save_path: str = R"{log_dir}/pruned_model",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_save_path = model_save_path
+        self.calib_set = calib_set
+        self.max_block_size = max_block_size
+        self.n_blocks_for_stat = n_blocks_for_stat
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.seed = seed
+        self.num_preserved_experts = num_preserved_experts
+
+    def run(self, modelpool: fb.modelpool.CausalLMPool):
+        """
+        Args:
+            modelpool (fb.modelpool.CausalLMPool): The model pool to run the algorithm on.
+                Example Config: config/modelpool/CausalLMPool/mixtral-8x7b.yaml
+        """
+        self.modelpool = modelpool
+        # set random seed
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        # parse model_save_path
+        self.model_save_path = self.model_save_path.format(log_dir=self.log_dir)
+
+        with self.profile("load model"):
+            model = modelpool.load_pretrained_or_first_model()
+            tokenizer = modelpool.load_tokenizer()
+
+        # Load the calibration data
+        with self.profile("load calibration data"):
+            calib_loader = build_calib_loader(
+                self.calib_set,
+                tokenizer=tokenizer,
+                max_block_size=self.max_block_size,
+                n_blocks_for_stat=self.n_blocks_for_stat,
+                batch_size=self.batch_size,
+                num_workers=self.num_workers,
+                seed=self.seed,
+            )
+
+        with self.profile("prune model"):
+            model, info = progressive_pruning(
+                model,
+                calib_loader,
+                r=self.num_preserved_experts,
+            )
+
+        if self.model_save_path is not None:
+            with self.profile("save model"):
+                modelpool.save_model(
+                    model,
+                    path=self.model_save_path,
+                    tokenizer=tokenizer,
+                )
+                torch.save(info, os.path.join(self.log_dir, "pruning_info.pt"))
+
+        self.print_profile_summary()
+        return model