Code for the paper "Merge, Then Compress: Demystify Efficient SMoE with Hints from Its Routing Policy"
import os

os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'

mcsmoe/__init__.py
UNITES-Lab/MC-SMoE
from .flops_counter import calculate_flops
from .utils import bytes_to_string
from .utils import flops_to_string
from .utils import generate_transformer_input
from .utils import macs_to_string
from .utils import number_to_string
from .utils import params_to_string

mcsmoe/calflops/__init__.py
UNITES-Lab/MC-SMoE
from functools import partial

from .pytorch_ops import MODULE_HOOK_MAPPING
from .pytorch_ops import _patch_functionals
from .pytorch_ops import _patch_tensor_methods
from .pytorch_ops import _reload_functionals
from .pytorch_ops import _reload_tensor_methods
from .utils import flops_to_string
from .utils import get_module_flops
from .utils import get_module_macs
from .utils import macs_to_string
from .utils import number_to_string
from .utils import params_to_string

DEFAULT_PRECISION = 2
module_flop_count = []
module_mac_count = []
old_functions = {}


class CalFlopsPipline(object):
    """The Pipline of calculating FLOPs（number of estimated floating-point operations） and Parameters of each module in a PyTorch model.
    The pipline is calculating the forward(and alson include back propagation) pass of a PyTorch model and prints the model graph with the calculated static attached to each module. 
    It can easily get only final resulst of FLOPs about model, and also can be showed how flops and parameters are spent in the model and which modules or layers could be the bottleneck in detailed.
    """

    def __init__(self, model, include_backPropagation, compute_bp_factor, is_sparse):
        """Init Pipline of Calculating the FLOPs about model.

        Args:
            model (pytorch model): The model must be a pytorh model now.
            compute_fwd_factor (float): Defaults to 2.0. According to https://epochai.org/blog/backward-forward-FLOP-ratio
        """

        self.model = model
        self.include_backPropagation = include_backPropagation  # Whether the calculation results include model backpropagation
        self.compute_bp_factor = compute_bp_factor  # Backpropagation takes twice as much computation as forward propagation
        self.pipline_started = False  # The flag of calculating FLOPs pipline started
        self.func_patched = False  # The flag of wheather calculating functional are patched
        self.is_sparse = is_sparse  # Whether to exclude sparse matrix flops

    def start_flops_calculate(self, ignore_list=None):
        """Starts the pipline of calculating FLOPs.

        Extra attributes are added recursively to all the modules and the calculate torch.nn.functionals are monkey patched.

        Args:
            ignore_list (list, optional): the list of modules to ignore while Piplining. Defaults to None.
        """

        self.reset_flops_calculate()
        _patch_functionals(old_functions, module_flop_count, module_mac_count)
        _patch_tensor_methods(old_functions, module_flop_count, module_mac_count)

        def register_module_hooks(module, ignore_list):
            if ignore_list and type(module) in ignore_list:
                return

            # if computing the flops of a module directly
            if type(module) in MODULE_HOOK_MAPPING:
                if not hasattr(module, "__flops_handle__"):
                    module.__flops_handle__ = module.register_forward_hook(MODULE_HOOK_MAPPING[type(module)])
                return

            # if computing the flops of the functionals in a module
            def pre_hook(module, input):
                module_flop_count.append([])
                module_mac_count.append([])

            if not hasattr(module, "__pre_hook_handle__"):
                module.__pre_hook_handle__ = module.register_forward_pre_hook(pre_hook)

            def post_hook(module, input, output):
                if module_flop_count:
                    module.__flops__ += sum([elem[1] for elem in module_flop_count[-1]])
                    module_flop_count.pop()
                    module.__macs__ += sum([elem[1] for elem in module_mac_count[-1]])
                    module_mac_count.pop()

            if not hasattr(module, "__post_hook_handle__"):
                module.__post_hook_handle__ = module.register_forward_hook(post_hook)

        self.model.apply(partial(register_module_hooks, ignore_list=ignore_list))
        self.pipline_started = True
        self.func_patched = True

    def stop_flops_calculate(self):
        """Stop the pipline of calculating FLOPs.

        All torch.nn.functionals are restored to their originals.
        """
        if self.pipline_started and self.func_patched:
            _reload_functionals(old_functions)
            _reload_tensor_methods(old_functions)
            self.func_patched = False

        def remove_calculate_attrs(module):
            if hasattr(module, "__pre_hook_handle__"):
                module.__pre_hook_handle__.remove()
                del module.__pre_hook_handle__
            if hasattr(module, "__post_hook_handle__"):
                module.__post_hook_handle__.remove()
                del module.__post_hook_handle__
            if hasattr(module, "__flops_handle__"):
                module.__flops_handle__.remove()
                del module.__flops_handle__

        self.model.apply(remove_calculate_attrs)

    def reset_flops_calculate(self):
        """Resets the pipline of calculating FLOPs.

        Adds or resets the extra attributes, include flops、macs、params.
        """

        def add_or_reset_attrs(module):
            module.__flops__ = 0
            module.__macs__ = 0
            module.__params__ = sum(
                p.count_nonzero().item() for p in module.parameters() if p.requires_grad
            ) if self.is_sparse else sum(
                p.numel() for p in module.parameters() if p.requires_grad)
            # just calculate parameter need training.

        self.model.apply(add_or_reset_attrs)

    def end_flops_calculate(self):
        """Ends the pipline of calculating FLOPs.

        The added attributes and handles are removed recursively on all the modules.
        """
        if not self.pipline_started:
            return
        self.stop_flops_calculate()
        self.pipline_started = False

        def remove_calculate_attrs(module):
            if hasattr(module, "__flops__"):
                del module.__flops__
            if hasattr(module, "__macs__"):
                del module.__macs__
            if hasattr(module, "__params__"):
                del module.__params__

        self.model.apply(remove_calculate_attrs)

    def get_total_flops(self, as_string=False):
        """Returns the total flops of the model.

        Args:
            as_string (bool, optional): whether to output the flops as string. Defaults to False.

        Returns:
            The number of multiply-accumulate operations of the model forward pass.
        """
        total_flops = get_module_flops(self.model, is_sparse=self.is_sparse)
        return number_to_string(total_flops) if as_string else total_flops

    def get_total_macs(self, as_string=False):
        """Returns the total MACs of the model.

        Args:
            as_string (bool, optional): whether to output the flops as string. Defaults to False.

        Returns:
            The number of multiply-accumulate operations of the model forward pass.
        """
        total_macs = get_module_macs(self.model, is_sparse=self.is_sparse)
        return macs_to_string(total_macs) if as_string else total_macs

    def get_total_params(self, as_string=False):
        """Returns the total number of parameters stored per rank.

        Args:
            as_string (bool, optional): whether to output the parameters as string. Defaults to False.
            is_sparse (bool, optional): whether to output the parameters as string. Defaults to False.

        Returns:
            The total number of parameters stored per rank.
        """
        total_params = self.model.__params__
        return params_to_string(total_params) if as_string else total_params

    def print_return_model_pipline(self, units=None, precision=DEFAULT_PRECISION, print_detailed=True,
                                   print_results=True):
        """Prints the model graph with the calculateing pipline attached to each module.

        Args:
            module_depth (int, optional): The depth of the model to which to print the aggregated module information. When set to -1, it prints information from the top to the innermost modules (the maximum depth).
            top_modules (int, optional): Limits the aggregated profile output to the number of top modules specified.
            print_detailed (bool, optional): Whether to print the detailed model profile.
        """
        if not self.pipline_started:
            return

        total_flops = self.get_total_flops()
        total_macs = self.get_total_macs()
        total_params = self.get_total_params()

        self.flops = total_flops
        self.macs = total_macs
        self.params = total_params

        prints = []
        prints.append(
            "\n------------------------------------- Calculate Flops Results -------------------------------------")

        prints.append("Notations:\n" +
                      "number of parameters (Params), number of multiply-accumulate operations(MACs),\n" +
                      "number of floating-point operations (FLOPs), floating-point operations per second (FLOPS),\n" +
                      "fwd FLOPs (model forward propagation FLOPs), bwd FLOPs (model backward propagation FLOPs),\n" +
                      "default model backpropagation takes %.2f times as much computation as forward propagation.\n" % self.compute_bp_factor)

        line_fmt = '{:<70}  {:<8}'
        prints.append(line_fmt.format('Total Training Params: ', params_to_string(total_params)))

        prints.append(line_fmt.format('fwd MACs: ', macs_to_string(total_macs, units=units,
                                                                   precision=precision)))
        prints.append(line_fmt.format('fwd FLOPs: ', flops_to_string(total_flops, units=units,
                                                                     precision=precision)))
        prints.append(line_fmt.format('fwd+bwd MACs: ', macs_to_string(total_macs * (1 + self.compute_bp_factor),
                                                                       units=units, precision=precision)))
        prints.append(line_fmt.format('fwd+bwd FLOPs: ', flops_to_string(total_flops * (1 + self.compute_bp_factor),
                                                                         units=units, precision=precision)))

        def flops_repr(module):
            params = module.__params__
            flops = get_module_flops(module)
            macs = get_module_macs(module)
            items = [
                "{} = {:g}% Params".format(
                    params_to_string(params),
                    round(100 * params / total_params, precision) if total_params else 0),
                "{} = {:g}% MACs".format(macs_to_string(macs),
                                         round(100 * macs / total_macs, precision) if total_macs else 0),
                "{} = {:g}% FLOPs".format(flops_to_string(flops),
                                          round(100 * macs / total_flops, precision) if total_flops else 0),
            ]
            original_extra_repr = module.original_extra_repr()
            if original_extra_repr:
                items.append(original_extra_repr)
            return ", ".join(items)

        def add_extra_repr(module):
            flops_extra_repr = flops_repr.__get__(module)
            if module.extra_repr != flops_extra_repr:
                module.original_extra_repr = module.extra_repr
                module.extra_repr = flops_extra_repr
                assert module.extra_repr != module.original_extra_repr

        def del_extra_repr(module):
            if hasattr(module, "original_extra_repr"):
                module.extra_repr = module.original_extra_repr
                del module.original_extra_repr

        self.model.apply(add_extra_repr)

        if print_detailed:
            prints.append(
                "\n-------------------------------- Detailed Calculated FLOPs Results --------------------------------")
            prints.append(
                "Each module caculated is listed after its name in the following order: \nparams, percentage of total params, MACs, percentage of total MACs, FLOPS, percentage of total FLOPs"
            )
            prints.append(
                "\nNote: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). \n They are not counted as submodules in calflops and not to be printed out. However they make up the difference between a parent's MACs and the sum of its submodules'.\n2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n"
            )
            prints.append(str(self.model))

        self.model.apply(del_extra_repr)

        prints.append(
            "---------------------------------------------------------------------------------------------------")

        return_print = ""
        for line in prints:
            if print_results:
                print(line)
            return_print += line + "\n"
        return return_print

    def print_model_pipline(self, units=None, precision=DEFAULT_PRECISION, print_detailed=True):
        """Prints the model graph with the calculateing pipline attached to each module.

        Args:
            module_depth (int, optional): The depth of the model to which to print the aggregated module information. When set to -1, it prints information from the top to the innermost modules (the maximum depth).
            top_modules (int, optional): Limits the aggregated profile output to the number of top modules specified.
            print_detailed (bool, optional): Whether to print the detailed model profile.
        """
        if not self.pipline_started:
            return

        total_flops = self.get_total_flops()
        total_macs = self.get_total_macs()
        total_params = self.get_total_params()

        self.flops = total_flops
        self.macs = total_macs
        self.params = total_params

        print("\n------------------------------------- Calculate Flops Results -------------------------------------")

        print("Notations:\n"
              "number of parameters (Params), number of multiply-accumulate operations(MACs),\n"
              "number of floating-point operations (FLOPs), floating-point operations per second (FLOPS),\n"
              "fwd FLOPs (model forward propagation FLOPs), bwd FLOPs (model backward propagation FLOPs),\n"
              "default model backpropagation takes %.2f times as much computation as forward propagation.\n" % self.compute_bp_factor)

        line_fmt = '{:<70}  {:<8}'

        print(line_fmt.format('Total Training Params: ', params_to_string(total_params)))

        print(line_fmt.format('fwd MACs: ', macs_to_string(total_macs, units=units,
                                                           precision=precision)))
        print(line_fmt.format('fwd FLOPs: ', flops_to_string(total_flops, units=units,
                                                             precision=precision)))
        print(line_fmt.format('fwd+bwd MACs: ', macs_to_string(total_macs * (1 + self.compute_bp_factor),
                                                               units=units, precision=precision)))
        print(line_fmt.format('fwd+bwd FLOPs: ', flops_to_string(total_flops * (1 + self.compute_bp_factor),
                                                                 units=units, precision=precision)))

        def flops_repr(module):
            params = module.__params__
            flops = get_module_flops(module)
            macs = get_module_macs(module)
            items = [
                "{} = {:g}% Params".format(
                    params_to_string(params),
                    round(100 * params / total_params, precision) if total_params else 0),
                "{} = {:g}% MACs".format(macs_to_string(macs),
                                         round(100 * macs / total_macs, precision) if total_macs else 0),
                "{} = {:g}% FLOPs".format(flops_to_string(flops),
                                          round(100 * macs / total_flops, precision) if total_flops else 0),
            ]
            original_extra_repr = module.original_extra_repr()
            if original_extra_repr:
                items.append(original_extra_repr)
            return ", ".join(items)

        def add_extra_repr(module):
            flops_extra_repr = flops_repr.__get__(module)
            if module.extra_repr != flops_extra_repr:
                module.original_extra_repr = module.extra_repr
                module.extra_repr = flops_extra_repr
                assert module.extra_repr != module.original_extra_repr

        def del_extra_repr(module):
            if hasattr(module, "original_extra_repr"):
                module.extra_repr = module.original_extra_repr
                del module.original_extra_repr

        self.model.apply(add_extra_repr)

        if print_detailed:
            print(
                "\n-------------------------------- Detailed Calculated FLOPs Results --------------------------------")
            print(
                "Each module caculated is listed after its name in the following order: \nparams, percentage of total params, MACs, percentage of total MACs, FLOPS, percentage of total FLOPs"
            )
            print(
                "\nNote: 1. A module can have torch.nn.module or torch.nn.functional to compute logits (e.g. CrossEntropyLoss). \n They are not counted as submodules in calflops and not to be printed out. However they make up the difference between a parent's MACs and the sum of its submodules'.\n2. Number of floating-point operations is a theoretical estimation, thus FLOPS computed using that could be larger than the maximum system throughput.\n"
            )
            print(self.model)

        self.model.apply(del_extra_repr)

        print("---------------------------------------------------------------------------------------------------")

mcsmoe/calflops/calculate_pipline.py
UNITES-Lab/MC-SMoE
import torch
import torch.nn as nn

from .calculate_pipline import CalFlopsPipline
from .utils import flops_to_string
from .utils import generate_transformer_input
from .utils import macs_to_string
from .utils import params_to_string


def calculate_flops(model,
                    input_shape=None,
                    transformer_tokenizer=None,
                    args=[],
                    kwargs={},
                    forward_mode="forward",
                    include_backPropagation=False,
                    compute_bp_factor=2.0,
                    print_results=True,
                    print_detailed=True,
                    output_as_string=True,
                    output_precision=2,
                    output_unit=None,
                    ignore_modules=None,
                    is_sparse=False):
    """Returns the total floating-point operations, MACs, and parameters of a model.

    Args:
        model ([torch.nn.Module]): The model of input must be a PyTorch model.
        input_shape (tuple, optional): Input shape to the model. If args and kwargs is empty, the model takes a tensor with this shape as the only positional argument. Default to [].
        transformers_tokenizer (None, optional): Transforemrs Toekenizer must be special if model type is transformers and args、kwargs is empty. Default to None
        args (list, optional): list of positional arguments to the model, such as bert input args is [input_ids, token_type_ids, attention_mask]. Default to []
        kwargs (dict, optional): dictionary of keyword arguments to the model, such as bert input kwargs is {'input_ids': ..., 'token_type_ids':..., 'attention_mask':...}. Default to {}
        forward_mode (str, optional): To determine the mode of model inference, Default to 'forward'. And use 'generate' if model inference uses model.generate().
        include_backPropagation (bool, optional): Decides whether the final return FLOPs computation includes the computation for backpropagation.
        compute_bp_factor (float, optional): The model backpropagation is a multiple of the forward propagation computation. Default to 2.
        print_results (bool, optional): Whether to print the model profile. Defaults to True.
        print_detailed (bool, optional): Whether to print the detailed model profile. Defaults to True.
        output_as_string (bool, optional): Whether to print the output as string. Defaults to True.
        output_precision (int, optional) : Output holds the number of decimal places if output_as_string is True. Default to 2.
        output_unit (str, optional): The unit used to output the result value, such as T, G, M, and K. Default is None, that is the unit of the output decide on value.
        ignore_modules ([type], optional): the list of modules to ignore during profiling. Defaults to None.
        is_sparse (bool, optional): Whether to exclude sparse matrix flops. Defaults to False.

    Example:
    .. code-block:: python
    from calflops import calculate_flops

    # Deep Learning Model, such as alexnet.
    from torchvision import models

    model = models.alexnet()
    batch_size = 1
    flops, macs, params = calculate_flops(model=model, 
                                          input_shape=(batch_size, 3, 224, 224),
                                          output_as_string=True,
                                          output_precision=4)
    print("Alexnet FLOPs:%s   MACs:%s   Params:%s \n" %(flops, macs, params))
    #Alexnet FLOPs:1.4297 GFLOPS   MACs:714.188 MMACs   Params:61.1008 M 

    # Transformers Model, such as bert.
    from transformers import AutoModel
    from transformers import AutoTokenizer
    batch_size = 1
    max_seq_length = 128
    model_name = "hfl/chinese-roberta-wwm-ext/"
    model_save = "../pretrain_models/" + model_name
    model = AutoModel.from_pretrained(model_save)
    tokenizer = AutoTokenizer.from_pretrained(model_save)
    flops, macs, params = calculate_flops(model=model, 
                                          input_shape=(batch_size, max_seq_length),
                                          transformer_tokenizer=tokenizer)
    print("Bert(hfl/chinese-roberta-wwm-ext) FLOPs:%s   MACs:%s   Params:%s \n" %(flops, macs, params))
    #Bert(hfl/chinese-roberta-wwm-ext) FLOPs:22.36 GFLOPS   MACs:11.17 GMACs   Params:102.27 M 

    # Large Languase Model, such as llama2-7b.
    from transformers import LlamaTokenizer
    from transformers import LlamaForCausalLM
    batch_size = 1
    max_seq_length = 128
    model_name = "llama2_hf_7B"
    model_save = "../model/" + model_name
    model = LlamaForCausalLM.from_pretrained(model_save)
    tokenizer = LlamaTokenizer.from_pretrained(model_save)
    flops, macs, params = calculate_flops(model=model,
                                          input_shape=(batch_size, max_seq_length),
                                          transformer_tokenizer=tokenizer)
    print("Llama2(7B) FLOPs:%s   MACs:%s   Params:%s \n" %(flops, macs, params))
    #Llama2(7B) FLOPs:1.7 TFLOPS   MACs:850.00 GMACs   Params:6.74 B 

    Returns:
        The number of floating-point operations, multiply-accumulate operations (MACs), and parameters in the model.
    """

    assert isinstance(model, nn.Module), "model must be a PyTorch module"
    # assert transformers_tokenizer and auto_generate_transformers_input and "transformers" in str(type(model)), "The model must be a transformers model if args of auto_generate_transformers_input is True and transformers_tokenizer is not None"
    model.eval()

    is_transformer = True if "transformers" in str(type(model)) else False

    calculate_flops_pipline = CalFlopsPipline(model=model,
                                              include_backPropagation=include_backPropagation,
                                              compute_bp_factor=compute_bp_factor,
                                              is_sparse=is_sparse)
    calculate_flops_pipline.start_flops_calculate(ignore_list=ignore_modules)

    device = next(model.parameters()).device
    model = model.to(device)

    if input_shape is not None:
        assert len(args) == 0 and len(
            kwargs) == 0, "args and kwargs must be empty value if input_shape is not None, then will be generate random input by inpust_shape"
        assert type(input_shape) is tuple, "input_shape must be a tuple"
        assert len(input_shape) >= 1, "input_shape must have at least one element"

        if transformer_tokenizer is None:  # model is not transformers model
            assert is_transformer is False, "the model is must not transformer model if input_shape is not None and transformer_tokenizer is None"
            try:
                input = torch.ones(()).new_empty(
                    (*input_shape,),
                    dtype=next(model.parameters()).dtype,
                    device=device,
                )
            except StopIteration:
                input = torch.ones(()).new_empty((*input_shape,))
            args = [input]
        else:
            assert len(
                input_shape) == 2, "the format of input_shape must be (batch_size, seq_len) if model is transformers model and auto_generate_transformers_input if True"
            kwargs = generate_transformer_input(input_shape=input_shape,
                                                model_tokenizer=transformer_tokenizer,
                                                device=device)
    else:
        assert transformer_tokenizer or (len(args) > 0 or len(
            kwargs) > 0), "input_shape or args or kwargs one of there parameters must specified if auto_generate_input is False"
        if transformer_tokenizer:
            kwargs = generate_transformer_input(input_shape=None,
                                                model_tokenizer=transformer_tokenizer,
                                                device=device)

    if kwargs:
        for key, value in kwargs.items():
            kwargs[key] = value.to(device)

        if forward_mode == 'forward':
            _ = model(*args, **kwargs)
        if forward_mode == 'generate':
            _ = model.generate(*args, **kwargs)
    else:
        for index in range(len(args)):
            args[index] = args[index].to(device)

        if forward_mode == 'forward':
            _ = model(*args)
        if forward_mode == 'generate':
            _ = model.generate(*args)

    flops = calculate_flops_pipline.get_total_flops()
    macs = calculate_flops_pipline.get_total_macs()
    params = calculate_flops_pipline.get_total_params()

    if print_results:
        return_print = calculate_flops_pipline.print_model_pipline(units=output_unit,
                                                                   precision=output_precision,
                                                                   print_detailed=print_detailed)

    calculate_flops_pipline.end_flops_calculate()

    if include_backPropagation:
        flops = flops * (1 + compute_bp_factor)
        macs = macs * (1 + compute_bp_factor)

    if output_as_string:
        return flops_to_string(flops, units=output_unit, precision=output_precision), \
            macs_to_string(macs, units=output_unit, precision=output_precision), \
            params_to_string(params, units=output_unit, precision=output_precision)

    return flops, macs, params

mcsmoe/calflops/flops_counter.py
UNITES-Lab/MC-SMoE
from collections import OrderedDict
from typing import List
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

Tensor = torch.Tensor


def _prod(dims):
    p = 1
    for v in dims:
        p *= v
    return p


def _linear_flops_compute(input, weight, bias=None):
    out_features = weight.shape[0]
    macs = input.numel() * out_features
    return 2 * macs, macs


# Activation just calculate FLOPs， MACs is 0
def _relu_flops_compute(input, inplace=False):
    return input.numel(), 0


def _prelu_flops_compute(input: Tensor, weight: Tensor):
    return input.numel(), 0


def _elu_flops_compute(input: Tensor, alpha: float = 1.0, inplace: bool = False):
    return input.numel(), 0


def _leaky_relu_flops_compute(input: Tensor, negative_slope: float = 0.01, inplace: bool = False):
    return input.numel(), 0


def _relu6_flops_compute(input: Tensor, inplace: bool = False):
    return input.numel(), 0


def _silu_flops_compute(input: Tensor, inplace: bool = False):
    return input.numel(), 0


def _gelu_flops_compute(input, **kwargs):
    return input.numel(), 0


def _pool_flops_compute(input,
                        kernel_size,
                        stride=None,
                        padding=0,
                        dilation=None,
                        ceil_mode=False,
                        count_include_pad=True,
                        divisor_override=None,
                        return_indices=None):
    return input.numel(), 0


def _conv_flops_compute(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    assert weight.shape[1] * groups == input.shape[1]

    batch_size = input.shape[0]
    in_channels = input.shape[1]
    out_channels = weight.shape[0]
    kernel_dims = list(weight.shape[2:])
    input_dims = list(input.shape[2:])

    length = len(input_dims)

    strides = stride if type(stride) is tuple else (stride,) * length
    dilations = dilation if type(dilation) is tuple else (dilation,) * length
    if isinstance(padding, str):
        if padding == 'valid':
            paddings = (0,) * length
        elif padding == 'same':
            paddings = ()
            for d, k in zip(dilations, kernel_dims):
                total_padding = d * (k - 1)
                paddings += (total_padding // 2,)
    elif isinstance(padding, tuple):
        paddings = padding
    else:
        paddings = (padding,) * length

    output_dims = []
    for idx, input_dim in enumerate(input_dims):
        output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] *
                                                       (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1
        output_dims.append(output_dim)

    filters_per_channel = out_channels // groups
    conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel
    active_elements_count = batch_size * int(_prod(output_dims))
    overall_conv_macs = conv_per_position_macs * active_elements_count
    overall_conv_flops = 2 * overall_conv_macs

    bias_flops = 0
    if bias is not None:
        bias_flops = out_channels * active_elements_count

    return int(overall_conv_flops + bias_flops), int(overall_conv_macs)


def _conv_trans_flops_compute(
        input,
        weight,
        bias=None,
        stride=1,
        padding=0,
        output_padding=0,
        groups=1,
        dilation=1,
):
    batch_size = input.shape[0]
    in_channels = input.shape[1]
    out_channels = weight.shape[1]
    kernel_dims = list(weight.shape[2:])
    input_dims = list(input.shape[2:])

    length = len(input_dims)

    paddings = padding if type(padding) is tuple else (padding,) * length
    strides = stride if type(stride) is tuple else (stride,) * length
    dilations = dilation if type(dilation) is tuple else (dilation,) * length

    output_dims = []
    for idx, input_dim in enumerate(input_dims):
        output_dim = (input_dim + 2 * paddings[idx] - (dilations[idx] *
                                                       (kernel_dims[idx] - 1) + 1)) // strides[idx] + 1
        output_dims.append(output_dim)

    paddings = padding if type(padding) is tuple else (padding, padding)
    strides = stride if type(stride) is tuple else (stride, stride)
    dilations = dilation if type(dilation) is tuple else (dilation, dilation)

    filters_per_channel = out_channels // groups
    conv_per_position_macs = int(_prod(kernel_dims)) * in_channels * filters_per_channel
    active_elements_count = batch_size * int(_prod(input_dims))
    overall_conv_macs = conv_per_position_macs * active_elements_count
    overall_conv_flops = 2 * overall_conv_macs

    bias_flops = 0
    if bias is not None:
        bias_flops = out_channels * batch_size * int(_prod(output_dims))

    return int(overall_conv_flops + bias_flops), int(overall_conv_macs)


def _batch_norm_flops_compute(
        input,
        running_mean,
        running_var,
        weight=None,
        bias=None,
        training=False,
        momentum=0.1,
        eps=1e-05,
):
    has_affine = weight is not None
    if training:
        # estimation
        return input.numel() * (5 if has_affine else 4), 0
    flops = input.numel() * (2 if has_affine else 1)
    return flops, 0


def _layer_norm_flops_compute(
        input: Tensor,
        normalized_shape: List[int],
        weight: Optional[Tensor] = None,
        bias: Optional[Tensor] = None,
        eps: float = 1e-5,
):
    has_affine = weight is not None
    # estimation
    return input.numel() * (5 if has_affine else 4), 0


def _group_norm_flops_compute(input: Tensor,
                              num_groups: int,
                              weight: Optional[Tensor] = None,
                              bias: Optional[Tensor] = None,
                              eps: float = 1e-5):
    has_affine = weight is not None
    # estimation
    return input.numel() * (5 if has_affine else 4), 0


def _instance_norm_flops_compute(
        input: Tensor,
        running_mean: Optional[Tensor] = None,
        running_var: Optional[Tensor] = None,
        weight: Optional[Tensor] = None,
        bias: Optional[Tensor] = None,
        use_input_stats: bool = True,
        momentum: float = 0.1,
        eps: float = 1e-5,
):
    has_affine = weight is not None
    # estimation
    return input.numel() * (5 if has_affine else 4), 0


def _upsample_flops_compute(*args, **kwargs):
    input = args[0]
    size = kwargs.get('size', None)
    if size is None and len(args) > 1:
        size = args[1]

    if size is not None:
        if isinstance(size, tuple) or isinstance(size, list):
            return int(_prod(size)), 0
        else:
            return int(size), 0

    scale_factor = kwargs.get('scale_factor', None)
    if scale_factor is None and len(args) > 2:
        scale_factor = args[2]
    assert scale_factor is not None, "either size or scale_factor should be defined"

    flops = input.numel()
    if isinstance(scale_factor, tuple) and len(scale_factor) == len(input):
        flops * int(_prod(scale_factor))
    else:
        flops * scale_factor ** len(input)
    return flops, 0


def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None):
    return input.numel(), 0


def _embedding_flops_compute(
        input,
        weight,
        padding_idx=None,
        max_norm=None,
        norm_type=2.0,
        scale_grad_by_freq=False,
        sparse=False,
):
    return 0, 0


def _dropout_flops_compute(input, p=0.5, training=True, inplace=False):
    return 0, 0


def _matmul_flops_compute(input, other, *, out=None):
    """
    Count flops for the matmul operation.
    """
    macs = _prod(input.shape) * other.shape[-1]
    return 2 * macs, macs


def _addmm_flops_compute(input, mat1, mat2, *, beta=1, alpha=1, out=None):
    """
    Count flops for the addmm operation.
    """
    macs = _prod(mat1.shape) * mat2.shape[-1]
    return 2 * macs + _prod(input.shape), macs


def _einsum_flops_compute(equation, *operands):
    """
    Count flops for the einsum operation.
    """
    equation = equation.replace(" ", "")
    input_shapes = [o.shape for o in operands]

    # Re-map equation so that same equation with different alphabet
    # representations will look the same.
    letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys()
    mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)}
    equation = equation.translate(mapping)

    np_arrs = [np.zeros(s) for s in input_shapes]
    optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1]
    for line in optim.split("\n"):
        if "optimized flop" in line.lower():
            flop = int(float(line.split(":")[-1]))
            return flop, 0
    raise NotImplementedError("Unsupported einsum operation.")


def _tensor_addmm_flops_compute(self, mat1, mat2, *, beta=1, alpha=1, out=None):
    """
    Count flops for the tensor addmm operation.
    """
    macs = _prod(mat1.shape) * mat2.shape[-1]
    return 2 * macs + _prod(self.shape), macs


def _mul_flops_compute(input, other, *, out=None):
    return _elementwise_flops_compute(input, other)


def _add_flops_compute(input, other, *, alpha=1, out=None):
    return _elementwise_flops_compute(input, other)


def _elementwise_flops_compute(input, other):
    if not torch.is_tensor(input):
        if torch.is_tensor(other):
            return _prod(other.shape), 0
        else:
            return 1, 0
    elif not torch.is_tensor(other):
        return _prod(input.shape), 0
    else:
        dim_input = len(input.shape)
        dim_other = len(other.shape)
        max_dim = max(dim_input, dim_other)

        final_shape = []
        for i in range(max_dim):
            in_i = input.shape[i] if i < dim_input else 1
            ot_i = other.shape[i] if i < dim_other else 1
            if in_i > ot_i:
                final_shape.append(in_i)
            else:
                final_shape.append(ot_i)
        flops = _prod(final_shape)
        return flops, 0


def wrapFunc(func, funcFlopCompute, old_functions, module_flop_count, module_mac_count):
    oldFunc = func
    name = func.__str__
    old_functions[name] = oldFunc

    def newFunc(*args, **kwds):
        flops, macs = funcFlopCompute(*args, **kwds)
        if module_flop_count:
            module_flop_count[-1].append((name, flops))
        if module_mac_count and macs:
            module_mac_count[-1].append((name, macs))
        return oldFunc(*args, **kwds)

    newFunc.__str__ = func.__str__

    return newFunc


def _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size):
    gates_size = w_ih.shape[0]
    # matrix matrix mult ih state and internal state
    flops += 2 * w_ih.shape[0] * w_ih.shape[1] - gates_size
    # matrix matrix mult hh state and internal state
    flops += 2 * w_hh.shape[0] * w_hh.shape[1] - gates_size
    if isinstance(rnn_module, (nn.RNN, nn.RNNCell)):
        # add both operations
        flops += rnn_module.hidden_size
    elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)):
        # hadamard of r
        flops += rnn_module.hidden_size
        # adding operations from both states
        flops += rnn_module.hidden_size * 3
        # last two hadamard _product and add
        flops += rnn_module.hidden_size * 3
    elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)):
        # adding operations from both states
        flops += rnn_module.hidden_size * 4
        # two hadamard _product and add for C state
        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
        # final hadamard
        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
    return flops


def _rnn_forward_hook(rnn_module, input, output):
    flops = 0
    # input is a tuple containing a sequence to process and (optionally) hidden state
    inp = input[0]
    batch_size = inp.shape[0]
    seq_length = inp.shape[1]
    num_layers = rnn_module.num_layers

    for i in range(num_layers):
        w_ih = rnn_module.__getattr__("weight_ih_l" + str(i))
        w_hh = rnn_module.__getattr__("weight_hh_l" + str(i))
        if i == 0:
            input_size = rnn_module.input_size
        else:
            input_size = rnn_module.hidden_size
        flops = _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size)
        if rnn_module.bias:
            b_ih = rnn_module.__getattr__("bias_ih_l" + str(i))
            b_hh = rnn_module.__getattr__("bias_hh_l" + str(i))
            flops += b_ih.shape[0] + b_hh.shape[0]

    flops *= batch_size
    flops *= seq_length
    if rnn_module.bidirectional:
        flops *= 2
    rnn_module.__flops__ += int(flops)


def _rnn_cell_forward_hook(rnn_cell_module, input, output):
    flops = 0
    inp = input[0]
    batch_size = inp.shape[0]
    w_ih = rnn_cell_module.__getattr__("weight_ih")
    w_hh = rnn_cell_module.__getattr__("weight_hh")
    input_size = inp.shape[1]
    flops = _rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size)
    if rnn_cell_module.bias:
        b_ih = rnn_cell_module.__getattr__("bias_ih")
        b_hh = rnn_cell_module.__getattr__("bias_hh")
        flops += b_ih.shape[0] + b_hh.shape[0]

    flops *= batch_size
    rnn_cell_module.__flops__ += int(flops)


MODULE_HOOK_MAPPING = {
    # RNN
    nn.RNN: _rnn_forward_hook,
    nn.GRU: _rnn_forward_hook,
    nn.LSTM: _rnn_forward_hook,
    nn.RNNCell: _rnn_cell_forward_hook,
    nn.LSTMCell: _rnn_cell_forward_hook,
    nn.GRUCell: _rnn_cell_forward_hook,
}


def _patch_functionals(old_functions, module_flop_count, module_mac_count):
    # FC
    F.linear = wrapFunc(F.linear, _linear_flops_compute, old_functions, module_flop_count, module_mac_count)
    # convolutions
    F.conv1d = wrapFunc(F.conv1d, _conv_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.conv2d = wrapFunc(F.conv2d, _conv_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.conv3d = wrapFunc(F.conv3d, _conv_flops_compute, old_functions, module_flop_count, module_mac_count)

    # conv transposed
    F.conv_transpose1d = wrapFunc(F.conv_transpose1d, _conv_trans_flops_compute, old_functions, module_flop_count,
                                  module_mac_count)
    F.conv_transpose2d = wrapFunc(F.conv_transpose2d, _conv_trans_flops_compute, old_functions, module_flop_count,
                                  module_mac_count)
    F.conv_transpose3d = wrapFunc(F.conv_transpose3d, _conv_trans_flops_compute, old_functions, module_flop_count,
                                  module_mac_count)

    # activations
    F.relu = wrapFunc(F.relu, _relu_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.prelu = wrapFunc(F.prelu, _prelu_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.elu = wrapFunc(F.elu, _elu_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.leaky_relu = wrapFunc(F.leaky_relu, _leaky_relu_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.relu6 = wrapFunc(F.relu6, _relu6_flops_compute, old_functions, module_flop_count, module_mac_count)
    if hasattr(F, "silu"):
        F.silu = wrapFunc(F.silu, _silu_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.gelu = wrapFunc(F.gelu, _gelu_flops_compute, old_functions, module_flop_count, module_mac_count)

    # Normalizations
    F.batch_norm = wrapFunc(F.batch_norm, _batch_norm_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.layer_norm = wrapFunc(F.layer_norm, _layer_norm_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.instance_norm = wrapFunc(F.instance_norm, _instance_norm_flops_compute, old_functions, module_flop_count,
                               module_mac_count)
    F.group_norm = wrapFunc(F.group_norm, _group_norm_flops_compute, old_functions, module_flop_count, module_mac_count)

    # poolings
    F.avg_pool1d = wrapFunc(F.avg_pool1d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.avg_pool2d = wrapFunc(F.avg_pool2d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.avg_pool3d = wrapFunc(F.avg_pool3d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.max_pool1d = wrapFunc(F.max_pool1d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.max_pool2d = wrapFunc(F.max_pool2d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.max_pool3d = wrapFunc(F.max_pool3d, _pool_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.adaptive_avg_pool1d = wrapFunc(F.adaptive_avg_pool1d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)
    F.adaptive_avg_pool2d = wrapFunc(F.adaptive_avg_pool2d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)
    F.adaptive_avg_pool3d = wrapFunc(F.adaptive_avg_pool3d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)
    F.adaptive_max_pool1d = wrapFunc(F.adaptive_max_pool1d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)
    F.adaptive_max_pool2d = wrapFunc(F.adaptive_max_pool2d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)
    F.adaptive_max_pool3d = wrapFunc(F.adaptive_max_pool3d, _pool_flops_compute, old_functions, module_flop_count,
                                     module_mac_count)

    # upsample
    F.upsample = wrapFunc(F.upsample, _upsample_flops_compute, old_functions, module_flop_count, module_mac_count)
    F.interpolate = wrapFunc(F.interpolate, _upsample_flops_compute, old_functions, module_flop_count, module_mac_count)

    # softmax
    F.softmax = wrapFunc(F.softmax, _softmax_flops_compute, old_functions, module_flop_count, module_mac_count)

    # embedding
    F.embedding = wrapFunc(F.embedding, _embedding_flops_compute, old_functions, module_flop_count, module_mac_count)


def _patch_tensor_methods(old_functions, module_flop_count, module_mac_count):
    torch.matmul = wrapFunc(torch.matmul, _matmul_flops_compute, old_functions, module_flop_count, module_mac_count)
    torch.Tensor.matmul = wrapFunc(torch.Tensor.matmul, _matmul_flops_compute, old_functions, module_flop_count,
                                   module_mac_count)
    # torch.mm = wrapFunc(torch.mm, _matmul_flops_compute, old_functions, module_flop_count, module_mac_count)
    # torch.Tensor.mm = wrapFunc(torch.Tensor.mm, _matmul_flops_compute, old_functions, module_flop_count, module_mac_count)
    # torch.bmm = wrapFunc(torch.bmm, _matmul_flops_compute, old_functions, module_flop_count, module_mac_count)
    # torch.Tensor.bmm = wrapFunc(torch.Tensor.bmm, _matmul_flops_compute, old_functions, module_flop_count, module_mac_count)

    torch.addmm = wrapFunc(torch.addmm, _addmm_flops_compute, old_functions, module_flop_count, module_mac_count)
    torch.Tensor.addmm = wrapFunc(torch.Tensor.addmm, _tensor_addmm_flops_compute, old_functions, module_flop_count,
                                  module_mac_count)

    torch.mul = wrapFunc(torch.mul, _mul_flops_compute, old_functions, module_flop_count, module_mac_count)
    torch.Tensor.mul = wrapFunc(torch.Tensor.mul, _mul_flops_compute, old_functions, module_flop_count,
                                module_mac_count)

    torch.add = wrapFunc(torch.add, _add_flops_compute, old_functions, module_flop_count, module_mac_count)
    torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute, old_functions, module_flop_count,
                                module_mac_count)

    torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute, old_functions, module_flop_count, module_mac_count)

    torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute, old_functions, module_flop_count,
                             module_mac_count)


def _reload_functionals(old_functions):
    # torch.nn.functional does not support importlib.reload()
    F.linear = old_functions[F.linear.__str__]
    F.conv1d = old_functions[F.conv1d.__str__]
    F.conv2d = old_functions[F.conv2d.__str__]
    F.conv3d = old_functions[F.conv3d.__str__]
    F.conv_transpose1d = old_functions[F.conv_transpose1d.__str__]
    F.conv_transpose2d = old_functions[F.conv_transpose2d.__str__]
    F.conv_transpose3d = old_functions[F.conv_transpose3d.__str__]
    F.relu = old_functions[F.relu.__str__]
    F.prelu = old_functions[F.prelu.__str__]
    F.elu = old_functions[F.elu.__str__]
    F.leaky_relu = old_functions[F.leaky_relu.__str__]
    F.relu6 = old_functions[F.relu6.__str__]
    if hasattr(F, "silu"):
        F.silu = old_functions[F.silu.__str__]
    F.gelu = old_functions[F.gelu.__str__]
    F.batch_norm = old_functions[F.batch_norm.__str__]
    F.layer_norm = old_functions[F.layer_norm.__str__]
    F.instance_norm = old_functions[F.instance_norm.__str__]
    F.group_norm = old_functions[F.group_norm.__str__]
    F.avg_pool1d = old_functions[F.avg_pool1d.__str__]
    F.avg_pool2d = old_functions[F.avg_pool2d.__str__]
    F.avg_pool3d = old_functions[F.avg_pool3d.__str__]
    F.max_pool1d = old_functions[F.max_pool1d.__str__]
    F.max_pool2d = old_functions[F.max_pool2d.__str__]
    F.max_pool3d = old_functions[F.max_pool3d.__str__]
    F.adaptive_avg_pool1d = old_functions[F.adaptive_avg_pool1d.__str__]
    F.adaptive_avg_pool2d = old_functions[F.adaptive_avg_pool2d.__str__]
    F.adaptive_avg_pool3d = old_functions[F.adaptive_avg_pool3d.__str__]
    F.adaptive_max_pool1d = old_functions[F.adaptive_max_pool1d.__str__]
    F.adaptive_max_pool2d = old_functions[F.adaptive_max_pool2d.__str__]
    F.adaptive_max_pool3d = old_functions[F.adaptive_max_pool3d.__str__]
    F.upsample = old_functions[F.upsample.__str__]
    F.interpolate = old_functions[F.interpolate.__str__]
    F.softmax = old_functions[F.softmax.__str__]
    F.embedding = old_functions[F.embedding.__str__]


def _reload_tensor_methods(old_functions):
    torch.matmul = old_functions[torch.matmul.__str__]
    torch.Tensor.matmul = old_functions[torch.Tensor.matmul.__str__]
    # torch.mm = old_functions[torch.mm.__str__]
    # torch.Tensor.mm = old_functions[torch.Tensor.mm.__str__]
    # torch.bmm = old_functions[torch.matmul.__str__]
    # torch.Tensor.bmm = old_functions[torch.Tensor.bmm.__str__]
    torch.addmm = old_functions[torch.addmm.__str__]
    torch.Tensor.addmm = old_functions[torch.Tensor.addmm.__str__]
    torch.mul = old_functions[torch.mul.__str__]
    torch.Tensor.mul = old_functions[torch.Tensor.mul.__str__]
    torch.add = old_functions[torch.add.__str__]
    torch.Tensor.add = old_functions[torch.Tensor.add.__str__]
    torch.einsum = old_functions[torch.einsum.__str__]
    torch.baddbmm = old_functions[torch.baddbmm.__str__]

mcsmoe/calflops/pytorch_ops.py
UNITES-Lab/MC-SMoE
import importlib

import torch

DEFAULT_PRECISION = 2


def generate_transformer_input(model_tokenizer, input_shape, device):
    """Automatically generates data in the form of transformes model input format.
    
    Args:
        input_shape (tuple):transformers model input shape: (batch_size, seq_len).
        tokenizer (transformer.model.tokenization): transformers model tokenization.tokenizer.

    Returns:
        dict: data format of transformers model input, it is a dict contain 'input_ids', 'attention_mask', sometime contain 'token_type_ids'.
    """

    if input_shape is None:
        input_shape = [1, 128]  # defautl (batch_size=1, seq_len=128)

    max_length = input_shape[1]
    model_input_ids = []
    model_attention_mask = []
    model_token_type_ids = []
    model_position_ids = []

    inp_seq = ""
    for _ in range(input_shape[0]):
        inputs = model_tokenizer.encode_plus(
            inp_seq,
            add_special_tokens=True,
            truncation_strategy='longest_first',
        )
        origin_length = len(inputs["input_ids"])
        padding_length = max_length - origin_length

        for key in inputs.keys():
            if key == "input_ids":
                input_ids = inputs["input_ids"]
                pad_token = model_tokenizer.pad_token_id if model_tokenizer.pad_token_id else 0
                input_ids = input_ids + ([pad_token] * padding_length)
                assert len(input_ids) == max_length, "len(input_ids) must equal max_length"
                model_input_ids.append(input_ids)
            elif key == "attention_mask":
                attention_mask = [1] * origin_length
                attention_mask = attention_mask + ([0] * padding_length)
                assert len(attention_mask) == max_length, "len(attention_mask) must equal max_length"
                model_attention_mask.append(attention_mask)
            elif key == "token_type_ids":
                token_type_ids = inputs['token_type_ids']
                pad_token_segment_id = 0
                token_type_ids = token_type_ids + ([pad_token_segment_id] * padding_length)
                assert len(token_type_ids) == max_length, "len(token_type_ids) must equal max_length"
                model_token_type_ids.append(token_type_ids)
            elif key == "position_ids":  # chatglm2 use position id
                position_ids = inputs['position_ids']
                for i in range(origin_length, max_length):
                    position_ids.append(i)
                assert len(position_ids) == max_length, "len(position_ids) must equal max_length"
                model_position_ids.append(position_ids)

    # Batch size input_shape[0], sequence length input_shape[128]
    inputs = {}
    if len(model_input_ids) > 0:
        inputs.update({"input_ids": torch.tensor(model_input_ids).to(device)})
    if len(model_attention_mask) > 0:
        inputs.update({"attention_mask": torch.tensor(model_attention_mask).to(device)})
    if len(model_token_type_ids) > 0:
        inputs.update({'token_type_ids': torch.tensor(model_token_type_ids).to(device)})
    if len(model_position_ids) > 0:
        inputs.update({'position_ids': torch.tensor(model_position_ids).to(device)})

    return inputs


def number_to_string(num, units=None, precision=DEFAULT_PRECISION):
    if units is None:
        if num >= 1e12:
            magnitude, units = 1e12, "T"
        elif num >= 1e9:
            magnitude, units = 1e9, "G"
        elif num >= 1e6:
            magnitude, units = 1e6, "M"
        elif num >= 1e3:
            magnitude, units = 1e3, "K"
        elif num >= 1 or num == 0:
            magnitude, units = 1, ""
        elif num >= 1e-3:
            magnitude, units = 1e-3, "m"
        else:
            magnitude, units = 1e-6, "u"
    else:
        if units == "T":
            magnitude = 1e12
        elif units == "G":
            magnitude = 1e9
        elif units == "M":
            magnitude = 1e6
        elif units == "K":
            magnitude = 1e3
        elif units == "m":
            magnitude = 1e-3
        elif units == "u":
            magnitude = 1e-6
        else:
            magnitude = 1
    return f"{round(num / magnitude, precision):g} {units}"


def macs_to_string(macs, units=None, precision=DEFAULT_PRECISION):
    """Converts macs in numeric form to string form.

    Args:
        macs (int): Calculate the results of the model macs in numerical form.
        units (str, optional): The unit of macs after conversion to string representation, such as TMACs、GMACs、MMACs、KMACs
        precision (int, optional): The number of digits of the result is preserved. Defaults to DEFAULT_PRECISION.

    Returns:
        string: The string representation of macs.
    """
    return f"{number_to_string(macs, units=units, precision=precision)}MACs"


def flops_to_string(flops, units=None, precision=DEFAULT_PRECISION):
    """Converts flops in numeric form to string form.

    Args:
        flops (int): Calculate the results of the model flops in numerical form.
        units (str, optional): The unit of flops after conversion to string representation, such as TFLOPs,GFLOPs,MFLOPs,KFLOPs.
        precision (int, optional): The number of digits of the result is preserved. Defaults to DEFAULT_PRECISION.

    Returns:
        string: The string representation of flops.
    """
    return f"{number_to_string(flops, units=units, precision=precision)}FLOPS"


def bytes_to_string(b, units=None, precision=DEFAULT_PRECISION):
    """Converts bytes in numeric form to string form.

    Args:
        b (int): Calculate the results of the bytes in numerical form.
        units (str, optional): The unit of bytes after conversion to string representation, such as TB,GB,MB,KB.
        precision (int, optional): The number of digits of the result is preserved. Defaults to DEFAULT_PRECISION.

    Returns:
        string: The string representation of bytes.
    """
    return f"{number_to_string(b, units=units, precision=precision)}B"


def params_to_string(params_num, units=None, precision=DEFAULT_PRECISION):
    """Converts params in numeric form to string form.

    Args:
        params_num (int): Calculate the results of the model param in numerical form.
        units (str, optional): The unit of params after conversion to string representation.
        precision (int, optional): The number of digits of the result is preserved. Defaults to DEFAULT_PRECISION.

    Returns:
        string: The string representation of params.
    """
    units = units.replace("B", "G") if units else units
    return number_to_string(params_num, units=units, precision=precision).replace("G", "B").strip()


def get_module_flops(module, is_sparse=False):
    """Recursively compute the FLOP s of the model

    Args:
        module (pytorch module): model format must be pytorch
        is_sparse (bool): Whether to exclude sparse weight.

    Returns:
        int: The sum of the entire model flops
    """
    sum_flops = module.__flops__ * sum(
        p.count_nonzero().item() for p in module.parameters() if p.requires_grad
    ) / (1e-8 + sum(p.numel() for p in module.parameters() if p.requires_grad)) if is_sparse else module.__flops__
    # iterate over immediate children modules
    for child in module.children():
        sum_flops += get_module_flops(child, is_sparse=is_sparse)
    return sum_flops


def get_module_macs(module, is_sparse=False):
    """Recursively compute the macs s of the model

    Args:
        module (pytorch module): model format must be pytorch
        is_sparse (bool): Whether to exclude sparse weight.

    Returns:
        int: The sum of the entire model macs
    """
    sum_macs = module.__macs__ * sum(
        p.count_nonzero().item() for p in module.parameters() if p.requires_grad
    ) / (1e-8 + sum(p.numel() for p in module.parameters() if p.requires_grad)) if is_sparse else module.__macs__
    # iterate over immediate children modules
    for child in module.children():
        sum_macs += get_module_macs(child, is_sparse=is_sparse)
    return sum_macs


def convert_bytes(size):
    "Converts `size` from bytes to the largest possible unit"
    for x in ["bytes", "KB", "MB", "GB", "TB"]:
        if size < 1024.0:
            return f"{round(size, 2)} {x}"
        size /= 1024.0

    return f"{round(size, 2)} PB"


def _is_package_available(pkg_name):
    # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
    package_exists = importlib.util.find_spec(pkg_name) is not None
    if package_exists:
        try:
            _ = importlib.metadata.metadata(pkg_name)
            return True
        except importlib.metadata.PackageNotFoundError:
            return False

mcsmoe/calflops/utils.py
UNITES-Lab/MC-SMoE
from .data_collator import (
    DataCollatorForSeq2Seq,
    DataCollatorForLanguageModeling
)
from .evaluation import (
    get_classification_label_index_and_token_ids,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    gather_predictions_references_by_casual_lm_loss
)
from .preprocessors import (
    TASK_MAPPING_DATASET_ARGUMENTS,
    TASK_MAPPING_PROMPT_KEY,
    get_label_mapping_id,
    Seq2SeqDataPreProcessor,
    Seq2SeqZeroShotDataPreProcessor,
    CasualZeroShotDataPreProcessor,
    CasualOneShotDataPreProcessor,
    tokenize_seq2seq,
    tokenize_seq2seq_zero_shot,
    tokenize_seq2se2_to_casual_lm,
    tokenize_casual_generation,
    tokenize_casual_zero_shot,
    keep_only_supporting_facts_in_context_for_hotpotqa,
    build_index_for_dataset
)

mcsmoe/data/__init__.py
UNITES-Lab/MC-SMoE
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
from transformers.data.data_collator import DataCollatorMixin
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
from transformers.utils import PaddingStrategy


@dataclass
class DataCollatorForSeq2Seq:
    """
    Data collator that will dynamically pad the inputs received, as well as the labels.

    Args:
        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
            The tokenizer used for encoding the data.
        model ([`PreTrainedModel`]):
            The model that is being trained. If set and has the *prepare_decoder_input_ids_from_labels*, use it to
            prepare the *decoder_input_ids*

            This is useful when using *label_smoothing* to avoid calculating loss twice.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
            among:

            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
              sequence is provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
        max_length (`int`, *optional*):
            Maximum length of the returned list and optionally padding length (see above).
        pad_to_multiple_of (`int`, *optional*):
            If set will pad the sequence to a multiple of the provided value.

            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
        label_pad_token_id (`int`, *optional*, defaults to -100):
            The id to use when padding the labels (-100 will be automatically ignored by PyTorch loss functions).
        return_tensors (`str`):
            The type of Tensor to return. Allowable values are "np", "pt" and "tf".
    """

    tokenizer: PreTrainedTokenizerBase
    model: Optional[Any] = None
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    label_pad_token_id: int = -100
    return_tensors: str = "pt"
    keys_to_ignore: Optional[List[str]] = None

    def __call__(self, features, return_tensors=None):
        if return_tensors is None:
            return_tensors = self.return_tensors
        labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
        # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
        # same length to return tensors.
        features_ignored = {}
        if self.keys_to_ignore is None:
            self.keys_to_ignore = []
        for key in self.keys_to_ignore:
            if key in features[0].keys():
                features_ignored[key] = [feature.pop(key) for feature in features]
        if labels is not None:
            max_label_length = max(len(l) for l in labels)
            if self.pad_to_multiple_of is not None:
                max_label_length = (
                        (max_label_length + self.pad_to_multiple_of - 1)
                        // self.pad_to_multiple_of
                        * self.pad_to_multiple_of
                )

            padding_side = self.tokenizer.padding_side
            for feature in features:
                remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
                if isinstance(feature["labels"], list):
                    feature["labels"] = (
                        feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
                    )
                elif padding_side == "right":
                    feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
                else:
                    feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)

        features = self.tokenizer.pad(
            features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=return_tensors,
        )

        # prepare decoder_input_ids
        if (
                labels is not None
                and self.model is not None
                and hasattr(self.model, "prepare_decoder_input_ids_from_labels")
        ):
            decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(labels=features["labels"])
            features["decoder_input_ids"] = decoder_input_ids

        return {**features, **features_ignored}


@dataclass
class DataCollatorForLanguageModeling(DataCollatorMixin):
    """
    Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they
    are not all of the same length.

    Args:
        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
            The tokenizer used for encoding the data.
        mlm (`bool`, *optional*, defaults to `True`):
            Whether or not to use masked language modeling. If set to `False`, the labels are the same as the inputs
            with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for non-masked
            tokens and the value to predict for the masked token.
        mlm_probability (`float`, *optional*, defaults to 0.15):
            The probability with which to (randomly) mask tokens in the input, when `mlm` is set to `True`.
        pad_to_multiple_of (`int`, *optional*):
            If set will pad the sequence to a multiple of the provided value.
        return_tensors (`str`):
            The type of Tensor to return. Allowable values are "np", "pt" and "tf".

    <Tip>

    For best performance, this data collator should be used with a dataset having items that are dictionaries or
    BatchEncoding, with the `"special_tokens_mask"` key, as returned by a [`PreTrainedTokenizer`] or a
    [`PreTrainedTokenizerFast`] with the argument `return_special_tokens_mask=True`.

    </Tip>"""

    tokenizer: PreTrainedTokenizerBase
    mlm: bool = False
    mlm_probability: float = 0.15
    pad_to_multiple_of: Optional[int] = None
    tf_experimental_compile: bool = False
    return_tensors: str = "pt"
    keys_to_ignore: Optional[List[str]] = None

    def __post_init__(self):
        if self.mlm and self.tokenizer.mask_token is None:
            raise ValueError(
                "This tokenizer does not have a mask token which is necessary for masked language modeling. "
                "You should pass `mlm=False` to train on causal language modeling instead."
            )

    def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
        # Handle dict or lists with proper padding and conversion to tensor.
        examples_ignored = {}
        if self.keys_to_ignore is None:
            self.keys_to_ignore = []
        for key in self.keys_to_ignore:
            if key in examples[0].keys():
                examples_ignored[key] = [sample.pop(key) for sample in examples]
        if "references" in examples[0].keys():
            # pad reference
            max_reference_length = max(len(sample['references']) for sample in examples)
            padding_side = self.tokenizer.padding_side
            for sample in examples:
                ref = sample["references"]
                remainder = [self.tokenizer.pad_token_id] * (max_reference_length - len(ref))
                sample["references"] = (
                    ref + remainder if padding_side == "right" else remainder + ref
                )
        if isinstance(examples[0], Mapping):
            batch = self.tokenizer.pad(examples, return_tensors="pt", pad_to_multiple_of=self.pad_to_multiple_of)
        else:
            batch = {
                "input_ids": _torch_collate_batch(examples, self.tokenizer, pad_to_multiple_of=self.pad_to_multiple_of)
            }

        # If special token mask has been preprocessed, pop it from the dict.
        special_tokens_mask = batch.pop("special_tokens_mask", None)
        if self.mlm:
            batch["input_ids"], batch["labels"] = self.torch_mask_tokens(
                batch["input_ids"], special_tokens_mask=special_tokens_mask
            )
        else:
            labels = batch["input_ids"].clone()
            if self.tokenizer.pad_token_id is not None:
                labels[labels == self.tokenizer.pad_token_id] = -100
            batch["labels"] = labels
        return {**batch, **examples_ignored}

    def torch_mask_tokens(self, inputs: Any, special_tokens_mask: Optional[Any] = None) -> Tuple[Any, Any]:
        """
        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
        """
        import torch

        labels = inputs.clone()
        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
        probability_matrix = torch.full(labels.shape, self.mlm_probability)
        if special_tokens_mask is None:
            special_tokens_mask = [
                self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
            ]
            special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
        else:
            special_tokens_mask = special_tokens_mask.bool()

        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # We only compute loss on masked tokens

        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)

        # 10% of the time, we replace masked input tokens with random word
        indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
        random_words = torch.randint(len(self.tokenizer), labels.shape, dtype=torch.long)
        inputs[indices_random] = random_words[indices_random]

        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
        return inputs, labels

mcsmoe/data/data_collator.py
UNITES-Lab/MC-SMoE
import random
from typing import Callable, Union, List, Optional, Dict

import torch
from datasets.arrow_dataset import Dataset
from evaluate import load
from promptsource.templates import DatasetTemplates
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

from .preprocessors import TASK_MAPPING_DATASET_ARGUMENTS, TASK_MAPPING_PROMPT_KEY, get_label_mapping_id

__all__ = [
    "get_evaluate_fn",
    "get_classification_label_index_and_token_ids",
    "EXTRA_KEYS_FOR_EVAL",
]

EXTRA_KEYS_FOR_EVAL = ["id", "idx"]


def get_evaluate_fn(
        task: str,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        raw_eval_dataset: Optional[Dataset] = None
) -> Callable:
    if task in ["squad", "squad_v2", "hotpotqa"]:
        return get_squad_evaluate_fn(tokenizer)
    elif task == "openbookqa":
        return get_openbookqa_evaluate_fn(tokenizer)
    elif task == "copa":
        return get_copa_evaluate_fn(tokenizer, raw_eval_dataset)
    elif task == "multirc":
        return get_multirc_evaluate_fn(tokenizer)
    elif task == "stsb":
        return get_stsb_evaluate_fn(tokenizer)
    else:
        # including other GLUE tasks, WinoGrande, WikiQA
        return get_cls_evaluate_fn(task, tokenizer)


def get_classification_label_index_and_token_ids(
        task: str,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
) -> tuple:
    prompt = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[task])[TASK_MAPPING_PROMPT_KEY[task]]
    if task in ["sst2", "mrpc"]:
        classification_token_index = 0
        choices = " ".join(prompt.get_fixed_answer_choices_list())
        classification_label_token_ids = tokenizer.encode(choices, add_special_tokens=False)
    else:
        return None, None
    return classification_token_index, classification_label_token_ids


def get_cls_evaluate_fn(
        task: str,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
) -> Callable:
    """
    Get the evaluate function for GLUE tasks.
    """

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], labels: torch.Tensor,
                    **kwargs):
        try:
            metric = load(*TASK_MAPPING_DATASET_ARGUMENTS[task])
        except FileNotFoundError:
            print(f"[Evaluation warning] No metric found for task {task}, using accuracy instead.")
            metric = load("accuracy")
        label_mapping_id = get_label_mapping_id(task)
        eos_token_id = tokenizer.eos_token_id
        predictions = [p[:p.index(eos_token_id) + 1] if eos_token_id in p else p for
                       p in predictions]
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        not_even_wrong = 0
        not_even_wrong_predictions = []
        for p in predictions:
            if p not in label_mapping_id:
                not_even_wrong += 1
                not_even_wrong_predictions.append(p)
        if not_even_wrong > 0:
            print(
                f"[Evaluation warning] {not_even_wrong} among total {len(predictions)} predictions are not in the label mapping.")
            print(f"[Evaluation warning] Some of the not even wrong predictions are: {not_even_wrong_predictions[:16]}")
        print(predictions[:32])
        predictions = [label_mapping_id[p] if p in label_mapping_id else random.choice(list(label_mapping_id.values()))
                       for p in predictions]
        references = labels.long().masked_fill(labels == -100, tokenizer.pad_token_id).tolist()
        references = tokenizer.batch_decode(references, skip_special_tokens=True)
        print(references[:32])
        references = [label_mapping_id[r]
                      for r in references]
        return metric.compute(predictions=predictions, references=references)

    return evaluate_fn


def get_stsb_evaluate_fn(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
) -> Callable:
    """
    Get the evaluate function for GLUE tasks.
    """

    def _is_float(s):
        try:
            float(s)
            return True
        except ValueError:
            return False

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], labels: torch.Tensor,
                    **kwargs):
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS["stsb"])
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        references = labels.long().masked_fill(labels == -100, tokenizer.pad_token_id).tolist()
        references = tokenizer.batch_decode(references, skip_special_tokens=True)
        not_even_wrong = 0
        not_even_wrong_predictions = []
        not_even_wrong_references = []
        for i, p in enumerate(predictions):
            if not _is_float(p):
                not_even_wrong += 1
                not_even_wrong_predictions.append(p)
                not_even_wrong_references.append(references[i])
        if not_even_wrong > 0:
            print(
                f"[Evaluation warning] {not_even_wrong} among total {len(predictions)} predictions are not in the label mapping.")
        predictions = [float(p) if _is_float(p) else round(random.uniform(0, 25)) / 5 for p in predictions]

        references = [float(r) for r in references]
        return metric.compute(predictions=predictions, references=references)

    return evaluate_fn


def get_multirc_evaluate_fn(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) -> Callable:
    """
    Get the evaluate function for SuperGLUE-MultiRC task.
    """

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], labels: torch.Tensor,
                    ids: List[Dict[str, int]], **kwargs):
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS['multirc'])
        label_mapping_id = get_label_mapping_id('multirc')
        # cut up tokens after the first [EOS] token, predictions is a list of token id list
        eos_token_id = tokenizer.eos_token_id
        predictions = [p[:p.index(eos_token_id) + 1] if eos_token_id in p else p for
                       p in predictions]
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        not_even_wrong = 0
        not_even_wrong_preds = []
        for p in predictions:
            if p not in label_mapping_id:
                not_even_wrong += 1
                not_even_wrong_preds.append(p)
        if not_even_wrong > 0:
            print(
                f"[Evaluation warning] {not_even_wrong} among total {len(predictions)} predictions are not in the label mapping.")
            print(f"[Evaluation warning] Some of the not even wrong predictions are: {not_even_wrong_preds[:16]}")
        predictions = [label_mapping_id[p] if p in label_mapping_id else random.choice([0, 1])
                       for p in predictions]
        references = labels.long().masked_fill(labels == -100, tokenizer.pad_token_id).tolist()
        references = tokenizer.batch_decode(references, skip_special_tokens=True)
        predictions = [
            {'prediction': p, 'idx': id} for p, id in zip(predictions, ids)
        ]
        references = [label_mapping_id[r] for r in references]
        results = metric.compute(predictions=predictions, references=references)
        return {
            'f1_a': results['f1_a'],
            'f1_m': results['f1_m'],
            'exact_match': results['exact_match']
        }

    return evaluate_fn


def get_squad_evaluate_fn(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) -> Callable:
    """
    Get the evaluate function for SQuAD tasks.
    """

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], labels: torch.Tensor,
                    ids: List[str], **kwargs):
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS['squad'])
        # cut up tokens after the first [EOS] token, predictions is a list of token id list
        eos_token_id = tokenizer.eos_token_id
        predictions = [p[:p.index(eos_token_id) + 1] if eos_token_id in p else p for
                       p in predictions]
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        references = labels.long().masked_fill(labels == -100, tokenizer.pad_token_id).tolist()
        references = tokenizer.batch_decode(references, skip_special_tokens=True)
        for i in range(8):
            print(predictions[i], references[i])
        predictions = [
            {'prediction_text': p, 'id': id} for p, id in zip(predictions, ids)
        ]
        references = [
            # answer_start is not used in the evaluation, so fake it
            {'answers': {'text': [reference], 'answer_start': [2333]}, 'id': id} for reference, id in
            zip(references, ids)
        ]
        return metric.compute(predictions=predictions, references=references)

    return evaluate_fn


def get_openbookqa_evaluate_fn(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]) -> Callable:
    """
    Get the evaluate function for OpenBookQA tasks.
    """

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], labels: torch.Tensor,
                    **kwargs):
        metric = load("accuracy")
        label_mapping_id = get_label_mapping_id('openbookqa')
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        predictions = [label_mapping_id[p] if p in label_mapping_id else -1
                       for p in predictions]
        references = labels.long().masked_fill(labels == -100, tokenizer.pad_token_id).tolist()
        references = tokenizer.batch_decode(references, skip_special_tokens=True)
        references = [label_mapping_id[reference]
                      for reference in references]
        return metric.compute(predictions=predictions, references=references)

    return evaluate_fn


def get_copa_evaluate_fn(tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
                         raw_eval_dataset: Dataset) -> Callable:
    """
    Get the evaluate function for SuperGLUE-COPA task.
    """
    id2choices = {
        item['idx']: [item['choice1'], item['choice2']] for item in raw_eval_dataset
    }
    id2references = {
        item['idx']: item['label'] for item in raw_eval_dataset
    }

    def evaluate_fn(predictions: Union[torch.Tensor, List[int], List[torch.LongTensor]], ids: List[int], **kwargs):
        metric = load('super_glue', 'copa')
        # cut up tokens after the first [EOS] token, predictions is a list of token id list
        eos_token_id = tokenizer.eos_token_id
        predictions = [p[:p.index(eos_token_id) + 1] if eos_token_id in p else p for
                       p in predictions]
        predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)  # string lengths batch
        predictions = [id2choices[idx].index(p) if p in id2choices[idx] else -1 for idx, p in zip(ids, predictions)]
        references = [id2references[idx] for idx in ids]
        return metric.compute(predictions=predictions, references=references)

    return evaluate_fn


def gather_predictions_references_by_casual_lm_loss(
        ids_list: List[int],
        answer_ids_list: List[int],
        choice_ids_list: List[int],
        losses_list: List[float],
) -> Dict[str, List[int]]:
    assert len(ids_list) == len(answer_ids_list) == len(choice_ids_list) == len(losses_list)
    num_choices = max(choice_ids_list) + 1
    idx_choice_idx_to_loss = {}
    idx_to_answer_idx = {}
    for idx, choice_idx, answer_idx, loss in zip(ids_list, choice_ids_list, answer_ids_list, losses_list):
        idx_choice_idx_to_loss[f"{idx}-{choice_idx}"] = loss
        if idx not in idx_to_answer_idx:
            idx_to_answer_idx[idx] = answer_idx
        else:
            assert idx_to_answer_idx[idx] == answer_idx
    predictions = []
    references = []
    for idx in sorted(idx_to_answer_idx.keys()):
        idx_loss_list = []
        for choice_idx in range(num_choices):
            idx_loss_list.append(idx_choice_idx_to_loss[f"{idx}-{choice_idx}"])
        best_choice_idx = idx_loss_list.index(min(idx_loss_list))
        predictions.append(best_choice_idx)
        references.append(idx_to_answer_idx[idx])
    return {"predictions": predictions, "references": references}

mcsmoe/data/evaluation.py
UNITES-Lab/MC-SMoE
"""
Functionalities for pre-processing and post-processing of GLUE datasets.
"""
from typing import Dict, List, Union, Any, Optional

import torch
from datasets.arrow_dataset import Dataset
from promptsource.templates import DatasetTemplates
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast

TASK_MAPPING_DATASET_ARGUMENTS = {
    "cola": ["glue", "cola"],
    "stsb": ["glue", "stsb"],
    "rte": ["glue", "rte"],
    "mnli": ["glue", "mnli"],
    "sst2": ["glue", "sst2"],
    "qqp": ["glue", "qqp"],
    "qnli": ["glue", "qnli"],
    "mrpc": ["glue", "mrpc"],
    "multirc": ["super_glue", "multirc"],
    "openbookqa": ["openbookqa", "main"],
    "squad": ["squad"],
    "squad_v2": ["squad_v2"],
    "copa": ["super_glue", "copa"],
    "winogrande": ["winogrande", "winogrande_l"],
    "wikiqa": ["wiki_qa"],
    "triviaqa": ["trivia_qa", "unfiltered"],
    "hotpotqa": ["hotpot_qa", "distractor"],
    "hellaswag": ["hellaswag"]
}

TASK_MAPPING_PROMPT_KEY = {
    "multirc": "found_this_answer",
    # "openbookqa": "pick_using_id",
    "openbookqa": "pick_answer_with_options",
    "sst2": "positive negative after",
    "mrpc": "paraphrase",
    "rte": "does the claim… follow the fact…",
    "squad": "answer_given_context_and_question",
    "squad_v2": "Questions with Context",
    # "copa": "cause_effect",
    "copa": "best_option",
    "qqp": "answer",
    "cola": "Make sense yes no",
    "stsb": "examples",
    "qnli": "based only on",
    "winogrande": "True or False",
    "wikiqa": "Decide_good_answer",
    "hotpotqa": "generate_answer_affirmative",
    "mnli": "GPT-3 style",
    "hellaswag": "Appropriate continuation - Yes or No",
}


def get_label_mapping_id(task: str) -> Dict[str, int]:
    """
    Examples
    --------
    >>> get_label_mapping_id("multirc")
    {'No': 0, 'Yes': 1}
    """
    prompt = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[task])[TASK_MAPPING_PROMPT_KEY[task]]
    choices = prompt.get_fixed_answer_choices_list()
    if task == "winogrande":
        # to minuscule for more balanced `input_ids` length
        choices = [choice.lower() for choice in choices]
    return {choice: i for i, choice in enumerate(choices)}


class Seq2SeqDataPreProcessor:
    """
    Examples
    --------
    >>> from datasets import load_dataset
    >>> proc = Seq2SeqDataPreProcessor("multirc")
    >>> dataset = load_dataset("super_glue", "multirc", split="train[:4]")
    >>> proc(dataset[0]).keys()
    dict_keys(['inputs', 'targets'])
    >>> len(proc(dataset[:2])['inputs'])
    2
    """

    def __init__(self, benchmark: str, keep_specific_keys: List[str] = None):
        self.benchmark = benchmark
        available_prompts = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[self.benchmark])
        self.prompt = available_prompts[TASK_MAPPING_PROMPT_KEY[self.benchmark]]
        self.keep_specific_keys = keep_specific_keys if keep_specific_keys else []

    def __call__(self, examples: Dict[str, Union[List, Any]], batched: Optional[bool] = True) -> Dict[str, List]:
        first_key = list(examples.keys())[0]
        if isinstance(examples[first_key], list) or batched:
            batch_size = len(examples["label"]) if "label" in examples else len(examples[first_key])
            ret = {'inputs': [], 'targets': []}
            for i in range(batch_size):
                result = self.prompt.apply({k: v[i] for k, v in examples.items()})
                ret['inputs'].append(result[0])
                if self.benchmark == "winogrande":
                    ret['targets'].append(result[1].lower())
                else:
                    ret['targets'].append(result[1])
        else:
            result = self.prompt.apply(examples)
            ret = {
                'inputs': result[0],
                'targets': result[1] if self.benchmark != "winogrande" else result[1].lower()
            }
        for key in examples:
            if key not in ret and key in self.keep_specific_keys:
                ret[key] = examples[key]
        return ret


class Seq2SeqZeroShotDataPreProcessor:
    """
    Examples
    --------
    >>> from datasets import load_dataset
    >>> proc = Seq2SeqZeroShotDataPreProcessor("winogrande")
    >>> dataset = load_dataset("winogrande", "winogrande_l", split="train[:4]")
    >>> proc(dataset[:2], batched=True).keys()
    dict_keys(['inputs', 'candidates', 'answer_ids'])
    >>> len(proc(dataset[:2])['inputs'])
    2
    """

    def __init__(self, benchmark: str, keep_specific_keys: List[str] = None):
        self.benchmark = benchmark
        available_prompts = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[self.benchmark])
        self.prompt = available_prompts[TASK_MAPPING_PROMPT_KEY[self.benchmark]]
        self.keep_specific_keys = keep_specific_keys if keep_specific_keys else []

    def __call__(self, examples: Dict[str, Union[List, Any]], batched: Optional[bool] = True) -> Dict[str, List]:
        first_key = list(examples.keys())[0]
        if isinstance(examples[first_key], list) or batched:
            batch_size = len(examples["label"]) if "label" in examples else len(examples[first_key])
            ret = {'inputs': [], 'candidates': [], 'answer_ids': []}
            for i in range(batch_size):
                sample = {k: v[i] for k, v in examples.items()}
                prompt, answer = self.prompt.apply(sample)
                choices_list = self.prompt.get_answer_choices_list(sample)
                answer_idx = choices_list.index(answer)
                ret['inputs'].append(prompt)
                ret['candidates'].append(choices_list)
                ret['answer_ids'].append(answer_idx)
        else:
            raise NotImplementedError("Seq2SeqZeroShotDataPreProcessor only supports batched=True")
        for key in examples:
            if key not in ret and key in self.keep_specific_keys:
                ret[key] = examples[key]
        return ret


class CasualZeroShotDataPreProcessor:
    """
    Examples
    --------
    >>> from datasets import load_dataset
    >>> proc = CasualZeroShotDataPreProcessor("winogrande")
    >>> dataset = load_dataset("winogrande", "winogrande_l", split="train[:4]")
    >>> proc(dataset[:2], batched=True).keys()
    dict_keys(['candidates', 'answer_ids'])
    >>> len(proc(dataset[:2], batched=True)['answer_ids'])
    2
    """

    def __init__(self, benchmark: str, keep_specific_keys: List[str] = None,
                 prompt_key: Optional[str] = None):
        self.benchmark = benchmark
        available_prompts = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[self.benchmark])
        self.prompt = available_prompts[TASK_MAPPING_PROMPT_KEY[self.benchmark]] if prompt_key is None \
            else available_prompts[prompt_key]
        self.keep_specific_keys = keep_specific_keys if keep_specific_keys else []

    def __call__(self, examples: Dict[str, Union[List, Any]], batched: Optional[bool] = True) -> Dict[str, List]:
        assert "idx" in examples, "Call build_index_for_dataset() first."
        if isinstance(examples["idx"], list) or batched:
            batch_size = len(examples["label"]) if "label" in examples else len(examples["idx"])
            ret = {'text': [], 'answer_idx': [], "choice_idx": [], "idx": []}
            for key in examples:
                if key not in ret and key in self.keep_specific_keys and key != "idx":
                    ret[key] = []
            for i in range(batch_size):
                sample = {k: v[i] for k, v in examples.items()}
                prompt, answer = self.prompt.apply(sample)
                choices_list = self.prompt.get_answer_choices_list(sample)
                answer_idx = choices_list.index(answer)
                for j, choice in enumerate(choices_list):
                    ret['text'].append(prompt + "\n" + choice)
                    ret['answer_idx'].append(answer_idx)
                    ret['choice_idx'].append(j)
                    ret['idx'].append(sample['idx'])
                    for key in sample:
                        if key not in ret and key in self.keep_specific_keys and key != "idx":
                            ret[key].append(sample[key])
        else:
            raise NotImplementedError("ZeroShotCasualDataPreProcessor only supports batched=True")
        return ret


class CasualOneShotDataPreProcessor:
    """
    Examples
    --------
    >>> from datasets import load_dataset
    >>> proc = CasualZeroShotDataPreProcessor("winogrande")
    >>> dataset = load_dataset("winogrande", "winogrande_l", split="train[:4]")
    >>> proc(dataset[:2], batched=True).keys()
    dict_keys(['candidates', 'answer_ids'])
    >>> len(proc(dataset[:2], batched=True)['answer_ids'])
    2
    """

    def __init__(
            self,
            benchmark: str,
            example: Dict[str, Any],
            keep_specific_keys: List[str] = None,
            prompt_key: Optional[str] = None
    ):
        self.benchmark = benchmark
        available_prompts = DatasetTemplates(*TASK_MAPPING_DATASET_ARGUMENTS[self.benchmark])
        self.prompt = available_prompts[TASK_MAPPING_PROMPT_KEY[self.benchmark]] if prompt_key is None \
            else available_prompts[prompt_key]
        self.keep_specific_keys = keep_specific_keys if keep_specific_keys else []
        self.example = "\n".join(self.prompt.apply(example))

    def __call__(self, examples: Dict[str, Union[List, Any]], batched: Optional[bool] = True) -> Dict[str, List]:
        assert "idx" in examples, "Call build_index_for_dataset() first."
        if isinstance(examples["idx"], list) or batched:
            batch_size = len(examples["label"]) if "label" in examples else len(examples["idx"])
            ret = {'text': [], 'answer_idx': [], "choice_idx": [], "idx": []}
            for key in examples:
                if key not in ret and key in self.keep_specific_keys and key != "idx":
                    ret[key] = []
            for i in range(batch_size):
                sample = {k: v[i] for k, v in examples.items()}
                prompt, answer = self.prompt.apply(sample)
                choices_list = self.prompt.get_answer_choices_list(sample)
                answer_idx = choices_list.index(answer)
                for j, choice in enumerate(choices_list):
                    ret['text'].append(self.example + "\n\n" + prompt + "\n" + choice)
                    ret['answer_idx'].append(answer_idx)
                    ret['choice_idx'].append(j)
                    ret['idx'].append(sample['idx'])
                    for key in sample:
                        if key not in ret and key in self.keep_specific_keys and key != "idx":
                            ret[key].append(sample[key])
        else:
            raise NotImplementedError("ZeroShotCasualDataPreProcessor only supports batched=True")
        return ret


def keep_only_supporting_facts_in_context_for_hotpotqa(examples: Dict[str, Any]):
    """ This is for fxxking long context in HotpotQA. Now keep only supporting facts in context. ^^ """
    new_context = {
        'title': [],
        'sentences': []
    }
    sup_facts = examples['supporting_facts']
    for title, sent_ids in zip(sup_facts['title'], sup_facts['sent_id']):
        vanilla_index = examples['context']['title'].index(title)
        vanilla_sentences = examples['context']['sentences'][vanilla_index]
        if len(vanilla_sentences) <= sent_ids:
            continue
        if title not in new_context['title']:
            new_context['title'].append(title)
            new_context['sentences'].append([vanilla_sentences[sent_ids]])
        else:
            new_context['sentences'][new_context['title'].index(title)].append(
                vanilla_sentences[sent_ids])
    examples['context'] = new_context
    return examples


def tokenize_seq2seq(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        batch: Dict[str, List],
        keep_other_keys=False
) -> Dict[str, List]:
    inputs = tokenizer(batch.pop("inputs"), truncation=True, return_attention_mask=True)
    targets = tokenizer(batch.pop("targets"), truncation=True, padding=False, return_attention_mask=False)
    labels = targets["input_ids"]
    # Replace pad_token_id 0 to -100 in labels
    labels = [[-100 if token == tokenizer.pad_token_id else token for token in label] for label in labels]
    ret = {"input_ids": inputs["input_ids"],
           "attention_mask": inputs["attention_mask"],
           "labels": labels}
    # This is for some dataset evaluation like "idx" in MultiRC, "id" in SQuAD
    if keep_other_keys:
        for key in batch:
            ret[key] = batch[key]
    return ret


def tokenize_seq2se2_to_casual_lm(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        batch: Dict[str, List],
) -> Dict[str, List]:
    inputs_list = batch.pop("inputs")  # of shape (batch_size)
    targets_list = batch.pop("targets")  # of shape (batch_size)
    model_inputs = tokenizer(
        [_replace_new_line_with_eos(inp + "\n" + tar, tokenizer.eos_token) for inp, tar in
         zip(inputs_list, targets_list)],
        truncation=True, return_attention_mask=True
    )
    ret = {
        "input_ids": model_inputs["input_ids"],
        "attention_mask": model_inputs["attention_mask"],
    }
    return ret


def tokenize_casual_generation(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        batch: Dict[str, List],
        for_eval: Optional[bool] = False,
) -> Dict[str, List]:
    """
    For training dataset:
        Training samples are supposed to be input to model(**batch) for loss calculation.
        - input_ids: "inputs" + "\n" + "targets" + "eos_token"

    For evaluation dataset:
        Evaluation samples are supposed to be input to model.generate() for prediction.
        - inputs: "inputs" + "\n"
        - references: "targets" + "eos_token"

    """
    if for_eval:
        inputs = []
        references = []
        for inp, tar in zip(batch.pop("inputs"), batch.pop("targets")):
            inp += "\n"
            inputs.append(inp)
            references.append(tar + tokenizer.eos_token)
        inputs = tokenizer(inputs, truncation=True, return_attention_mask=True)
        references = tokenizer(references, truncation=True, padding=False, return_attention_mask=False)
        ret = {"input_ids": inputs["input_ids"],
               "attention_mask": inputs["attention_mask"],
               "references": references["input_ids"]}
        if for_eval:
            for key in batch:
                ret[key] = batch[key]
        return ret
    else:
        inputs_targets = []
        for inp, tar in zip(batch.pop("inputs"), batch.pop("targets")):
            inp += "\n"
            inputs_targets.append(inp + tar + tokenizer.eos_token)
        inputs_targets = tokenizer(inputs_targets, truncation=True, padding=False, return_attention_mask=True)
        ret = {"input_ids": inputs_targets["input_ids"],
               "attention_mask": inputs_targets["attention_mask"]}
        return ret


def tokenize_seq2seq_zero_shot(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        batch: Dict[str, List],
) -> Dict[str, List[Union[Dict[str, List[int]], int]]]:
    """
    Returns
    -------
    Dict:
        - "candidates": List[Dict]
            - "input_ids"
            - "attention_mask"
        - "answer_ids": List[int]
    """
    inputs_list = batch.pop("inputs")  # of shape (batch_size)
    candidates_list = batch.pop("candidates")  # of shape (batch_size, num_choices)
    answer_ids_list = batch.pop("answer_ids")  # of shape (batch_size)
    batch_size = len(inputs_list)
    num_choices = len(candidates_list[0])
    candidate_texts = [candidates_list[i // num_choices][i % num_choices] for i in range(batch_size * num_choices)]
    ret = {"input_ids": [],
           "attention_mask": [],
           "candidates": [],
           "answer_ids": answer_ids_list}
    for i in range(batch_size):
        inputs = tokenizer(
            inputs_list[i],
            truncation=True, return_attention_mask=True, return_tensors='pt'
        )
        candidates = tokenizer(
            candidate_texts[i * num_choices: (i + 1) * num_choices],
            truncation=True, padding=True, return_attention_mask=False, return_tensors=None
        )['input_ids']
        candidates = [[-100 if token == tokenizer.pad_token_id else token for token in candidate]
                      for candidate in candidates]
        ret["candidates"].append(torch.tensor(candidates))
        ret["input_ids"].append(inputs['input_ids'].repeat(num_choices, 1))
        ret["attention_mask"].append(inputs['attention_mask'].repeat(num_choices, 1))
    for key in batch:
        # keep other keys
        ret[key] = batch[key]
    return ret


def _replace_new_line_with_eos(text: str, eos_token: str) -> str:
    return text.replace("\n", eos_token)


def tokenize_casual_zero_shot(
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
        batch: Dict[str, List],
        replace_new_line_with_eos: Optional[bool] = True,
        keep_input_only: Optional[bool] = False,
) -> Dict[str, List[Union[Dict[str, List[int]], int]]]:
    texts = batch.pop("text")
    if replace_new_line_with_eos:
        texts = [_replace_new_line_with_eos(text, tokenizer.eos_token) + tokenizer.eos_token for text in texts]
    else:
        texts = [text + tokenizer.eos_token for text in texts]
    texts = tokenizer(texts, truncation=True, return_attention_mask=True)
    ret = {
        "input_ids": texts["input_ids"],
        "attention_mask": texts["attention_mask"],
    }
    if not keep_input_only:
        for key in batch:
            ret[key] = batch[key]
    return ret


def convert_hotpot_to_squad_format(
        json_dict, gold_paras_only=True, combine_context=True
):
    new_dict = {"data": []}
    count = 0
    for example in json_dict:
        raw_contexts = example["context"]
        if gold_paras_only:
            support = {
                para_title: line_num
                for para_title, line_num in example["supporting_facts"].items()
            }
            raw_contexts = [lst for lst in raw_contexts.items() if lst[0] in support]
        contexts = ["".join(lst[1]) for lst in raw_contexts]
        if combine_context:
            contexts = [" ".join(contexts)]
        answer = example["answer"]
        for context in contexts:
            context = add_yes_no(context)
            answer_start = context.index(answer) if answer in context else -1
            new_dict["data"].append(
                create_para_dict(
                    create_example_dict(
                        context=context,
                        answer_start=answer_start,
                        answer=answer,
                        id=str(count),  # SquadExample.__repr__ only accepts type==str
                        is_impossible=(answer_start == -1),
                        question=example["question"],
                    )
                )
            )
            count += 1
    return new_dict


def build_index_for_dataset(
        dataset: Dataset
):
    """ add a key 'idx' to each example in dataset """
    if "idx" in dataset.column_names:
        return dataset
    ids = list(range(len(dataset)))
    dataset = dataset.add_column("idx", ids)
    return dataset

mcsmoe/data/preprocessors.py
UNITES-Lab/MC-SMoE
import math
import os
import os.path
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import T5TokenizerFast, get_scheduler, SwitchTransformersForConditionalGeneration

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.utils.training_utils import freeze_switch_routers_for_finetuning

logger = get_logger(__name__)
logger.setLevel(20)


@dataclass
class SwitchFineTuningConfig:
    task: Optional[str] = "cola"
    num_epochs: Optional[int] = 5
    batch_size: Optional[int] = 32
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 16
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1


def finetune_on_downstream_task(
        checkpoint: Optional[str] = None,
        output_dir: Optional[str] = None,
        task: Optional[str] = "cola",
        num_experts: Optional[int] = 32,
        no_eval_until_epochs: Optional[int] = 0,
        num_eval_steps: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        num_epochs: Optional[int] = 10,
        report: Optional[bool] = True,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 5e-4,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 16,
        save_each_epoch: Optional[bool] = False,
        preprocessing_num_workers: Optional[int] = None,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    if task not in TASK_MAPPING_DATASET_ARGUMENTS:
        raise ValueError(f"task must be one of {list(TASK_MAPPING_DATASET_ARGUMENTS.keys())}")
    if task not in output_dir:
        output_dir = os.path.join(output_dir, task)

    training_config = SwitchFineTuningConfig(
        task=task,
        num_epochs=num_epochs,
        adam_epsilon=1e-8,
        adam_betas=(0.9, 0.98),
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        dropout=0.1,
    )

    checkpoint = f"google/switch-base-{num_experts}" if checkpoint is None else checkpoint
    low_rank_factor_postfix = "-vanilla"

    model = SwitchTransformersForConditionalGeneration.from_pretrained(checkpoint)
    model = freeze_switch_routers_for_finetuning(model)

    tokenizer = T5TokenizerFast.from_pretrained(f"google/switch-base-{num_experts}")

    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        print(f"Number of parameters in Switch Transformers: {model.num_parameters()}")
        if report:
            wandb.init(project="mc-smoe",
                       config={**model.config.__dict__, **training_config.__dict__},
                       name=f"switch-{num_experts}-{low_rank_factor_postfix}-{task}-{per_device_train_batch_size}-bs-{learning_rate}-lr")

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names,
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names,
        )

        tokenized_train_dataset = train_dataset.map(
            lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
            num_proc=preprocessing_num_workers,
            batched=True,
            remove_columns=train_dataset.column_names,
            load_from_cache_file=False
        )
        tokenized_eval_dataset = eval_dataset.map(
            lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
            num_proc=preprocessing_num_workers,
            batched=True,
            remove_columns=eval_dataset.column_names,
            load_from_cache_file=False
        )

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )

    # optimizer
    no_decay = ["bias", "layer_norm.weight", "LayerNorm", "layernorm", "layer_norm", "ln"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_parameters,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset['validation']
    )

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        # model.train()
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                log_dict = {"train_loss": loss.item(),
                            "epoch": completed_steps / num_update_steps_per_epoch,
                            "learning_rate": lr_scheduler.get_lr()}
                wandb.log(log_dict, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)

                # model.train()

            if completed_steps >= max_train_steps:
                break
        if save_each_epoch:
            save_dir = os.path.join(output_dir, f"epoch-{epoch + 1}")
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(save_dir,
                                            is_main_process=accelerator.is_local_main_process,
                                            save_function=accelerator.save)
            if accelerator.is_local_main_process:
                tokenizer.save_pretrained(save_dir)

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(finetune_on_downstream_task)

mcsmoe/finetune-switch-transformers.py
UNITES-Lab/MC-SMoE
import os
import os.path
from typing import Optional, Union, List

import torch
from accelerate.logging import get_logger
from datasets import load_dataset
from evaluate import load
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import GPT2TokenizerFast
from transformers import logging as hf_logging

from mcsmoe.data import (
    CasualZeroShotDataPreProcessor,
    tokenize_casual_zero_shot,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForLanguageModeling,
    build_index_for_dataset,
    gather_predictions_references_by_casual_lm_loss
)
from mcsmoe.merging.utils import load_merged_fsgpt_moe_from_checkpoint
from mcsmoe.pruning.losparse import (
    LoSparseConfig,
    substitute_fsgpt_moe_expert_weights,
    LoSparsePrunerForFSGPT
)

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_error()


def sanitize_pruning_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def post_merging_losparse(
        checkpoint: str,
        output_dir: str,
        task: str = "mrpc",
        final_threshold: Optional[float] = 0.01,
        low_rank_factor: Optional[int] = 16,
        warmup_steps: Optional[int] = 4,
        samples_batch_size: Optional[int] = 32,
        eval_batch_size: Optional[int] = 32,
        compression_layers: Optional[Union[str, List, int]] = None,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, f"rank-{low_rank_factor}")
        os.makedirs(output_dir, exist_ok=True)

    compression_layers = sanitize_pruning_layers(compression_layers)
    try:
        tokenizer = GPT2TokenizerFast.from_pretrained(checkpoint)
    except OSError:
        tokenizer = GPT2TokenizerFast.from_pretrained("Phando/fairseq-moe-15b-bf16")
    config = LoSparseConfig(
        final_threshold=final_threshold,
        warmup_steps=warmup_steps
    )
    # load merge
    model = load_merged_fsgpt_moe_from_checkpoint(
        checkpoint=checkpoint,
    )
    d_model = model.config.hidden_size
    d_ff = model.config.intermediate_size
    parameter_ratio = (d_model + d_ff) * low_rank_factor / (d_model * d_ff)
    model = substitute_fsgpt_moe_expert_weights(
        model=model,
        parameter_ratio=parameter_ratio,
        has_sparse=True,
        use_svd=True,
        compression_layers=compression_layers
    )

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
        keys_to_ignore=["answer_idx", "choice_idx", "idx"]
    )
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset['train'].select(range(512))
    dataset = build_index_for_dataset(dataset)
    dataset = dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_dataset = dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x, keep_input_only=True),
        num_proc=8,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    dataloader = DataLoader(
        tokenized_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=samples_batch_size,
        num_workers=8
    )

    eval_dataset = raw_dataset["validation"]
    eval_dataset = build_index_for_dataset(eval_dataset)
    eval_dataset = eval_dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x),
        num_proc=8,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        batch_size=eval_batch_size,
        num_workers=4,
        collate_fn=data_collator
    )

    pruner = LoSparsePrunerForFSGPT(
        config=config,
        total_step=len(dataloader),
        mask_param_names=['sparse']
    )

    complete_steps = 0
    model.eval()
    for step, batch in enumerate(tqdm(dataloader, desc="LoSparse Pruning")):
        with torch.autocast("cuda"):
            # fake training, no optim
            for key in batch:
                if key in ["answer_idx", "choice_idx", "idx"]:
                    batch.pop(key)
            batch = {k: v.cuda() for k, v in batch.items()}
            outputs = model(**batch)
            outputs.loss.backward()
            threshold, mask_threshold = pruner.update_and_pruning(
                named_parameters_caller=model.named_parameters,
                global_step=complete_steps
            )
            model.zero_grad()
            print(
                f"Step {complete_steps} threshold: {threshold} mask_threshold: {mask_threshold} num-param: {sum([p.count_nonzero() for p in model.parameters()])}")
        complete_steps += 1
    print(f"LoSparse pruned model num params: {model.num_parameters()}")
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)

    # eval
    ids_list = []
    answer_ids_list = []
    choice_ids_list = []
    losses_list = []
    try:
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    except FileNotFoundError:
        print(f"[Evaluation warning] No metric found for task {task}, using accuracy instead.")
        metric = load("accuracy")
    for eval_step, eval_batch in tqdm(enumerate(eval_dataloader),
                                      total=len(eval_dataloader), desc="Evaluating"):
        ids_list += eval_batch.pop("idx")
        answer_ids_list += eval_batch.pop("answer_idx")
        choice_ids_list += eval_batch.pop("choice_idx")
        eval_batch = {k: v.cuda() for k, v in eval_batch.items()}
        with torch.no_grad():
            with torch.autocast("cuda"):
                logits = model(**eval_batch).logits
            # eval probs
            loss_fct = torch.nn.CrossEntropyLoss()
            labels = eval_batch["labels"]
            shift_labels = labels.new_zeros(labels.shape)
            shift_labels[:, :-1] = labels[:, 1:].clone()
            shift_labels[:, -1] = model.config.pad_token_id
            losses = torch.stack([
                loss_fct(logits[i], shift_labels[i].to(logits.device)) for i in range(logits.shape[0])
            ])
            losses_list += losses.tolist()
    predictions_references = gather_predictions_references_by_casual_lm_loss(
        ids_list=ids_list,
        answer_ids_list=answer_ids_list,
        choice_ids_list=choice_ids_list,
        losses_list=losses_list,
    )
    eval_res = metric.compute(predictions=predictions_references["predictions"],
                              references=predictions_references["references"])
    print(f"Task-specific {task} evaluation result: {eval_res}")


if __name__ == "__main__":
    Fire(post_merging_losparse)

mcsmoe/fsgpt-zero-shot-losparse.py
UNITES-Lab/MC-SMoE
import os
import os.path
from typing import Optional, Union, List

import torch
from accelerate.logging import get_logger
from datasets import load_dataset
from evaluate import load
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import GPT2TokenizerFast
from transformers import logging as hf_logging

from mcsmoe.data import (
    CasualZeroShotDataPreProcessor,
    tokenize_casual_zero_shot,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForLanguageModeling,
    build_index_for_dataset,
    gather_predictions_references_by_casual_lm_loss
)
from mcsmoe.merging.grouping_fsgpt import (
    ExpertUsageFrequencyTrackerForFSGPT
)
from mcsmoe.models import (
    FSGPTMoEForCausalLM,
    FSGPTMoEConfig
)
from mcsmoe.pruning import (
    TaskSpecificExpertPrunerForFSGPT,
    transform_fsgpt_model_for_task_specific_pruning
)

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_error()


def sanitize_pruning_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def forward_only_with_task_specific_pruning(
        checkpoint: str,
        output_dir: str,
        task: str = "mrpc",
        num_samples_for_pruning: Optional[int] = 512,
        samples_batch_size: Optional[int] = 32,
        eval_batch_size: Optional[int] = 32,
        threshold_beta: Optional[float] = 0.1,
        minimum_num_experts: Optional[int] = 8,
        pruning_layers: Optional[Union[str, List, int]] = None,
        num_pruning_times: Optional[int] = 4,
        debug: Optional[bool] = False,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, "task-specific-pruning")
        os.makedirs(output_dir, exist_ok=True)

    pruning_layers = sanitize_pruning_layers(pruning_layers)
    tokenizer = GPT2TokenizerFast.from_pretrained(checkpoint)
    if debug:
        config = FSGPTMoEConfig.from_pretrained(
            checkpoint,
            moe_eval_capacity_token_fraction=1.0,
        )
        model = FSGPTMoEForCausalLM(config).bfloat16().cuda()
        print("[Merging]Debug mode, using fake model")
    else:
        model = FSGPTMoEForCausalLM.from_pretrained(
            checkpoint,
            device_map='auto',
            torch_dtype=torch.bfloat16,
            moe_eval_capacity_token_fraction=1.0,
        )
    model = transform_fsgpt_model_for_task_specific_pruning(model)
    tokenizer = GPT2TokenizerFast.from_pretrained(checkpoint)
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
        keys_to_ignore=["answer_idx", "choice_idx", "idx"]
    )
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"].select(range(num_samples_for_pruning))
    dataset = build_index_for_dataset(dataset)
    dataset = dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_dataset = dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x, keep_input_only=True),
        num_proc=8,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    dataloader = DataLoader(
        tokenized_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=samples_batch_size,
        num_workers=8
    )

    eval_dataset = raw_dataset["validation"]
    eval_dataset = build_index_for_dataset(eval_dataset)
    eval_dataset = eval_dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x),
        num_proc=8,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        batch_size=eval_batch_size,
        num_workers=4,
        collate_fn=data_collator
    )

    usage_tracker = ExpertUsageFrequencyTrackerForFSGPT(
        named_parameters_caller=model.named_parameters,
        compute_every_n_steps=1,
        beta=0.6,
        device='cuda'
    )
    task_specifc = TaskSpecificExpertPrunerForFSGPT(
        model_config=model.config,
        threshold_beta=threshold_beta,
        minimum_num_experts=minimum_num_experts,
        pruning_layers=pruning_layers
    )

    complete_steps = 0
    model.eval()
    max_train_steps = len(dataloader)
    pruning_interval_steps = int(max_train_steps / num_pruning_times)
    for step, batch in enumerate(tqdm(dataloader, desc="Task-specific Pruning")):
        with torch.no_grad(), torch.autocast("cuda"):
            # fake training, no optim
            for key in batch:
                if key in ["answer_idx", "choice_idx", "idx"]:
                    batch.pop(key)
            batch = {k: v.cuda() for k, v in batch.items()}
            outputs = model(**batch)
            usage_tracker.step(
                model_outputs=outputs,
                global_step=complete_steps
            )
            if complete_steps % pruning_interval_steps == 0:
                task_specifc.execute_pruning(
                    model=model,
                    usage_frequency_state_dict=usage_tracker.usage_frequency_state_dict(),
                    prune_to_minimum=False
                )
        complete_steps += 1
    task_specifc.execute_pruning(
        model=model,
        usage_frequency_state_dict=usage_tracker.usage_frequency_state_dict(),
        prune_to_minimum=True
    )
    print(f"Task-specific pruned model num params: {model.num_parameters()}")

    ids_list = []
    answer_ids_list = []
    choice_ids_list = []
    losses_list = []
    try:
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    except FileNotFoundError:
        print(f"[Evaluation warning] No metric found for task {task}, using accuracy instead.")
        metric = load("accuracy")
    for eval_step, eval_batch in tqdm(enumerate(eval_dataloader),
                                      total=len(eval_dataloader), desc="Evaluating"):
        ids_list += eval_batch.pop("idx")
        answer_ids_list += eval_batch.pop("answer_idx")
        choice_ids_list += eval_batch.pop("choice_idx")
        eval_batch = {k: v.cuda() for k, v in eval_batch.items()}
        with torch.no_grad():
            with torch.autocast("cuda"):
                logits = model(**eval_batch).logits
            # eval probs
            loss_fct = torch.nn.CrossEntropyLoss()
            labels = eval_batch["labels"]
            shift_labels = labels.new_zeros(labels.shape)
            shift_labels[:, :-1] = labels[:, 1:].clone()
            shift_labels[:, -1] = model.config.pad_token_id
            losses = torch.stack([
                loss_fct(logits[i], shift_labels[i].to(logits.device)) for i in range(logits.shape[0])
            ])
            losses_list += losses.tolist()
    predictions_references = gather_predictions_references_by_casual_lm_loss(
        ids_list=ids_list,
        answer_ids_list=answer_ids_list,
        choice_ids_list=choice_ids_list,
        losses_list=losses_list,
    )
    eval_res = metric.compute(predictions=predictions_references["predictions"],
                              references=predictions_references["references"])
    print(f"Task-specific {task} evaluation result: {eval_res}")


if __name__ == "__main__":
    Fire(forward_only_with_task_specific_pruning)

mcsmoe/fsgpt-zero-shot-with-task-specific-pruning.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.merging.utils import load_merged_switch_transformers_from_checkpoint
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)
from mcsmoe.pruning.losparse import (
    LoSparseConfig,
    substitute_switch_transformers_expert_weights,
    LoSparsePruner
)
from mcsmoe.utils.training_utils import freeze_switch_routers_for_finetuning

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]) -> List[int]:
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


@dataclass
class SwitchDistillationConfig:
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.1
    hd_lambda: Optional[float] = 0.01
    hd_cos_sim: Optional[bool] = False


def distill_downstream_for_losparse(
        output_dir: Optional[str] = None,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        final_threshold: Optional[float] = 0.01,
        low_rank_factor: Optional[int] = 16,
        encoder_compression_layers: Optional[Union[str, List, int]] = None,
        decoder_compression_layers: Optional[Union[str, List, int]] = None,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-4,
        gradient_accumulation_steps: Optional[int] = 1,
        warmup_steps: Optional[int] = 16,
        num_epochs: Optional[int] = 10,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, f"rank-{low_rank_factor}")
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        raise ValueError("student_checkpoint must be specified")

    if encoder_compression_layers is None:
        encoder_compression_layers = list(range(3, 12, 2))  # TODO

    training_config = SwitchDistillationConfig(
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    config = LoSparseConfig(
        final_threshold=final_threshold,
        warmup_steps=warmup_steps
    )
    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = load_merged_switch_transformers_from_checkpoint(
        checkpoint=student_checkpoint
    ) if "merged" in student_checkpoint else HFSwitch.from_pretrained(student_checkpoint)

    d_model = student_model.config.hidden_size
    d_ff = student_model.config.d_ff
    parameter_ratio = (d_model + d_ff) * low_rank_factor / (d_model * d_ff)
    student_model = substitute_switch_transformers_expert_weights(
        student_model,
        parameter_ratio=parameter_ratio,
        has_sparse=True,
        use_svd=True,
        encoder_compression_layers=encoder_compression_layers,
        decoder_compression_layers=decoder_compression_layers
    )
    student_model = freeze_switch_routers_for_finetuning(student_model)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=0.0,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"losparse-{task}-rank-{str(low_rank_factor)}-sparse-{str(final_threshold)}"
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )

    if accelerator.is_local_main_process:
        print(f"Number of training examples: {len(tokenized_train_dataset)}")
        print(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )

    no_weight_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    pruner = LoSparsePruner(
        config=config,
        total_step=max_train_steps,
        mask_param_names=['sparse']
    )

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                task_loss = outputs.task_loss
                accelerator.backward(loss)
                optimizer.step()

                if accelerator.sync_gradients:
                    threshold, mask_threshold = pruner.update_and_pruning(
                        named_parameters_caller=model.module.student.named_parameters, global_step=completed_steps
                    )

                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_threshold": threshold,
                           "train_mask_threshold": mask_threshold,
                           "num_parameters": sum(p.count_nonzero() for p in model.module.student.parameters()),
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                # Notice: Always save the latest as best for pruning
                accelerator.wait_for_everyone()
                unwrapped_model = accelerator.unwrap_model(model)
                unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                                is_main_process=accelerator.is_local_main_process,
                                                save_function=accelerator.save)
                if accelerator.is_local_main_process:
                    tokenizer.save_pretrained(os.path.join(output_dir, "latest"))
                    print("Latest model saved")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)
    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(distill_downstream_for_losparse)

mcsmoe/losparse-downstream.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union, Dict

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_classification_label_index_and_token_ids,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups,
    compute_experts_empirical_fisher_matrix
    # merge_by_groups_decoder_only as merge_by_groups,
)
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)
from mcsmoe.utils.sparsity import compute_weight_stable_rank

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def save_stable_rank_to(
        core_experts: Dict[str, List[int]], state_dict: Dict[str, torch.Tensor], save_dir: str, save_name: str
):
    num_groups = len(core_experts[list(core_experts.keys())[0]])
    wi_stable_rank_dict = {key: torch.zeros(num_groups) for key in core_experts.keys()}
    wo_stable_rank_dict = {key: torch.zeros(num_groups) for key in core_experts.keys()}
    for mlp_name in tqdm(core_experts, desc="Computing stable rank"):
        for i, core_idx in enumerate(core_experts[mlp_name]):
            name = f"{mlp_name}.experts.expert_{core_idx}.wi.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wi_stable_rank_dict[mlp_name][i] = stable_rank
            name = f"{mlp_name}.experts.expert_{core_idx}.wo.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wo_stable_rank_dict[mlp_name][i] = stable_rank
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    torch.save(wi_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wi.pt"))
    torch.save(wo_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wo.pt"))


@dataclass
class SwitchDistillationConfig:
    max_steps: Optional[int] = 10000
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    mlm_probability: Optional[float] = 0.15
    mean_noise_span_length: Optional[int] = 3
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.1
    hd_lambda: Optional[float] = 0.01
    hd_cos_sim: Optional[bool] = False
    next_rank_for_norm: Optional[int] = 384


def merge_and_distill_downstream_for_recover(
        output_dir: Optional[str] = None,
        save_stable_rank: Optional[bool] = False,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        fisher_state_dict_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        merging_strategy: Optional[str] = "average",
        num_samples_for_merging: Optional[int] = 32,
        similarity_base: Optional[str] = "weight",
        similarity_fn: Optional[str] = "cosine",
        num_groups: Optional[int] = 16,
        globally_group: Optional[bool] = False,
        permute_when_merge: Optional[bool] = False,
        encoder_merging_layers: Optional[Union[str, List, int]] = None,
        decoder_merging_layers: Optional[Union[str, List, int]] = None,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-3,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 100,
        num_epochs: Optional[int] = 10,
        max_train_steps: Optional[int] = 20000,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, merging_strategy, similarity_base)
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        student_checkpoint = teacher_checkpoint

    encoder_merging_layers = sanitize_merging_layers(encoder_merging_layers)
    decoder_merging_layers = sanitize_merging_layers(decoder_merging_layers)

    training_config = SwitchDistillationConfig(
        max_steps=max_train_steps,
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        mlm_probability=0.15,
        mean_noise_span_length=3,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = HFSwitch.from_pretrained(student_checkpoint)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=training_config.reg_lambda,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"merge-{task}-global-{merging_strategy}-{similarity_base}" if globally_group else (
            f"merge-{task}-uniform-{merging_strategy}-{similarity_base}")
        if permute_when_merge:
            run_name += "-permuted"
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    # random select a subset of training dataset for merging
    dataset_for_merging = tokenized_train_dataset.shuffle(seed=2333).select(range(num_samples_for_merging))

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    dataloader_for_merging = DataLoader(
        dataset_for_merging,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=num_samples_for_merging,
        num_workers=4
    )

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )

    if merging_strategy == "fisher" and fisher_state_dict_checkpoint is not None:
        print(f"[Merging]Loading fisher state dict from {fisher_state_dict_checkpoint}")
        fisher_state_dict = torch.load(fisher_state_dict_checkpoint)
    elif merging_strategy == "fisher":
        classification_token_index, classification_label_token_ids = get_classification_label_index_and_token_ids(
            tokenizer=tokenizer,
            task=task
        )
        fisher_dataloader = DataLoader(
            dataset_for_merging,
            shuffle=False,
            collate_fn=data_collator,
            batch_size=1,
            num_workers=4
        )
        fisher_state_dict = compute_experts_empirical_fisher_matrix(
            model=model.student,
            dataloader=fisher_dataloader,
            no_fill=False
        )["experts_fisher_state_dict"]
    else:
        fisher_state_dict = None

    print(f"Merging with {merging_strategy} strategy")
    print(f"Similarities are computed based on {similarity_base} and {similarity_fn}")
    print(f"Number of groups: {num_groups}")
    print(f"Merge encoder layers: {encoder_merging_layers}")
    print(f"Merge decoder layers: {decoder_merging_layers}")

    grouper = ExpertsGrouperForSwitch(
        config=model.student.config,
        similarity_fn=similarity_fn,
        similarity_base=similarity_base
    )
    grouper.compute_all_similarities(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    grouper.compute_all_usages(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    grouper.group_experts_into_clusters_by_spectral_clustering(
        num_groups=num_groups
    )
    print(grouper.group_state_dict())
    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"before-merging-{num_groups}"
        )
    model.student = merge_by_groups(
        model=model.student,
        grouper=grouper,
        strategy=merging_strategy,
        fisher_state_dict=fisher_state_dict,
        encoder_merging_layers=encoder_merging_layers,
        decoder_merging_layers=decoder_merging_layers,
        permute=permute_when_merge,
    )

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-merging-{num_groups}"
        )

    print(f"Number of parameters before merging: {model.teacher.num_parameters()}")
    print(f"Number of parameters after merging: {model.student.num_parameters()}")

    no_weight_decay = ["bias", "LayerNorm", "mat_u", "mat_v"]  # no weight decay for orthogonalization matrices
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                hd_loss = outputs.hd_loss
                task_loss = outputs.task_loss
                norm_loss = outputs.norm_loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_hd_loss": hd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_norm_loss": norm_loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task == "squad":
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)
                # model.student.train()

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=unwrapped_model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-training-{num_groups}"
        )

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(merge_and_distill_downstream_for_recover)

mcsmoe/merge-and-kd-for-recover-downstream.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union, Dict

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_classification_label_index_and_token_ids,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups,
    # merge_by_groups_decoder_only as merge_by_groups,
)
from mcsmoe.merging.fisher import (
    compute_experts_exact_fisher_matrix_for_classification,
    compute_experts_empirical_fisher_matrix
)
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)
from mcsmoe.utils.sparsity import compute_weight_stable_rank

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def save_stable_rank_to(
        core_experts: Dict[str, List[int]], state_dict: Dict[str, torch.Tensor], save_dir: str, save_name: str
):
    wi_stable_rank_dict = {key: [] for key in core_experts.keys()}
    wo_stable_rank_dict = {key: [] for key in core_experts.keys()}
    for mlp_name in tqdm(core_experts, desc="Computing stable rank"):
        for i, core_idx in enumerate(core_experts[mlp_name]):
            name = f"{mlp_name}.experts.expert_{core_idx}.wi.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wi_stable_rank_dict[mlp_name].append(stable_rank)
            name = f"{mlp_name}.experts.expert_{core_idx}.wo.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wo_stable_rank_dict[mlp_name].append(stable_rank)
        wi_stable_rank_dict[mlp_name] = torch.tensor(wi_stable_rank_dict[mlp_name])
        wo_stable_rank_dict[mlp_name] = torch.tensor(wo_stable_rank_dict[mlp_name])
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    torch.save(wi_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wi.pt"))
    torch.save(wo_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wo.pt"))


@dataclass
class SwitchDistillationConfig:
    max_steps: Optional[int] = 10000
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    mlm_probability: Optional[float] = 0.15
    mean_noise_span_length: Optional[int] = 3
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.1
    hd_lambda: Optional[float] = 0.01
    hd_cos_sim: Optional[bool] = False
    next_rank_for_norm: Optional[int] = 384


def merge_and_distill_downstream_for_recover(
        output_dir: Optional[str] = None,
        save_stable_rank: Optional[bool] = False,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        fisher_state_dict_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        merging_strategy: Optional[str] = "average",
        num_samples_for_merging: Optional[int] = 32,
        similarity_base: Optional[str] = "weight",
        similarity_fn: Optional[str] = "cosine",
        num_groups: Optional[int] = 16,
        globally_group: Optional[bool] = False,
        permute_when_merge: Optional[bool] = False,
        encoder_merging_layers: Optional[Union[str, List, int]] = None,
        decoder_merging_layers: Optional[Union[str, List, int]] = None,
        exact_fisher: Optional[bool] = False,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-3,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 100,
        num_epochs: Optional[int] = 10,
        max_train_steps: Optional[int] = 20000,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, merging_strategy, similarity_base)
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        student_checkpoint = teacher_checkpoint
    if task in ['squad', 'openbookqs']:
        exact_fisher = False
        print(f"[WARNING]exact_fisher is set to False for {task} task")

    encoder_merging_layers = sanitize_merging_layers(encoder_merging_layers)
    decoder_merging_layers = sanitize_merging_layers(decoder_merging_layers)

    training_config = SwitchDistillationConfig(
        max_steps=max_train_steps,
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        mlm_probability=0.15,
        mean_noise_span_length=3,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = HFSwitch.from_pretrained(student_checkpoint)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=training_config.reg_lambda,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"merge-{task}-global-dense-{merging_strategy}-{similarity_base}" if globally_group else (
            f"merge-{task}-uniform-dense-{merging_strategy}-{similarity_base}")
        if permute_when_merge:
            run_name += "-permuted"
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    # random select a subset of training dataset for merging
    dataset_for_merging = tokenized_train_dataset.shuffle(seed=2333).select(range(num_samples_for_merging))

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    dataloader_for_merging = DataLoader(
        dataset_for_merging,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=num_samples_for_merging,
        num_workers=4
    )

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )

    if merging_strategy == "fisher" and fisher_state_dict_checkpoint is not None:
        print(f"[Merging]Loading fisher state dict from {fisher_state_dict_checkpoint}")
        fisher_state_dict = torch.load(fisher_state_dict_checkpoint)
    elif merging_strategy == "fisher":
        classification_token_index, classification_label_token_ids = get_classification_label_index_and_token_ids(
            tokenizer=tokenizer,
            task=task
        )
        fisher_dataloader = DataLoader(
            dataset_for_merging,
            shuffle=False,
            collate_fn=data_collator,
            batch_size=1,
            num_workers=4
        )
        fisher_state_dict = compute_experts_exact_fisher_matrix_for_classification(
            model=model.student,
            dataloader=fisher_dataloader,
            classification_token_index=classification_token_index,
            classification_label_token_ids=classification_label_token_ids,
            no_fill=False
        )["experts_fisher_state_dict"] if exact_fisher else compute_experts_empirical_fisher_matrix(
            model=model.student,
            dataloader=fisher_dataloader,
            no_fill=False
        )["experts_fisher_state_dict"]
    else:
        fisher_state_dict = None

    print(f"Merging with {merging_strategy} strategy")
    print(f"Similarities are computed based on {similarity_base} and {similarity_fn}")
    print(f"Number of groups: {num_groups}")
    print(f"Merge encoder layers: {encoder_merging_layers}")
    print(f"Merge decoder layers: {decoder_merging_layers}")

    grouper = ExpertsGrouperForSwitch(
        config=model.student.config,
        similarity_fn=similarity_fn,
        similarity_base=similarity_base
    )
    grouper.compute_all_similarities(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    grouper.compute_all_usages(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=num_groups,
        merging_encoder_layers=encoder_merging_layers,
        merging_decoder_layers=decoder_merging_layers,
    ) if globally_group else grouper.group_experts_into_clusters_by_routing_guided(
        num_groups=num_groups,
    )
    print(grouper.group_state_dict())
    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"before-merging-{num_groups}"
        )
    torch.save(grouper.group_state_dict(), os.path.join(output_dir, "group_state_dict.pt"))
    if globally_group:
        torch.save(core_experts, os.path.join(output_dir, "core_experts.pt"))
    model.student = merge_by_groups(
        model=model.student,
        grouper=grouper,
        strategy=merging_strategy,
        fisher_state_dict=fisher_state_dict,
        encoder_merging_layers=encoder_merging_layers,
        decoder_merging_layers=decoder_merging_layers,
        permute=permute_when_merge,
    )

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-merging-{num_groups}"
        )

    print(f"Number of parameters before merging: {model.teacher.num_parameters()}")
    print(f"Number of parameters after merging: {model.student.num_parameters()}")

    no_weight_decay = ["bias", "LayerNorm", "mat_u", "mat_v"]  # no weight decay for orthogonalization matrices
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                hd_loss = outputs.hd_loss
                task_loss = outputs.task_loss
                norm_loss = outputs.norm_loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_hd_loss": hd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_norm_loss": norm_loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task == "squad":
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)
                # model.student.train()

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=unwrapped_model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-training-{num_groups}"
        )

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(merge_and_distill_downstream_for_recover)

mcsmoe/merge-by-averaging.py
UNITES-Lab/MC-SMoE
import os
import resource
from typing import Optional, Union, List

import torch
import transformers.utils.logging as hf_logging
from accelerate.utils import set_seed
from datasets import load_dataset
from evaluate import load
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    AutoTokenizer
)

from mcsmoe.data import (
    CasualZeroShotDataPreProcessor,
    tokenize_casual_zero_shot,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForLanguageModeling,
    build_index_for_dataset,
    gather_predictions_references_by_casual_lm_loss
)
from mcsmoe.merging.grouping_fsgpt import (
    ExpertsGrouperForFSGPT,
    merge_by_groups_with_usage_frequency_weighting,
)
from mcsmoe.models import FSGPTMoEForCausalLM, FSGPTMoEConfig

set_seed(42)
hf_logging.set_verbosity_info()
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))


def _sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def merge_fsgpt_by_usage_frequency_weighting(
        output_dir: Optional[str] = None,
        checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        num_samples_for_merging: Optional[int] = 512,
        samples_batch_size: Optional[int] = 32,
        eval_batch_size: Optional[int] = 32,
        similarity_base: Optional[str] = "router-logits",
        similarity_fn: Optional[str] = "cosine",
        num_groups: Optional[int] = 16,
        merging_layers: Optional[Union[str, List, int]] = None,
        debug: Optional[bool] = False,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, similarity_base)
        os.makedirs(output_dir, exist_ok=True)
    merging_layers = _sanitize_merging_layers(merging_layers)

    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    if debug:
        config = FSGPTMoEConfig.from_pretrained(
            checkpoint,
            moe_eval_capacity_token_fraction=1.0,
        )
        model = FSGPTMoEForCausalLM(config).bfloat16().cuda()
        print("[Merging]Debug mode, using fake model")
    else:
        model = FSGPTMoEForCausalLM.from_pretrained(
            checkpoint,
            device_map='auto',
            torch_dtype=torch.bfloat16,
            moe_eval_capacity_token_fraction=1.0,
        )
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
        keys_to_ignore=["answer_idx", "choice_idx", "idx"]
    )
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"].select(range(min(num_samples_for_merging, len(raw_dataset["train"]))))
    dataset = build_index_for_dataset(dataset)
    dataset = dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_dataset = dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x, keep_input_only=True),
        num_proc=8,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    dataloader = DataLoader(
        tokenized_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=samples_batch_size,
        num_workers=8
    )

    grouper = ExpertsGrouperForFSGPT(
        config=model.config,
        similarity_base=similarity_base,
        similarity_fn=similarity_fn,
    )
    grouper.compute_all_similarities(
        model=model,
        dataloader=dataloader,
    )
    grouper.compute_all_usages(
        model=model,
        dataloader=dataloader,
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=num_groups,
        merging_layers=merging_layers,
    )
    print(f"[Merging]Core experts: {core_experts}")
    torch.save(grouper.group_state_dict(), os.path.join(output_dir, "group_state_dict.pt"))
    torch.save(core_experts, os.path.join(output_dir, "core_experts.pt"))
    print(f"[Merging]Number of parameters before merging: {model.num_parameters()}")
    model = merge_by_groups_with_usage_frequency_weighting(
        model=model,
        grouper=grouper,
        merging_layers=merging_layers,
    )
    print(f"[Merging]Number of parameters after merging: {model.num_parameters()}")
    model.save_pretrained(output_dir)

    # ========= evaluate downstream zero-shot =========
    eval_dataset = raw_dataset["validation"]
    eval_dataset = build_index_for_dataset(eval_dataset)
    eval_dataset = eval_dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x),
        num_proc=8,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        batch_size=eval_batch_size,
        num_workers=4,
        collate_fn=data_collator
    )

    model.eval()
    ids_list = []
    answer_ids_list = []
    choice_ids_list = []
    losses_list = []
    try:
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    except FileNotFoundError:
        print(f"[Evaluation warning] No metric found for task {task}, using accuracy instead.")
        metric = load("accuracy")
    for eval_step, eval_batch in tqdm(enumerate(eval_dataloader),
                                      total=len(eval_dataloader), desc="Evaluating"):
        ids_list += eval_batch.pop("idx")
        answer_ids_list += eval_batch.pop("answer_idx")
        choice_ids_list += eval_batch.pop("choice_idx")
        eval_batch = {k: v.cuda() for k, v in eval_batch.items()}
        with torch.no_grad():
            with torch.autocast("cuda"):
                logits = model(**eval_batch).logits
            # eval probs
            loss_fct = torch.nn.CrossEntropyLoss()
            labels = eval_batch["labels"]
            shift_labels = labels.new_zeros(labels.shape)
            shift_labels[:, :-1] = labels[:, 1:].clone()
            shift_labels[:, -1] = model.config.pad_token_id
            losses = torch.stack([
                loss_fct(logits[i], shift_labels[i].to(logits.device)) for i in range(logits.shape[0])
            ])
            losses_list += losses.tolist()
    predictions_references = gather_predictions_references_by_casual_lm_loss(
        ids_list=ids_list,
        answer_ids_list=answer_ids_list,
        choice_ids_list=choice_ids_list,
        losses_list=losses_list,
    )
    eval_res = metric.compute(predictions=predictions_references["predictions"],
                              references=predictions_references["references"])
    print(f"{task} evaluation result: {eval_res}")


if __name__ == '__main__':
    Fire(merge_fsgpt_by_usage_frequency_weighting)

mcsmoe/merge-fsgpt-by-usage-frequency-weighted.py
UNITES-Lab/MC-SMoE
from .fisher import (
    compute_experts_exact_fisher_matrix_for_classification,
    compute_experts_empirical_fisher_matrix,
    fisher_merge_layer_experts,
    ExpertFisherTracker
)
from .grouping import (
    ExpertsGrouperForSwitch,
    merge_by_groups,
    LEGAL_SIMILARITY_BASES,
    ExpertUsageFrequencyTracker,
    prune_non_core_experts_by_groups,
    merge_by_groups_with_usage_frequency_weighting,
    merge_non_core_by_groups_with_usage_frequency_weighting,
    merge_all_non_core_to_one_with_usage_frequency_weighting,
    merge_by_groups_within_and_across_models,
    merge_non_core_by_groups_within_and_across_models,
    merge_all_non_core_to_one_within_and_across_models
)
from .utils import load_merged_switch_transformers_from_checkpoint

mcsmoe/merging/__init__.py
UNITES-Lab/MC-SMoE
from collections import OrderedDict
from copy import deepcopy
from typing import Dict, List, Any, Tuple, Iterator, Callable, Optional

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import PreTrainedModel
from transformers.models.switch_transformers.modeling_switch_transformers import (
    SwitchTransformersDenseActDense as HFDenseActDense,
    SwitchTransformersSparseMLP,
    SwitchTransformersConfig
)

__all__ = [
    "compute_experts_exact_fisher_matrix_for_classification",
    "compute_experts_empirical_fisher_matrix",
    "fisher_merge_layer_experts",
    "ExpertFisherTracker",
]


def compute_experts_exact_fisher_matrix_for_classification(
        model: PreTrainedModel,
        dataloader: DataLoader,
        classification_token_index: int,
        classification_label_token_ids: List[int],
        no_fill: bool,
) -> Dict[str, Any]:
    """
    Given a model and a dataloader, compute the Fisher information matrix of the model.
    Note:
        1. Here we only estimate the diagonal of the Fisher matrix
        2. Since T5 is a seq2seq model, we specially design the prompt by keeping only one classification token
            different in the output sequence. Typical examples are "yes" vs "no", "A" vs "B" vs "C", etc.
        3. Thus, the classification probability of the model is the softmax of the logits of the prediction at the
            classification token index.
    """
    if no_fill:
        return {}
    if dataloader.batch_size != 1:
        raise ValueError("The batch size of the dataloader must be 1")
    if classification_label_token_ids is None:
        raise ValueError("The classification label token ids list must be provided")

    model.eval()
    model.cuda()
    num_labels = len(classification_label_token_ids)
    print(
        f"Computing the Fisher information matrix of the model with {num_labels} labels and {len(dataloader)} samples"
    )
    experts_fisher_state_dict = {}
    for name, param in model.named_parameters():
        if 'expert' in name:
            experts_fisher_state_dict[name] = torch.zeros_like(param.data, requires_grad=False, device='cuda')

    for batch in tqdm(dataloader, desc="[Fisher]Running through the dataloader"):
        # forward pass
        model.zero_grad()
        batch = {k: v.to('cuda') for k, v in batch.items()}
        logits = model(**batch).logits  # shape (1, seq_len, vocab_size)
        logits = logits.squeeze(dim=0)[classification_token_index, classification_label_token_ids]  # shape (num_labels)
        log_probs = torch.log_softmax(logits, dim=-1)
        probs = torch.softmax(logits, dim=-1)

        # compute the diagonal of the Fisher matrix
        for i in range(num_labels):
            log_probs[i].backward(retain_graph=True)
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if 'expert_' in name:
                        experts_fisher_state_dict[name] += ((param.grad.data ** 2) * probs[i])
            model.zero_grad()

    experts_fisher_state_dict = {name: fisher.cpu() / len(dataloader)
                                 for name, fisher in experts_fisher_state_dict.items()}
    model.cpu()
    return {"experts_fisher_state_dict": experts_fisher_state_dict}


def compute_experts_empirical_fisher_matrix(
        model: PreTrainedModel,
        dataloader: DataLoader,
        no_fill: bool,
) -> Dict[str, Any]:
    """
    Similar to `compute_experts_fisher_matrix_for_classification`, but we compute the empirical Fisher matrix.
    This is useful for NLG/QA tasks which have an exponential explosion of the number of classes.
    """
    if no_fill:
        return {}
    if dataloader.batch_size != 1:
        raise ValueError("The batch size of the dataloader must be 1")

    model.eval()
    model.cuda()
    print(
        f"Computing the Empirical Fisher information matrix of the model with {len(dataloader)} samples"
    )
    experts_fisher_state_dict = {}
    for name, param in model.named_parameters():
        if 'expert' in name:
            experts_fisher_state_dict[name] = torch.zeros_like(param.data, requires_grad=False, device='cuda')

    for batch in tqdm(dataloader, desc=f"[Empirical Fisher]Running through {len(dataloader)} samples"):
        # forward pass
        model.zero_grad()
        batch = {k: v.to('cuda') for k, v in batch.items()}
        outputs = model(**batch)
        outputs.loss.backward(retain_graph=True)
        with torch.no_grad():
            for name, param in model.named_parameters():
                if 'expert_' in name:
                    experts_fisher_state_dict[name] += (param.grad.data ** 2)

    experts_fisher_state_dict = {name: fisher.cpu() / len(dataloader)
                                 for name, fisher in experts_fisher_state_dict.items()}
    model.cpu()
    return {"experts_fisher_state_dict": experts_fisher_state_dict}


def fisher_merge_layer_experts(
        sparse_mlp: SwitchTransformersSparseMLP,
        config: SwitchTransformersConfig,
        name_prefix: str,
        experts_fisher_state_dict: Dict[str, torch.Tensor] = None,
        no_fill: bool = False,
) -> HFDenseActDense:
    """
    Merge experts in a SwitchTransformersSparseMLP into one MLP of SwitchTransformersDenseActDense
     by fisher information matrix, requires the modules to retain gradients
    """
    dense = HFDenseActDense(config)
    if no_fill:
        return dense

    wi_weight_list = []
    wo_weight_list = []
    wi_fisher_list = []
    wo_fisher_list = []
    num_experts = config.num_experts
    for expert_idx in range(num_experts):
        wi_fisher_list.append(experts_fisher_state_dict[f"{name_prefix}.expert_{expert_idx}.wi.weight"])
        wo_fisher_list.append(experts_fisher_state_dict[f"{name_prefix}.expert_{expert_idx}.wo.weight"])
        wi_weight_list.append(sparse_mlp.experts[f"expert_{expert_idx}"].wi.weight)
        wo_weight_list.append(sparse_mlp.experts[f"expert_{expert_idx}"].wo.weight)
    wi_fisher_list = torch.stack(wi_fisher_list, dim=0)
    wo_fisher_list = torch.stack(wo_fisher_list, dim=0)
    wi_weight_list = torch.stack(wi_weight_list, dim=0)
    wo_weight_list = torch.stack(wo_weight_list, dim=0)

    # nan check
    wi_weight = torch.sum(wi_weight_list * wi_fisher_list, dim=0) / (torch.sum(wi_fisher_list, dim=0) + 1e-6)
    wo_weight = torch.sum(wo_weight_list * wo_fisher_list, dim=0) / (torch.sum(wo_fisher_list, dim=0) + 1e-6)

    with torch.no_grad():
        dense.wi.weight.copy_(wi_weight)
        dense.wo.weight.copy_(wo_weight)
    return dense


class ExpertFisherTracker(object):
    """
    A class to track the Empirical Fisher information matrix of the experts during training,
        using exponential moving average


    Examples
    --------
    # 1. define the tracker
    >>> tracker = ExpertFisherTracker(named_parameters_caller=model.named_parameters,
    >>>                               begin_step=0,
    >>>                               end_step=len(train_dataloader),
    >>>                               compute_fisher_every_n_steps=1,
    >>>                               beta=0.9,
    >>>                               device='cpu')
    # --- inside the training loop ---
    # 2. compute the fisher matrix at the step `complete_steps`
    >>> fisher_state_dict = compute_experts_empirical_fisher_matrix(model, fisher_dataloader)
    >>> tracker.step(fisher_state_dict, complete_steps)
    """

    def __init__(
            self,
            named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]],
            begin_step: int,
            end_step: int,
            compute_fisher_every_n_steps: Optional[int] = 1,
            beta: Optional[float] = 0.9,
            device: Optional[str] = 'cpu',
    ):
        self.beta = beta
        self.begin_step = begin_step
        self.end_step = end_step
        self.device = device
        self.compute_fisher_every_n_steps = compute_fisher_every_n_steps
        self.exp_experts_fisher_dict = None
        self.last_error = -1

        self._init_fisher_state(named_parameters_caller)

    def fisher_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self.exp_experts_fisher_dict)

    def _init_fisher_state(self, named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]]):
        self.exp_experts_fisher_dict = OrderedDict()
        for name, param in named_parameters_caller():
            if "expert_" in name:
                self.exp_experts_fisher_dict[name] = torch.zeros_like(param.data,
                                                                      requires_grad=False, device=self.device)

    def _update_exp_fisher_state(self, fisher_state_dict: Dict[str, torch.Tensor]):
        is_first_step = self.last_error < 0
        fisher_state_dict = {k: v.to(self.device) for k, v in fisher_state_dict.items()}
        self.last_error = sum(
            [torch.sum(torch.abs(self.exp_experts_fisher_dict[name] - fisher_state_dict[name])).item()
             for name in self.exp_experts_fisher_dict.keys()]
        )
        for name in self.exp_experts_fisher_dict.keys():
            if is_first_step:
                self.exp_experts_fisher_dict[name] = fisher_state_dict[name].detach()
            else:
                self.exp_experts_fisher_dict[name] = self.beta * self.exp_experts_fisher_dict[name] + \
                                                     (1 - self.beta) * fisher_state_dict[name]

    def step(self, fisher_state_dict: Dict[str, torch.Tensor], global_step: int) -> float:
        if (self.begin_step <= global_step <= self.end_step) and (
                global_step % self.compute_fisher_every_n_steps == 0):
            self._update_exp_fisher_state(fisher_state_dict)

        return self.last_error

mcsmoe/merging/fisher.py
UNITES-Lab/MC-SMoE
import os
from collections import OrderedDict
from copy import deepcopy
from typing import Dict, Optional, List, Union, Callable, Iterator, Tuple

import torch
from scipy.optimize import linear_sum_assignment
from sklearn.cluster import SpectralClustering
from torch.nn import functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    SwitchTransformersConfig,
    SwitchTransformersForConditionalGeneration,
    SwitchTransformersSparseMLP,
    SwitchTransformersTop1Router,
    PretrainedConfig
)
from transformers.modeling_outputs import Seq2SeqMoEOutput
from transformers.models.switch_transformers.modeling_switch_transformers import SwitchTransformersDenseActDense

from .permutation import (
    permute_switch_mlp_dense_expert_,
    compute_switch_permutation_by_weight_matching,
    compute_switch_permutation_by_activation_matching,
    merge_switch_mlp_by_activation_matching_within_and_across_models,
    merge_switch_mlp_by_weight_matching_within_and_across_models
)
from .utils import generate_random_group_labels
from ..utils.constants import FP32_EPS

__all__ = [
    'ExpertsGrouperForSwitch',
    'merge_by_groups',
    'LEGAL_SIMILARITY_BASES',
    'SIMILARITY_MAPPING_FUNCTION',
    'ExpertUsageFrequencyTracker',
    'prune_non_core_experts_by_groups',
    'merge_by_groups_with_usage_frequency_weighting',
    'merge_non_core_by_groups_with_usage_frequency_weighting',
    'merge_all_non_core_to_one_with_usage_frequency_weighting',
    'merge_by_groups_within_and_across_models',
    'merge_non_core_by_groups_within_and_across_models',
    'merge_all_non_core_to_one_within_and_across_models',
]

SIMILARITY_MAPPING_FUNCTION = {
    "cosine": lambda x, y: (F.cosine_similarity(x, y, dim=-1, eps=FP32_EPS) + 1) / 2,
    "mse": lambda x, y: 1 / (1 + 0.1 * torch.log(F.mse_loss(x, y, reduction="sum"))),
}
LEGAL_SIMILARITY_BASES = ["weight", "feature", "feature.abs", "weight-feature", "gradient", "weight-gradient",
                          "router-logits", "router-weight", "router-weight-feature", "mse", "random",
                          "feature-correlation.lsa", "feature-correlation.max"]


class ExpertsGrouperForSwitch(object):
    def __init__(
            self,
            config: Union[SwitchTransformersConfig, PretrainedConfig],
            similarity_fn: str = "cosine",
            similarity_base: str = "weight",
    ):
        if similarity_fn not in SIMILARITY_MAPPING_FUNCTION:
            raise ValueError(
                f"[Merging]similarity_fn should be one of {SIMILARITY_MAPPING_FUNCTION.keys()}, got {similarity_fn} instead.")
        if similarity_base not in LEGAL_SIMILARITY_BASES:
            raise ValueError(
                f"[Merging]similarity_base should be one of {LEGAL_SIMILARITY_BASES}, got {similarity_base} instead.")

        self.num_experts = config.num_experts
        self.d_model = config.d_model
        self.sparse_layer_indices = list(range(1, config.num_layers, config.encoder_sparse_step))
        self.similarity_fn = SIMILARITY_MAPPING_FUNCTION[similarity_fn]
        self.similarity_base = similarity_base
        self._group_state_dict = None
        self._similarity_state_dict = None
        self._usage_frequency_state_dict = None
        self.reset_all()

    def reset_all(self):
        if self.similarity_base == "mse":
            self.similarity_fn = SIMILARITY_MAPPING_FUNCTION["mse"]
            print("[Merging]Set similarity_fn to mse for mse similarity_base.")
        self._group_state_dict = dict()
        self._similarity_state_dict = dict()
        self._usage_frequency_state_dict = dict()
        # Similarity range: [0, 2]
        for layer_idx in self.sparse_layer_indices:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            self._group_state_dict[encoder_mlp_name] = torch.arange(self.num_experts,
                                                                    device="cuda")
            self._group_state_dict[decoder_mlp_name] = torch.arange(self.num_experts,
                                                                    device="cuda")
            self._similarity_state_dict[encoder_mlp_name] = torch.zeros(
                (self.num_experts, self.num_experts), device="cuda"
            ) + torch.eye(self.num_experts, device="cuda")
            self._similarity_state_dict[decoder_mlp_name] = torch.zeros(
                (self.num_experts, self.num_experts), device="cuda"
            ) + torch.eye(self.num_experts, device="cuda")
            self._usage_frequency_state_dict[encoder_mlp_name] = torch.zeros(self.num_experts, device="cuda")
            self._usage_frequency_state_dict[decoder_mlp_name] = torch.zeros(self.num_experts, device="cuda")

    def similarity_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self._similarity_state_dict)

    def group_state_dict(self) -> Dict[str, torch.LongTensor]:
        return deepcopy(self._group_state_dict)

    def usage_frequency_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self._usage_frequency_state_dict)

    def save_similarity(self, mlp_name: str, i: int, j: int, similarity: float):
        self._similarity_state_dict[mlp_name][i, j] = similarity
        self._similarity_state_dict[mlp_name][j, i] = similarity

    def get_similarity(self, mlp_name: str, i: int, j: int) -> float:
        return self._similarity_state_dict[mlp_name][i, j].item()

    def get_similarity_matrix(self, mlp_name: str) -> torch.Tensor:
        return deepcopy(self._similarity_state_dict[mlp_name])

    def save_group_state_dict(self, save_dir: str):
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        torch.save(self._group_state_dict, os.path.join(save_dir, "group_state_dict.pt"))

    def load_group_state_dict(self, load_dir: str):
        self._group_state_dict = torch.load(os.path.join(load_dir, "group_state_dict.pt"))

    @staticmethod
    def defend_core_experts_fisher(
            fisher_state_dict: Dict[str, torch.Tensor],
            core_experts: Dict[str, List[int]],
    ) -> Dict[str, torch.Tensor]:
        """TODO: This is a unfinished solution to avoid zero fisher in core experts."""
        for mlp_name, core_expert_list in core_experts.items():
            for i in core_expert_list:
                expert_name = f"{mlp_name}.experts.expert_{i}.wi.weight"
                if fisher_state_dict[expert_name].abs().sum() < FP32_EPS:
                    fisher_state_dict[expert_name] = torch.ones_like(fisher_state_dict[expert_name]) * FP32_EPS
                expert_name = f"{mlp_name}.experts.expert_{i}.wo.weight"
                if fisher_state_dict[expert_name].abs().sum() < FP32_EPS:
                    fisher_state_dict[expert_name] = torch.ones_like(fisher_state_dict[expert_name]) * FP32_EPS
        return fisher_state_dict

    def _assign_num_groups_per_layer(
            self,
            average_num_groups: int,
            merging_encoder_layers: List[int],
            merging_decoder_layers: List[int],
    ) -> Dict[str, int]:
        num_grouping_layers = len(merging_encoder_layers) + len(merging_decoder_layers)
        total_num_groups = average_num_groups * num_grouping_layers + self.num_experts * (
                len(self.sparse_layer_indices) * 2 - num_grouping_layers
        )
        all_usage_frequency = []
        usage_frequency_dict = deepcopy(self._usage_frequency_state_dict)
        for i, layer_idx in enumerate(self.sparse_layer_indices):
            # 1. Experts in the excluded layers are always not merged.
            if layer_idx not in merging_encoder_layers:
                usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"] = torch.ones_like(
                    usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"]
                )
            if layer_idx not in merging_decoder_layers:
                usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"] = torch.ones_like(
                    usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"]
                )
            # 2. Each layer must have at least one group, set the most used expert in a layer to frequency 1.
            max_usage_index = torch.argmax(usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"])
            usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"][max_usage_index] = 1.0
            max_usage_index = torch.argmax(usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"])
            usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"][max_usage_index] = 1.0
            # 3. Collect all usage frequency.
            all_usage_frequency.append(usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"])
            all_usage_frequency.append(usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"])
        all_usage_frequency = torch.cat(all_usage_frequency, dim=0)
        sorted_usage_frequency, sorted_indices = torch.sort(all_usage_frequency, descending=True)
        num_groups_per_layer = dict()

        # Note: When threshold is 0.0, the actual number of groups is smaller than total_num_groups.
        if average_num_groups == self.num_experts:
            total_num_groups = total_num_groups - 1
        frequency_threshold = sorted_usage_frequency[total_num_groups]

        print(f"[Merging]Frequency threshold: {frequency_threshold}")
        if frequency_threshold == 1.0:
            raise ValueError("[Merging]The number of groups is too large, please reduce the number of groups.")
        for i, layer_idx in enumerate(self.sparse_layer_indices):
            num_groups_per_layer[f"encoder.block.{layer_idx}.layer.1.mlp"] = torch.sum(
                (usage_frequency_dict[f"encoder.block.{layer_idx}.layer.1.mlp"] > frequency_threshold).long()
            ).item()
            num_groups_per_layer[f"decoder.block.{layer_idx}.layer.2.mlp"] = torch.sum(
                (usage_frequency_dict[f"decoder.block.{layer_idx}.layer.2.mlp"] > frequency_threshold).long()
            ).item()
        return num_groups_per_layer

    def group_experts_randomly(
            self,
            num_groups: int,
    ):
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Randomly merging experts into {num_groups} clusters"):
            # 1. Encoder
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            group_labels = generate_random_group_labels(self.num_experts, num_groups)
            self._group_state_dict[encoder_mlp_name] = group_labels
            # 2. Decoder
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            group_labels = generate_random_group_labels(self.num_experts, num_groups)
            self._group_state_dict[decoder_mlp_name] = group_labels

    def group_experts_into_clusters_by_routing_guided_globally(
            self,
            average_num_groups: int,
            merging_encoder_layers: List[int],
            merging_decoder_layers: List[int],
            encoder_group_capacity: Optional[int] = None,
            decoder_group_capacity: Optional[int] = None,
    ) -> Dict[str, List[int]]:
        """
        Globally group experts into clusters by routing-guided clustering, each layer will have different number of
         clusters. The total number of clusters is determined by average_num_groups.

        Parameters
        ----------
        average_num_groups: int
            The average number of clusters for all layers.
        merging_encoder_layers: List[int]
            The layers of encoder that are excluded from merging.
        merging_decoder_layers: List[int]
            The layers of decoder that are excluded from merging.
        encoder_group_capacity: Optional[int]
            The maximum number of experts in each group in Encoder layers. If None, the number of experts in each group is not limited.
        decoder_group_capacity: Optional[int]
            The maximum number of experts in each group in Decoder layers. If None, the number of experts in each group is not limited.

        Returns
        -------
        core_experts: Dict[str, List[int]]
            The core experts of each cluster
        """
        # By default, the first layer of encoder is excluded.
        # 1. Assign num_groups respectively for each layer according to average_num_groups
        encoder_group_capacity = encoder_group_capacity if encoder_group_capacity is not None else self.num_experts
        decoder_group_capacity = decoder_group_capacity if decoder_group_capacity is not None else self.num_experts
        num_groups_per_layer = self._assign_num_groups_per_layer(
            average_num_groups, merging_encoder_layers, merging_decoder_layers
        )
        print(f"[Merging]Number of groups of each layer: {num_groups_per_layer}")
        # 2. Group experts into clusters for each layer
        core_experts = dict()
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Globally routing-guided clustering experts into average {average_num_groups} clusters"):
            # 1. Encoder
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            num_groups = num_groups_per_layer[encoder_mlp_name]
            group_member_count = torch.zeros(num_groups)
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[encoder_mlp_name], descending=True)
            # 1.1 Assign top-K most-used experts with label 0 to K-1 respectively
            core_expert_indices = indices_sorted_by_usage[:num_groups]
            core_experts[encoder_mlp_name] = core_expert_indices.tolist()
            for i in range(num_groups):
                self._group_state_dict[encoder_mlp_name][core_expert_indices[i]] = i
                group_member_count[i] += 1
            # 1.2 Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(encoder_mlp_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = core_expert_indices[
                    torch.argmax(similarity_matrix[i, core_expert_indices])
                ]
                most_similar_group_label = self._group_state_dict[encoder_mlp_name][most_similar_core]
                self._group_state_dict[encoder_mlp_name][i] = most_similar_group_label
                group_member_count[most_similar_group_label] += 1
                if group_member_count[self._group_state_dict[encoder_mlp_name][i]] >= encoder_group_capacity:
                    if len(core_expert_indices) == 1:
                        raise ValueError(
                            f"[Merging]The number of groups at Encoder layer {layer_idx} is too small!"
                        )
                    # Kick out the filled group as well as its core, by pop the core from core_experts
                    core_index = torch.argmax(similarity_matrix[i, core_expert_indices])
                    core_expert_indices = torch.cat(
                        [core_expert_indices[:core_index], core_expert_indices[core_index + 1:]]
                    )

            # 2. Decoder
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            num_groups = num_groups_per_layer[decoder_mlp_name]
            group_member_count = torch.zeros(num_groups)
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[decoder_mlp_name], descending=True)
            # 2.1 Assign top-K most-used experts with label 0 to K-1 respectively
            core_expert_indices = indices_sorted_by_usage[:num_groups]
            core_experts[decoder_mlp_name] = core_expert_indices.tolist()
            for i in range(num_groups):
                self._group_state_dict[decoder_mlp_name][indices_sorted_by_usage[i]] = i
                group_member_count[i] += 1
            # 2.2 Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(decoder_mlp_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = core_expert_indices[
                    torch.argmax(similarity_matrix[i, core_expert_indices])
                ]
                most_similar_group_label = self._group_state_dict[decoder_mlp_name][most_similar_core]
                self._group_state_dict[decoder_mlp_name][i] = most_similar_group_label
                if group_member_count[self._group_state_dict[decoder_mlp_name][i]] >= decoder_group_capacity:
                    if len(core_expert_indices) == 1:
                        raise ValueError(
                            f"[Merging]The number of groups at Decoder layer {layer_idx} is too small!"
                        )
                    # Kick out the filled group as well as its core, by pop the core from core_experts
                    core_index = torch.argmax(similarity_matrix[i, core_expert_indices])
                    core_expert_indices = torch.cat(
                        [core_expert_indices[:core_index], core_expert_indices[core_index + 1:]]
                    )
        return core_experts

    def group_experts_into_clusters_by_routing_guided(
            self,
            num_groups: int,
    ) -> Dict[str, List[int]]:
        """
        1. Select top-K most-used experts, assign them with label 0 to K-1 respectively.
            These K experts are the cores of each cluster
        2. For each left unassigned expert, assign it to the cluster with the most similar core.

        Returns
        -------
        core_experts: Dict[str, List[int]]
            The core experts of each cluster
        """
        core_experts = {}
        if num_groups > self.num_experts:
            raise ValueError(
                f"[Merging]num_groups should be less than or equal to {self.num_experts}, got {num_groups} instead.")

        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Routing-guided clustering experts into {num_groups} clusters"):
            # 1. Encoder
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[encoder_mlp_name], descending=True)
            # 1.1 Assign top-K most-used experts with label 0 to K-1 respectively
            core_experts[encoder_mlp_name] = indices_sorted_by_usage[:num_groups].tolist()
            for i in range(num_groups):
                self._group_state_dict[encoder_mlp_name][indices_sorted_by_usage[i]] = i
            # 1.2 Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(encoder_mlp_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = indices_sorted_by_usage[
                    torch.argmax(similarity_matrix[i, indices_sorted_by_usage[:num_groups]])
                ]
                self._group_state_dict[encoder_mlp_name][i] = self._group_state_dict[encoder_mlp_name][
                    most_similar_core]

            # 2. Decoder
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[decoder_mlp_name], descending=True)
            # 2.1 Assign top-K most-used experts with label 0 to K-1 respectively
            core_experts[decoder_mlp_name] = indices_sorted_by_usage[:num_groups].tolist()
            for i in range(num_groups):
                self._group_state_dict[decoder_mlp_name][indices_sorted_by_usage[i]] = i
            # 2.2 Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(decoder_mlp_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = indices_sorted_by_usage[
                    torch.argmax(similarity_matrix[i, indices_sorted_by_usage[:num_groups]])
                ]
                self._group_state_dict[decoder_mlp_name][i] = self._group_state_dict[decoder_mlp_name][
                    most_similar_core]
        return core_experts

    def group_experts_into_clusters_by_spectral_clustering(
            self,
            num_groups: int
    ):
        """
        Group experts by spectral clustering.
        """
        if num_groups > self.num_experts:
            raise ValueError(
                f"[Merging]num_groups should be less than or equal to {self.num_experts}, got {num_groups} instead.")
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Spectral clustering experts into {num_groups} clusters..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"

            # 1. Encoder layer_idx MLP
            # 1.1 Compute grouping labels
            similarity_matrix = self.get_similarity_matrix(encoder_mlp_name).cpu().numpy()
            clustering = SpectralClustering(
                n_clusters=num_groups,
                affinity="precomputed",
                assign_labels="kmeans",
                random_state=0
            ).fit(similarity_matrix)
            self._group_state_dict[encoder_mlp_name] = torch.tensor(clustering.labels_, device="cuda")

            # 2. Decoder layer_idx MLP
            similarity_matrix = self.get_similarity_matrix(decoder_mlp_name).cpu().numpy()
            clustering = SpectralClustering(
                n_clusters=num_groups,
                affinity="precomputed",
                assign_labels="kmeans",
                random_state=0
            ).fit(similarity_matrix)
            self._group_state_dict[decoder_mlp_name] = torch.tensor(clustering.labels_, device="cuda")

    def compute_all_usages(
            self,
            model: SwitchTransformersForConditionalGeneration,
            batch: Dict[str, torch.Tensor],
            mini_batch_size: Optional[int] = 128,
    ):
        model.cuda()
        model.eval()
        total_batch_size = batch["input_ids"].shape[0]
        if mini_batch_size > total_batch_size:
            mini_batch_size = total_batch_size
        num_batches = total_batch_size // mini_batch_size
        for i in tqdm(range(num_batches), desc="[Merging]Computing all usages..."):
            with torch.no_grad():
                mini_batch = {k: v[i * mini_batch_size: (i + 1) * mini_batch_size] for k, v in batch.items()}
                mini_batch = {k: v.cuda() for k, v in mini_batch.items()}
                outputs = model(**mini_batch)
                for layer_idx in self.sparse_layer_indices:
                    encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                    encoder_router_expert_index = outputs.encoder_router_logits[layer_idx][1].reshape(-1)
                    for idx in encoder_router_expert_index:
                        self._usage_frequency_state_dict[encoder_mlp_name][idx] += 1
                    decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                    decoder_router_expert_index = outputs.decoder_router_logits[layer_idx][1].reshape(-1)
                    for idx in decoder_router_expert_index:
                        self._usage_frequency_state_dict[decoder_mlp_name][idx] += 1
        self._usage_frequency_state_dict = {
            k: v / torch.sum(v) for k, v in self._usage_frequency_state_dict.items()
        }

    def reverse_all_similarities(self):
        print("[Merging]Reversing all similarities...")
        for key in self._similarity_state_dict.keys():
            self._similarity_state_dict[key] = 1 - self._similarity_state_dict[key]

    def compute_all_similarities(
            self,
            model: SwitchTransformersForConditionalGeneration,
            batch: Dict[str, torch.Tensor] = None
    ):
        if self.similarity_base not in ["weight", "router-weight"] and batch is None:
            raise ValueError(
                "[Merging]batch should be provided when similarity_base is not 'weight' or 'router-weight'")

        model = model.cuda()
        model = model.eval()
        if self.similarity_base == "weight":
            self._compute_all_similarities_by_weight(model.state_dict())
        elif self.similarity_base == 'router-weight':
            self._compute_all_similarities_by_router_weight(model.state_dict())
        elif self.similarity_base == 'router-logits':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_router_logits(model, batch)
        elif self.similarity_base == 'gradient':
            model = self._accumulate_gradient(model, batch)
            self._compute_all_similarities_by_gradient(model)
            model.zero_grad()
        elif self.similarity_base == 'weight-gradient':
            model = self._accumulate_gradient(model, batch)
            self._compute_all_similarities_by_weight_gradient(model)
            model.zero_grad()
        elif self.similarity_base == 'feature':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_feature(model, batch)
        elif self.similarity_base == 'feature.abs':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_feature_abs(model, batch)
        elif self.similarity_base == 'weight-feature':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_weight_feature(model, batch)
        elif self.similarity_base == 'router-weight-feature':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_router_weight_feature(model, batch)
        elif self.similarity_base == 'mse':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_mse(model, batch)
        elif self.similarity_base == 'random':
            self._compute_all_similarities_by_random()
        elif self.similarity_base == 'feature-correlation.lsa':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_average_feature_correlation(model, batch, lsa=True)
        elif self.similarity_base == 'feature-correlation.max':
            batch = {k: v.cuda() for k, v in batch.items()}
            self._compute_all_similarities_by_average_feature_correlation(model, batch, lsa=False)
        else:
            raise NotImplementedError

    def _compute_all_similarities_by_weight(self, state_dict: Dict[str, torch.Tensor]):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by weight..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            for i in range(self.num_experts):
                for j in range(i + 1, self.num_experts):
                    # 1. Encoder
                    i_flat = torch.cat(
                        [state_dict[f"{encoder_mlp_name}.experts.expert_{i}.wi.weight"].flatten(),
                         state_dict[f"{encoder_mlp_name}.experts.expert_{i}.wo.weight"].flatten()],
                        dim=0
                    )
                    j_flat = torch.cat(
                        [state_dict[f"{encoder_mlp_name}.experts.expert_{j}.wi.weight"].flatten(),
                         state_dict[f"{encoder_mlp_name}.experts.expert_{j}.wo.weight"].flatten()],
                        dim=0
                    )
                    similarity = self.similarity_fn(i_flat, j_flat)
                    self.save_similarity(encoder_mlp_name, i, j, similarity)

                    # 2. Decoder
                    i_flat = torch.cat(
                        [state_dict[f"{decoder_mlp_name}.experts.expert_{i}.wi.weight"].flatten(),
                         state_dict[f"{decoder_mlp_name}.experts.expert_{i}.wo.weight"].flatten()],
                        dim=0
                    )
                    j_flat = torch.cat(
                        [state_dict[f"{decoder_mlp_name}.experts.expert_{j}.wi.weight"].flatten(),
                         state_dict[f"{decoder_mlp_name}.experts.expert_{j}.wo.weight"].flatten()],
                        dim=0
                    )
                    similarity = self.similarity_fn(i_flat, j_flat)
                    self.save_similarity(decoder_mlp_name, i, j, similarity)

    @staticmethod
    def _accumulate_gradient(
            model: SwitchTransformersForConditionalGeneration,
            batch: Dict[str, torch.Tensor],
            mini_batch_size: Optional[int] = 16,
    ) -> SwitchTransformersForConditionalGeneration:
        batch_size = batch["input_ids"].shape[0]
        num_batches = batch_size // mini_batch_size
        for i in tqdm(range(num_batches), desc="[Merging]Accumulating gradient..."):
            mini_batch = {k: v[i * mini_batch_size: (i + 1) * mini_batch_size] for k, v in batch.items()}
            mini_batch = {k: v.cuda() for k, v in mini_batch.items()}
            loss = model(**mini_batch).loss
            loss /= num_batches
            loss.backward(retain_graph=True)
        return model

    def _compute_all_similarities_by_gradient(
            self, model: SwitchTransformersForConditionalGeneration
    ):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by gradient..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
            decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_flat = torch.cat(
                            [encoder_mlp.experts[f"expert_{i}"].wi.weight.grad.data.flatten(),
                             encoder_mlp.experts[f"expert_{i}"].wo.weight.grad.data.flatten()],
                            dim=0
                        )
                        j_flat = torch.cat(
                            [encoder_mlp.experts[f"expert_{j}"].wi.weight.grad.data.flatten(),
                             encoder_mlp.experts[f"expert_{j}"].wo.weight.grad.data.flatten()],
                            dim=0
                        )
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_flat = torch.cat(
                            [decoder_mlp.experts[f"expert_{i}"].wi.weight.grad.data.flatten(),
                             decoder_mlp.experts[f"expert_{i}"].wo.weight.grad.data.flatten()],
                            dim=0
                        )
                        j_flat = torch.cat(
                            [decoder_mlp.experts[f"expert_{j}"].wi.weight.grad.data.flatten(),
                             decoder_mlp.experts[f"expert_{j}"].wo.weight.grad.data.flatten()],
                            dim=0
                        )
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)

    def _compute_all_similarities_by_weight_gradient(
            self, model: SwitchTransformersForConditionalGeneration
    ):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by weight*gradient..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
            decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_flat = torch.cat(
                            [encoder_mlp.experts[f"expert_{i}"].wi.weight.grad.data.flatten(),
                             encoder_mlp.experts[f"expert_{i}"].wo.weight.grad.data.flatten()],
                            dim=0
                        ) * torch.cat(
                            [encoder_mlp.experts[f"expert_{i}"].wi.weight.data.flatten(),
                             encoder_mlp.experts[f"expert_{i}"].wo.weight.data.flatten()],
                            dim=0
                        )
                        j_flat = torch.cat(
                            [encoder_mlp.experts[f"expert_{j}"].wi.weight.grad.data.flatten(),
                             encoder_mlp.experts[f"expert_{j}"].wo.weight.grad.data.flatten()],
                            dim=0
                        ) * torch.cat(
                            [encoder_mlp.experts[f"expert_{j}"].wi.weight.data.flatten(),
                             encoder_mlp.experts[f"expert_{j}"].wo.weight.data.flatten()],
                            dim=0
                        )
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_flat = torch.cat(
                            [decoder_mlp.experts[f"expert_{i}"].wi.weight.grad.data.flatten(),
                             decoder_mlp.experts[f"expert_{i}"].wo.weight.grad.data.flatten()],
                            dim=0
                        ) * torch.cat(
                            [decoder_mlp.experts[f"expert_{i}"].wi.weight.data.flatten(),
                             decoder_mlp.experts[f"expert_{i}"].wo.weight.data.flatten()],
                            dim=0
                        )
                        j_flat = torch.cat(
                            [decoder_mlp.experts[f"expert_{j}"].wi.weight.grad.data.flatten(),
                             decoder_mlp.experts[f"expert_{j}"].wo.weight.grad.data.flatten()],
                            dim=0
                        ) * torch.cat(
                            [decoder_mlp.experts[f"expert_{j}"].wi.weight.data.flatten(),
                             decoder_mlp.experts[f"expert_{j}"].wo.weight.data.flatten()],
                            dim=0
                        )
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)

    def _compute_all_similarities_by_router_weight(
            self, state_dict: Dict[str, torch.Tensor]
    ):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by router rows..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            for i in range(self.num_experts):
                for j in range(i + 1, self.num_experts):
                    # 1. Encoder
                    i_flat = state_dict[f"{encoder_mlp_name}.router.classifier.weight"][i]
                    j_flat = state_dict[f"{encoder_mlp_name}.router.classifier.weight"][j]
                    similarity = self.similarity_fn(i_flat, j_flat)
                    self.save_similarity(encoder_mlp_name, i, j, similarity)

                    # 2. Decoder
                    i_flat = state_dict[f"{decoder_mlp_name}.router.classifier.weight"][i]
                    j_flat = state_dict[f"{decoder_mlp_name}.router.classifier.weight"][j]
                    similarity = self.similarity_fn(i_flat, j_flat)
                    self.save_similarity(decoder_mlp_name, i, j, similarity)

    def _compute_all_similarities_by_router_logits(
            self, model: SwitchTransformersForConditionalGeneration, batch: Dict[str, torch.Tensor]
    ):
        with torch.no_grad():
            outputs = model(**batch)
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by router logits..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            encoder_router_logits = outputs.encoder_router_logits[layer_idx][0].reshape(-1, self.num_experts)
            decoder_router_logits = outputs.decoder_router_logits[layer_idx][0].reshape(-1, self.num_experts)
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_flat = encoder_router_logits[:, i].flatten()
                        j_flat = encoder_router_logits[:, j].flatten()
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_flat = decoder_router_logits[:, i].flatten()
                        j_flat = decoder_router_logits[:, j].flatten()
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)

    def _compute_all_similarities_by_router_weight_feature(
            self, model: SwitchTransformersForConditionalGeneration, batch: Dict[str, torch.Tensor]
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name)
            ))
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name)
            ))
        with torch.no_grad():
            model(**batch)
        state_dict = model.state_dict()
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc="[Merging]Computing similarities by router weight*feature..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            # features of shape (batch_size*sequence_length, hidden_size)
            encoder_features = self.activations[encoder_mlp_name]
            decoder_features = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        features = torch.linalg.norm(
                            encoder_features,
                            dim=0
                        )
                        i_flat = state_dict[f"{encoder_mlp_name}.router.classifier.weight"][i] * features
                        j_flat = state_dict[f"{encoder_mlp_name}.router.classifier.weight"][j] * features
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        features = torch.linalg.norm(
                            decoder_features,
                            dim=0
                        )
                        i_flat = state_dict[f"{decoder_mlp_name}.router.classifier.weight"][i] * features
                        j_flat = state_dict[f"{decoder_mlp_name}.router.classifier.weight"][j] * features
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)
        for handle in handles:
            handle.remove()

    def _get_mlp_activation(self, name):
        def hook(model, input, output):
            self.activations[name] = input[0].detach().reshape(
                -1, self.d_model)  # of shape (batch_size*sequence_length, hidden_size)

        return hook

    def _compute_all_similarities_by_feature(
            self, model: SwitchTransformersForConditionalGeneration, batch: Dict[str, torch.Tensor]
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name)
            ))
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name)
            ))
        with torch.no_grad():
            outputs = model(**batch)
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by feature..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            # router index of shape (batch_size*sequence_length)
            encoder_router_expert_index = outputs.encoder_router_logits[layer_idx][1].reshape(-1)
            decoder_router_expert_index = outputs.decoder_router_logits[layer_idx][1].reshape(-1)
            # features of shape (batch_size*sequence_length, hidden_size)
            encoder_features = self.activations[encoder_mlp_name]
            decoder_features = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_flat = torch.mean(
                            encoder_features[encoder_router_expert_index == i],
                            dim=0
                        )
                        if torch.isnan(i_flat).any():
                            i_flat = torch.zeros(self.d_model, device='cuda')
                        j_flat = torch.mean(
                            encoder_features[encoder_router_expert_index == j],
                            dim=0
                        )
                        if torch.isnan(j_flat).any():
                            j_flat = torch.zeros(self.d_model, device='cuda')
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_flat = torch.mean(
                            decoder_features[decoder_router_expert_index == i],
                            dim=0
                        )
                        if torch.isnan(i_flat).any():
                            i_flat = torch.zeros(self.d_model, device='cuda')
                        j_flat = torch.mean(
                            decoder_features[decoder_router_expert_index == j],
                            dim=0
                        )
                        if torch.isnan(j_flat).any():
                            j_flat = torch.zeros(self.d_model, device='cuda')
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)
        for handle in handles:
            handle.remove()

    def _compute_all_similarities_by_feature_abs(
            self, model: SwitchTransformersForConditionalGeneration, batch: Dict[str, torch.Tensor]
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name))
            )
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name))
            )
        with torch.no_grad():
            outputs = model(**batch)
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by feature.abs..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            # router index of shape (batch_size*sequence_length)
            encoder_router_expert_index = outputs.encoder_router_logits[layer_idx][1].reshape(-1)
            decoder_router_expert_index = outputs.decoder_router_logits[layer_idx][1].reshape(-1)
            # features of shape (batch_size*sequence_length, hidden_size)
            encoder_features = self.activations[encoder_mlp_name]
            decoder_features = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_flat = torch.mean(
                            encoder_features[encoder_router_expert_index == i].abs(),
                            dim=0
                        )
                        if torch.isnan(i_flat).any():
                            i_flat = torch.zeros(self.d_model, device='cuda')
                        j_flat = torch.mean(
                            encoder_features[encoder_router_expert_index == j].abs(),
                            dim=0
                        )
                        if torch.isnan(j_flat).any():
                            j_flat = torch.zeros(self.d_model, device='cuda')
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_flat = torch.mean(
                            decoder_features[decoder_router_expert_index == i].abs(),
                            dim=0
                        )
                        if torch.isnan(i_flat).any():
                            i_flat = torch.zeros(self.d_model, device='cuda')
                        j_flat = torch.mean(
                            decoder_features[decoder_router_expert_index == j].abs(),
                            dim=0
                        )
                        if torch.isnan(j_flat).any():
                            j_flat = torch.zeros(self.d_model, device='cuda')
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)
        for handle in handles:
            handle.remove()

    def _compute_all_similarities_by_weight_feature(
            self, model: SwitchTransformersForConditionalGeneration, batch: Dict[str, torch.Tensor]
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name)
            ))
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name)
            ))
        with torch.no_grad():
            outputs = model(**batch)
        state_dict = model.state_dict()
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by weight*feature..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            # router index of shape (batch_size*sequence_length)
            encoder_router_expert_index = outputs.encoder_router_logits[layer_idx][1].reshape(-1)
            decoder_router_expert_index = outputs.decoder_router_logits[layer_idx][1].reshape(-1)
            # features of shape (batch_size*sequence_length, hidden_size)
            encoder_features = self.activations[encoder_mlp_name]
            decoder_features = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # 1. Encoder
                        i_feature = torch.linalg.norm(
                            encoder_features[encoder_router_expert_index == i],
                            dim=0
                        )
                        if torch.isnan(i_feature).any():
                            i_feature = torch.zeros(self.d_model, device='cuda')
                        i_flat = (state_dict[f"{encoder_mlp_name}.experts.expert_{i}.wi.weight"] * i_feature).flatten()
                        j_feature = torch.linalg.norm(
                            encoder_features[encoder_router_expert_index == j],
                            dim=0
                        )
                        if torch.isnan(j_feature).any():
                            j_feature = torch.zeros(self.d_model, device='cuda')
                        j_flat = (state_dict[f"{encoder_mlp_name}.experts.expert_{j}.wi.weight"] * j_feature).flatten()
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        i_feature = torch.linalg.norm(
                            decoder_features[decoder_router_expert_index == i],
                            dim=0
                        )
                        if torch.isnan(i_feature).any():
                            i_feature = torch.zeros(self.d_model, device='cuda')
                        i_flat = (state_dict[f"{decoder_mlp_name}.experts.expert_{i}.wi.weight"] * i_feature).flatten()
                        j_feature = torch.linalg.norm(
                            decoder_features[decoder_router_expert_index == j],
                            dim=0
                        )
                        if torch.isnan(j_feature).any():
                            j_feature = torch.zeros(self.d_model, device='cuda')
                        j_flat = (state_dict[f"{decoder_mlp_name}.experts.expert_{j}.wi.weight"] * j_feature).flatten()
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)

        for handle in handles:
            handle.remove()

    def _compute_all_similarities_by_mse(
            self,
            model: SwitchTransformersForConditionalGeneration,
            batch: Dict[str, torch.Tensor]
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name)
            ))
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name)
            ))
        with torch.no_grad():
            outputs = model(**batch)
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by MSE..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            # router index of shape (batch_size*sequence_length)
            encoder_router_expert_index = outputs.encoder_router_logits[layer_idx][1].reshape(-1)
            decoder_router_expert_index = outputs.decoder_router_logits[layer_idx][1].reshape(-1)
            # features of shape (batch_size*sequence_length, hidden_size)
            encoder_features = self.activations[encoder_mlp_name]
            decoder_features = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    encoder_i = model.encoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{i}"]
                    encoder_i_inputs = encoder_features[encoder_router_expert_index == i]
                    encoder_i_outputs = encoder_i(encoder_i_inputs)
                    decoder_i = model.decoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{i}"]
                    decoder_i_inputs = decoder_features[decoder_router_expert_index == i]
                    decoder_i_outputs = decoder_i(decoder_i_inputs)
                    for j in range(i + 1, self.num_experts):
                        encoder_j = model.encoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{j}"]
                        encoder_j_inputs = encoder_features[encoder_router_expert_index == j]
                        encoder_j_outputs = encoder_j(encoder_j_inputs)
                        merged_new_mlp = SwitchTransformersDenseActDense(model.config).cuda()

                        # 1. Encoder
                        merged_new_mlp.wi.weight.copy_((encoder_i.wi.weight.data + encoder_j.wi.weight.data) / 2)
                        merged_new_mlp.wo.weight.copy_((encoder_i.wo.weight.data + encoder_j.wo.weight.data) / 2)
                        merged_new_outputs_from_i = merged_new_mlp(encoder_i_inputs)
                        merged_new_outputs_from_j = merged_new_mlp(encoder_j_inputs)
                        if len(encoder_i_inputs) == 0 and len(encoder_j_inputs) == 0:
                            # When both empty, similarity is 1
                            similarity = 1.0
                        else:
                            merged_new_outputs = torch.cat(
                                [merged_new_outputs_from_i, merged_new_outputs_from_j], dim=0
                            )
                            encoder_outputs = torch.cat([encoder_i_outputs, encoder_j_outputs], dim=0)
                            similarity = self.similarity_fn(merged_new_outputs, encoder_outputs)
                        self.save_similarity(encoder_mlp_name, i, j, similarity)

                        # 2. Decoder
                        decoder_j = model.decoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{j}"]
                        decoder_j_inputs = decoder_features[decoder_router_expert_index == j]
                        decoder_j_outputs = decoder_j(decoder_j_inputs)
                        merged_new_mlp.wi.weight.copy_((decoder_i.wi.weight.data + decoder_j.wi.weight.data) / 2)
                        merged_new_mlp.wo.weight.copy_((decoder_i.wo.weight.data + decoder_j.wo.weight.data) / 2)
                        merged_new_outputs_from_i = merged_new_mlp(decoder_i_inputs)
                        merged_new_outputs_from_j = merged_new_mlp(decoder_j_inputs)
                        if len(decoder_i_inputs) == 0 and len(decoder_j_inputs) == 0:
                            # When both empty, similarity is 1
                            similarity = 1.0
                        else:
                            merged_new_outputs = torch.cat(
                                [merged_new_outputs_from_i, merged_new_outputs_from_j], dim=0
                            )
                            decoder_outputs = torch.cat(
                                [decoder_i_outputs, decoder_j_outputs], dim=0
                            )
                            similarity = self.similarity_fn(merged_new_outputs, decoder_outputs)
                        self.save_similarity(decoder_mlp_name, i, j, similarity)
        for handle in handles:
            handle.remove()

    def _compute_all_similarities_by_random(self):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities randomly..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            for i in range(self.num_experts):
                for j in range(i + 1, self.num_experts):
                    self.save_similarity(encoder_mlp_name, i, j, torch.rand(1).item())
                    self.save_similarity(decoder_mlp_name, i, j, torch.rand(1).item())

    @staticmethod
    def _compute_similarity_by_average_correlation(
            first_mlp: SwitchTransformersSparseMLP,
            second_mlp: SwitchTransformersSparseMLP,
            forwarded_hidden_states: torch.Tensor,
            lsa: bool
    ) -> float:
        if len(forwarded_hidden_states) <= 1:
            # TODO: What to do when both experts are never selected?
            # For now, return 1.0 to just merge them
            return 1.0

        first_activations = []
        second_activations = []

        def _first_hook(model, input, output):
            first_activations.append(input[0].detach().reshape(-1, input[0].shape[-1]))

        def _second_hook(model, input, output):
            second_activations.append(input[0].detach().reshape(-1, input[0].shape[-1]))

        first_handle = first_mlp.wo.register_forward_hook(_first_hook)
        second_handle = second_mlp.wo.register_forward_hook(_second_hook)
        with torch.no_grad():
            first_mlp(forwarded_hidden_states)
            second_mlp(forwarded_hidden_states)
        first_handle.remove()
        second_handle.remove()
        first_activations = torch.cat(first_activations, dim=0)
        second_activations = torch.cat(second_activations, dim=0)

        # Compute correlation
        mean_first = torch.mean(first_activations, dim=0, keepdim=True)  # (1, d_ff)
        mean_second = torch.mean(second_activations, dim=0, keepdim=True)  # (1, d_ff)
        std_first = torch.std(first_activations, dim=0, keepdim=True)  # (1, d_ff)
        std_second = torch.std(second_activations, dim=0, keepdim=True)  # (1, d_ff)
        covar = torch.mm(
            (first_activations - mean_first).t(),
            (second_activations - mean_second)
        ) / (len(forwarded_hidden_states) - 1)  # (d_ff, d_ff)
        corr = covar / (std_first.t() * std_second + FP32_EPS)  # (d_ff, d_ff)

        # Compute linear sum assignment to find maximum correlation
        if lsa:
            corr = corr.cpu().numpy()
            row_ind, col_ind = linear_sum_assignment(corr, maximize=True)
            # Normalize to [0, 1]
            return (corr[row_ind, col_ind] + 1).sum() / corr.shape[0] / 2
        else:
            return (corr.max(dim=0).values + 1).sum() / corr.shape[0] / 2

    def _compute_all_similarities_by_average_feature_correlation(
            self,
            model: SwitchTransformersForConditionalGeneration,
            batch: Dict[str, torch.Tensor],
            lsa: bool
    ):
        self.activations = {}
        handles = []
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Registering forward hook..."):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(encoder_mlp_name)
            ))
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                self._get_mlp_activation(decoder_mlp_name)
            ))
        model = model.eval()
        encoder_router_indices = {l: None for l in self.sparse_layer_indices}
        decoder_router_indices = {l: None for l in self.sparse_layer_indices}
        with torch.no_grad():
            outputs = model(**batch)
            for layer_idx in self.sparse_layer_indices:
                encoder_router_indices[layer_idx] = outputs.encoder_router_logits[layer_idx][1]
                decoder_router_indices[layer_idx] = outputs.decoder_router_logits[layer_idx][1]
        del outputs
        for handle in handles:
            handle.remove()
        torch.cuda.empty_cache()

        progress_bar = tqdm(
            range(len(self.sparse_layer_indices) * self.num_experts * (self.num_experts - 1) // 2),
            desc="[Merging]Computing similarities by average feature correlation..."
        )

        for layer_idx in self.sparse_layer_indices:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            encoder_router_expert_index = encoder_router_indices[layer_idx].reshape(-1)
            decoder_router_expert_index = decoder_router_indices[layer_idx].reshape(-1)
            encoder_hidden_states = self.activations[encoder_mlp_name]
            decoder_hidden_states = self.activations[decoder_mlp_name]
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        similarity = self._compute_similarity_by_average_correlation(
                            model.encoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{i}"],
                            model.encoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{j}"],
                            torch.cat([
                                encoder_hidden_states[encoder_router_expert_index == i],
                                encoder_hidden_states[encoder_router_expert_index == j]
                            ], dim=0),
                            lsa=lsa
                        )
                        self.save_similarity(encoder_mlp_name, i, j, similarity)
                        similarity = self._compute_similarity_by_average_correlation(
                            model.decoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{i}"],
                            model.decoder.block[layer_idx].layer[-1].mlp.experts[f"expert_{j}"],
                            torch.cat([
                                decoder_hidden_states[decoder_router_expert_index == i],
                                decoder_hidden_states[decoder_router_expert_index == j]
                            ], dim=0),
                            lsa=lsa
                        )
                        self.save_similarity(decoder_mlp_name, i, j, similarity)
                        progress_bar.update(1)


def _merge_mlp_experts_by_averaging(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        permute: bool,
        permute_strategy: str,
        forwarded_hidden_states: Optional[Tuple[torch.Tensor]] = None,
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        if permute and permute_strategy == "weight-matching":
            for expert_idx in expert_indices[1:]:
                perm = compute_switch_permutation_by_weight_matching(
                    reference_mlp=mlp.experts[f"expert_{expert_indices[0]}"],
                    target_mlp=mlp.experts[f"expert_{expert_idx}"],
                    include_wo=True
                )
                mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                    mlp.experts[f"expert_{expert_idx}"], perm
                )
        elif permute and permute_strategy == "activation-matching":
            group_forwarded_hidden_states = torch.cat([
                forwarded_hidden_states[expert_idx] for expert_idx in expert_indices
            ], dim=0)
            for expert_idx in expert_indices[1:]:
                perm = compute_switch_permutation_by_activation_matching(
                    reference_mlp=mlp.experts[f"expert_{expert_indices[0]}"],
                    target_mlp=mlp.experts[f"expert_{expert_idx}"],
                    forwarded_hidden_states=group_forwarded_hidden_states,
                )
                mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                    mlp.experts[f"expert_{expert_idx}"], perm
                )
        elif permute:
            raise ValueError(f"Unknown permute strategy: {permute_strategy}")

        with torch.no_grad():
            wi_weight = torch.mean(
                torch.stack([mlp.experts[f"expert_{expert_idx}"].wi.weight for expert_idx in expert_indices]),
                dim=0
            )
            wo_weight = torch.mean(
                torch.stack([mlp.experts[f"expert_{expert_idx}"].wo.weight for expert_idx in expert_indices]),
                dim=0
            )
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(wi_weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(wo_weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]
    return mlp


def _merge_mlp_experts_by_fisher_weighted_averaging(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        name_prefix: str,
        experts_fisher_state_dict: Dict[str, torch.Tensor],
        permute: Optional[bool] = False
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        if permute:
            for expert_idx in expert_indices[1:]:
                perm = compute_switch_permutation_by_weight_matching(
                    reference_mlp=mlp.experts[f"expert_{expert_indices[0]}"],
                    target_mlp=mlp.experts[f"expert_{expert_idx}"],
                    include_wo=True
                )
                mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                    mlp.experts[f"expert_{expert_idx}"], perm
                )
        with torch.no_grad():
            wi_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wi.weight for expert_idx in expert_indices], dim=0
            )
            wo_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wo.weight for expert_idx in expert_indices], dim=0
            )
            wi_fisher_list = torch.stack(
                [experts_fisher_state_dict[f"{name_prefix}.expert_{expert_idx}.wi.weight"]
                 for expert_idx in expert_indices], dim=0
            )
            wo_fisher_list = torch.stack(
                [experts_fisher_state_dict[f"{name_prefix}.expert_{expert_idx}.wo.weight"]
                 for expert_idx in expert_indices], dim=0
            )

            wi_weight = torch.sum(wi_weight_list * wi_fisher_list, dim=0) / (
                    torch.sum(wi_fisher_list, dim=0) + FP32_EPS)
            wo_weight = torch.sum(wo_weight_list * wo_fisher_list, dim=0) / (
                    torch.sum(wo_fisher_list, dim=0) + FP32_EPS)

            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(wi_weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(wo_weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]
    return mlp


def merge_by_groups(
        model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        strategy: str = "average",
        fisher_state_dict: Optional[Dict[str, torch.Tensor]] = None,
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
        permute: Optional[bool] = False,
        permute_strategy: Optional[str] = "weight-matching",
        dataloader: Optional[DataLoader] = None,
) -> SwitchTransformersForConditionalGeneration:
    """
    Merge experts by groups, strategies include:
        1. "average": simply average the weights of experts in the same group.
        2. "fisher": average the weights of experts in the same group weighted by (diagnonal) fisher information.
    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to merge experts.
    grouper: ExpertsGrouperForSwitch
        The grouper to group experts, supposed to have been called `grouper.compute_all_usages()` and
            one of `grouper.group_experts()` (i.e. have grouped labels).
    fisher_state_dict: Optional[Dict[str, torch.Tensor]]
        The fisher state dict, if None, can only use (vanilla )averaging strategy, else use fisher averaging strategy.
    strategy: str
        The strategy to merge experts, one of ["replace", "drop"].
    encoder_merging_layers: Optional[List[int]]
        The encoder layers to merge experts, if None, merge all layers.
    decoder_merging_layers: Optional[List[int]]
        The decoder layers to merge experts, if None, merge all layers.
    permute: Optional[bool]
        Whether to permute the experts in the same group.
    permute_strategy: Optional[str]
        The strategy to permute experts, one of ["weight-matching", "random"].
    dataloader: Optional[DataLoader]
        The dataloader to compute activations, only used when `strategy` is "activation-matching".
    """
    forwarded_hidden_states = dict()
    router_indices = dict()
    if permute_strategy == "activation-matching":
        model.eval().cuda()
        handles = []

        def _get_activation_hook(name):
            def hook(module, input, output):
                forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

            return hook

        for layer_idx in tqdm(
                grouper.sparse_layer_indices,
                desc=f"[Merging]Registering forward hook..."
        ):
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            forwarded_hidden_states[encoder_mlp_name] = []
            forwarded_hidden_states[decoder_mlp_name] = []
            handles.append(model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                _get_activation_hook(encoder_mlp_name))
            )
            handles.append(model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
                _get_activation_hook(decoder_mlp_name))
            )

        # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
        router_indices = {name: [] for name in forwarded_hidden_states.keys()}
        with torch.no_grad():
            for batch in tqdm(dataloader, desc="[Merging]Computing activations..."):
                batch = {k: v.cuda() for k, v in batch.items()}
                outputs = model(**batch)
                for layer_idx in grouper.sparse_layer_indices:
                    router_indices[f"encoder.block.{layer_idx}.layer.1.mlp"].append(
                        outputs.encoder_router_logits[layer_idx][1].reshape(-1)
                    )
                    router_indices[f"decoder.block.{layer_idx}.layer.2.mlp"].append(
                        outputs.decoder_router_logits[layer_idx][1].reshape(-1)
                    )

        for handle in handles:
            handle.remove()

    num_experts = grouper.num_experts

    if permute:
        print(f"[Merging]Permutation is enabled, will permute experts by {permute_strategy} in the same group.")
    if strategy == "average":
        for layer_idx in tqdm(grouper.sparse_layer_indices,
                              desc="[Merging]Merging experts with averaging..."):
            if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
                encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
                if permute_strategy == "activation-matching":
                    layer_forwarded_hidden_states = tuple()
                    for expert_idx in range(num_experts):
                        layer_forwarded_hidden_states += (
                            torch.cat(
                                [forwarded_hidden_states[encoder_mlp_name][i][
                                     router_indices[encoder_mlp_name][i] == expert_idx]
                                 for i in range(len(dataloader))], dim=0),
                        )
                    model.encoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_averaging(
                        mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                        group_labels=encoder_group_labels,
                        permute=permute,
                        permute_strategy=permute_strategy,
                        forwarded_hidden_states=layer_forwarded_hidden_states
                    )
                else:
                    model.encoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_averaging(
                        mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                        group_labels=encoder_group_labels,
                        permute=permute,
                        permute_strategy=permute_strategy,
                    )
            if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
                decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
                if permute_strategy == "activation-matching":
                    layer_forwarded_hidden_states = tuple()
                    for expert_idx in range(num_experts):
                        layer_forwarded_hidden_states += (
                            torch.cat(
                                [forwarded_hidden_states[decoder_mlp_name][i][
                                     router_indices[decoder_mlp_name][i] == expert_idx]
                                 for i in range(len(dataloader))], dim=0),
                        )
                    model.decoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_averaging(
                        mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                        group_labels=decoder_group_labels,
                        permute=permute,
                        permute_strategy=permute_strategy,
                        forwarded_hidden_states=layer_forwarded_hidden_states
                    )
                else:
                    model.decoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_averaging(
                        mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                        group_labels=decoder_group_labels,
                        permute=permute,
                        permute_strategy=permute_strategy,
                    )
    elif strategy == "fisher":
        if fisher_state_dict is None:
            raise ValueError("[Merging]Fisher averaging strategy requires fisher_state_dict")
        fisher_state_dict = {k: v.cuda() for k, v in fisher_state_dict.items()}
        for layer_idx in tqdm(grouper.sparse_layer_indices, desc="[Merging]Merging experts with fisher..."):
            if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
                encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                encoder_group_labels = grouper._group_state_dict[encoder_mlp_name]
                model.encoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_fisher_weighted_averaging(
                    mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                    group_labels=encoder_group_labels,
                    name_prefix=encoder_mlp_name + ".experts",
                    experts_fisher_state_dict=fisher_state_dict,
                    permute=permute
                )
            if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
                decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                decoder_group_labels = grouper._group_state_dict[decoder_mlp_name]
                model.decoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_fisher_weighted_averaging(
                    mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                    group_labels=decoder_group_labels,
                    name_prefix=decoder_mlp_name + ".experts",
                    experts_fisher_state_dict=fisher_state_dict,
                    permute=permute
                )
    return model


def _prune_mlp_experts_by_replacing_non_core(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        core_expert_indices: List[int],
) -> SwitchTransformersSparseMLP:
    """
    Prune non-core experts by replacing them with core experts
    """
    for core_idx in core_expert_indices:
        with torch.no_grad():
            experts_indices = torch.where(group_labels == group_labels[core_idx])[0]
            core_expert = mlp.experts[f"expert_{core_idx}"]
            for expert_idx in experts_indices:
                if expert_idx == core_idx:
                    continue
                mlp.experts[f"expert_{expert_idx}"] = core_expert

    return mlp


def _prune_mlp_experts_by_dropping_non_core(
        mlp: SwitchTransformersSparseMLP,
        core_expert_indices: List[int],
) -> SwitchTransformersSparseMLP:
    non_core_expert_mask = torch.ones(len(mlp.experts), dtype=torch.bool)
    non_core_expert_mask[core_expert_indices] = False
    non_core_expert_indices = torch.where(non_core_expert_mask)[0]
    mask_value = -1e6

    def _custom_forward(self, hidden_states: torch.Tensor) -> Tuple:
        _, router_logits = self._compute_router_probabilities(hidden_states)

        router_logits = router_logits.masked_fill(non_core_expert_mask.to(router_logits.device), mask_value)
        router_probs = F.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)

        expert_index = torch.argmax(router_probs, dim=-1)
        expert_index = F.one_hot(expert_index, num_classes=self.num_experts)
        # Since experts are pruned, no need to do max-capacity dropping
        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
        return expert_index, router_probs, router_logits

    mlp.router.forward = _custom_forward.__get__(mlp.router, SwitchTransformersTop1Router)

    for expert_idx in non_core_expert_indices:
        # fake expert, this will raise error if used by mistake
        mlp.experts[f"expert_{expert_idx}"] = torch.nn.Linear(mlp.router.classifier.in_features, 1)

    return mlp


def prune_non_core_experts_by_groups(
        model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        core_experts: Dict[str, List[int]],
        strategy: str = "replace",
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
) -> SwitchTransformersForConditionalGeneration:
    """
    Simply prune non-core experts by replacing them with core experts, or dropping them.
    (The difference between replacing and dropping is whether to route to the pruned experts and redirect to core experts)

    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to merge experts
    grouper: ExpertsGrouperForSwitch
        The grouper to group experts, supposed to have been called `grouper.compute_all_usages()` and
            one of `grouper.group_experts()` (i.e. have grouped labels)
    core_experts: Dict[str, List[int]]
        The core experts dict, normally returned by `grouper.group_experts_into_clusters_by_routing_guided_globally()`
    strategy: str
        The strategy to merge experts, one of ["replace", "drop"]
    encoder_merging_layers: Optional[List[int]]
        The encoder layers to merge experts, if None, merge all layers
    decoder_merging_layers: Optional[List[int]]
        The decoder layers to merge experts, if None, merge all layers
    """
    if strategy == "replace":
        for layer_idx in tqdm(grouper.sparse_layer_indices, desc="[Merging]Pruning non-core experts by replacing..."):
            if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
                encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
                model.encoder.block[layer_idx].layer[-1].mlp = _prune_mlp_experts_by_replacing_non_core(
                    mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                    group_labels=encoder_group_labels,
                    core_expert_indices=core_experts[encoder_mlp_name]
                )
            if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
                decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
                model.decoder.block[layer_idx].layer[-1].mlp = _prune_mlp_experts_by_replacing_non_core(
                    mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                    group_labels=decoder_group_labels,
                    core_expert_indices=core_experts[decoder_mlp_name]
                )
    elif strategy == "drop":
        for layer_idx in tqdm(grouper.sparse_layer_indices, desc="[Merging]Pruning non-core experts by dropping..."):
            if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
                encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                model.encoder.block[layer_idx].layer[-1].mlp = _prune_mlp_experts_by_dropping_non_core(
                    mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                    core_expert_indices=core_experts[encoder_mlp_name]
                )
            if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
                decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                model.decoder.block[layer_idx].layer[-1].mlp = _prune_mlp_experts_by_dropping_non_core(
                    mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                    core_expert_indices=core_experts[decoder_mlp_name]
                )
    return model


def _merge_mlp_experts_by_usage_frequency_weighting(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        usage_frequencies: torch.Tensor,
        permute: bool,
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        if permute:
            for expert_idx in expert_indices[1:]:
                perm = compute_switch_permutation_by_weight_matching(
                    reference_mlp=mlp.experts[f"expert_{expert_indices[0]}"],
                    target_mlp=mlp.experts[f"expert_{expert_idx}"],
                    include_wo=True
                )
                mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                    mlp.experts[f"expert_{expert_idx}"], perm
                )
        with torch.no_grad():
            wi_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wi.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            wo_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wo.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            wi_weight = torch.sum(wi_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + FP32_EPS)
            wo_weight = torch.sum(wo_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + FP32_EPS)
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(wi_weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(wo_weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]
    return mlp


def _merge_mlp_experts_by_usage_frequency_weighting_within_and_across_models(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        usage_frequencies: torch.Tensor
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        with torch.no_grad():
            merged_mlp = merge_switch_mlp_by_weight_matching_within_and_across_models(
                mlp_list=[mlp.experts[f"expert_{expert_idx}"] for expert_idx in expert_indices],
                average_coefs=usage_frequencies[expert_indices].tolist()
            )
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(merged_mlp.wi.weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(merged_mlp.wo.weight)
            del merged_mlp
            torch.cuda.empty_cache()
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]
    return mlp


def merge_by_groups_with_usage_frequency_weighting(
        model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        strategy: str = "normal",
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
        permute: Optional[bool] = False,
        within_and_across_models: Optional[bool] = False,
) -> SwitchTransformersForConditionalGeneration:
    """
    Merge experts by usage-frequency-weighted averaging, strategies include:
        1. normal: merge experts in each group by usage-frequency-weighted averaging.
        2. reversed: reverse usage frequencies by 1 - usage_frequency and merge experts in each group by
                        usage-frequency-weighted averaging.
        3. random: randomly initialize usage frequencies and merge experts in each group by
                        usage-frequency-weighted averaging.

    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to merge experts
    grouper: ExpertsGrouperForSwitch
        The grouper to group experts, supposed to have been called `grouper.compute_all_usages()` and
            one of `grouper.group_experts()` (i.e. have grouped labels)
    strategy: str
        The strategy to merge experts, one of ["normal", "reversed", "random"]
    encoder_merging_layers: Optional[List[int]]
        The encoder layers to merge experts, if None, merge all layers
    decoder_merging_layers: Optional[List[int]]
        The decoder layers to merge experts, if None, merge all layers
    permute: Optional[bool]
        Whether to permute the experts in the same group, only availabel when `within_and_across_models` is False.
    within_and_across_models: Optional[bool]
        Whether to merge experts within and across models.
    """
    if permute:
        print("[Merging]Permutation is enabled, will permute experts in the same group.")
    usage_frequency_dict = grouper.usage_frequency_state_dict()
    if strategy == "reversed":
        for key, value in usage_frequency_dict.items():
            usage_frequency_dict[key] = 1 - value
    elif strategy == "random":
        for key, value in usage_frequency_dict.items():
            usage_frequency_dict[key] = torch.rand_like(value)
    elif strategy != "normal":
        raise ValueError(f"[Merging]Unknown strategy {strategy}")

    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging experts with {strategy} usage-frequency-weighted averaging..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
            usage_frequencies = usage_frequency_dict[encoder_mlp_name]
            model.encoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_usage_frequency_weighting(
                mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                group_labels=encoder_group_labels,
                usage_frequencies=usage_frequencies,
                permute=permute
            ) if not within_and_across_models else _merge_mlp_experts_by_usage_frequency_weighting_within_and_across_models(
                mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                group_labels=encoder_group_labels,
                usage_frequencies=usage_frequencies,
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
            usage_frequencies = usage_frequency_dict[decoder_mlp_name]
            model.decoder.block[layer_idx].layer[-1].mlp = _merge_mlp_experts_by_usage_frequency_weighting(
                mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                group_labels=decoder_group_labels,
                usage_frequencies=usage_frequencies,
                permute=permute
            ) if not within_and_across_models else _merge_mlp_experts_by_usage_frequency_weighting_within_and_across_models(
                mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                group_labels=decoder_group_labels,
                usage_frequencies=usage_frequencies,
            )
    return model


def _merge_mlp_experts_within_and_across_models(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        forwarded_hidden_states: Tuple[torch.Tensor],
        dominant_alone: bool,
        core_expert_indices: Optional[List[int]] = None,
        usage_frequencies: Optional[torch.Tensor] = None,
) -> SwitchTransformersSparseMLP:
    """
    Merge grouped experts within and across models.

    Parameters
    ----------
    mlp: SwitchTransformersSparseMLP
        The mlp to merge experts.
    group_labels: torch.LongTensor
        The group labels of experts.
    forwarded_hidden_states: Tuple[torch.Tensor]
        The forwarded hidden states of each expert, should be of length num_experts
    dominant_alone: bool
        Whether to merge the dominant expert alone.
        If True, the merging process in a group will be done in two steps:
            1. Merge all experts except the dominant one.
            2. Merge the dominant expert with the merged expert in step 1.
    core_expert_indices: List[int]

    Returns
    -------
    mlp: SwitchTransformersSparseMLP
        The merged mlp.
    """
    if dominant_alone and core_expert_indices is None:
        raise ValueError("[Merging]dominant_alone is True, but core_expert_indices is None")

    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        with torch.no_grad():
            if dominant_alone:
                group_core_expert_indices = torch.stack([
                    idx for idx in expert_indices if idx in core_expert_indices])
                to_skip = False
                if len(group_core_expert_indices) == len(expert_indices):
                    merged_expert = mlp.experts[f"expert_{expert_indices[0]}"]
                    to_skip = True
                elif usage_frequencies is not None and len(group_core_expert_indices) == 1:
                    non_core_usage_sum = torch.sum(
                        usage_frequencies[[expert_idx.item() for expert_idx in
                                           expert_indices if expert_idx not in group_core_expert_indices]]).item()
                    if non_core_usage_sum == 0:
                        merged_expert = mlp.experts[f"expert_{group_core_expert_indices[0]}"]
                        to_skip = True
                    else:
                        to_skip = False
                if not to_skip:
                    # Stage 1: merge all experts except the dominant one
                    group_forwarded_hidden_states = torch.cat([
                        forwarded_hidden_states[expert_idx] for expert_idx in expert_indices if
                        expert_idx not in group_core_expert_indices
                    ], dim=0)
                    if usage_frequencies is not None:
                        non_core_usages = usage_frequencies[[expert_idx.item() for expert_idx in expert_indices if
                                                             expert_idx not in group_core_expert_indices]]
                    merged_expert = merge_switch_mlp_by_activation_matching_within_and_across_models(
                        mlp_list=[mlp.experts[f"expert_{expert_idx}"] for expert_idx in expert_indices if
                                  expert_idx not in group_core_expert_indices],
                        forwarded_hidden_states=group_forwarded_hidden_states,
                        average_coefs=non_core_usages.tolist() if usage_frequencies is not None else None
                    )
                    # Stage 2: merge the dominant expert with the merged expert in stage 1
                    group_forwarded_hidden_states = torch.cat([
                        forwarded_hidden_states[expert_idx] for expert_idx in expert_indices
                    ], dim=0)
                    if usage_frequencies is not None:
                        core_usages = usage_frequencies[group_core_expert_indices]
                        non_core_usage_sum = torch.sum(non_core_usages).item()
                    merged_expert = merge_switch_mlp_by_activation_matching_within_and_across_models(
                        mlp_list=[merged_expert] + [mlp.experts[f"expert_{expert_idx}"] for expert_idx in
                                                    group_core_expert_indices],
                        forwarded_hidden_states=group_forwarded_hidden_states,
                        average_coefs=[non_core_usage_sum] + core_usages.tolist(
                        ) if usage_frequencies is not None else None
                    )
            else:
                # Merge all experts in the group
                group_forwarded_hidden_states = torch.cat([
                    forwarded_hidden_states[expert_idx] for expert_idx in expert_indices
                ], dim=0)
                merged_expert = merge_switch_mlp_by_activation_matching_within_and_across_models(
                    mlp_list=[mlp.experts[f"expert_{expert_idx}"] for expert_idx in expert_indices],
                    forwarded_hidden_states=group_forwarded_hidden_states,
                    average_coefs=usage_frequencies[expert_indices].tolist() if usage_frequencies is not None else None
                )
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(merged_expert.wi.weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(merged_expert.wo.weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]

    return mlp


def merge_by_groups_within_and_across_models(
        switch_model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        dataloader: DataLoader,
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
        dominant_alone: Optional[bool] = False,
        core_experts: Optional[Dict[str, List[int]]] = None,
        usage_weighted: Optional[bool] = False,
) -> SwitchTransformersForConditionalGeneration:
    # {name: values}, values  will be of shape (len(dataloader), batch_size * seq_len, d_ff)
    forwarded_hidden_states = dict()

    usage_frequencies = grouper.usage_frequency_state_dict()

    switch_model.eval().cuda()
    handles = []

    def _get_activation_hook(name):
        def hook(module, input, output):
            forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

        return hook

    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Registering forward hook..."
    ):
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
        forwarded_hidden_states[encoder_mlp_name] = []
        forwarded_hidden_states[decoder_mlp_name] = []
        handles.append(switch_model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(encoder_mlp_name))
        )
        handles.append(switch_model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(decoder_mlp_name))
        )

    # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
    router_indices = {name: [] for name in forwarded_hidden_states.keys()}
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="[Merging]Computing activations..."):
            batch = {k: v.cuda() for k, v in batch.items()}
            outputs = switch_model(**batch)
            for layer_idx in grouper.sparse_layer_indices:
                router_indices[f"encoder.block.{layer_idx}.layer.1.mlp"].append(
                    outputs.encoder_router_logits[layer_idx][1].reshape(-1)
                )
                router_indices[f"decoder.block.{layer_idx}.layer.2.mlp"].append(
                    outputs.decoder_router_logits[layer_idx][1].reshape(-1)
                )

    for handle in handles:
        handle.remove()

    num_experts = grouper.num_experts
    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging by groups within and across experts..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[encoder_mlp_name][i][router_indices[encoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.encoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_experts_within_and_across_models(
                mlp=switch_model.encoder.block[layer_idx].layer[-1].mlp,
                group_labels=encoder_group_labels,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                dominant_alone=dominant_alone,
                core_expert_indices=core_experts[encoder_mlp_name] if core_experts is not None else None,
                usage_frequencies=usage_frequencies[encoder_mlp_name] if usage_weighted else None,
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[decoder_mlp_name][i][router_indices[decoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.decoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_experts_within_and_across_models(
                mlp=switch_model.decoder.block[layer_idx].layer[-1].mlp,
                group_labels=decoder_group_labels,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                dominant_alone=dominant_alone,
                core_expert_indices=core_experts[decoder_mlp_name] if core_experts is not None else None,
                usage_frequencies=usage_frequencies[decoder_mlp_name] if usage_weighted else None,
            )

    del forwarded_hidden_states, router_indices
    torch.cuda.empty_cache()
    return switch_model


def _merge_mlp_non_core_experts_by_usage_frequency_weighting(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        usage_frequencies: torch.Tensor,
        core_experts: List[int],
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        # Only merge non-core experts
        expert_indices = torch.tensor(
            [expert_idx.item() for expert_idx in expert_indices if expert_idx not in core_experts]
        )
        if len(expert_indices) == 0:
            continue
        with torch.no_grad():
            wi_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wi.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            wo_weight_list = torch.stack(
                [mlp.experts[f"expert_{expert_idx}"].wo.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            wi_weight = torch.sum(wi_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + FP32_EPS)
            wo_weight = torch.sum(wo_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + FP32_EPS)
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(wi_weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(wo_weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]
    return mlp


def merge_non_core_by_groups_with_usage_frequency_weighting(
        model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        core_experts: Dict[str, List[int]],
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
) -> SwitchTransformersForConditionalGeneration:
    """
    Similar to `merge_by_groups_with_usage_frequency_weighting`, but only merge non-core experts.

    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to merge experts
    grouper: ExpertsGrouperForSwitch
        The grouper to group experts, supposed to have been called `grouper.compute_all_usages()` and
            one of `grouper.group_experts()` (i.e. have grouped labels)
    core_experts: Dict[str, List[int]]
        The core experts dict, normally returned by `grouper.group_experts_into_clusters_by_routing_guided_globally()`
    encoder_merging_layers: Optional[List[int]]
        The encoder layers to merge experts, if None, merge all layers
    decoder_merging_layers: Optional[List[int]]
        The decoder layers to merge experts, if None, merge all layers

    Returns
    -------
    model: SwitchTransformersForConditionalGeneration
        The model with merged non-core experts
    """
    usage_frequency_dict = grouper.usage_frequency_state_dict()
    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging only non-core experts with usage-frequency-weighted averaging..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
            usage_frequencies = usage_frequency_dict[encoder_mlp_name]
            model.encoder.block[layer_idx].layer[-1].mlp = _merge_mlp_non_core_experts_by_usage_frequency_weighting(
                mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                group_labels=encoder_group_labels,
                usage_frequencies=usage_frequencies,
                core_experts=core_experts[encoder_mlp_name]
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
            usage_frequencies = usage_frequency_dict[decoder_mlp_name]
            model.decoder.block[layer_idx].layer[-1].mlp = _merge_mlp_non_core_experts_by_usage_frequency_weighting(
                mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                group_labels=decoder_group_labels,
                usage_frequencies=usage_frequencies,
                core_experts=core_experts[decoder_mlp_name]
            )
    return model


def _merge_mlp_non_core_experts_within_and_across_models(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.LongTensor,
        forwarded_hidden_states: Tuple[torch.Tensor],
        core_expert_indices: Optional[List[int]] = None,
) -> SwitchTransformersSparseMLP:
    """
    Merge grouped experts within and across models, only merge non-core experts. (N groups 2N experts)

    Parameters
    ----------
    mlp: SwitchTransformersSparseMLP
        The mlp to merge experts.
    group_labels: torch.LongTensor
        The group labels of experts.
    forwarded_hidden_states: Tuple[torch.Tensor]
        The forwarded hidden states of each expert, should be of length num_experts
    core_expert_indices: List[int]

    Returns
    -------
    mlp: SwitchTransformersSparseMLP
        The merged mlp.
    """
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        # Only merge non-core experts
        expert_indices = torch.tensor(
            [expert_idx.item() for expert_idx in expert_indices if expert_idx not in core_expert_indices]
        )
        with torch.no_grad():
            if len(expert_indices) == 0:
                continue
            # Merge all experts except the dominant one
            group_forwarded_hidden_states = torch.cat([
                forwarded_hidden_states[expert_idx] for expert_idx in expert_indices
            ], dim=0)
            merged_expert = merge_switch_mlp_by_activation_matching_within_and_across_models(
                mlp_list=[mlp.experts[f"expert_{expert_idx}"] for expert_idx in expert_indices],
                forwarded_hidden_states=group_forwarded_hidden_states,
            )
            mlp.experts[f"expert_{expert_indices[0]}"].wi.weight.copy_(merged_expert.wi.weight)
            mlp.experts[f"expert_{expert_indices[0]}"].wo.weight.copy_(merged_expert.wo.weight)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{expert_indices[0]}"]

    return mlp


def merge_non_core_by_groups_within_and_across_models(
        switch_model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        dataloader: DataLoader,
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
        core_experts: Optional[Dict[str, List[int]]] = None,
) -> SwitchTransformersForConditionalGeneration:
    # {name: values}, values  will be of shape (len(dataloader), batch_size * seq_len, d_ff)
    forwarded_hidden_states = dict()

    switch_model.eval()
    handles = []

    def _get_activation_hook(name):
        def hook(module, input, output):
            forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

        return hook

    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Registering forward hook..."
    ):
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
        forwarded_hidden_states[encoder_mlp_name] = []
        forwarded_hidden_states[decoder_mlp_name] = []
        handles.append(switch_model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(encoder_mlp_name))
        )
        handles.append(switch_model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(decoder_mlp_name))
        )

    # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
    router_indices = {name: [] for name in forwarded_hidden_states.keys()}
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="[Merging]Computing activations..."):
            batch = {k: v.cuda() for k, v in batch.items()}
            outputs = switch_model(**batch)
            for layer_idx in grouper.sparse_layer_indices:
                router_indices[f"encoder.block.{layer_idx}.layer.1.mlp"].append(
                    outputs.encoder_router_logits[layer_idx][1].reshape(-1)
                )
                router_indices[f"decoder.block.{layer_idx}.layer.2.mlp"].append(
                    outputs.decoder_router_logits[layer_idx][1].reshape(-1)
                )

    for handle in handles:
        handle.remove()

    num_experts = grouper.num_experts
    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging non core by groups within and across experts..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            encoder_group_labels = grouper.group_state_dict()[encoder_mlp_name]
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[encoder_mlp_name][i][router_indices[encoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.encoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_non_core_experts_within_and_across_models(
                mlp=switch_model.encoder.block[layer_idx].layer[-1].mlp,
                group_labels=encoder_group_labels,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                core_expert_indices=core_experts[encoder_mlp_name],
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            decoder_group_labels = grouper.group_state_dict()[decoder_mlp_name]
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[decoder_mlp_name][i][router_indices[decoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.decoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_non_core_experts_within_and_across_models(
                mlp=switch_model.decoder.block[layer_idx].layer[-1].mlp,
                group_labels=decoder_group_labels,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                core_expert_indices=core_experts[decoder_mlp_name],
            )

    del forwarded_hidden_states, router_indices
    torch.cuda.empty_cache()
    return switch_model


def _merge_mlp_all_non_core_experts_to_one_with_frequency_weighting(
        mlp: SwitchTransformersSparseMLP,
        usage_frequencies: torch.Tensor,
        core_experts: List[int],
) -> SwitchTransformersSparseMLP:
    non_core_expert_indices = torch.tensor(
        [expert_idx for expert_idx in range(len(mlp.experts)) if expert_idx not in core_experts]
    )
    if len(non_core_expert_indices) == 0:
        return mlp
    with torch.no_grad():
        wi_weight_list = torch.stack(
            [mlp.experts[f"expert_{expert_idx}"].wi.weight * usage_frequencies[expert_idx] for expert_idx in
             non_core_expert_indices], dim=0
        )
        wo_weight_list = torch.stack(
            [mlp.experts[f"expert_{expert_idx}"].wo.weight * usage_frequencies[expert_idx] for expert_idx in
             non_core_expert_indices], dim=0
        )
        wi_weight = torch.sum(wi_weight_list, dim=0) / (
                torch.sum(usage_frequencies[non_core_expert_indices], dim=0) + FP32_EPS)
        wo_weight = torch.sum(wo_weight_list, dim=0) / (
                torch.sum(usage_frequencies[non_core_expert_indices], dim=0) + FP32_EPS)
        mlp.experts[f"expert_{non_core_expert_indices[0]}"].wi.weight.copy_(wi_weight)
        mlp.experts[f"expert_{non_core_expert_indices[0]}"].wo.weight.copy_(wo_weight)
        for expert_idx in non_core_expert_indices[1:]:
            # Binding merged experts to the first of them
            mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{non_core_expert_indices[0]}"]

    return mlp


def merge_all_non_core_to_one_with_usage_frequency_weighting(
        model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        core_experts: Dict[str, List[int]],
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
) -> SwitchTransformersForConditionalGeneration:
    """
    Merge all non-core experts to one, and leave core experts alone

    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to merge experts
    grouper: ExpertsGrouperForSwitch
        The grouper to group experts, supposed to have been called `grouper.compute_all_usages()` only
            (i.e. no need to have grouped labels)
    core_experts: Dict[str, List[int]]
        The core experts dict, normally returned by `grouper.group_experts_into_clusters_by_routing_guided_globally()`
    encoder_merging_layers: Optional[List[int]]
        The encoder layers to merge experts, if None, merge all layers
    decoder_merging_layers: Optional[List[int]]
        The decoder layers to merge experts, if None, merge all layers

    Returns
    -------
    model: SwitchTransformersForConditionalGeneration
        The model with merged non-core experts
    """
    usage_frequency_dict = grouper.usage_frequency_state_dict()
    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging all non-core experts to one with usage-frequency-weighted averaging..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            usage_frequencies = usage_frequency_dict[encoder_mlp_name]
            model.encoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_all_non_core_experts_to_one_with_frequency_weighting(
                mlp=model.encoder.block[layer_idx].layer[-1].mlp,
                usage_frequencies=usage_frequencies,
                core_experts=core_experts[encoder_mlp_name]
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            usage_frequencies = usage_frequency_dict[decoder_mlp_name]
            model.decoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_all_non_core_experts_to_one_with_frequency_weighting(
                mlp=model.decoder.block[layer_idx].layer[-1].mlp,
                usage_frequencies=usage_frequencies,
                core_experts=core_experts[decoder_mlp_name]
            )
    return model


def _merge_mlp_all_non_core_experts_to_one_within_and_across_models(
        mlp: SwitchTransformersSparseMLP,
        forwarded_hidden_states: Tuple[torch.Tensor],
        core_experts: List[int],
) -> SwitchTransformersSparseMLP:
    non_core_expert_indices = torch.tensor(
        [expert_idx for expert_idx in range(len(mlp.experts)) if expert_idx not in core_experts]
    )
    if len(non_core_expert_indices) == 0:
        return mlp
    with torch.no_grad():
        non_core_forwarded_hidden_states = torch.cat([
            forwarded_hidden_states[expert_idx] for expert_idx in non_core_expert_indices
        ], dim=0)
        merged_expert = merge_switch_mlp_by_activation_matching_within_and_across_models(
            mlp_list=[mlp.experts[f"expert_{expert_idx}"] for expert_idx in non_core_expert_indices],
            forwarded_hidden_states=non_core_forwarded_hidden_states,
        )
        mlp.experts[f"expert_{non_core_expert_indices[0]}"].wi.weight.copy_(merged_expert.wi.weight)
        mlp.experts[f"expert_{non_core_expert_indices[0]}"].wo.weight.copy_(merged_expert.wo.weight)
        for expert_idx in non_core_expert_indices[1:]:
            # Binding merged experts to the first of them
            mlp.experts[f"expert_{expert_idx}"] = mlp.experts[f"expert_{non_core_expert_indices[0]}"]

    return mlp


def merge_all_non_core_to_one_within_and_across_models(
        switch_model: SwitchTransformersForConditionalGeneration,
        grouper: ExpertsGrouperForSwitch,
        dataloader: DataLoader,
        core_experts: Dict[str, List[int]],
        encoder_merging_layers: Optional[List[int]] = None,
        decoder_merging_layers: Optional[List[int]] = None,
) -> SwitchTransformersForConditionalGeneration:
    # {name: values}, values  will be of shape (len(dataloader), batch_size * seq_len, d_ff)
    forwarded_hidden_states = dict()

    switch_model.eval()
    handles = []

    def _get_activation_hook(name):
        def hook(module, input, output):
            forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

        return hook

    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Registering forward hook..."
    ):
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
        forwarded_hidden_states[encoder_mlp_name] = []
        forwarded_hidden_states[decoder_mlp_name] = []
        handles.append(switch_model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(encoder_mlp_name))
        )
        handles.append(switch_model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(decoder_mlp_name))
        )

    # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
    router_indices = {name: [] for name in forwarded_hidden_states.keys()}
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="[Merging]Computing activations..."):
            batch = {k: v.cuda() for k, v in batch.items()}
            outputs = switch_model(**batch)
            for layer_idx in grouper.sparse_layer_indices:
                router_indices[f"encoder.block.{layer_idx}.layer.1.mlp"].append(
                    outputs.encoder_router_logits[layer_idx][1].reshape(-1)
                )
                router_indices[f"decoder.block.{layer_idx}.layer.2.mlp"].append(
                    outputs.decoder_router_logits[layer_idx][1].reshape(-1)
                )

    for handle in handles:
        handle.remove()

    num_experts = grouper.num_experts

    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging all non-core experts to one within and across experts..."
    ):
        if encoder_merging_layers is None or layer_idx in encoder_merging_layers:
            encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[encoder_mlp_name][i][router_indices[encoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.encoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_all_non_core_experts_to_one_within_and_across_models(
                mlp=switch_model.encoder.block[layer_idx].layer[-1].mlp,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                core_experts=core_experts[encoder_mlp_name],
            )
        if decoder_merging_layers is None or layer_idx in decoder_merging_layers:
            decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            layer_forwarded_hidden_states = tuple()
            for expert_idx in range(num_experts):
                layer_forwarded_hidden_states += (
                    torch.cat(
                        [forwarded_hidden_states[decoder_mlp_name][i][router_indices[decoder_mlp_name][i] == expert_idx]
                         for i in range(len(dataloader))], dim=0),
                )
            switch_model.decoder.block[layer_idx].layer[
                -1].mlp = _merge_mlp_all_non_core_experts_to_one_within_and_across_models(
                mlp=switch_model.decoder.block[layer_idx].layer[-1].mlp,
                forwarded_hidden_states=layer_forwarded_hidden_states,
                core_experts=core_experts[decoder_mlp_name],
            )

    return switch_model


class ExpertUsageFrequencyTracker(object):
    """
    A class to track the usage frequencies of experts in the model during the training process.
    """

    def __init__(
            self,
            named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]],
            beta: Optional[float] = 0.9,
            compute_every_n_steps: Optional[int] = 10,
            device: Optional[str] = 'cpu',
    ):
        """

        Parameters
        ----------
        named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]]
            Normally, it is `model.named_parameters`
        beta: Optional[float]
            The beta parameter in exponential moving average
        compute_every_n_steps: Optional[int]
            Compute usage EMA every n steps
        device: Optional[str]
            The device to store the usage frequency state dict
        """
        self.beta = beta
        self.compute_fisher_every_n_steps = compute_every_n_steps
        self.device = device
        self.exp_expert_usage_frequency_dict = None
        self.last_error = -1

        self._init_usage_frequency_dict(named_parameters_caller)
        self.num_layers = None

    def usage_frequency_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self.exp_expert_usage_frequency_dict)

    def _init_usage_frequency_dict(self,
                                   named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]]):
        self.exp_expert_usage_frequency_dict = OrderedDict()
        for name, _ in named_parameters_caller():
            if "router.classifier" in name:
                mlp_name = name.split(".router")[0]
                self.exp_expert_usage_frequency_dict[mlp_name] = []

        for name, _ in named_parameters_caller():
            if "expert_" in name and "wi" in name:
                mlp_name = name.split(".experts.expert_")[0]
                self.exp_expert_usage_frequency_dict[mlp_name].append(1.0)

        for key, value in self.exp_expert_usage_frequency_dict.items():
            value = torch.tensor(value, device=self.device)
            # init as average used frequency
            self.exp_expert_usage_frequency_dict[key] = value / torch.sum(value)
        self.num_layers = len(self.exp_expert_usage_frequency_dict) // 2

    def _update_exp_usage_frequency_state(
            self,
            model_outputs: Union[List[Seq2SeqMoEOutput], Seq2SeqMoEOutput]
    ):
        usage_frequency_state_dict = OrderedDict()
        is_first_step = self.last_error < 0
        for name in self.exp_expert_usage_frequency_dict.keys():
            usage_frequency_state_dict[name] = torch.zeros_like(self.exp_expert_usage_frequency_dict[name])
            # `name` is like "encoder.block.1.layer.1.mlp" or "decoder.block.3.layer.2.mlp" or etc.,
            # Capture the layer index from the name
            layer_idx = int(name.split(".block.")[1].split(".layer.")[0])
            is_encoder = "encoder" in name
            if is_encoder:
                if isinstance(model_outputs, list):
                    router_expert_index = [op.encoder_router_logits[layer_idx][1].reshape(-1) for op in model_outputs]
                    router_expert_index = torch.concat(router_expert_index)
                else:
                    router_expert_index = model_outputs.encoder_router_logits[layer_idx][1].reshape(-1)
            else:
                if isinstance(model_outputs, list):
                    router_expert_index = [op.decoder_router_logits[layer_idx][1].reshape(-1) for op in model_outputs]
                    router_expert_index = torch.concat(router_expert_index)
                else:
                    router_expert_index = model_outputs.decoder_router_logits[layer_idx][1].reshape(-1)
            for idx in router_expert_index:
                usage_frequency_state_dict[name][idx] += 1
            usage_frequency_state_dict[name] = usage_frequency_state_dict[name] / torch.sum(
                usage_frequency_state_dict[name])
        self.last_error = sum(
            [torch.sum(torch.abs(usage_frequency_state_dict[name] - self.exp_expert_usage_frequency_dict[name]))
             for name in self.exp_expert_usage_frequency_dict.keys()]
        )
        for name in self.exp_expert_usage_frequency_dict.keys():
            if is_first_step:
                self.exp_expert_usage_frequency_dict[name] = usage_frequency_state_dict[name]
            else:
                self.exp_expert_usage_frequency_dict[name] = self.beta * self.exp_expert_usage_frequency_dict[
                    name] + (1 - self.beta) * usage_frequency_state_dict[name]
        # normalize
        for name in self.exp_expert_usage_frequency_dict.keys():
            self.exp_expert_usage_frequency_dict[name] = self.exp_expert_usage_frequency_dict[name] / torch.sum(
                self.exp_expert_usage_frequency_dict[name])

    def step(self, model_outputs: Union[List[Seq2SeqMoEOutput], Seq2SeqMoEOutput], global_step: int) -> float:
        """
        Update the usage frequency state of experts and return the last error of the state which is L1 norm

        Parameters
        ----------
        model_outputs: Union[List[Seq2SeqMoEOutput], Seq2SeqMoEOutput]
            The outputs of the model on dataset for merging
        global_step: int
            The global step of the training process

        Returns
        -------
        float
            The last error of the state which is L1 norm
        """
        if global_step % self.compute_fisher_every_n_steps == 0:
            self._update_exp_usage_frequency_state(model_outputs)
        return self.last_error

mcsmoe/merging/grouping.py
UNITES-Lab/MC-SMoE
import os
from collections import OrderedDict
from copy import deepcopy
from typing import Dict, Optional, List, Union, Tuple, Callable, Iterator

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    PretrainedConfig
)
from transformers.modeling_outputs import MoECausalLMOutputWithPast

from mcsmoe.merging.permutation import (
    permute_fsgpt_ffn_dense_expert_,
    compute_fsgpt_permutation_by_weight_matching,
    compute_switch_permutation_by_activation_matching
)
from mcsmoe.models.fsgpt_moe import (
    FSGPTMoEConfig,
    FSGPTMoEForCausalLM,
    FSGPTMoESparseMLP,
    FSGPTMoETop2Router
)
from mcsmoe.utils.constants import BF16_EPS
from .grouping import SIMILARITY_MAPPING_FUNCTION, LEGAL_SIMILARITY_BASES
from .utils import generate_random_group_labels


class ExpertsGrouperForFSGPT(object):
    def __init__(
            self,
            config: Union[FSGPTMoEConfig, PretrainedConfig],
            similarity_fn: str = "cosine",
            similarity_base: str = "router-logits",
    ):
        if similarity_fn not in SIMILARITY_MAPPING_FUNCTION:
            raise ValueError(
                f"[Merging]similarity_fn should be one of {SIMILARITY_MAPPING_FUNCTION.keys()}, got {similarity_fn} instead.")
        if similarity_base not in LEGAL_SIMILARITY_BASES:
            raise ValueError(
                f"[Merging]similarity_base should be one of {LEGAL_SIMILARITY_BASES}, got {similarity_base} instead.")

        self.num_experts = config.num_experts
        self.hidden_size = config.hidden_size
        self.sparse_layer_indices = list(range(1, config.num_layers, config.sparse_step))
        self.similarity_fn = SIMILARITY_MAPPING_FUNCTION[similarity_fn]
        self.similarity_base = similarity_base
        self._group_state_dict = None
        self._similarity_state_dict = None
        self._usage_frequency_state_dict = None
        self.reset_all()

    def reset_all(self):
        if self.similarity_base == "mse":
            self.similarity_fn = SIMILARITY_MAPPING_FUNCTION["mse"]
            print("[Merging]Set similarity_fn to mse for mse similarity_base.")
        self._group_state_dict = dict()
        self._similarity_state_dict = dict()
        self._usage_frequency_state_dict = dict()
        # Similarity range: [0, 1]
        for layer_idx in self.sparse_layer_indices:
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            self._group_state_dict[ffn_name] = torch.arange(self.num_experts,
                                                            device="cuda")
            self._similarity_state_dict[ffn_name] = torch.zeros(
                (self.num_experts, self.num_experts), device="cuda"
            ) + torch.eye(self.num_experts, device="cuda")
            self._usage_frequency_state_dict[ffn_name] = torch.ones(self.num_experts, device="cuda") / self.num_experts

    def similarity_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self._similarity_state_dict)

    def group_state_dict(self) -> Dict[str, torch.LongTensor]:
        return deepcopy(self._group_state_dict)

    def usage_frequency_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self._usage_frequency_state_dict)

    def save_similarity(self, ffn_name: str, i: int, j: int, similarity: float):
        self._similarity_state_dict[ffn_name][i, j] = similarity
        self._similarity_state_dict[ffn_name][j, i] = similarity

    def get_similarity(self, ffn_name: str, i: int, j: int) -> float:
        return self._similarity_state_dict[ffn_name][i, j].item()

    def get_similarity_matrix(self, ffn_name: str) -> torch.Tensor:
        return deepcopy(self._similarity_state_dict[ffn_name])

    def save_group_state_dict(self, save_dir: str):
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        torch.save(self._group_state_dict, os.path.join(save_dir, "group_state_dict.pt"))

    def load_group_state_dict(self, load_dir: str):
        self._group_state_dict = torch.load(os.path.join(load_dir, "group_state_dict.pt"))

    def _assign_num_groups_per_layer(
            self,
            average_num_groups: int,
            merging_layers: List[int],
    ) -> Dict[str, int]:
        num_grouping_layers = len(merging_layers)
        total_num_groups = average_num_groups * num_grouping_layers + self.num_experts * (
                len(self.sparse_layer_indices) - num_grouping_layers
        )
        all_usage_frequency = []
        usage_frequency_dict = deepcopy(self._usage_frequency_state_dict)
        for i, layer_idx in enumerate(self.sparse_layer_indices):
            # 1. Experts in the excluded layers are always not merged.
            if layer_idx not in merging_layers:
                usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"] = torch.ones_like(
                    usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"]
                )
            # 2. Each layer must have at least one group, set the most used expert in a layer to frequency 1.
            max_usage_index = torch.argmax(usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"])
            usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"][max_usage_index] = 1.0
            # 3. Collect all usage frequency.
            all_usage_frequency.append(usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"])
        all_usage_frequency = torch.cat(all_usage_frequency, dim=0)
        sorted_usage_frequency, sorted_indices = torch.sort(all_usage_frequency, descending=True)
        num_groups_per_layer = dict()

        # Note: When threshold is 0.0, the actual number of groups is smaller than total_num_groups.
        if average_num_groups == self.num_experts:
            total_num_groups = total_num_groups - 1
        frequency_threshold = sorted_usage_frequency[total_num_groups]

        print(f"[Merging]Frequency threshold: {frequency_threshold}")
        if frequency_threshold == 1.0:
            raise ValueError("[Merging]The number of groups is too large, please reduce the number of groups.")
        for i, layer_idx in enumerate(self.sparse_layer_indices):
            num_groups_per_layer[f"decoder.layers.{layer_idx}.ffn"] = torch.sum(
                (usage_frequency_dict[f"decoder.layers.{layer_idx}.ffn"] > frequency_threshold).long()
            ).item()
        return num_groups_per_layer

    def group_experts_into_clusters_by_routing_guided_globally(
            self,
            average_num_groups: int,
            merging_layers: List[int],
            group_capacity: Optional[int] = None,
    ) -> Dict[str, List[int]]:
        """
        Globally group experts into clusters by routing-guided clustering, each layer will have different number of
         clusters. The total number of clusters is determined by average_num_groups.

        Parameters
        ----------
        average_num_groups: int
            The average number of clusters for all layers.
        merging_layers: List[int]
            The layers that are excluded from merging.
        group_capacity: Optional[int]
            The maximum number of experts in each group. If None, the number of experts in each group is not limited.

        Returns
        -------
        core_experts: Dict[str, List[int]]
            The core experts of each cluster
        """
        # 1. Assign num_groups respectively for each layer according to average_num_groups
        group_capacity = group_capacity if group_capacity is not None else self.num_experts
        num_groups_per_layer = self._assign_num_groups_per_layer(
            average_num_groups, merging_layers
        )
        print(f"[Merging]Number of groups of each layer: {num_groups_per_layer}")
        # 2. Group experts into clusters for each layer
        core_experts = dict()
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Globally routing-guided clustering experts into average {average_num_groups} clusters"):
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            num_groups = num_groups_per_layer[ffn_name]
            group_member_count = torch.zeros(num_groups)
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[ffn_name], descending=True)
            # 1. Assign top-K most-used experts with label 0 to K-1 respectively
            core_expert_indices = indices_sorted_by_usage[:num_groups]
            core_experts[ffn_name] = core_expert_indices.tolist()
            for i in range(num_groups):
                self._group_state_dict[ffn_name][indices_sorted_by_usage[i]] = i
                group_member_count[i] += 1
            # 2. Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(ffn_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = core_expert_indices[
                    torch.argmax(similarity_matrix[i, core_expert_indices])
                ]
                most_similar_group_label = self._group_state_dict[ffn_name][most_similar_core]
                self._group_state_dict[ffn_name][i] = most_similar_group_label
                if group_member_count[self._group_state_dict[ffn_name][i]] >= group_capacity:
                    if len(core_expert_indices) == 1:
                        raise ValueError(
                            f"[Merging]The number of groups at Decoder layer {layer_idx} is too small!"
                        )
                    # Kick out the filled group as well as its core, by pop the core from core_experts
                    core_index = torch.argmax(similarity_matrix[i, core_expert_indices])
                    core_expert_indices = torch.cat(
                        [core_expert_indices[:core_index], core_expert_indices[core_index + 1:]]
                    )
        return core_experts

    def group_experts_into_clusters_by_routing_guided(
            self,
            num_groups: int,
    ) -> Dict[str, List[int]]:
        """
        1. Select top-K most-used experts, assign them with label 0 to K-1 respectively.
            These K experts are the cores of each cluster
        2. For each left unassigned expert, assign it to the cluster with the most similar core.

        Returns
        -------
        core_experts: Dict[str, List[int]]
            The core experts of each cluster
        """
        core_experts = {}
        if num_groups > self.num_experts:
            raise ValueError(
                f"[Merging]num_groups should be less than or equal to {self.num_experts}, got {num_groups} instead.")

        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Routing-guided clustering experts into {num_groups} clusters"):
            # Decoder
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            indices_sorted_by_usage = torch.argsort(self._usage_frequency_state_dict[ffn_name], descending=True)
            # 1. Assign top-K most-used experts with label 0 to K-1 respectively
            core_experts[ffn_name] = indices_sorted_by_usage[:num_groups].tolist()
            for i in range(num_groups):
                self._group_state_dict[ffn_name][indices_sorted_by_usage[i]] = i
            # 2. Assign left unassigned experts to the cluster with the most similar core
            similarity_matrix = self.get_similarity_matrix(ffn_name)
            for i in range(num_groups, self.num_experts):
                # Find the most similar core
                most_similar_core = indices_sorted_by_usage[
                    torch.argmax(similarity_matrix[i, indices_sorted_by_usage[:num_groups]])
                ]
                self._group_state_dict[ffn_name][i] = self._group_state_dict[ffn_name][
                    most_similar_core]
        return core_experts

    def compute_all_usages(
            self,
            model: FSGPTMoEForCausalLM,
            dataloader: DataLoader,
    ):
        model.eval()
        # reset all usage frequency to 0
        self._usage_frequency_state_dict = {
            k: torch.zeros_like(v) for k, v in self._usage_frequency_state_dict.items()
        }
        for batch in tqdm(dataloader, desc="[Merging]Computing all usages..."):
            with torch.no_grad():
                with torch.autocast("cuda"):
                    batch = {k: v.cuda() for k, v in batch.items()}
                    outputs = model(**batch)
                    for layer_idx in self.sparse_layer_indices:
                        ffn_name = f"decoder.layers.{layer_idx}.ffn"
                        decoder_router_probs = outputs.router_logits[layer_idx][1]
                        decoder_router_expert_index = torch.argsort(decoder_router_probs)[:, -2:].reshape(-1)
                        for idx in decoder_router_expert_index:
                            self._usage_frequency_state_dict[ffn_name][idx] += 1
        self._usage_frequency_state_dict = {
            k: v / torch.sum(v) for k, v in self._usage_frequency_state_dict.items()
        }

    def compute_all_similarities(
            self,
            model: FSGPTMoEForCausalLM,
            dataloader: DataLoader = None,
    ):
        if self.similarity_base not in ["weight", "router-weight"] and dataloader is None:
            raise ValueError(
                "[Merging]Dataloader should be provided when similarity_base is not 'weight' or 'router-weight'")

        model = model.eval()
        if self.similarity_base == 'router-logits':
            self._compute_all_similarities_by_router_logits(model, dataloader)
        elif self.similarity_base == 'random':
            self._compute_all_similarities_by_random()
        else:
            raise NotImplementedError

    def _compute_all_similarities_by_random(self):
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities randomly..."):
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            for i in range(self.num_experts):
                for j in range(i + 1, self.num_experts):
                    self.save_similarity(ffn_name, i, j, torch.rand(1).item())

    def _compute_all_similarities_by_router_logits(
            self, model: FSGPTMoEForCausalLM, dataloader: DataLoader
    ):
        router_logits_outputs_list = []
        for batch in dataloader:
            with torch.no_grad():
                with torch.autocast("cuda"):
                    batch = {k: v.cuda() for k, v in batch.items()}
                    outputs = model(**batch)
                    router_logits_outputs_list.append(outputs.router_logits)
        for layer_idx in tqdm(self.sparse_layer_indices, desc="[Merging]Computing similarities by router logits..."):
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            router_logits_list = [
                router_logits[layer_idx][0].reshape(
                    -1, self.num_experts) for router_logits in router_logits_outputs_list
            ]
            router_logits = torch.cat(router_logits_list, dim=0)
            with torch.no_grad():
                for i in range(self.num_experts):
                    for j in range(i + 1, self.num_experts):
                        # Decoder
                        i_flat = router_logits[:, i].flatten()
                        j_flat = router_logits[:, j].flatten()
                        similarity = self.similarity_fn(i_flat, j_flat)
                        self.save_similarity(ffn_name, i, j, similarity)

    def group_experts_randomly(
            self,
            num_groups: int,
    ):
        for layer_idx in tqdm(self.sparse_layer_indices,
                              desc=f"Randomly merging experts into {num_groups} clusters"):
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            group_labels = generate_random_group_labels(self.num_experts, num_groups)
            self._group_state_dict[ffn_name] = group_labels


def _merge_ffn_experts_by_usage_frequency_weighting(
        ffn: FSGPTMoESparseMLP,
        group_labels: torch.LongTensor,
        usage_frequencies: torch.Tensor,
) -> FSGPTMoESparseMLP:
    assert len(group_labels) == len(usage_frequencies) == len(ffn.experts)
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        with torch.no_grad():
            fc1_weight_list = torch.stack(
                [ffn.experts[f"expert_{expert_idx}"].fc1.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            fc2_weight_list = torch.stack(
                [ffn.experts[f"expert_{expert_idx}"].fc2.weight * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            fc1_bias_list = torch.stack(
                [ffn.experts[f"expert_{expert_idx}"].fc1.bias * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            fc2_bias_list = torch.stack(
                [ffn.experts[f"expert_{expert_idx}"].fc2.bias * usage_frequencies[expert_idx] for expert_idx in
                 expert_indices], dim=0
            )
            fc1_weight = torch.sum(fc1_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + BF16_EPS)
            fc1_bias = torch.sum(fc1_bias_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + BF16_EPS)
            fc2_weight = torch.sum(fc2_weight_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + BF16_EPS)
            fc2_bias = torch.sum(fc2_bias_list, dim=0) / (
                    torch.sum(usage_frequencies[expert_indices], dim=0) + BF16_EPS)
            ffn.experts[f"expert_{expert_indices[0]}"].fc1.weight.copy_(fc1_weight)
            ffn.experts[f"expert_{expert_indices[0]}"].fc1.bias.copy_(fc1_bias)
            ffn.experts[f"expert_{expert_indices[0]}"].fc2.weight.copy_(fc2_weight)
            ffn.experts[f"expert_{expert_indices[0]}"].fc2.bias.copy_(fc2_bias)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                ffn.experts[f"expert_{expert_idx}"] = ffn.experts[f"expert_{expert_indices[0]}"]
    return ffn


def merge_by_groups_with_usage_frequency_weighting(
        model: FSGPTMoEForCausalLM,
        grouper: ExpertsGrouperForFSGPT,
        merging_layers: Optional[List[int]] = None,
) -> FSGPTMoEForCausalLM:
    usage_frequency_dict = grouper.usage_frequency_state_dict()
    for layer_idx in tqdm(
            grouper.sparse_layer_indices,
            desc=f"[Merging]Merging experts with usage-frequency-weighted averaging..."
    ):
        if merging_layers is None or layer_idx in merging_layers:
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            group_labels = grouper.group_state_dict()[ffn_name]
            usage_frequencies = usage_frequency_dict[ffn_name]
            model.decoder.layers[layer_idx].ffn = _merge_ffn_experts_by_usage_frequency_weighting(
                ffn=model.decoder.layers[layer_idx].ffn,
                group_labels=group_labels,
                usage_frequencies=usage_frequencies,
            )
    return model


def _merge_ffn_experts_by_averaging(
        ffn: FSGPTMoESparseMLP,
        group_labels: torch.LongTensor,
        permute: bool,
        permute_strategy: str,
        forwarded_hidden_states: Optional[Tuple[torch.Tensor]] = None,
) -> FSGPTMoESparseMLP:
    for label in group_labels.unique():
        expert_indices = torch.where(group_labels == label)[0]
        if permute and permute_strategy == "weight-matching":
            for expert_idx in expert_indices[1:]:
                perm = compute_fsgpt_permutation_by_weight_matching(
                    reference_ffn=ffn.experts[f"expert_{expert_indices[0]}"],
                    target_ffn=ffn.experts[f"expert_{expert_idx}"],
                    include_wo=True
                )
                ffn.experts[f"expert_{expert_idx}"] = permute_fsgpt_ffn_dense_expert_(
                    ffn.experts[f"expert_{expert_idx}"], perm
                )
        elif permute and permute_strategy == "activation-matching":
            group_forwarded_hidden_states = torch.cat([
                forwarded_hidden_states[expert_idx] for expert_idx in expert_indices
            ], dim=0)
            for expert_idx in expert_indices[1:]:
                perm = compute_switch_permutation_by_activation_matching(
                    reference_mlp=ffn.experts[f"expert_{expert_indices[0]}"],
                    target_mlp=ffn.experts[f"expert_{expert_idx}"],
                    forwarded_hidden_states=group_forwarded_hidden_states,
                )
                ffn.experts[f"expert_{expert_idx}"] = permute_fsgpt_ffn_dense_expert_(
                    ffn.experts[f"expert_{expert_idx}"], perm
                )
        elif permute:
            raise ValueError(f"Unknown permute strategy: {permute_strategy}")

        with torch.no_grad():
            fc1_weight = torch.mean(
                torch.stack([ffn.experts[f"expert_{expert_idx}"].fc1.weight for expert_idx in expert_indices]),
                dim=0
            )
            fc2_weight = torch.mean(
                torch.stack([ffn.experts[f"expert_{expert_idx}"].fc2.weight for expert_idx in expert_indices]),
                dim=0
            )
            fc1_bias = torch.mean(
                torch.stack([ffn.experts[f"expert_{expert_idx}"].fc1.bias for expert_idx in expert_indices]),
                dim=0
            )
            fc2_bias = torch.mean(
                torch.stack([ffn.experts[f"expert_{expert_idx}"].fc2.bias for expert_idx in expert_indices]),
                dim=0
            )
            ffn.experts[f"expert_{expert_indices[0]}"].fc1.weight.copy_(fc1_weight)
            ffn.experts[f"expert_{expert_indices[0]}"].fc2.weight.copy_(fc2_weight)
            ffn.experts[f"expert_{expert_indices[0]}"].fc1.bias.copy_(fc1_bias)
            ffn.experts[f"expert_{expert_indices[0]}"].fc2.bias.copy_(fc2_bias)
            for expert_idx in expert_indices[1:]:
                # Binding merged experts to the first of them
                ffn.experts[f"expert_{expert_idx}"] = ffn.experts[f"expert_{expert_indices[0]}"]
    return ffn


def merge_by_groups(
        model: FSGPTMoEForCausalLM,
        grouper: ExpertsGrouperForFSGPT,
        merging_layers: Optional[List[int]] = None,
        permute: Optional[bool] = False,
        permute_strategy: Optional[str] = "weight-matching",
        dataloader: Optional[DataLoader] = None,
) -> FSGPTMoEForCausalLM:
    forwarded_hidden_states = dict()
    router_indices = dict()
    if permute_strategy == "activation-matching":
        model.eval().cuda()
        handles = []

        def _get_activation_hook(name):
            def hook(module, input, output):
                forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

            return hook

        for layer_idx in tqdm(
                grouper.sparse_layer_indices,
                desc=f"[Merging]Registering forward hook..."
        ):
            fff_name = f"decoder.layers.{layer_idx}.ffn"
            forwarded_hidden_states[fff_name] = []
            handles.append(model.decoder.layers[layer_idx].ffn.register_forward_hook(
                _get_activation_hook(fff_name))
            )

        # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
        router_indices = {name: [] for name in forwarded_hidden_states.keys()}
        with torch.no_grad(), torch.autocast("cuda"):
            for batch in tqdm(dataloader, desc="[Merging]Computing activations..."):
                batch = {k: v.cuda() for k, v in batch.items()}
                outputs = model(**batch)
                for layer_idx in grouper.sparse_layer_indices:
                    router_indices[f"decoder.layers.{layer_idx}.ffn"].append(
                        outputs.router_logits[layer_idx][2].reshape(-1)
                    )

        for handle in handles:
            handle.remove()

    num_experts = grouper.num_experts

    if permute:
        print(f"[Merging]Permutation is enabled, will permute experts by {permute_strategy} in the same group.")
    for layer_idx in tqdm(grouper.sparse_layer_indices,
                          desc="[Merging]Merging experts with averaging..."):
        if merging_layers is None or layer_idx in merging_layers:
            fff_name = f"decoder.layers.{layer_idx}.ffn"
            group_labels = grouper.group_state_dict()[fff_name]
            if permute_strategy == "activation-matching":
                layer_forwarded_hidden_states = tuple()
                for expert_idx in range(num_experts):
                    layer_forwarded_hidden_states += (
                        torch.cat(
                            [forwarded_hidden_states[fff_name][i][
                                 router_indices[fff_name][i] == expert_idx]
                             for i in range(len(dataloader))], dim=0),
                    )
                model.decoder.layers[layer_idx].ffn = _merge_ffn_experts_by_averaging(
                    ffn=model.decoder.layers[layer_idx].ffn,
                    group_labels=group_labels,
                    permute=permute,
                    permute_strategy=permute_strategy,
                    forwarded_hidden_states=layer_forwarded_hidden_states
                )
            else:
                model.decoder.layers[layer_idx].ffn = _merge_ffn_experts_by_averaging(
                    ffn=model.decoder.layers[layer_idx].ffn,
                    group_labels=group_labels,
                    permute=permute,
                    permute_strategy=permute_strategy,
                )
    return model


def _prune_ffn_experts_by_dropping_non_core(
        ffn: FSGPTMoESparseMLP,
        core_expert_indices: List[int],
) -> FSGPTMoESparseMLP:
    non_core_expert_mask = torch.ones(len(ffn.experts), dtype=torch.bool)
    non_core_expert_mask[core_expert_indices] = False
    non_core_expert_indices = torch.where(non_core_expert_mask)[0]
    bf16_mask_value = -1e4

    def _custom_forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.LongTensor] = None) -> Tuple:
        self.input_dtype = hidden_states.dtype
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.reshape((batch_size * sequence_length), hidden_dim)
        hidden_states = hidden_states.to(self.dtype)
        self._cast_classifier()
        router_logits = self.classifier(hidden_states)

        router_logits = router_logits.masked_fill(non_core_expert_mask.to(router_logits.device), bf16_mask_value)

        top_1_mask, router_probs = self.route_tokens(router_logits, self.input_dtype, padding_mask)
        return top_1_mask, router_probs, router_logits

    ffn.router.forward = _custom_forward.__get__(ffn.router, FSGPTMoETop2Router)

    for expert_idx in non_core_expert_indices:
        # fake expert, this will raise error if used by mistake
        ffn.experts[f"expert_{expert_idx}"] = torch.nn.Linear(ffn.router.classifier.in_features, 1)

    return ffn


def prune_non_core_experts_by_groups(
        model: FSGPTMoEForCausalLM,
        grouper: ExpertsGrouperForFSGPT,
        core_experts: Dict[str, List[int]],
        merging_layers: Optional[List[int]] = None,
) -> FSGPTMoEForCausalLM:
    for layer_idx in tqdm(grouper.sparse_layer_indices, desc="[Merging]Pruning non-core experts by dropping..."):
        if merging_layers is None or layer_idx in merging_layers:
            ffn_name = f"decoder.layers.{layer_idx}.ffn"
            model.decoder.layers[layer_idx].ffn = _prune_ffn_experts_by_dropping_non_core(
                ffn=model.decoder.layers[layer_idx].ffn,
                core_expert_indices=core_experts[ffn_name]
            )
    return model


class ExpertUsageFrequencyTrackerForFSGPT(object):
    """
    A class to track the usage frequencies of experts in the model during the training process.
    """

    def __init__(
            self,
            named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]],
            beta: Optional[float] = 0.9,
            compute_every_n_steps: Optional[int] = 10,
            device: Optional[str] = 'cpu',
    ):
        """

        Parameters
        ----------
        named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]]
            Normally, it is `model.named_parameters`
        beta: Optional[float]
            The beta parameter in exponential moving average
        compute_every_n_steps: Optional[int]
            Compute usage EMA every n steps
        device: Optional[str]
            The device to store the usage frequency state dict
        """
        self.beta = beta
        self.compute_fisher_every_n_steps = compute_every_n_steps
        self.device = device
        self.exp_expert_usage_frequency_dict = None
        self.last_error = -1

        self._init_usage_frequency_dict(named_parameters_caller)
        self.num_layers = None

    def usage_frequency_state_dict(self) -> Dict[str, torch.Tensor]:
        return deepcopy(self.exp_expert_usage_frequency_dict)

    def _init_usage_frequency_dict(self,
                                   named_parameters_caller: Callable[[], Iterator[Tuple[str, torch.nn.Parameter]]]):
        self.exp_expert_usage_frequency_dict = OrderedDict()
        for name, _ in named_parameters_caller():
            if "router.classifier" in name:
                mlp_name = name.split(".router")[0]
                self.exp_expert_usage_frequency_dict[mlp_name] = []

        for name, _ in named_parameters_caller():
            if "expert_" in name and "fc1.weight" in name:
                mlp_name = name.split(".experts.expert_")[0]
                self.exp_expert_usage_frequency_dict[mlp_name].append(1.0)

        for key, value in self.exp_expert_usage_frequency_dict.items():
            value = torch.tensor(value, device=self.device)
            # init as average used frequency
            self.exp_expert_usage_frequency_dict[key] = value / torch.sum(value)
        self.num_layers = len(self.exp_expert_usage_frequency_dict) // 2

    def _update_exp_usage_frequency_state(
            self,
            model_outputs: Union[List[MoECausalLMOutputWithPast], MoECausalLMOutputWithPast]
    ):
        usage_frequency_state_dict = OrderedDict()
        is_first_step = self.last_error < 0
        for name in self.exp_expert_usage_frequency_dict.keys():
            usage_frequency_state_dict[name] = torch.zeros_like(self.exp_expert_usage_frequency_dict[name])
            # Capture the layer index from the name
            layer_idx = int(name.split(".layers.")[1].split(".ffn")[0])
            if isinstance(model_outputs, list):
                router_expert_index = [
                    torch.argsort(outputs.router_logits[layer_idx][1])[:, -2:].reshape(-1) for outputs in model_outputs]
                router_expert_index = torch.concat(router_expert_index)
            else:
                router_expert_index = torch.argsort(model_outputs.router_logits[layer_idx][1])[:, -2:].reshape(-1)
            for idx in router_expert_index:
                usage_frequency_state_dict[name][idx] += 1
            usage_frequency_state_dict[name] = usage_frequency_state_dict[name] / torch.sum(
                usage_frequency_state_dict[name])
        self.last_error = sum(
            [torch.sum(torch.abs(usage_frequency_state_dict[name] - self.exp_expert_usage_frequency_dict[name]))
             for name in self.exp_expert_usage_frequency_dict.keys()]
        )
        for name in self.exp_expert_usage_frequency_dict.keys():
            if is_first_step:
                self.exp_expert_usage_frequency_dict[name] = usage_frequency_state_dict[name]
            else:
                self.exp_expert_usage_frequency_dict[name] = self.beta * self.exp_expert_usage_frequency_dict[
                    name] + (1 - self.beta) * usage_frequency_state_dict[name]
        # normalize
        for name in self.exp_expert_usage_frequency_dict.keys():
            self.exp_expert_usage_frequency_dict[name] = self.exp_expert_usage_frequency_dict[name] / torch.sum(
                self.exp_expert_usage_frequency_dict[name])

    def step(self, model_outputs: Union[List[MoECausalLMOutputWithPast], MoECausalLMOutputWithPast],
             global_step: int) -> float:
        """
        Update the usage frequency state of experts and return the last error of the state which is L1 norm

        Parameters
        ----------
        model_outputs: Union[List[MoECausalLMOutputWithPast], MoECausalLMOutputWithPast]
            The outputs of the model on dataset for merging
        global_step: int
            The global step of the training process

        Returns
        -------
        float
            The last error of the state which is L1 norm
        """
        if global_step % self.compute_fisher_every_n_steps == 0:
            self._update_exp_usage_frequency_state(model_outputs)
        return self.last_error

mcsmoe/merging/grouping_fsgpt.py
UNITES-Lab/MC-SMoE
from copy import deepcopy
from typing import Optional, List

import torch
from scipy.optimize import linear_sum_assignment
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers.models.switch_transformers.modeling_switch_transformers import (
    SwitchTransformersDenseActDense,
    SwitchTransformersForConditionalGeneration
)

from mcsmoe.models.fsgpt_moe import (
    FSGPTMoEDenseActDense,
    FSGPTMoEForCausalLM,
    FSGPTMoEConfig
)
from mcsmoe.utils.constants import FP32_EPS

__all__ = [
    "permute_switch_mlp_dense_expert_",
    "compute_switch_permutation_by_weight_matching",
    "align_switch_permutation_for_all_experts_by_weight_matching",
    "align_switch_permutation_for_all_experts_by_activation_matching",
    "align_fsgpt_permutation_for_all_experts_by_weight_matching",
    "compute_switch_permutation_by_activation_matching",
    "merge_switch_mlp_by_activation_matching_within_and_across_models",
    "merge_switch_mlp_by_weight_matching_within_and_across_models",
    "permute_fsgpt_ffn_dense_expert_",
    "compute_fsgpt_permutation_by_weight_matching",
]


def permute_switch_mlp_dense_expert_(
        dense_mlp: SwitchTransformersDenseActDense,
        perm: torch.Tensor,
) -> SwitchTransformersDenseActDense:
    """
    Permute the weights of MLP according to the given permutation, the behavior of the MLP is still the same as before.

    The original weights of the MLP are Wi and Wo, the permutation matrix is P, the input is x, then the output of the
    MLP is Wo @ act(Wi @ x^T), where act is the activation function. After the permutation, the output of the MLP is
    (Wo @ P) @ act((P @ Wi) @ x^T), which is the same as the original output.

    So the permutation of the MLP is (Wi, Wo) -> (P @ Wi, Wo @ P).

    Parameters
    ----------
    dense_mlp: SwitchTransformersDenseActDense
        The MLP to be permuted.
    perm: torch.Tensor, of shape (d_ff, )
        The permutation vector (not permutation matrix) to be applied to the MLP.

    Returns
    -------
    mlp: SwitchTransformersDenseActDense
        The permuted MLP, from in-place operation.

    Examples
    --------
    >>> from transformers import SwitchTransformersConfig
    >>> mlp = SwitchTransformersDenseActDense(SwitchTransformersConfig(d_model=16, d_ff=32, dropout_rate=0))
    >>> dummy_input = torch.randn(4, 16)
    >>> permuted_mlp = permute_switch_mlp_dense_expert_(deepcopy(mlp), torch.randperm(32))
    >>> torch.allclose(mlp(dummy_input), permuted_mlp(dummy_input))
    True
    """
    d_ff = dense_mlp.wi.out_features

    # Check the permutation vector
    if perm.shape != (d_ff,):
        raise ValueError(f"The shape of the permutation vector should be (d_ff, ), but got {perm.shape}.")
    if not torch.allclose(perm.sort()[0], torch.arange(d_ff, device=perm.device)):
        raise ValueError("The permutation vector should be a permutation.")

    # Permute the weights of the MLP
    with torch.no_grad():
        dense_mlp.wi.weight.data = dense_mlp.wi.weight.data[perm, :]
        dense_mlp.wo.weight.data = dense_mlp.wo.weight.data[:, perm]

    return dense_mlp


def compute_switch_permutation_by_weight_matching(
        reference_mlp: SwitchTransformersDenseActDense,
        target_mlp: SwitchTransformersDenseActDense,
        include_wo: bool,
) -> torch.Tensor:
    """
    Compute the permutation vector by weight match that can permute the weights of the target MLP
     to match the weights of the reference MLP.

    Specifically, find a permutation of the target MLP such that the summation of L2 error between each row of
     reference MLP and target MLP is minimal.

    This can be formulated as linear sum assignment problem:
        argmin{||vec(Wr) - vec(Wt @ P)||^2} = argmax{vec(Wr) * vec(Wt @ P)}

    Parameters
    ----------
    reference_mlp: SwitchTransformersDenseActDense
        The reference MLP.
    target_mlp: SwitchTransformersDenseActDense
        The target MLP.
    include_wo: bool, default False
        Whether to include the (each column) weights of the second layer of MLP in the weight matching.

    Returns
    -------
    perm: torch.Tensor, of shape (d_ff, )
        The permutation vector that can permute the weights of the target MLP to match the weights of the reference MLP.

    Examples
    --------
    >>> from transformers import SwitchTransformersConfig
    >>> mlp = SwitchTransformersDenseActDense(SwitchTransformersConfig(d_model=16, d_ff=32, dropout_rate=0))
    >>> perm = compute_switch_permutation_by_weight_matching(mlp, mlp)
    >>> torch.allclose(perm, torch.arange(32))
    True
    """
    with torch.no_grad():
        lsa_cost_matrix = torch.mm(
            reference_mlp.wi.weight.data, target_mlp.wi.weight.data.t()
        )
        if include_wo:
            lsa_cost_matrix += torch.mm(
                reference_mlp.wo.weight.data.t(), target_mlp.wo.weight.data
            )
    _, perm = linear_sum_assignment(lsa_cost_matrix.cpu().numpy(), maximize=True)
    return torch.from_numpy(perm).to(lsa_cost_matrix.device)


def align_switch_permutation_for_all_experts_by_weight_matching(
        switch_model: SwitchTransformersForConditionalGeneration,
        include_wo: bool,
) -> SwitchTransformersForConditionalGeneration:
    """
    Align the permutation of all experts in the switch model by weight matching.

    Parameters
    ----------
    switch_model: SwitchTransformersForConditionalGeneration
        The switch model to be aligned.
    include_wo: bool, default False
        Whether to include the (each column) weights of the second layer of MLP in the weight matching.

    Returns
    -------
    switch_model: SwitchTransformersForConditionalGeneration
        The aligned switch model, from in-place operation.
    """
    config = switch_model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.encoder_sparse_step))
    num_experts = config.num_experts
    for layer_idx in tqdm(sparse_layer_indices,
                          desc=f"[Permutation]Aligning permutation {'with' if include_wo else 'without'} Wo"):
        encoder_mlp = switch_model.encoder.block[layer_idx].layer[-1].mlp
        decoder_mlp = switch_model.decoder.block[layer_idx].layer[-1].mlp
        for expert_idx in range(1, num_experts):
            # Permute encoder
            perm = compute_switch_permutation_by_weight_matching(
                encoder_mlp.experts["expert_0"],
                encoder_mlp.experts[f"expert_{expert_idx}"],
                include_wo=include_wo,
            )
            encoder_mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                encoder_mlp.experts[f"expert_{expert_idx}"], perm
            )
            # Permute decoder
            perm = compute_switch_permutation_by_weight_matching(
                decoder_mlp.experts["expert_0"],
                decoder_mlp.experts[f"expert_{expert_idx}"],
                include_wo=include_wo,
            )
            decoder_mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                decoder_mlp.experts[f"expert_{expert_idx}"], perm
            )

    return switch_model


def permute_fsgpt_ffn_dense_expert_(
        dense_ffn: FSGPTMoEDenseActDense,
        perm: torch.Tensor,
) -> FSGPTMoEDenseActDense:
    """
    Similar to permute_switch_mlp_dense_expert, but for fsgpt

    Examples
    --------
    >>> ffn = FSGPTMoEDenseActDense(FSGPTMoEConfig(hidden_size=8), intermediate_size=32).eval()
    >>> perm = torch.randperm(32)
    >>> permuted_ffn = permute_fsgpt_ffn_dense_expert_(deepcopy(ffn), perm)
    >>> dummy_input = torch.arange(32).reshape(4,8).float()
    >>> torch.allclose(ffn(dummy_input), permuted_ffn(dummy_input))
    True
    """
    hidden_size = dense_ffn.fc1.out_features

    # Check the permutation vector
    if perm.shape != (hidden_size,):
        raise ValueError(f"The shape of the permutation vector should be (d_ff, ), but got {perm.shape}.")
    if not torch.allclose(perm.sort()[0], torch.arange(hidden_size, device=perm.device)):
        raise ValueError("The permutation vector should be a permutation.")

    # Permute the weights of the FFN
    with torch.no_grad():
        dense_ffn.fc1.weight.data = dense_ffn.fc1.weight.data[perm, :]
        dense_ffn.fc1.bias.data = dense_ffn.fc1.bias.data[perm]
        dense_ffn.fc2.weight.data = dense_ffn.fc2.weight.data[:, perm]

    return dense_ffn


def compute_fsgpt_permutation_by_weight_matching(
        reference_ffn: FSGPTMoEDenseActDense,
        target_ffn: FSGPTMoEDenseActDense,
        include_wo: Optional[bool] = True,
) -> torch.Tensor:
    """
    Similar to compute_switch_permutation_by_weight_matching, but for fsgpt

    Will transform BF16 to FP16 for computation
    Examples
    --------
    >>> ffn = FSGPTMoEDenseActDense(FSGPTMoEConfig(hidden_size=8), intermediate_size=32)
    >>> perm = compute_fsgpt_permutation_by_weight_matching(ffn, ffn)
    >>> torch.allclose(perm, torch.arange(32))
    True
    """
    with torch.no_grad():
        lsa_cost_matrix = torch.mm(
            reference_ffn.fc1.weight.data.float(), target_ffn.fc1.weight.data.t().float()
        )
        if include_wo:
            lsa_cost_matrix += torch.mm(
                reference_ffn.fc2.weight.data.t().float(), target_ffn.fc2.weight.data.float()
            )
    _, perm = linear_sum_assignment(lsa_cost_matrix.cpu().numpy(), maximize=True)
    return torch.from_numpy(perm).to(lsa_cost_matrix.device)


def align_fsgpt_permutation_for_all_experts_by_weight_matching(
        fsgpt_model: FSGPTMoEForCausalLM,
        include_wo: bool = True,
) -> FSGPTMoEForCausalLM:
    config = fsgpt_model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.sparse_step))
    num_experts = config.num_experts
    progress_bar = tqdm(range(len(sparse_layer_indices) * (num_experts - 1)),
                        desc=f"[Permutation]Aligning permutation {'with' if include_wo else 'without'} Wo")
    for layer_idx in sparse_layer_indices:
        ffn = fsgpt_model.decoder.layers[layer_idx].ffn
        for expert_idx in range(1, num_experts):
            perm = compute_fsgpt_permutation_by_weight_matching(
                ffn.experts["expert_0"],
                ffn.experts[f"expert_{expert_idx}"],
                include_wo=include_wo,
            )
            ffn.experts[f"expert_{expert_idx}"] = permute_fsgpt_ffn_dense_expert_(
                ffn.experts[f"expert_{expert_idx}"], perm
            )
            progress_bar.update(1)

    return fsgpt_model


def compute_switch_permutation_by_activation_matching(
        reference_mlp: SwitchTransformersDenseActDense,
        target_mlp: SwitchTransformersDenseActDense,
        forwarded_hidden_states: torch.Tensor,
        mini_batch_size: Optional[int] = None,
) -> torch.Tensor:
    """
    In particular, the forwarded_hidden_states is supposed to be gathered from tokens routed to both
    reference_mlp and target_mlp, and the permutation is computed by matching the activation of the
    forwarded_hidden_states.

    Parameters
    ----------
    reference_mlp: SwitchTransformersDenseActDense
        The reference MLP.
    target_mlp: SwitchTransformersDenseActDense
        The target MLP.
    forwarded_hidden_states: torch.Tensor, of shape (batch_size * seq_len, d_ff)
        The hidden states that are forwarded to both reference_mlp and target_mlp.
    mini_batch_size: int, default None
        The mini batch size for the activation matching. If None, the mini batch size is set to be the batch size.

    Returns
    -------
    perm: torch.Tensor, of shape (d_ff, )
        The permutation vector that can permute the weights of the target MLP to match the weights of the reference MLP.

    Examples
    --------
    >>> from transformers import SwitchTransformersConfig
    >>> from copy import deepcopy
    >>> mlp = SwitchTransformersDenseActDense(SwitchTransformersConfig(d_model=16, d_ff=32, dropout_rate=0))
    >>> perm = compute_switch_permutation_by_activation_matching(mlp, deepcopy(mlp), torch.randn(64, 16))
    >>> torch.allclose(perm, torch.arange(32))
    True
    """
    if len(forwarded_hidden_states) == 0 or len(forwarded_hidden_states) == 1:
        return torch.arange(reference_mlp.wi.out_features, device=forwarded_hidden_states.device)

    if forwarded_hidden_states.shape[-1] != reference_mlp.wi.in_features:
        raise ValueError(
            f"The last dimension of forwarded_hidden_states should be {reference_mlp.wi.in_features}, "
            f"but got {forwarded_hidden_states.shape[-1]}."
        )
    if mini_batch_size is None:
        mini_batch_size = forwarded_hidden_states.shape[0]
    reference_activations = []
    target_activations = []

    def _ref_activation_hook(module, input, output):
        reference_activations.append(input[0].detach().reshape(-1, input[0].shape[-1]))

    def _target_activation_hook(module, input, output):
        target_activations.append(input[0].detach().reshape(-1, input[0].shape[-1]))

    reference_handle = reference_mlp.wo.register_forward_hook(_ref_activation_hook)
    target_handle = target_mlp.wo.register_forward_hook(_target_activation_hook)
    with torch.no_grad():
        for i in range(0, forwarded_hidden_states.shape[0], mini_batch_size):
            reference_mlp(forwarded_hidden_states[i:i + mini_batch_size])
            target_mlp(forwarded_hidden_states[i:i + mini_batch_size])

    reference_activations = torch.cat(reference_activations, dim=0)  # (batch_size * seq_len, d_ff)
    target_activations = torch.cat(target_activations, dim=0)  # (batch_size * seq_len, d_ff)

    # Compute the correlation matrix as the cost matrix
    mean_ref = reference_activations.mean(dim=0, keepdim=True)  # (1, d_ff)
    mean_target = target_activations.mean(dim=0, keepdim=True)  # (1, d_ff)
    std_ref = reference_activations.std(dim=0, keepdim=True)  # (1, d_ff)
    std_target = target_activations.std(dim=0, keepdim=True)  # (1, d_ff)
    covar = torch.mm(
        (reference_activations - mean_ref).t(),
        (target_activations - mean_target)
    ) / (reference_activations.shape[0] - 1)  # (d_ff, d_ff)
    cost_matrix = covar / (std_ref.t() * std_target + FP32_EPS)  # (d_ff, d_ff)

    _, perm = linear_sum_assignment(cost_matrix.cpu().numpy(), maximize=True)
    reference_handle.remove()
    target_handle.remove()
    return torch.from_numpy(perm).to(cost_matrix.device)


def align_switch_permutation_for_all_experts_by_activation_matching(
        switch_model: SwitchTransformersForConditionalGeneration,
        dataloader: DataLoader,
) -> SwitchTransformersForConditionalGeneration:
    """
    Align the permutation of all experts in the switch model by activation matching.

    In practice, the hidden states of the tokens routed to both reference_mlp and target_mlp are gathered
        from the dataloader, and the permutation is computed by matching the activation of the hidden states.

    Parameters
    ----------
    switch_model: SwitchTransformersForConditionalGeneration
        The switch model to be aligned.
    dataloader: DataLoader
        The dataloader to be used to gather the hidden states for activation matching.

    Returns
    -------
    switch_model: SwitchTransformersForConditionalGeneration
        The aligned switch model, from in-place operation.

    """
    # {name: values}, values  will be of shape (len(dataloader), batch_size * seq_len, d_ff)
    forwarded_hidden_states = dict()

    switch_model.eval()
    config = switch_model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.encoder_sparse_step))

    # Register the activation hook for all experts
    handles = []

    def _get_activation_hook(name):
        def hook(module, input, output):
            forwarded_hidden_states[name].append(input[0].detach().reshape(-1, input[0].shape[-1]))

        return hook

    for layer_idx in tqdm(sparse_layer_indices, desc="[Permutation]Registering forward hook..."):
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
        forwarded_hidden_states[encoder_mlp_name] = []
        forwarded_hidden_states[decoder_mlp_name] = []
        handles.append(switch_model.encoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(encoder_mlp_name))
        )
        handles.append(switch_model.decoder.block[layer_idx].layer[-1].mlp.register_forward_hook(
            _get_activation_hook(decoder_mlp_name))
        )

    # {name: values}, values will be of shape (len(dataloader), batch_size * seq_len)
    router_indices = {name: [] for name in forwarded_hidden_states.keys()}
    with torch.no_grad():
        for batch in tqdm(dataloader, desc="[Permutation]Computing activations..."):
            batch = {k: v.cuda() for k, v in batch.items()}
            output = switch_model(**batch)
            for layer_idx in sparse_layer_indices:
                router_indices[f"encoder.block.{layer_idx}.layer.1.mlp"].append(
                    output.encoder_router_logits[layer_idx][1].reshape(-1)
                )
                router_indices[f"decoder.block.{layer_idx}.layer.2.mlp"].append(
                    output.decoder_router_logits[layer_idx][1].reshape(-1)
                )

    # Compute the permutation for all experts
    num_experts = config.num_experts
    progress_bar = tqdm(range(len(sparse_layer_indices) * (num_experts - 1)),
                        desc="[Permutation]Aligning permutation by activation matching")
    for layer_idx in sparse_layer_indices:
        encoder_mlp = switch_model.encoder.block[layer_idx].layer[-1].mlp
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        decoder_mlp = switch_model.decoder.block[layer_idx].layer[-1].mlp
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"

        expert0_encoder_hidden_states = torch.cat([
            forwarded_hidden_states[encoder_mlp_name][i][router_indices[encoder_mlp_name][i] == 0] for i in
            range(len(forwarded_hidden_states[encoder_mlp_name]))
        ])
        expert0_decoder_hidden_states = torch.cat([
            forwarded_hidden_states[decoder_mlp_name][i][router_indices[decoder_mlp_name][i] == 0] for i in
            range(len(forwarded_hidden_states[decoder_mlp_name]))
        ])
        for expert_idx in range(1, num_experts):
            # Permute encoder
            expert_encoder_hidden_states = torch.cat([
                forwarded_hidden_states[encoder_mlp_name][i][router_indices[encoder_mlp_name][i] == expert_idx]
                for i in range(len(forwarded_hidden_states[encoder_mlp_name]))
            ])
            perm = compute_switch_permutation_by_activation_matching(
                reference_mlp=encoder_mlp.experts["expert_0"],
                target_mlp=encoder_mlp.experts[f"expert_{expert_idx}"],
                forwarded_hidden_states=torch.cat([expert0_encoder_hidden_states, expert_encoder_hidden_states], dim=0)
            )
            encoder_mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                encoder_mlp.experts[f"expert_{expert_idx}"], perm
            )
            # Permute decoder
            expert_decoder_hidden_states = torch.cat([
                forwarded_hidden_states[decoder_mlp_name][i][router_indices[decoder_mlp_name][i] == expert_idx]
                for i in range(len(forwarded_hidden_states[decoder_mlp_name]))
            ])
            perm = compute_switch_permutation_by_activation_matching(
                reference_mlp=decoder_mlp.experts["expert_0"],
                target_mlp=decoder_mlp.experts[f"expert_{expert_idx}"],
                forwarded_hidden_states=torch.cat([expert0_decoder_hidden_states, expert_decoder_hidden_states], dim=0)
            )
            decoder_mlp.experts[f"expert_{expert_idx}"] = permute_switch_mlp_dense_expert_(
                decoder_mlp.experts[f"expert_{expert_idx}"], perm
            )

            progress_bar.update(1)

    for handle in handles:
        handle.remove()
    return switch_model


def merge_switch_mlp_by_activation_matching_within_and_across_models(
        mlp_list: List[SwitchTransformersDenseActDense],
        forwarded_hidden_states: torch.Tensor,
        mini_batch_size: Optional[int] = None,
        alpha_for_repeated_merging: Optional[float] = 0.1,
        average_coefs: Optional[List[float]] = None,
) -> SwitchTransformersDenseActDense:
    """
    Merge the MLPs in the mlp_list by activation matching. Specifically, the MLPs in the mlp_list are merged
    by matching the activation of the forwarded_hidden_states, within and across models.


    Parameters
    ----------
    mlp_list: List[SwitchTransformersDenseActDense]
        The MLPs to be merged.
    forwarded_hidden_states: torch.Tensor, of shape (batch_size * seq_len, d_ff)
        The hidden states that are forwarded to all MLPs in the mlp_list.
    mini_batch_size: int, default None
        The mini batch size for the activation matching. If None, the mini batch size is set to be the batch size.
    alpha_for_repeated_merging: float, default 0.1
        The alpha for repeated merging, which is used to update the correlation matrix after each merging.
        In particular, the correlation between the merged feature and the other features is set to be the minimum
            of the old features' correlation weighted.
    average_coefs: Optional[List[float]], default None
        The average coefficients for the features in the MLPs. If None, the average coefficients are set to be 1.0.

    Returns
    -------
    merged_mlp: SwitchTransformersDenseActDense
        The merged MLP.

    Examples
    --------
    >>> from transformers import SwitchTransformersConfig
    >>> from copy import deepcopy
    >>> mlp = SwitchTransformersDenseActDense(SwitchTransformersConfig(d_model=16, d_ff=32, dropout_rate=0))
    >>> mlp_merged = merge_switch_mlp_by_activation_matching_within_and_across_models([mlp, deepcopy(mlp)], torch.randn(64, 16))
    >>> dummy_input = torch.randn(4, 16)
    >>> torch.allclose(mlp_merged(dummy_input), mlp(dummy_input))
    True
    """
    mlp_list = [mlp.eval() for mlp in mlp_list]
    concat_mlp = deepcopy(mlp_list[0])
    d_ff, d_model = concat_mlp.wi.out_features, concat_mlp.wi.in_features
    if average_coefs is None:
        average_coefs = [1.0] * len(mlp_list) * d_ff
    elif len(average_coefs) == len(mlp_list):
        average_coefs = [coef for coef in average_coefs for _ in range(d_ff)]
    elif len(average_coefs) != len(mlp_list) * d_ff:
        raise ValueError(
            f"The length of average_coefs should be either {len(mlp_list)} or {len(mlp_list) * d_ff}, "
            f"but got {len(average_coefs)}."
        )
    num_mlp = len(mlp_list)
    if len(forwarded_hidden_states) == 0 or len(forwarded_hidden_states) == 1:
        return concat_mlp
    if mini_batch_size is None:
        mini_batch_size = forwarded_hidden_states.shape[0]

    mlp_all_wi = torch.cat([
        mlp.wi.weight.data for mlp in mlp_list
    ], dim=0)
    mlp_all_wo = torch.cat([
        mlp.wo.weight.data for mlp in mlp_list
    ], dim=1)
    concat_mlp.wi = torch.nn.Linear(d_model, d_ff * num_mlp, bias=False)
    concat_mlp.wo = torch.nn.Linear(d_ff * num_mlp, d_model, bias=False)
    with torch.no_grad():
        concat_mlp.wi.weight.data = mlp_all_wi  # (d_ff * num_mlp, d_model)
        concat_mlp.wo.weight.data = mlp_all_wo  # (d_model, d_ff * num_mlp)
    concat_mlp = concat_mlp.eval().to(forwarded_hidden_states.device)

    activations = []

    def _activation_hook(module, input, output):
        activations.append(input[0].detach().reshape(-1, input[0].shape[-1]))

    handle = concat_mlp.wo.register_forward_hook(_activation_hook)
    with torch.no_grad():
        for i in range(0, forwarded_hidden_states.shape[0], mini_batch_size):
            concat_mlp(forwarded_hidden_states[i:i + mini_batch_size])

    activations = torch.cat(activations, dim=0)  # (batch_size * seq_len, d_ff * num_mlp)

    # Initialize the correlation matrix
    mean = activations.mean(dim=0, keepdim=True)  # (1, d_ff * num_mlp)
    std = activations.std(dim=0, keepdim=True)  # (1, d_ff * num_mlp)
    covar = torch.mm(
        (activations - mean).t(),
        (activations - mean)
    ) / (activations.shape[0] - 1)  # (d_ff * num_mlp, d_ff * num_mlp)
    corr_matrix = covar / (std.t() * std + FP32_EPS)  # (d_ff * num_mlp, d_ff * num_mlp)

    del activations, covar, std, mean
    torch.cuda.empty_cache()

    corr_matrix[torch.arange(d_ff * num_mlp), torch.arange(d_ff * num_mlp)] = -1  # Remove self-correlation

    # Greedy Merging!
    while mlp_all_wi.shape[0] > d_ff:
        # Select the most correlated pair
        max_index = torch.argmax(corr_matrix)
        max_i, max_j = max_index // corr_matrix.shape[0], max_index % corr_matrix.shape[0]

        # Merge the most correlated pair, replace the first feature with the merged one
        i_coef, j_coef = average_coefs[max_i], average_coefs[max_j]
        mlp_all_wi[max_i] = (i_coef * mlp_all_wi[max_i] + j_coef * mlp_all_wi[max_j]) / (i_coef + j_coef + FP32_EPS)
        mlp_all_wo[:, max_i] = (i_coef * mlp_all_wo[:, max_i] + j_coef * mlp_all_wo[:, max_j]) / (
                i_coef + j_coef + FP32_EPS)

        # Remove the second feature
        mlp_all_wi = torch.cat([
            mlp_all_wi[:max_j],
            mlp_all_wi[max_j + 1:]
        ], dim=0)
        mlp_all_wo = torch.cat([
            mlp_all_wo[:, :max_j],
            mlp_all_wo[:, max_j + 1:]
        ], dim=1)

        # Update the correlation matrix
        updated_corr_vec = alpha_for_repeated_merging * torch.min(
            torch.stack([corr_matrix[max_i], corr_matrix[max_j]]), dim=0
        ).values
        corr_matrix[max_i] = updated_corr_vec
        corr_matrix[:, max_i] = updated_corr_vec
        corr_matrix[max_i, max_i] = -1  # Remove self-correlation

        # Remove the second feature from the correlation matrix
        corr_matrix = torch.cat([
            corr_matrix[:, :max_j],
            corr_matrix[:, max_j + 1:]
        ], dim=1)
        corr_matrix = torch.cat([
            corr_matrix[:max_j],
            corr_matrix[max_j + 1:]
        ], dim=0)

        # Update the average coefs
        average_coefs[max_i] += average_coefs[max_j]
        average_coefs = average_coefs[:max_j] + average_coefs[max_j + 1:]

    handle.remove()
    merged_mlp = deepcopy(mlp_list[0])
    with torch.no_grad():
        merged_mlp.wi.weight.data = mlp_all_wi
        merged_mlp.wo.weight.data = mlp_all_wo

    return merged_mlp


def merge_switch_mlp_by_weight_matching_within_and_across_models(
        mlp_list: List[SwitchTransformersDenseActDense],
        include_wo: Optional[bool] = True,
        average_coefs: Optional[List[float]] = None,
) -> SwitchTransformersDenseActDense:
    """
    Merge the MLPs in the mlp_list by weight matching. Specifically, the MLPs in the mlp_list are merged

    Parameters
    ----------
    mlp_list: List[SwitchTransformersDenseActDense]
        The MLPs to be merged.
    include_wo: bool, default True
        Whether to include the (each column) weights of the second layer of MLP in the weight matching.
    average_coefs: List[float], default None
        The average coefs for each MLP in the mlp_list. If None, the average coefs are set to be 1.0 for all MLPs.

    Returns
    -------
    merged_mlp: SwitchTransformersDenseActDense
        The merged MLP.

    Examples
    --------
    >>> from transformers import SwitchTransformersConfig
    >>> from copy import deepcopy
    >>> mlp = SwitchTransformersDenseActDense(SwitchTransformersConfig(d_model=16, d_ff=32, dropout_rate=0))
    >>> mlp_merged = merge_switch_mlp_by_activation_matching_within_and_across_models([mlp, deepcopy(mlp)], torch.randn(64, 16))
    >>> dummy_input = torch.randn(4, 16)
    >>> torch.allclose(mlp_merged(dummy_input), mlp(dummy_input))
    True
    """
    d_ff, d_model = mlp_list[0].wi.out_features, mlp_list[0].wi.in_features
    if average_coefs is None:
        average_coefs = [1.0] * len(mlp_list) * d_ff
    elif len(average_coefs) == len(mlp_list):
        average_coefs = [coef for coef in average_coefs for _ in range(d_ff)]
    elif len(average_coefs) != len(mlp_list) * d_ff:
        raise ValueError(
            f"The length of average_coefs should be either {len(mlp_list)} or {len(mlp_list) * d_ff}, "
            f"but got {len(average_coefs)}."
        )
    concat_wi = torch.cat([
        mlp.wi.weight.data for mlp in mlp_list
    ], dim=0)  # of shape (d_ff * num_mlp, d_model)
    concat_wo = torch.cat([
        mlp.wo.weight.data.t() for mlp in mlp_list
    ], dim=0)  # of shape (d_ff * num_mlp, d_model)
    if include_wo:
        cost_matrix = torch.mm(concat_wi, concat_wi.t()) + torch.mm(concat_wo, concat_wo.t())
    else:
        cost_matrix = torch.mm(concat_wi, concat_wi.t())

    # Remove self-correlation
    dtype_min = torch.finfo(cost_matrix.dtype).min
    cost_matrix[torch.arange(d_ff * len(mlp_list)), torch.arange(d_ff * len(mlp_list))] = dtype_min

    # Greedy Merging!
    while concat_wi.shape[0] > d_ff:
        # Select the most correlated pair
        max_index = torch.argmax(cost_matrix)
        max_i, max_j = max_index // cost_matrix.shape[0], max_index % cost_matrix.shape[0]

        # Merge the most correlated pair, replace the first feature with the merged one
        i_coef = average_coefs[max_i]
        j_coef = average_coefs[max_j]
        concat_wi[max_i] = (i_coef * concat_wi[max_i] + j_coef * concat_wi[max_j]) / (i_coef + j_coef + FP32_EPS)
        concat_wo[max_i] = (i_coef * concat_wo[max_i] + j_coef * concat_wo[max_j]) / (i_coef + j_coef + FP32_EPS)

        # Update the cost matrix
        updated_cost_vec = torch.mm(
            concat_wi[max_i].reshape(1, -1),
            concat_wi.t()
        ) + torch.mm(
            concat_wo[max_i].reshape(1, -1),
            concat_wo.t()
        )
        cost_matrix[max_i] = updated_cost_vec
        cost_matrix[:, max_i] = updated_cost_vec
        cost_matrix[max_i, max_i] = dtype_min  # Remove self-correlation

        # Remove the second feature
        concat_wi = torch.cat([
            concat_wi[:max_j],
            concat_wi[max_j + 1:]
        ], dim=0)
        concat_wo = torch.cat([
            concat_wo[:max_j],
            concat_wo[max_j + 1:]
        ], dim=0)

        # Remove the second feature from the correlation matrix
        cost_matrix = torch.cat([
            cost_matrix[:, :max_j],
            cost_matrix[:, max_j + 1:]
        ], dim=1)
        cost_matrix = torch.cat([
            cost_matrix[:max_j],
            cost_matrix[max_j + 1:]
        ], dim=0)

        # Update the average coefs
        average_coefs[max_i] += average_coefs[max_j]
        average_coefs = average_coefs[:max_j] + average_coefs[max_j + 1:]

    merged_mlp = deepcopy(mlp_list[0])
    with torch.no_grad():
        merged_mlp.wi.weight.data = concat_wi
        merged_mlp.wo.weight.data = concat_wo.t()

    return merged_mlp

mcsmoe/merging/permutation.py
UNITES-Lab/MC-SMoE
import os
import random
from typing import Optional

import torch
from tqdm import tqdm
from transformers import (
    SwitchTransformersForConditionalGeneration,
    SwitchTransformersSparseMLP
)

from mcsmoe.models.fsgpt_moe import (
    FSGPTMoESparseMLP,
    FSGPTMoEForCausalLM
)


def _merge_switch_expert_module(
        mlp: SwitchTransformersSparseMLP,
        group_labels: torch.Tensor,
) -> SwitchTransformersSparseMLP:
    for label in group_labels.unique():
        indices = torch.where(group_labels == label)[0]
        if len(indices) == 1:
            continue
        # merge experts
        for idx in indices[1:]:
            mlp.experts[f"expert_{idx}"] = mlp.experts[f"expert_{indices[0]}"]
            # Despite only one expert within a group has the right weights,
            # this is supposed to be correct since the nn.ModuleDict always keep the largest index as the true expert.
            # (i.e. The equal sign is not that normal function here.)
            # In this way, as long as we assign the largest index in the end, we are good.
    return mlp


def load_merged_switch_transformers_from_checkpoint(
        checkpoint: str,
        group_state_checkpoint: Optional[str] = None,
        sanity_check: Optional[bool] = False,
) -> SwitchTransformersForConditionalGeneration:
    if group_state_checkpoint is None:
        group_state_checkpoint = os.path.join(checkpoint, "group_state_dict.pt")
    model = SwitchTransformersForConditionalGeneration.from_pretrained(checkpoint)
    state_dict = model.state_dict()
    group_state_dict = torch.load(group_state_checkpoint, map_location="cpu")

    # merge experts
    config = model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.encoder_sparse_step))
    for layer_idx in tqdm(sparse_layer_indices, desc="Merging expert modules"):
        # Merge encoder
        encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
        encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
        group_labels = group_state_dict[encoder_mlp_name]
        model.encoder.block[layer_idx].layer[-1].mlp = _merge_switch_expert_module(encoder_mlp, group_labels)
        # Merge decoder
        decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
        decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
        group_labels = group_state_dict[decoder_mlp_name]
        model.decoder.block[layer_idx].layer[-1].mlp = _merge_switch_expert_module(decoder_mlp, group_labels)

    if sanity_check:
        for name, param in tqdm(state_dict.items(), desc="Sanity checking"):
            if not torch.allclose(param, model.state_dict()[name]):
                raise ValueError(f"Parameter {name} is not merged correctly.")

    model.group_state_dict = group_state_dict
    return model


def _merge_fsgpt_expert_module(
        ffn: FSGPTMoESparseMLP,
        group_labels: torch.Tensor,
) -> FSGPTMoESparseMLP:
    for label in group_labels.unique():
        indices = torch.where(group_labels == label)[0]
        if len(indices) == 1:
            continue
        # merge experts
        for idx in indices[1:]:
            ffn.experts[f"expert_{idx}"] = ffn.experts[f"expert_{indices[0]}"]
            # Despite only one expert within a group has the right weights,
            # this is supposed to be correct since the nn.ModuleDict always keep the largest index as the true expert.
            # (i.e. The equal sign is not that normal function here.)
            # In this way, as long as we assign the largest index in the end, we are good.
    return ffn


def load_merged_fsgpt_moe_from_checkpoint(
        checkpoint: str,
        group_state_checkpoint: Optional[str] = None,
        sanity_check: Optional[bool] = False,
) -> FSGPTMoEForCausalLM:
    if group_state_checkpoint is None:
        group_state_checkpoint = os.path.join(checkpoint, "group_state_dict.pt")
    model = FSGPTMoEForCausalLM.from_pretrained(
        checkpoint,
        device_map="auto",
        torch_dtype=torch.bfloat16,
    )
    state_dict = model.state_dict() if sanity_check else None
    group_state_dict = torch.load(group_state_checkpoint, map_location="cpu")

    # merge experts
    config = model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.sparse_step))
    for layer_idx in tqdm(sparse_layer_indices, desc="Merging expert modules"):
        ffn_name = f"decoder.layers.{layer_idx}.ffn"
        ffn = model.decoder.layers[layer_idx].ffn
        group_labels = group_state_dict[ffn_name]
        model.decoder.layers[layer_idx].ffn = _merge_fsgpt_expert_module(ffn, group_labels)

    if sanity_check:
        for name, param in tqdm(state_dict.items(), desc="Sanity checking"):
            if not torch.allclose(param, model.state_dict()[name]):
                raise ValueError(f"Parameter {name} is not merged correctly.")

    model.group_state_dict = group_state_dict
    return model


def generate_random_group_labels(
        num_experts: int,
        num_groups: int,
) -> torch.Tensor:
    """
    Assign random group labels to experts, with each group has at least one expert.

    Examples
    --------
    >>> generate_random_group_labels(10, 3).unique().sort()[0]
    tensor([0, 1, 2])
    """
    group_labels = torch.zeros(num_experts, dtype=torch.long)
    for i in range(num_groups):
        group_labels[i] = i
    for i in range(num_groups, num_experts):
        group_labels[i] = random.randint(0, num_groups - 1)
    return group_labels[torch.randperm(num_experts)]

mcsmoe/merging/utils.py
UNITES-Lab/MC-SMoE
from .fsgpt import (
    FSGPTForCausalLM,
    FSGPTConfig
)
from .fsgpt_moe import (
    FSGPTMoEForCausalLM,
    FSGPTMoEConfig,
)
from .switch_transformers import SwitchTransformersWrapperForDistillation

mcsmoe/models/__init__.py
UNITES-Lab/MC-SMoE
"""
Reproduce the Fairseq GPT model.

Description from Fairseq:
We differ from Brown et al. (2020) in two ways:
(1) we use only dense attention, while they alternate between dense and locally banded sparse attention;
(2) we train our models with sinusoidal positional embeddings, following Shortformer (Press et al., 2020).


"""
import math
from typing import Any, Optional
from typing import Tuple, Union, Dict

import torch
import torch.onnx.operators
from torch import Tensor, nn
from torch.nn import CrossEntropyLoss
from transformers import PreTrainedModel
from transformers.activations import ACT2FN
from transformers.configuration_utils import PretrainedConfig
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast
)


def rename_fairseq_state_dict(
        state_dict: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
    convert_state_dict = dict()
    for name, param in state_dict.items():
        if name == 'decoder.version':
            continue
        name = name.replace('decoder.output_projection', 'lm_head')
        name = name.replace('fc1', 'ffn.fc1')
        name = name.replace('fc2', 'ffn.fc2')
        convert_state_dict[name] = param

    return convert_state_dict


class FSGPTConfig(PretrainedConfig):
    def __init__(
            self,
            vocab_size=51200,
            max_position_embeddings=2048,
            hidden_size=768,
            num_layers=12,
            num_heads=12,
            intermediate_size=3072,
            activation_function="gelu",
            resid_dropout=0.1,
            embed_dropout=0.0,
            attention_dropout=0.1,
            classifier_dropout=0.1,
            use_cache=True,
            pad_token_id=1,
            eos_token_id=2,
            unk_token_id=3,
            no_scale_embedding=False,
            **kwargs,
    ):
        self.vocab_size = vocab_size
        self.max_position_embeddings = max_position_embeddings
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.intermediate_size = intermediate_size
        self.activation_function = activation_function
        self.resid_dropout = resid_dropout
        self.embed_dropout = embed_dropout
        self.attention_dropout = attention_dropout
        self.classifier_dropout = classifier_dropout
        self.use_cache = use_cache

        self.pad_token_id = pad_token_id
        self.eos_token_id = eos_token_id
        self.unk_token_id = unk_token_id

        # Fairseq compatibility
        self.no_scale_embedding = no_scale_embedding

        super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, unk_token_id=unk_token_id, **kwargs)


def make_positions(tensor, padding_idx: int, offset: int):
    """Replace non-padding symbols with their position numbers.

    Position numbers begin at padding_idx+1. Padding symbols are ignored.
    """
    # The series of casts and type-conversions here are carefully
    # balanced to both work with ONNX export and XLA. In particular XLA
    # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
    # how to handle the dtype kwarg in cumsum.
    mask = tensor.ne(padding_idx).int()
    return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + offset - 1


class FSGPTSinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length.

    Padding symbols are ignored.
    """

    def __init__(self, embedding_dim, padding_idx, num_positions=1024):
        super().__init__()
        self.offset = 2
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx if padding_idx is not None else 0
        self.weights = FSGPTSinusoidalPositionalEmbedding.get_embedding(
            num_positions + self.offset, embedding_dim, self.offset
        )
        self.onnx_trace = False
        self.register_buffer("_float_tensor", torch.FloatTensor(1))
        self.max_positions = int(1e5)

    def prepare_for_onnx_export_(self):
        self.onnx_trace = True

    @staticmethod
    def get_embedding(
            num_embeddings: int, embedding_dim: int, offset: Optional[int] = None
    ):
        """Build sinusoidal embeddings.

        This matches the implementation in tensor2tensor, but differs slightly
        from the description in Section 3.5 of "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(
            1
        ) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(
            num_embeddings, -1
        )
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if offset is not None:
            emb[offset - 1, :] = 0
        return emb

    def forward(
            self,
            input,
            incremental_state: Optional[Any] = None,
            timestep: Optional[Tensor] = None,
            positions: Optional[Any] = None,
    ):
        """Input is expected to be of size [bsz x seqlen]."""
        bspair = torch.onnx.operators.shape_as_tensor(input)
        bsz, seq_len = bspair[0], bspair[1]
        max_pos = self.offset + seq_len
        if self.weights is None or max_pos > self.weights.size(0):
            # recompute/expand embeddings if needed
            self.weights = FSGPTSinusoidalPositionalEmbedding.get_embedding(
                max_pos, self.embedding_dim, self.offset
            )
        self.weights = self.weights.to(self._float_tensor)

        if incremental_state is not None:
            # positions is the same for every token when decoding a single step
            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
            if self.onnx_trace:
                return (
                    self.weights.index_select(index=self.offset - 1 + pos, dim=0)
                    .unsqueeze(1)
                    .repeat(bsz, 1, 1)
                )
            return self.weights[self.offset - 1 + pos, :].expand(bsz, 1, -1)

        positions = make_positions(
            input, self.padding_idx, self.offset
        )
        if self.onnx_trace:
            flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
            embedding_shape = torch.cat(
                (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))
            )
            embeddings = torch.onnx.operators.reshape_from_tensor_shape(
                flat_embeddings, embedding_shape
            )
            return embeddings
        return (
            self.weights.index_select(0, positions.view(-1))
            .view(bsz, seq_len, -1)
            .detach()
        )


class FSGPTSelfAttention(nn.Module):
    def __init__(self, config: FSGPTConfig):
        super().__init__()

        max_positions = config.max_position_embeddings

        bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
            1, 1, max_positions, max_positions
        )

        self.register_buffer("bias", bias, persistent=False)
        self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)

        self.num_heads = config.num_heads
        self.embed_dim = config.hidden_size
        self.head_dim = self.embed_dim // self.num_heads

        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
            )

        self.attn_dropout = nn.Dropout(config.attention_dropout)
        self.resid_dropout = nn.Dropout(float(config.resid_dropout))

        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.query_scaled = math.sqrt(self.head_dim)

    @staticmethod
    def _split_heads(tensor, num_heads, attn_head_size):
        """
        Splits hidden_size dim into attn_head_size and num_heads
        """
        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
        tensor = tensor.view(new_shape)
        return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)

    @staticmethod
    def _merge_heads(tensor, num_heads, attn_head_size):
        """
        Merges attn_head_size dim and num_attn_heads dim into hidden_size
        """
        tensor = tensor.permute(0, 2, 1, 3).contiguous()
        new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
        return tensor.view(new_shape)

    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = query.to(torch.float32)
        key = key.to(torch.float32)

        attn_weights = torch.matmul(query, key.transpose(-1, -2))

        query_length, key_length = query.size(-2), key.size(-2)
        causal_mask = self.bias[:, :, key_length - query_length: key_length, :key_length]
        mask_value = torch.finfo(attn_weights.dtype).min
        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
        mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
        attn_weights = torch.where(causal_mask, attn_weights, mask_value)

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
        attn_weights = attn_weights.to(value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

        # Mask heads if we want to
        if head_mask is not None:
            attn_weights = attn_weights * head_mask

        attn_output = torch.matmul(attn_weights, value)

        return attn_output, attn_weights

    def forward(
            self,
            hidden_states,
            attention_mask=None,
            layer_past=None,
            head_mask=None,
            use_cache=False,
            output_attentions=False,
    ):
        query = self.q_proj(hidden_states) / self.query_scaled
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs  # a, present, (attentions)


class FSGPTMLP(nn.Module):
    def __init__(self, config: FSGPTConfig, intermediate_size: int):
        super().__init__()
        self.fc1 = nn.Linear(config.hidden_size, intermediate_size)
        self.fc2 = nn.Linear(intermediate_size, config.hidden_size)
        self.act = ACT2FN[config.activation_function]
        self.dropout = nn.Dropout(config.resid_dropout)

    def forward(self, hidden_states):
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.fc2(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states


class FSGPTBlock(nn.Module):
    def __init__(self, config: FSGPTConfig):
        super().__init__()
        hidden_size = config.hidden_size
        inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * hidden_size
        self.self_attn = FSGPTSelfAttention(config)
        self.self_attn_layer_norm = nn.LayerNorm(hidden_size)
        self.ffn = FSGPTMLP(config, inner_dim)
        self.final_layer_norm = nn.LayerNorm(hidden_size)

    def forward(
            self,
            hidden_states,
            layer_past=None,
            attention_mask=None,
            head_mask=None,
            use_cache=False,
            output_attentions=False,
    ):
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        attn_outputs = self.self_attn(
            hidden_states,
            layer_past=layer_past,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attn_output = attn_outputs[0]
        outputs = attn_outputs[1:]
        # residual connection
        hidden_states = attn_output + residual

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        feed_forward_hidden_states = self.ffn(hidden_states)
        # residual connection
        hidden_states = residual + feed_forward_hidden_states

        if use_cache:
            outputs = (hidden_states,) + outputs
        else:
            outputs = (hidden_states,) + outputs[1:]

        return outputs  # hidden_states, present, (attentions, cross_attentions)


class FSGPTPreTrainedModel(PreTrainedModel):
    config_class = FSGPTConfig
    base_model_prefix = "decoder"
    supports_gradient_checkpointing = True
    _no_split_modules = ["FSGPTBlock"]
    _skip_keys_device_placement = "past_key_values"

    def __init__(self, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, (nn.Linear,)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.hidden_size ** -0.5)
            if module.weight.data.shape[0] == self.config.vocab_size:
                module.weight.data[self.config.pad_token_id].zero_()
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.hidden_size ** -0.5)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, GPTNeoModel):
            module.gradient_checkpointing = value


class FSGPTModel(FSGPTPreTrainedModel):
    def __init__(self, config: FSGPTConfig):
        super().__init__(config)

        self.embed_dim = config.hidden_size
        self.embed_tokens = nn.Embedding(config.vocab_size, self.embed_dim, padding_idx=config.pad_token_id)
        self.embed_positions = FSGPTSinusoidalPositionalEmbedding(
            embedding_dim=self.embed_dim,
            padding_idx=config.pad_token_id,
            num_positions=config.max_position_embeddings,
        )
        self.drop = nn.Dropout(float(config.embed_dropout))

        self.embed_scale = (
            1.0 if config.no_scale_embedding else math.sqrt(self.embed_dim)
        )
        self.layers = nn.ModuleList([FSGPTBlock(config) for _ in range(config.num_layers)])
        self.layer_norm = nn.LayerNorm(self.embed_dim)

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.embed_tokens = new_embeddings

    def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            batch_size = input_ids.shape[0]
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size = inputs_embeds.shape[0]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

        if past_key_values is None:
            past_length = 0
            past_key_values = tuple([None] * len(self.layers))
        else:
            past_length = past_key_values[0][0].size(-2)

        # if position_ids is None:
        #     position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
        #     position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])

        # Attention mask.
        if attention_mask is not None:
            if batch_size <= 0:
                raise ValueError("batch_size has to be defined and > 0")
            attention_mask = attention_mask.view(batch_size, -1)
            # We create a 3D attention mask from a 2D tensor mask.
            # Sizes are [batch_size, 1, 1, to_seq_length]
            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
            # this attention mask is more simple than the triangular masking of causal attention
            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
            attention_mask = attention_mask[:, None, None, :]

            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
            # masked positions, this operation will create a tensor which is 0.0 for
            # positions we want to attend and the dtype's smallest value for masked positions.
            # Since we are adding it to the raw scores before the softmax, this is
            # effectively the same as removing these entirely.
            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x num_heads x N x N
        # head_mask has shape n_layer x batch x num_heads x N x N
        head_mask = self.get_head_mask(head_mask, self.config.num_layers)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        position_embeds = self.embed_positions(input_ids)
        hidden_states = inputs_embeds + position_embeds

        if token_type_ids is not None:
            token_type_embeds = self.embed_tokens(token_type_ids)
            hidden_states = hidden_states + token_type_embeds

        hidden_states = self.drop(hidden_states)

        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None
        for i, (block, layer_past) in enumerate(zip(self.layers, past_key_values)):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            outputs = block(
                hidden_states,
                layer_past=layer_past,
                attention_mask=attention_mask,
                head_mask=head_mask[i],
                use_cache=use_cache,
                output_attentions=output_attentions,
            )

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)

        hidden_states = self.layer_norm(hidden_states)

        hidden_states = hidden_states.view(output_shape)
        # Add last hidden state
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )


class FSGPTForCausalLM(FSGPTPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: FSGPTConfig):
        super().__init__(config)
        self.decoder = FSGPTModel(config)
        # Share input output embeddings
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.decoder.embed_tokens.weight

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
        token_type_ids = kwargs.get("token_type_ids", None)
        # only last token for inputs_ids if past is defined in kwargs
        if past_key_values:
            input_ids = input_ids[:, -1].unsqueeze(-1)
            if token_type_ids is not None:
                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)

        attention_mask = kwargs.get("attention_mask", None)
        position_ids = kwargs.get("position_ids", None)

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "position_ids": position_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        return model_inputs

    def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            labels: Optional[torch.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.decoder(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Compute loss in fp32 to match with mesh-tf version
            # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
            lm_logits = lm_logits.to(torch.float32)

            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            lm_logits = lm_logits.to(hidden_states.dtype)
            loss = loss.to(hidden_states.dtype)

        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=lm_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

    @staticmethod
    def _reorder_cache(
            past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
    ) -> Tuple[Tuple[torch.Tensor]]:
        """
        This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
        [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
        beam_idx at every generation step.
        """
        return tuple(
            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
            for layer_past in past_key_values
        )

mcsmoe/models/fsgpt.py
UNITES-Lab/MC-SMoE
"""
Reproduce the Fairseq GPT model.

Description from Fairseq:
We differ from Brown et al. (2020) in two ways:
(1) we use only dense attention, while they alternate between dense and locally banded sparse attention;
(2) we train our models with sinusoidal positional embeddings, following Shortformer (Press et al., 2020).


"""
import math
from typing import Any, Optional
from typing import Tuple, Union, Dict

import torch
import torch.onnx.operators
from torch import Tensor, nn
from torch.nn import CrossEntropyLoss
from transformers import PreTrainedModel
from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    MoEModelOutputWithPastAndCrossAttentions,
    MoECausalLMOutputWithPast
)

from .fsgpt import FSGPTConfig

__all__ = [
    "FSGPTMoEConfig",
    "FSGPTMoETop2Router",
    "FSGPTMoEDenseActDense",
    "FSGPTMoESparseMLP",
    "FSGPTMoEBlock",
    "FSGPTMoEModel",
    "FSGPTMoEForCausalLM",
]


def rename_fairseq_state_dict(
        state_dict: Dict[str, torch.Tensor],
) -> Dict[str, torch.Tensor]:
    convert_state_dict = dict()
    for name, param in state_dict.items():
        if name == 'decoder.version':
            continue
        name = name.replace('decoder.output_projection', 'lm_head')
        name = name.replace('fc1', 'ffn.fc1')
        name = name.replace('fc2', 'ffn.fc2')
        convert_state_dict[name] = param

    return convert_state_dict


class FSGPTMoEConfig(FSGPTConfig):
    def __init__(
            self,
            sparse_step=2,
            router_bias=False,
            num_experts=512,
            expert_capacity=None,
            router_ignore_padding_tokens=True,
            router_dtype="float32",
            second_expert_policy="all",
            normalize_router_prob_before_dropping=False,
            batch_prioritized_routing=False,
            moe_eval_capacity_token_fraction=-1.0,
            moe_token_dropout=0.2,
            output_router_logits=True,
            **kwargs,
    ):
        self.sparse_step = sparse_step
        self.router_bias = router_bias
        self.num_experts = num_experts
        self.expert_capacity = expert_capacity
        self.router_ignore_padding_tokens = router_ignore_padding_tokens
        self.router_dtype = router_dtype
        self.second_expert_policy = second_expert_policy
        self.normalize_router_prob_before_dropping = normalize_router_prob_before_dropping
        self.batch_prioritized_routing = batch_prioritized_routing
        self.moe_eval_capacity_token_fraction = moe_eval_capacity_token_fraction
        self.moe_token_dropout = moe_token_dropout
        self.output_router_logits = output_router_logits

        super().__init__(**kwargs)


def _make_causal_mask(
        input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
):
    """
    Make causal mask used for bi-directional self-attention.
    """
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
    mask_cond = torch.arange(mask.size(-1), device=device)
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)

    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


# Copied from transformers.models.bart.modeling_bart._expand_mask
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
    """
    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
    """
    bsz, src_len = mask.size()
    tgt_len = tgt_len if tgt_len is not None else src_len

    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

    inverted_mask = 1.0 - expanded_mask

    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)


def make_positions(tensor, padding_idx: int, offset: int):
    """Replace non-padding symbols with their position numbers.

    Position numbers begin at padding_idx+1. Padding symbols are ignored.
    """
    # The series of casts and type-conversions here are carefully
    # balanced to both work with ONNX export and XLA. In particular XLA
    # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
    # how to handle the dtype kwarg in cumsum.
    mask = tensor.ne(padding_idx).int()
    return (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + offset - 1


class FSGPTMoESinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length.

    Padding symbols are ignored.
    """

    def __init__(self, embedding_dim, padding_idx, num_positions=1024):
        super().__init__()
        self.offset = 2
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx if padding_idx is not None else 0
        self.weights = FSGPTMoESinusoidalPositionalEmbedding.get_embedding(
            num_positions + self.offset, embedding_dim, self.offset
        )
        self.onnx_trace = False
        self.register_buffer("_float_tensor", torch.FloatTensor(1))
        self.max_positions = int(1e5)

    def prepare_for_onnx_export_(self):
        self.onnx_trace = True

    @staticmethod
    def get_embedding(
            num_embeddings: int, embedding_dim: int, offset: Optional[int] = None
    ):
        """Build sinusoidal embeddings.

        This matches the implementation in tensor2tensor, but differs slightly
        from the description in Section 3.5 of "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(
            1
        ) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(
            num_embeddings, -1
        )
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if offset is not None:
            emb[offset - 1, :] = 0
        return emb

    def forward(
            self,
            input,
            incremental_state: Optional[Any] = None,
            timestep: Optional[Tensor] = None,
            positions: Optional[Any] = None,
    ):
        """Input is expected to be of size [bsz x seqlen]."""
        bspair = torch.onnx.operators.shape_as_tensor(input)
        bsz, seq_len = bspair[0], bspair[1]
        max_pos = self.offset + seq_len
        if self.weights is None or max_pos > self.weights.size(0):
            # recompute/expand embeddings if needed
            self.weights = FSGPTMoESinusoidalPositionalEmbedding.get_embedding(
                max_pos, self.embedding_dim, self.offset
            )
        self.weights = self.weights.to(self._float_tensor)

        if incremental_state is not None:
            # positions is the same for every token when decoding a single step
            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
            if self.onnx_trace:
                return (
                    self.weights.index_select(index=self.offset - 1 + pos, dim=0)
                    .unsqueeze(1)
                    .repeat(bsz, 1, 1)
                )
            return self.weights[self.offset - 1 + pos, :].expand(bsz, 1, -1)

        positions = make_positions(
            input, self.padding_idx, self.offset
        )
        if self.onnx_trace:
            flat_embeddings = self.weights.detach().index_select(0, positions.view(-1))
            embedding_shape = torch.cat(
                (bsz.view(1), seq_len.view(1), torch.tensor([-1], dtype=torch.long))
            )
            embeddings = torch.onnx.operators.reshape_from_tensor_shape(
                flat_embeddings, embedding_shape
            )
            return embeddings
        return (
            self.weights.index_select(0, positions.view(-1))
            .view(bsz, seq_len, -1)
            .detach()
        )


class FSGPTMoESelfAttention(nn.Module):
    def __init__(self, config: FSGPTMoEConfig):
        super().__init__()

        max_positions = config.max_position_embeddings

        bias = torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
            1, 1, max_positions, max_positions
        )

        self.register_buffer("bias", bias, persistent=False)
        self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)

        self.num_heads = config.num_heads
        self.embed_dim = config.hidden_size
        self.head_dim = self.embed_dim // self.num_heads

        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
            )

        self.attn_dropout = nn.Dropout(config.attention_dropout)
        self.resid_dropout = nn.Dropout(float(config.resid_dropout))

        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
        self.query_scaled = math.sqrt(self.head_dim)

    @staticmethod
    def _split_heads(tensor, num_heads, attn_head_size):
        """
        Splits hidden_size dim into attn_head_size and num_heads
        """
        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
        tensor = tensor.view(new_shape)
        return tensor.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)

    @staticmethod
    def _merge_heads(tensor, num_heads, attn_head_size):
        """
        Merges attn_head_size dim and num_attn_heads dim into hidden_size
        """
        tensor = tensor.permute(0, 2, 1, 3).contiguous()
        new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
        return tensor.view(new_shape)

    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        # Keep the attention weights computation in fp32 to avoid overflow issues
        query = query.to(torch.float32)
        key = key.to(torch.float32)

        attn_weights = torch.matmul(query, key.transpose(-1, -2))

        query_length, key_length = query.size(-2), key.size(-2)
        causal_mask = self.bias[:, :, key_length - query_length: key_length, :key_length]
        mask_value = torch.finfo(attn_weights.dtype).min
        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
        mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(attn_weights.device)
        attn_weights = torch.where(causal_mask, attn_weights, mask_value)

        if attention_mask is not None:
            # Apply the attention mask
            attn_weights = attn_weights + attention_mask

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
        attn_weights = attn_weights.to(value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

        # Mask heads if we want to
        if head_mask is not None:
            attn_weights = attn_weights * head_mask

        attn_output = torch.matmul(attn_weights, value)

        return attn_output, attn_weights

    def forward(
            self,
            hidden_states,
            attention_mask=None,
            layer_past=None,
            head_mask=None,
            use_cache=False,
            output_attentions=False,
    ):
        query = self.q_proj(hidden_states) / self.query_scaled
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

        if layer_past is not None:
            past_key = layer_past[0]
            past_value = layer_past[1]
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
        attn_output = self.out_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs  # a, present, (attentions)


class FSGPTMoETop2Router(nn.Module):
    """
    Adapted form hugginfgace NLLB-MoE

    Router using tokens choose top-2 experts assignment.

    This router uses the same mechanism as in NLLB-MoE from the fairseq repository. Items are sorted by router_probs
    and then routed to their choice of expert until the expert's expert_capacity is reached. **There is no guarantee
    that each token is processed by an expert**, or that each expert receives at least one token.

    The router combining weights are also returned to make sure that the states that are not updated will be masked.

    """

    def __init__(self, config: FSGPTMoEConfig):
        super().__init__()
        self.num_experts = config.num_experts
        self.expert_capacity = config.expert_capacity
        self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
        self.router_ignore_padding_tokens = config.router_ignore_padding_tokens
        self.dtype = getattr(torch, config.router_dtype)

        self.second_expert_policy = config.second_expert_policy
        self.normalize_router_prob_before_dropping = config.normalize_router_prob_before_dropping
        self.batch_prioritized_routing = config.batch_prioritized_routing
        self.moe_eval_capacity_token_fraction = config.moe_eval_capacity_token_fraction

    def _cast_classifier(self):
        r"""
        `bitsandbytes` `Linear8bitLt` layers does not support manual casting Therefore we need to check if they are an
        instance of the `Linear8bitLt` class by checking special attributes.
        """
        if not (hasattr(self.classifier, "SCB") or hasattr(self.classifier, "CB")):
            self.classifier = self.classifier.to(self.dtype)

    def normalize_router_probabilities(self, router_probs, top_1_mask, top_2_mask):
        top_1_max_probs = (router_probs * top_1_mask).sum(dim=1)
        top_2_max_probs = (router_probs * top_2_mask).sum(dim=1)
        denom_s = torch.clamp(top_1_max_probs + top_2_max_probs, min=torch.finfo(router_probs.dtype).eps)
        top_1_max_probs = top_1_max_probs / denom_s
        top_2_max_probs = top_2_max_probs / denom_s
        return top_1_max_probs, top_2_max_probs

    def route_tokens(
            self,
            router_logits: torch.Tensor,
            input_dtype: torch.dtype = torch.float32,
            padding_mask: Optional[torch.LongTensor] = None,
    ) -> Tuple:
        """
        Computes the `dispatch_mask` and the `dispatch_weights` for each experts. The masks are adapted to the expert
        capacity.
        """
        nb_tokens = router_logits.shape[0]
        # Apply Softmax and cast back to the original `dtype`
        router_probs = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(input_dtype)
        top_1_expert_index = torch.argmax(router_probs, dim=-1)
        top_1_mask = torch.nn.functional.one_hot(top_1_expert_index, num_classes=self.num_experts)

        if self.second_expert_policy == "sampling":
            gumbel = torch.distributions.gumbel.Gumbel(0, 1).rsample
            router_logits += gumbel(router_logits.shape).to(router_logits.device)

        # replace top_1_expert_index with min values
        logits_except_top_1 = router_logits.masked_fill(top_1_mask.bool(), float("-inf"))
        top_2_expert_index = torch.argmax(logits_except_top_1, dim=-1)
        top_2_mask = torch.nn.functional.one_hot(top_2_expert_index, num_classes=self.num_experts)

        if self.normalize_router_prob_before_dropping:
            top_1_max_probs, top_2_max_probs = self.normalize_router_probabilities(
                router_probs, top_1_mask, top_2_mask
            )

        if self.second_expert_policy == "random":
            top_2_max_probs = (router_probs * top_2_mask).sum(dim=1)
            sampled = (2 * top_2_max_probs) > torch.rand_like(top_2_max_probs.float())
            top_2_mask = top_2_mask * sampled.repeat(self.num_experts, 1).transpose(1, 0)

        if padding_mask is not None and not self.router_ignore_padding_tokens:
            if len(padding_mask.shape) == 4:
                # only get the last causal mask
                padding_mask = padding_mask[:, :, -1, :].reshape(-1)[-nb_tokens:]
            non_padding = ~padding_mask.bool()
            top_1_mask = top_1_mask * non_padding.unsqueeze(-1).to(top_1_mask.dtype)
            top_2_mask = top_2_mask * non_padding.unsqueeze(-1).to(top_1_mask.dtype)

        if self.batch_prioritized_routing:
            # sort tokens based on their routing probability
            # to make sure important tokens are routed, first
            importance_scores = -1 * router_probs.max(dim=1)[0]
            sorted_top_1_mask = top_1_mask[importance_scores.argsort(dim=0)]
            sorted_cumsum1 = (torch.cumsum(sorted_top_1_mask, dim=0) - 1) * sorted_top_1_mask
            locations1 = sorted_cumsum1[importance_scores.argsort(dim=0).argsort(dim=0)]

            sorted_top_2_mask = top_2_mask[importance_scores.argsort(dim=0)]
            sorted_cumsum2 = (torch.cumsum(sorted_top_2_mask, dim=0) - 1) * sorted_top_2_mask
            locations2 = sorted_cumsum2[importance_scores.argsort(dim=0).argsort(dim=0)]
            # Update 2nd's location by accounting for locations of 1st
            locations2 += torch.sum(top_1_mask, dim=0, keepdim=True)

        else:
            locations1 = torch.cumsum(top_1_mask, dim=0) - 1
            locations2 = torch.cumsum(top_2_mask, dim=0) - 1
            # Update 2nd's location by accounting for locations of 1st
            locations2 += torch.sum(top_1_mask, dim=0, keepdim=True)

        if not self.training and self.moe_eval_capacity_token_fraction > 0:
            self.expert_capacity = math.ceil(self.moe_eval_capacity_token_fraction * nb_tokens)
        else:
            capacity = 2 * math.ceil(nb_tokens / self.num_experts)
            self.expert_capacity = capacity if self.expert_capacity is None else self.expert_capacity

        # Remove locations outside capacity from ( cumsum < capacity = False will not be routed)
        top_1_mask = top_1_mask * torch.lt(locations1, self.expert_capacity)
        top_2_mask = top_2_mask * torch.lt(locations2, self.expert_capacity)

        if not self.normalize_router_prob_before_dropping:
            top_1_max_probs, top_2_max_probs = self.normalize_router_probabilities(
                router_probs, top_1_mask, top_2_mask
            )

        # Calculate combine_weights and dispatch_mask
        gates1 = top_1_max_probs[:, None] * top_1_mask
        gates2 = top_2_max_probs[:, None] * top_2_mask
        router_probs = gates1 + gates2

        return top_1_mask, router_probs

    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.LongTensor] = None) -> Tuple:
        r"""
        The hidden states are reshaped to simplify the computation of the router probabilities (combining weights for
        each experts.)

        Args:
            hidden_states (`torch.Tensor`):
                (batch_size, sequence_length, hidden_dim) from which router probabilities are computed.
        Returns:
            top_1_mask (`torch.Tensor` of shape (batch_size, sequence_length)):
                Index tensor of shape [batch_size, sequence_length] corresponding to the expert selected for each token
                using the top1 probabilities of the router.
            router_probabilities (`torch.Tensor` of shape (batch_size, sequence_length, nump_experts)):
                Tensor of shape (batch_size, sequence_length, num_experts) corresponding to the probabilities for each
                token and expert. Used for routing tokens to experts.
            router_logits (`torch.Tensor` of shape (batch_size, sequence_length))):
                Logits tensor of shape (batch_size, sequence_length, num_experts) corresponding to raw router logits.
                This is used later for computing router z-loss.
        """
        self.input_dtype = hidden_states.dtype
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.reshape((batch_size * sequence_length), hidden_dim)
        hidden_states = hidden_states.to(self.dtype)
        self._cast_classifier()
        router_logits = self.classifier(hidden_states)
        top_1_mask, router_probs = self.route_tokens(router_logits, self.input_dtype, padding_mask)
        return top_1_mask, router_probs, router_logits


class FSGPTMoEDenseActDense(nn.Module):
    def __init__(self, config: FSGPTMoEConfig, intermediate_size: int):
        super().__init__()
        self.fc1 = nn.Linear(config.hidden_size, intermediate_size)
        self.fc2 = nn.Linear(intermediate_size, config.hidden_size)
        self.act = ACT2FN[config.activation_function]
        self.dropout = nn.Dropout(config.resid_dropout)

    def forward(self, hidden_states):
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.dropout(hidden_states)
        if (
                isinstance(self.fc2.weight, torch.Tensor)
                and hidden_states.dtype != self.fc2.weight.dtype
                and self.fc2.weight.dtype != torch.int8
        ):
            hidden_states = hidden_states.to(self.fc2.weight.dtype)
        hidden_states = self.fc2(hidden_states)
        return hidden_states


class FSGPTMoESparseMLP(nn.Module):
    r"""
    Adapted from huggingface NLLB-MoE
    """

    def __init__(self, config: FSGPTMoEConfig, ffn_dim: int, expert_class: nn.Module = FSGPTMoEDenseActDense):
        super().__init__()
        self.router = FSGPTMoETop2Router(config)
        self.moe_token_dropout = config.moe_token_dropout
        self.token_dropout = nn.Dropout(self.moe_token_dropout)
        self.num_experts = config.num_experts

        self.experts = nn.ModuleDict()
        for idx in range(self.num_experts):
            self.experts[f"expert_{idx}"] = expert_class(config, ffn_dim)

    def forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.Tensor] = False):
        r"""
        The goal of this forward pass is to have the same number of operation as the equivalent `NllbMoeDenseActDense`
        (mlp) layer. This means that all of the hidden states should be processed at most twice ( since we are using a
        top_2 gating mecanism). This means that we keep the complexity to O(batch_size x sequence_length x hidden_dim)
        instead of O(num_experts x batch_size x sequence_length x hidden_dim).

        1- Get the `router_probs` from the `router`. The shape of the `router_mask` is `(batch_size X sequence_length,
        num_expert)` and corresponds to the boolean version of the `router_probs`. The inputs are masked using the
        `router_mask`.

        2- Dispatch the hidden_states to its associated experts. The router probabilities are used to weight the
        contribution of each experts when updating the masked hidden states.

        Args:
            hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`):
                The hidden states
            padding_mask (`torch.Tensor`, *optional*, defaults to `False`):
                Attention mask. Can be in the causal form or not.

        Returns:
            hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_dim)`):
                Updated hidden states
            router_logits (`torch.Tensor` of shape `(batch_size, sequence_length, num_experts)`):
                Needed for computing the loss

        """
        batch_size, sequence_length, hidden_dim = hidden_states.shape

        top_1_mask, router_probs, router_logits = self.router(hidden_states, padding_mask)
        router_mask = router_probs.bool()
        hidden_states = hidden_states.reshape((batch_size * sequence_length), hidden_dim)
        masked_hidden_states = torch.einsum("bm,be->ebm", hidden_states, router_mask)
        for idx, expert in enumerate(self.experts.values()):
            token_indices = router_mask[:, idx]
            combining_weights = router_probs[token_indices, idx]
            if isinstance(expert, torch.nn.Linear):
                # After pruning,there will be fake experts with no parameters which are nn.Linear
                # need to move them to the same device as the hidden states
                expert = expert.to(masked_hidden_states.device)
            expert_output = expert(masked_hidden_states[idx, token_indices])
            if self.moe_token_dropout > 0:
                if self.training:
                    expert_output = self.token_dropout(expert_output)
                else:
                    expert_output *= 1 - self.moe_token_dropout
            masked_hidden_states[idx, token_indices] = torch.einsum("b,be->be", combining_weights, expert_output)
        hidden_states = masked_hidden_states.sum(dim=0).reshape(batch_size, sequence_length, hidden_dim)

        top_1_expert_index = torch.argmax(top_1_mask, dim=-1)
        return hidden_states, (router_logits, router_probs, top_1_expert_index)


class FSGPTMoEBlock(nn.Module):
    def __init__(self, config: FSGPTMoEConfig, is_sparse: bool = False):
        super().__init__()
        hidden_size = config.hidden_size
        inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * hidden_size
        self.is_sparse = is_sparse
        self.self_attn = FSGPTMoESelfAttention(config)
        self.self_attn_layer_norm = nn.LayerNorm(hidden_size)
        if self.is_sparse:
            self.ffn = FSGPTMoESparseMLP(config, inner_dim)
        else:
            self.ffn = FSGPTMoEDenseActDense(config, inner_dim)
        self.final_layer_norm = nn.LayerNorm(hidden_size)

    def forward(
            self,
            hidden_states,
            layer_past=None,
            attention_mask=None,
            head_mask=None,
            use_cache=False,
            output_attentions=False,
            output_router_logits=False,
    ):
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)
        attn_outputs = self.self_attn(
            hidden_states,
            layer_past=layer_past,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attn_output = attn_outputs[0]
        outputs = attn_outputs[1:]
        # residual connection
        hidden_states = attn_output + residual

        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)

        if self.is_sparse:
            hidden_states, router_states = self.ffn(hidden_states, attention_mask)
        else:
            hidden_states, router_states = self.ffn(hidden_states), None

        # residual connection
        hidden_states = residual + hidden_states

        if use_cache:
            outputs = (hidden_states,) + outputs
        else:
            outputs = (hidden_states,) + outputs[1:]

        if output_router_logits:
            outputs = outputs + (router_states,)

        return outputs  # hidden_states, present, (attentions, cross_attentions), router_logits


class FSGPTMoEPreTrainedModel(PreTrainedModel):
    config_class = FSGPTMoEConfig
    base_model_prefix = "decoder"
    supports_gradient_checkpointing = True
    _no_split_modules = ["FSGPTMoEBlock"]
    _skip_keys_device_placement = "past_key_values"

    def __init__(self, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)

    def _init_weights(self, module):
        """Initialize the weights."""
        if isinstance(module, (nn.Linear,)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.hidden_size ** -0.5)
            if module.weight.data.shape[0] == self.config.vocab_size:
                module.weight.data[self.config.pad_token_id].zero_()
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.hidden_size ** -0.5)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, GPTNeoModel):
            module.gradient_checkpointing = value


class FSGPTMoEModel(FSGPTMoEPreTrainedModel):
    def __init__(self, config: FSGPTMoEConfig):
        super().__init__(config)

        self.embed_dim = config.hidden_size
        self.embed_tokens = nn.Embedding(config.vocab_size, self.embed_dim, padding_idx=config.pad_token_id)
        self.embed_positions = FSGPTMoESinusoidalPositionalEmbedding(
            embedding_dim=self.embed_dim,
            padding_idx=config.pad_token_id,
            num_positions=config.max_position_embeddings,
        )
        self.drop = nn.Dropout(float(config.embed_dropout))

        self.embed_scale = (
            1.0 if config.no_scale_embedding else math.sqrt(self.embed_dim)
        )
        self.layers = nn.ModuleList()
        sparse_step = config.sparse_step
        for i in range(config.num_layers):
            is_sparse = (i + 1) % sparse_step == 0 if sparse_step > 0 else False
            self.layers.append(FSGPTMoEBlock(config, is_sparse=is_sparse))
        self.layer_norm = nn.LayerNorm(self.embed_dim)

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, new_embeddings):
        self.embed_tokens = new_embeddings

    def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            output_router_logits: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], MoEModelOutputWithPastAndCrossAttentions]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            batch_size = input_ids.shape[0]
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size = inputs_embeds.shape[0]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

        if past_key_values is None:
            past_length = 0
            past_key_values = tuple([None] * len(self.layers))
        else:
            past_length = past_key_values[0][0].size(-2)

        # if position_ids is None:
        #     position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
        #     position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])

        # Attention mask.
        if attention_mask is not None:
            if batch_size <= 0:
                raise ValueError("batch_size has to be defined and > 0")
            attention_mask = attention_mask.view(batch_size, -1)
            # We create a 3D attention mask from a 2D tensor mask.
            # Sizes are [batch_size, 1, 1, to_seq_length]
            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
            # this attention mask is more simple than the triangular masking of causal attention
            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
            attention_mask = attention_mask[:, None, None, :]

            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
            # masked positions, this operation will create a tensor which is 0.0 for
            # positions we want to attend and the dtype's smallest value for masked positions.
            # Since we are adding it to the raw scores before the softmax, this is
            # effectively the same as removing these entirely.
            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x num_heads x N x N
        # head_mask has shape n_layer x batch x num_heads x N x N
        head_mask = self.get_head_mask(head_mask, self.config.num_layers)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
        position_embeds = self.embed_positions(input_ids)
        hidden_states = inputs_embeds + position_embeds

        if token_type_ids is not None:
            token_type_embeds = self.embed_tokens(token_type_ids)
            hidden_states = hidden_states + token_type_embeds

        hidden_states = self.drop(hidden_states)

        output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_hidden_states = () if output_hidden_states else None
        all_router_probs = () if output_router_logits else None
        for i, (block, layer_past) in enumerate(zip(self.layers, past_key_values)):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            outputs = block(
                hidden_states,
                layer_past=layer_past,
                attention_mask=attention_mask,
                head_mask=head_mask[i],
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_router_logits=output_router_logits,
            )

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)

            if output_router_logits:
                all_router_probs += (outputs[-1],)

        hidden_states = self.layer_norm(hidden_states)

        hidden_states = hidden_states.view(output_shape)
        # Add last hidden state
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v for v in
                [hidden_states,
                 presents,
                 all_hidden_states,
                 all_self_attentions,
                 all_router_probs]
                if v is not None
            )

        return MoEModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            router_probs=all_router_probs
        )


class FSGPTMoEForCausalLM(FSGPTMoEPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: FSGPTMoEConfig):
        super().__init__(config)
        self.decoder = FSGPTMoEModel(config)
        # Share input output embeddings
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.lm_head.weight = self.decoder.embed_tokens.weight

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
        token_type_ids = kwargs.get("token_type_ids", None)
        # only last token for inputs_ids if past is defined in kwargs
        if past_key_values:
            input_ids = input_ids[:, -1].unsqueeze(-1)
            if token_type_ids is not None:
                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)

        attention_mask = kwargs.get("attention_mask", None)
        position_ids = kwargs.get("position_ids", None)

        if attention_mask is not None and position_ids is None:
            # create position_ids on the fly for batch generation
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "position_ids": position_ids,
                "attention_mask": attention_mask,
                "token_type_ids": token_type_ids,
            }
        )

        return model_inputs

    def forward(
            self,
            input_ids: Optional[torch.Tensor] = None,
            past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
            attention_mask: Optional[torch.Tensor] = None,
            token_type_ids: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.Tensor] = None,
            head_mask: Optional[torch.Tensor] = None,
            inputs_embeds: Optional[torch.Tensor] = None,
            labels: Optional[torch.Tensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            output_router_logits: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], MoECausalLMOutputWithPast]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        output_router_logits = (
            output_router_logits if output_router_logits is not None else self.config.output_router_logits
        )

        transformer_outputs = self.decoder(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            output_router_logits=output_router_logits,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Compute loss in fp32 to match with mesh-tf version
            # https://github.com/EleutherAI/gpt-neo/blob/89ce74164da2fb16179106f54e2269b5da8db333/models/gpt2/gpt2.py#L179
            lm_logits = lm_logits.to(torch.float32)

            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            lm_logits = lm_logits.to(hidden_states.dtype)
            loss = loss.to(hidden_states.dtype)

        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return MoECausalLMOutputWithPast(
            loss=loss,
            logits=lm_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
            router_logits=transformer_outputs.router_probs,
        )

    @staticmethod
    def _reorder_cache(
            past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
    ) -> Tuple[Tuple[torch.Tensor]]:
        """
        This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
        [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
        beam_idx at every generation step.
        """
        return tuple(
            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
            for layer_past in past_key_values
        )

mcsmoe/models/fsgpt_moe.py
UNITES-Lab/MC-SMoE
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from transformers.modeling_outputs import ModelOutput

__all__ = [
    "MoEModelOutputWithPastAndCrossAttentions",
    "MoECausalLMOutputWithPastAndCrossAttentions",
    "MoECasualLMDistillationOutputWithPastAndCrossAttention",
    "MoEDistillationOutput",
]


@dataclass
class MoEModelOutputWithPastAndCrossAttentions(ModelOutput):
    last_hidden_state: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    router_probs: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class MoECausalLMOutputWithPastAndCrossAttentions(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    z_loss: torch.FloatTensor = None
    aux_loss: torch.FloatTensor = None
    router_logits: Optional[Tuple[torch.FloatTensor]] = None


class MoECasualLMDistillationOutputWithPastAndCrossAttention(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    z_loss: torch.FloatTensor = None
    aux_loss: torch.FloatTensor = None
    router_logits: Optional[Tuple[torch.FloatTensor]] = None
    kd_loss: torch.FloatTensor = None
    lm_loss: torch.FloatTensor = None


class MoEDistillationOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    kd_loss: torch.FloatTensor = None
    task_loss: torch.FloatTensor = None
    hd_loss: torch.FloatTensor = None
    norm_loss: torch.FloatTensor = None
    sing_reg_loss: torch.FloatTensor = None
    orth_reg_loss: torch.FloatTensor = None
    student_logits: torch.FloatTensor = None
    teacher_logits: torch.FloatTensor = None

mcsmoe/models/modeling_outputs.py
UNITES-Lab/MC-SMoE
import torch
from torch import nn


class LoSparseLinear(nn.Module):
    """
    Examples
    --------
    >>> import torch
    >>> linear = LoSparseLinear(16, 32, 2, has_bias=False)
    >>> inp = torch.randn(2, 16)
    >>> out = linear(inp)
    >>> out.shape
    torch.Size([2, 32])
    """

    def __init__(self, in_feature, out_feature, reduced_rank, has_bias=True, has_sparse=True):
        super().__init__()
        self.in_feature = in_feature
        self.out_feature = out_feature
        self.reduced_rank = reduced_rank
        self.has_bias = has_bias
        self.has_sparse = has_sparse

        self.right = nn.Linear(in_feature, reduced_rank, bias=False)
        self.left = nn.Linear(reduced_rank, out_feature, bias=False)
        if self.has_sparse:
            self.sparse = nn.Linear(in_feature, out_feature, bias=False)

        if self.has_bias:
            self.bias = nn.Parameter(torch.zeros(out_feature, requires_grad=True))

        self.nonzero_idx = None
        self.sparse_weight_pruned = None
        self.SX = None
        self.SX_deberta = None  # Deberta will use Q and K again

    @property
    def weight(self):
        return self.left.weight @ self.right.weight

    def forward(self, x):
        """ Y = XW.T+B = X(LR+S).T+B = X(LR).T+XS.T+B """
        l_r_x = self.left(self.right(x))
        if self.has_sparse:
            if self.sparse_weight_pruned is not None:
                s_x_ = torch.matmul(x, self.sparse_weight_pruned.T)
                b_, l_, d_ = x.shape

                # restore y
                # keep record for the first forward
                if self.SX is None or self.SX_deberta is None:  # For QKV at the first time
                    out_feature, in_feature = self.sparse.weight.shape
                    device = x.device
                    if b_ != 1:
                        self.SX = torch.zeros(b_, l_, out_feature, device=device)
                        self.SX[..., self.nonzero_idx] = s_x_
                        y_ = l_r_x + self.SX + self.bias if self.has_bias else l_r_x + self.SX
                    else:  # For QK at the second time
                        self.SX_deberta = torch.zeros(b_, l_, out_feature, device=device)
                        self.SX_deberta[..., self.nonzero_idx] = s_x_
                        y_ = l_r_x + self.SX_deberta + self.bias if self.has_bias else l_r_x + self.SX_deberta

                # do not need to create new cuda memory
                else:
                    if b_ != 1:
                        self.SX[..., self.nonzero_idx] = s_x_
                        y_ = l_r_x + self.SX + self.bias if self.has_bias else l_r_x + self.SX
                    else:
                        self.SX_deberta[..., self.nonzero_idx] = s_x_
                        y_ = l_r_x + self.SX_deberta + self.bias if self.has_bias else l_r_x + self.SX_deberta
            else:
                s_x = self.sparse(x)
                y_ = l_r_x + s_x + self.bias if self.has_bias else l_r_x + s_x
        else:
            y_ = l_r_x + self.bias if self.has_bias else l_r_x
        return y_

    def initialize_weight(self, left_weight, right_weight, sparse_weight=None, bias=None):
        self.left.weight = nn.Parameter(left_weight, requires_grad=True)
        self.right.weight = nn.Parameter(right_weight, requires_grad=True)
        if self.has_sparse:
            self.sparse.weight = nn.Parameter(sparse_weight, requires_grad=True)
        if self.has_bias:
            self.bias = nn.Parameter(bias, requires_grad=True)

    def prune_sparse(self):
        self.nonzero_idx = torch.nonzero(self.sparse.weight.sum(dim=1)).flatten()
        # self.sparse_weight_pruned = self.sparse.weight[self.nonzero_idx, :]
        self.sparse_weight_pruned = nn.Parameter(self.sparse.weight[self.nonzero_idx, :])

mcsmoe/models/sparse.py
UNITES-Lab/MC-SMoE
""" Adapted from modeling_switch_transformers.py in HuggingFace transformers repository """

import copy
import math
import warnings
from dataclasses import dataclass
from typing import Dict, List, Any
from typing import Optional, Union
from typing import Tuple
from warnings import warn

import numpy as np
import torch
import torch.nn.functional as F
from torch import linalg as LA
from torch import nn
from torch.nn import CrossEntropyLoss
from torch.utils.checkpoint import checkpoint
from transformers import (
    AutoTokenizer,
    BatchEncoding,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
)
from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    MoEModelOutput,
    MoEModelOutputWithPastAndCrossAttentions,
    Seq2SeqMoEModelOutput,
    Seq2SeqMoEOutput
)
from transformers.modeling_utils import PreTrainedModel
from transformers.models.switch_transformers import SwitchTransformersConfig
from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
from transformers.utils import (
    DUMMY_INPUTS,
    DUMMY_MASK,
    is_torch_fx_proxy,
    logging,
)

from .modeling_outputs import MoEDistillationOutput

logger = logging.get_logger(__name__)


def compute_input_and_target_lengths(inputs_length, noise_density, mean_noise_span_length):
    """This function is copy of `random_spans_helper <https://github.com/google-research/text-to-text-transfer-transformer/blob/84f8bcc14b5f2c03de51bd3587609ba8f6bbd1cd/t5/data/preprocessors.py#L2466>`__ .

    [Copied from https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_t5_mlm_flax.py]
    Training parameters to avoid padding with random_spans_noise_mask.
    When training a model with random_spans_noise_mask, we would like to set the other
    training hyperparmeters in a way that avoids padding.
    This function helps us compute these hyperparameters.
    We assume that each noise span in the input is replaced by extra_tokens_per_span_inputs sentinel tokens,
    and each non-noise span in the targets is replaced by extra_tokens_per_span_targets sentinel tokens.
    This function tells us the required number of tokens in the raw example (for split_tokens())
    as well as the length of the encoded targets. Note that this function assumes
    the inputs and targets will have EOS appended and includes that in the reported length.

    Args:
        inputs_length: an integer - desired length of the tokenized inputs sequence
        noise_density: a float
        mean_noise_span_length: a float
    Returns:
        tokens_length: length of original text in tokens
        targets_length: an integer - length in tokens of encoded targets sequence

    >>> DataCollatorForT5MLM.compute_input_and_target_lengths(512, 0.15, 3.0)
    (568, 114)
    """

    def _tokens_length_to_inputs_length_targets_length(tokens_length):
        num_noise_tokens = int(round(tokens_length * noise_density))
        num_nonnoise_tokens = tokens_length - num_noise_tokens
        num_noise_spans = int(round(num_noise_tokens / mean_noise_span_length))
        # inputs contain all nonnoise tokens, sentinels for all noise spans
        # and one EOS token.
        _input_length = num_nonnoise_tokens + num_noise_spans + 1
        _output_length = num_noise_tokens + num_noise_spans + 1
        return _input_length, _output_length

    tokens_length = inputs_length

    while _tokens_length_to_inputs_length_targets_length(tokens_length + 1)[0] <= inputs_length:
        tokens_length += 1

    inputs_length, targets_length = _tokens_length_to_inputs_length_targets_length(tokens_length)

    # minor hack to get the targets length to be equal to inputs length
    # which is more likely to have been set to a nice round number.
    if noise_density == 0.5 and targets_length > inputs_length:
        tokens_length -= 1
        targets_length -= 1
    return tokens_length, targets_length

class SwitchTransformersWrapperForDistillation(nn.Module):
    EXCEPTED_CONFIG_KEYS = [
        "freeze_router",
        "sparse_encoder_layer_list",
        "sparse_decoder_layer_list",
        "expert_norm_coef",
        "router_z_loss_coef",
        "router_aux_loss_coef",
        "only_task_vectors_for_pruning",
        "keep_singular_values",
    ]

    def __init__(
            self,
            teacher=None,
            student=None,
            mlm_lambda: float = 1.0,
            kd_lambda: float = 1.0,
            kd_temperature: float = 1.0,
            model_max_length: int = 512,
            reg_lambda: float = 0.0,
            hd_lambda: float = 0.0,
            orth_reg_lambda: float = 0.0,
            sing_reg_lambda: float = 0.0,
            hd_cos_sim: bool = False,
            next_rank: int = None,
    ):
        super().__init__()
        self.kd_lambda = kd_lambda
        self.kd_temperature = kd_temperature
        self.model_max_length = model_max_length
        self.mlm_lambda = mlm_lambda
        self.reg_lambda = reg_lambda
        self.hd_cos_sim = hd_cos_sim
        self.hd_lambda = hd_lambda
        self.orth_reg_lambda = orth_reg_lambda
        self.sing_reg_lambda = sing_reg_lambda

        if next_rank is None and hasattr(student.config, "low_rank_factor"):
            self.next_rank = student.config.low_rank_factor
        elif next_rank is None:
            self.next_rank = student.config.d_model
        else:
            self.next_rank = next_rank

        self.teacher = teacher
        self.teacher.router_z_loss_coef = 0.0
        self.teacher.router_aux_loss_coef = 0.0

        self.student = student
        self.student.router_z_loss_coef = 0.0
        self.student.router_aux_loss_coef = 0.0

        for key, student_value in self.student.config.__dict__.items():
            if key not in self.EXCEPTED_CONFIG_KEYS and self.teacher.config.__dict__[key] != student_value:
                warn(f"Teacher and student have different config: {key}, "
                     f"at teacher: {self.teacher.config.__dict__[key]} and at student: {student_value}")

    def forward(
            self,
            *args,
            **kwargs
    ):
        # Shut down dropout and layer norm for accurate distillation
        self.teacher.eval()
        self.student.eval()

        kwargs["output_hidden_states"] = True

        with torch.no_grad():
            teacher_outputs = self.teacher(*args, **kwargs)
        teacher_encoder_hidden_states = teacher_outputs.encoder_hidden_states[1:]  # skip embedding hidden states
        teacher_decoder_hidden_states = teacher_outputs.decoder_hidden_states[1:]
        teacher_encoder_hidden_states = torch.stack(teacher_encoder_hidden_states, dim=0)
        teacher_decoder_hidden_states = torch.stack(teacher_decoder_hidden_states, dim=0)
        teacher_logits = teacher_outputs.logits

        student_outputs = self.student(*args, **kwargs)
        student_encoder_hidden_states = student_outputs.encoder_hidden_states[1:]
        student_decoder_hidden_states = student_outputs.decoder_hidden_states[1:]
        student_encoder_hidden_states = torch.stack(student_encoder_hidden_states, dim=0)
        student_decoder_hidden_states = torch.stack(student_decoder_hidden_states, dim=0)
        student_logits = student_outputs.logits

        # 1. kl divergence loss
        kd_criteria = nn.KLDivLoss(reduction="batchmean")
        kd_loss = kd_criteria(
            F.log_softmax(student_logits / self.kd_temperature, dim=-1),
            F.softmax(teacher_logits / self.kd_temperature, dim=-1)
        )
        kd_loss = kd_loss * self.kd_lambda

        # 2. language modeling loss
        task_loss = student_outputs.loss * self.mlm_lambda

        # 3. expert U & V norm loss
        device = task_loss.device
        if self.reg_lambda == 0:
            norm_loss = torch.tensor(0.0, device=device)
        elif self.next_rank != self.student.config.low_rank_factor:
            norm_loss = self.compute_u_v_reg_for_next_rank(device=device)
            norm_loss = norm_loss * self.reg_lambda
        else:
            norm_loss = torch.tensor(0.0, device=device)

        # 4. hidden states loss
        if self.hd_cos_sim:
            hd_criteria = nn.CosineSimilarity(dim=-1)
        else:
            hd_criteria = nn.MSELoss(reduction="none")

        encoder_hidden_states_loss = hd_criteria(student_encoder_hidden_states, teacher_encoder_hidden_states).mean()
        decoder_hidden_states_loss = hd_criteria(student_decoder_hidden_states, teacher_decoder_hidden_states).mean()
        hidden_states_loss = (encoder_hidden_states_loss + decoder_hidden_states_loss) / 2

        if self.hd_lambda == 0:
            hd_loss = torch.tensor(0.0, device=device)
        elif self.hd_cos_sim:
            # Cosine similarity is in [-1, 1], so we need to convert it to [0, 1]
            hd_loss = (1 - hidden_states_loss) * self.hd_lambda / 2
        else:
            hd_loss = torch.log(hidden_states_loss) * self.hd_lambda

        # 5. orthogonality regularization loss
        if self.orth_reg_lambda == 0:
            orth_reg_loss = torch.tensor(0.0, device=device)
        else:
            orth_reg_loss = self.compute_orthogonality_reg(device=device) * self.orth_reg_lambda

        # 6. singular value sparsity regularization loss
        if self.sing_reg_lambda == 0:
            sing_reg_loss = torch.tensor(0.0, device=device)
        else:
            sing_reg_loss = self.compute_singular_sparsity_reg(device=device) * self.sing_reg_lambda

        # total loss
        loss = kd_loss + task_loss + norm_loss + hd_loss + orth_reg_loss + sing_reg_loss

        return MoEDistillationOutput(
            loss=loss,
            student_logits=student_logits,
            kd_loss=kd_loss,
            hd_loss=hd_loss,
            task_loss=task_loss,
            norm_loss=norm_loss,
            sing_reg_loss=sing_reg_loss,
            orth_reg_loss=orth_reg_loss,
        )

    def save_pretrained(self, *args, **kwargs):
        self.student.save_pretrained(*args, **kwargs)

    def compute_u_v_reg_for_next_rank(self, device) -> torch.Tensor:
        """
        Intuitively, this computes the average singular values of components to be pruned next iter.

        Given next_rank, we compute the average singular values of components by
         U[:, next_rank:] @ V[:, next_rank:].T
        """
        reg_loss = torch.tensor(0.0, device=device)
        next_rank = self.next_rank
        for layer_idx in self.student.config.sparse_encoder_layer_list:
            encoder_mlp = self.student.encoder.block[layer_idx].layer[-1].mlp
            decoder_mlp = self.student.decoder.block[layer_idx].layer[-1].mlp
            for expert_idx in range(self.student.config.num_experts):
                u_norm_loss = LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wi_u,
                    ord=2,
                    dim=0
                )[next_rank:]
                v_norm_loss = LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wi_v,
                    ord=2,
                    dim=1
                )[next_rank:]
                reg_loss += (u_norm_loss * v_norm_loss).mean()

                u_norm_loss = LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wo_u,
                    ord=2,
                    dim=0
                )[next_rank:]
                v_norm_loss = LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wo_v,
                    ord=2,
                    dim=1
                )[next_rank:]
                reg_loss += (u_norm_loss * v_norm_loss).mean()

                u_norm_loss = LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wi_u,
                    ord=2,
                    dim=0
                )[next_rank:]
                v_norm_loss = LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wi_v,
                    ord=2,
                    dim=1
                )[next_rank:]
                reg_loss += (u_norm_loss * v_norm_loss).mean()

                u_norm_loss = LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wo_u,
                    ord=2,
                    dim=0
                )[next_rank:]
                v_norm_loss = LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wo_v,
                    ord=2,
                    dim=1
                )[next_rank:]
                reg_loss += (u_norm_loss * v_norm_loss).mean()

        reg_loss /= len(self.student.config.sparse_encoder_layer_list) * self.student.config.num_experts * 4

        return reg_loss

    def compute_singular_sparsity_reg(self, device) -> torch.Tensor:
        """
        Sparsity-inducing regularization loss for S matrix in USV decomposition.

        Given W=USV, we compute the sparsity-inducing regularization loss of W by
            reg_loss = ||W||_1 / ||W||_2
        """
        if not hasattr(self.student.config, "keep_singular_values") or not self.student.config.keep_singular_values:
            raise RuntimeError(f"Student model must be of form USV^T, but keep_singular_values is not set to True")

        reg_loss = torch.tensor(0.0, device=device)
        for layer_idx in self.student.config.sparse_encoder_layer_list:
            encoder_mlp = self.student.encoder.block[layer_idx].layer[-1].mlp
            decoder_mlp = self.student.decoder.block[layer_idx].layer[-1].mlp
            for expert_idx in range(self.student.config.num_experts):
                # Be careful with NaN
                reg_loss += LA.norm(encoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_s, ord=1) / (LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_s, ord=2) + 1e-5)
                reg_loss += LA.norm(encoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_s, ord=1) / (LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_s, ord=2) + 1e-5)
                reg_loss += LA.norm(decoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_s, ord=1) / (LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_s, ord=2) + 1e-5)
                reg_loss += LA.norm(decoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_s, ord=1) / (LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_s, ord=2) + 1e-5)

        reg_loss /= len(self.student.config.sparse_encoder_layer_list) * self.student.config.num_experts * 4

        return reg_loss

    def compute_orthogonality_reg(self, device) -> torch.Tensor:
        """
        Orthogonality inducer for U and V matrices in USV decomposition.

        Given W=USV, we compute the orthogonality-inducing regularization loss of U and V by
            reg_loss = ( ||U^T U - I||_F^2 + ||V^T V - I||_F^2 ) / 2

        """

        if not hasattr(self.student.config, "keep_singular_values") or not self.student.config.keep_singular_values:
            raise RuntimeError(f"Student model must be of form USV^T, but keep_singular_values is not set to True")

        reg_loss = torch.tensor(0.0, device=device)
        for layer_idx in self.student.config.sparse_encoder_layer_list:
            encoder_mlp = self.student.encoder.block[layer_idx].layer[-1].mlp
            decoder_mlp = self.student.decoder.block[layer_idx].layer[-1].mlp
            for expert_idx in range(self.student.config.num_experts):
                reg_loss += LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_u.T @ encoder_mlp.experts[
                        f"expert_{expert_idx}"].wi_mat_u - torch.eye(self.student.config.low_rank_factor,
                                                                     device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_v @ encoder_mlp.experts[
                        f"expert_{expert_idx}"].wi_mat_v.T - torch.eye(self.student.config.low_rank_factor,
                                                                       device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_u.T @ encoder_mlp.experts[
                        f"expert_{expert_idx}"].wo_mat_u - torch.eye(self.student.config.low_rank_factor,
                                                                     device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    encoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_v @ encoder_mlp.experts[
                        f"expert_{expert_idx}"].wo_mat_v.T - torch.eye(self.student.config.low_rank_factor,
                                                                       device=device, requires_grad=False),
                    ord="fro"
                )

                reg_loss += LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_u.T @ decoder_mlp.experts[
                        f"expert_{expert_idx}"].wi_mat_u - torch.eye(self.student.config.low_rank_factor,
                                                                     device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wi_mat_v @ decoder_mlp.experts[
                        f"expert_{expert_idx}"].wi_mat_v.T - torch.eye(self.student.config.low_rank_factor,
                                                                       device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_u.T @ decoder_mlp.experts[
                        f"expert_{expert_idx}"].wo_mat_u - torch.eye(self.student.config.low_rank_factor,
                                                                     device=device, requires_grad=False),
                    ord="fro"
                )
                reg_loss += LA.norm(
                    decoder_mlp.experts[f"expert_{expert_idx}"].wo_mat_v @ decoder_mlp.experts[
                        f"expert_{expert_idx}"].wo_mat_v.T - torch.eye(self.student.config.low_rank_factor,
                                                                       device=device, requires_grad=False),
                    ord="fro"
                )

        return reg_loss / (len(self.student.config.sparse_encoder_layer_list) * self.student.config.num_experts * 8)

mcsmoe/models/switch_transformers.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union, Dict

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups_with_usage_frequency_weighting
)
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)
from mcsmoe.utils.sparsity import compute_weight_stable_rank
from mcsmoe.utils.training_utils import freeze_switch_routers_for_finetuning

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def save_stable_rank_to(
        core_experts: Dict[str, List[int]], state_dict: Dict[str, torch.Tensor], save_dir: str, save_name: str
):
    wi_stable_rank_dict = {key: [] for key in core_experts.keys()}
    wo_stable_rank_dict = {key: [] for key in core_experts.keys()}
    for mlp_name in tqdm(core_experts, desc="Computing stable rank"):
        for i, core_idx in enumerate(core_experts[mlp_name]):
            name = f"{mlp_name}.experts.expert_{core_idx}.wi.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wi_stable_rank_dict[mlp_name].append(stable_rank)
            name = f"{mlp_name}.experts.expert_{core_idx}.wo.weight"
            stable_rank = compute_weight_stable_rank(state_dict[name].float())
            wo_stable_rank_dict[mlp_name].append(stable_rank)
        wi_stable_rank_dict[mlp_name] = torch.tensor(wi_stable_rank_dict[mlp_name])
        wo_stable_rank_dict[mlp_name] = torch.tensor(wo_stable_rank_dict[mlp_name])
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    torch.save(wi_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wi.pt"))
    torch.save(wo_stable_rank_dict, os.path.join(save_dir, f"{save_name}.wo.pt"))


@dataclass
class SwitchDistillationConfig:
    max_steps: Optional[int] = 10000
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    mlm_probability: Optional[float] = 0.15
    mean_noise_span_length: Optional[int] = 3
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.1
    hd_lambda: Optional[float] = 0.01
    hd_cos_sim: Optional[bool] = False
    next_rank_for_norm: Optional[int] = 384


def merge_and_distill_downstream_for_recover(
        output_dir: Optional[str] = None,
        save_stable_rank: Optional[bool] = False,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        merging_strategy: Optional[str] = "normal",
        num_samples_for_merging: Optional[int] = 32,
        similarity_base: Optional[str] = "router-logits",
        num_groups: Optional[int] = 16,
        reverse_similarity: Optional[bool] = False,
        globally_group: Optional[bool] = False,
        within_and_across_models: Optional[bool] = False,
        group_capacity: Optional[int] = None,
        permute_when_merge: Optional[bool] = False,
        encoder_merging_layers: Optional[Union[str, List, int]] = None,
        decoder_merging_layers: Optional[Union[str, List, int]] = None,
        exact_fisher: Optional[bool] = True,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-4,
        gradient_accumulation_steps: Optional[int] = 1,
        warmup_steps: Optional[int] = 16,
        num_epochs: Optional[int] = 10,
        max_train_steps: Optional[int] = 20000,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, merging_strategy, similarity_base)
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        student_checkpoint = teacher_checkpoint

    encoder_merging_layers = sanitize_merging_layers(encoder_merging_layers)
    decoder_merging_layers = sanitize_merging_layers(decoder_merging_layers)

    training_config = SwitchDistillationConfig(
        max_steps=max_train_steps,
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        mlm_probability=0.15,
        mean_noise_span_length=3,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = HFSwitch.from_pretrained(student_checkpoint)
    student_model = freeze_switch_routers_for_finetuning(student_model)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=training_config.reg_lambda,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"merge-by-usage-weighting-{task}-global-{similarity_base}-{merging_strategy}" if globally_group else (
            f"merge-by-usage-weighting-{task}-uniform-{similarity_base}-{merging_strategy}")
        if permute_when_merge:
            run_name += "-permuted"
        if reverse_similarity:
            run_name += "-revsim"
        if within_and_across_models:
            run_name += "-within-and-across-models"
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    # random select a subset of training dataset for merging
    dataset_for_merging = tokenized_train_dataset.shuffle(seed=2333).select(range(num_samples_for_merging))

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    dataloader_for_merging = DataLoader(
        dataset_for_merging,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=num_samples_for_merging,
        num_workers=3
    )

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=3
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=3
    )

    print(f"Merging experts by usage weighted averaging with {merging_strategy} usage")
    print(f"Similarities are computed based on {similarity_base}")
    print(f"Number of groups: {num_groups}")
    print(f"Merge encoder layers: {encoder_merging_layers}")
    print(f"Merge decoder layers: {decoder_merging_layers}")

    grouper = ExpertsGrouperForSwitch(
        config=model.student.config,
        similarity_base=similarity_base
    )
    grouper.compute_all_similarities(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    if reverse_similarity:
        grouper.reverse_all_similarities()
    grouper.compute_all_usages(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=num_groups,
        merging_encoder_layers=encoder_merging_layers,
        merging_decoder_layers=decoder_merging_layers,
        encoder_group_capacity=group_capacity,
    ) if globally_group else grouper.group_experts_into_clusters_by_routing_guided(
        num_groups=num_groups,
    )
    print(grouper.group_state_dict())
    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"before-merging-{num_groups}"
        )
    torch.save(grouper.group_state_dict(), os.path.join(output_dir, "group_state_dict.pt"))
    if globally_group:
        torch.save(core_experts, os.path.join(output_dir, "core_experts.pt"))
    model.student = merge_by_groups_with_usage_frequency_weighting(
        model=model.student,
        grouper=grouper,
        strategy=merging_strategy,
        encoder_merging_layers=encoder_merging_layers,
        decoder_merging_layers=decoder_merging_layers,
        permute=permute_when_merge,
        within_and_across_models=within_and_across_models,
    )

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-merging-{num_groups}"
        )

    print(f"Number of parameters before merging: {model.teacher.num_parameters()}")
    print(f"Number of parameters after merging: {model.student.num_parameters()}")

    no_weight_decay = ["bias", "LayerNorm"]  # no weight decay for orthogonalization matrices
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                hd_loss = outputs.hd_loss
                task_loss = outputs.task_loss
                norm_loss = outputs.norm_loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_hd_loss": hd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_norm_loss": norm_loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)
                # model.student.train()

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if save_stable_rank:
        save_stable_rank_to(
            core_experts=core_experts,
            state_dict=unwrapped_model.student.state_dict(),
            save_dir=f"./visualization/{task}/{os.path.basename(teacher_model.config._name_or_path)}/",
            save_name=f"after-training-{num_groups}"
        )

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(merge_and_distill_downstream_for_recover)

mcsmoe/msmoe-merging.py
UNITES-Lab/MC-SMoE
from typing import Optional

import torch
from accelerate.utils import set_seed
from fire import Fire
from transformers import (
    AutoTokenizer
)

from mcsmoe.merging.permutation import (
    align_fsgpt_permutation_for_all_experts_by_weight_matching
)
from mcsmoe.models import (
    FSGPTMoEForCausalLM
)

set_seed(2333)


def align_fsgpt_model_by_weight_matching_permutation(
        checkpoint: str,
        save_dir: str,
        include_wo: Optional[bool] = True,
):
    model = FSGPTMoEForCausalLM.from_pretrained(
        checkpoint,
        device_map='auto',
        torch_dtype=torch.bfloat16,
        moe_eval_capacity_token_fraction=1.0,
    )
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    model = align_fsgpt_permutation_for_all_experts_by_weight_matching(model, include_wo=include_wo).cpu()
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)


if __name__ == '__main__':
    Fire(align_fsgpt_model_by_weight_matching_permutation)

mcsmoe/permute-fsgpt-model.py
UNITES-Lab/MC-SMoE
import os
from typing import Optional

from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from transformers import (
    SwitchTransformersForConditionalGeneration as HFSwitch,
    AutoTokenizer
)

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq
)
from mcsmoe.merging.permutation import (
    align_switch_permutation_for_all_experts_by_weight_matching,
    align_switch_permutation_for_all_experts_by_activation_matching
)

set_seed(2333)


def align_switch_model_by_weight_matching_permutation(
        checkpoint: str,
        save_dir: str,
        include_wo: Optional[bool] = True,
):
    model = HFSwitch.from_pretrained(checkpoint).cuda()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    model = align_switch_permutation_for_all_experts_by_weight_matching(model, include_wo=include_wo).cpu()
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)


def align_switch_model_by_activation_matching_permutation(
        checkpoint: str,
        save_dir: str,
        task: str,
        num_samples: int,
        batch_size: Optional[int] = 128,
):
    if batch_size > num_samples:
        print(f"batch_size {batch_size} is larger than num_samples {num_samples}, "
              f"set batch_size to num_samples")
        batch_size = num_samples

    model = HFSwitch.from_pretrained(checkpoint).cuda()
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt')
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    train_dataset = raw_dataset["train"]
    train_dataset = train_dataset.map(
        Seq2SeqDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=train_dataset.column_names
    )
    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=8,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    dataset_for_permutation = tokenized_train_dataset.shuffle(seed=2333).select(range(num_samples))
    dataloader_for_permutation = DataLoader(
        dataset_for_permutation,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=batch_size,
        num_workers=3
    )
    model = align_switch_permutation_for_all_experts_by_activation_matching(model, dataloader_for_permutation).cpu()
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)


def main(
        checkpoint: str,
        save_dir: str,
        strategy: str,
        task: Optional[str] = None,
        num_samples: Optional[int] = None,
        batch_size: Optional[int] = 128,
        include_wo: Optional[bool] = True
):
    if os.path.isdir(save_dir):
        raise ValueError(f"save_dir {save_dir} already exists")
    if strategy == 'weight':
        align_switch_model_by_weight_matching_permutation(checkpoint, save_dir, include_wo)
    elif strategy == 'activation':
        align_switch_model_by_activation_matching_permutation(checkpoint, save_dir, task, num_samples, batch_size)
    else:
        raise ValueError(f"strategy {strategy} is not supported")


if __name__ == '__main__':
    Fire(main)

mcsmoe/permute-model.py
UNITES-Lab/MC-SMoE
import os
from typing import Optional, Dict, Union

import torch
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from transformers import (
    T5TokenizerFast,
    T5ForConditionalGeneration,
    SwitchTransformersForConditionalGeneration
)

from mcsmoe.calflops import calculate_flops
from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups_with_usage_frequency_weighting,
    load_merged_switch_transformers_from_checkpoint
)
from mcsmoe.pruning import substitute_switch_transformers_expert_weights

set_seed(233)


class ModelWrapper(torch.nn.Module):
    def __init__(self, model: Union[T5ForConditionalGeneration, SwitchTransformersForConditionalGeneration]):
        super().__init__()
        self.model = model

    def forward(
            self,
            input_ids: torch.LongTensor,
            attention_mask: torch.FloatTensor,
            labels: torch.LongTensor,
    ):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        ).logits


def prepare_dataloader(task: str, tokenizer: T5TokenizerFast, batch_size: int):
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"]
    dataset = dataset.map(
        Seq2SeqDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=6,
        remove_columns=dataset.column_names
    )
    dataset = dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=6,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt')
    dataloader = DataLoader(
        dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=batch_size,
        num_workers=8
    )
    return dataloader


def merge_switch_by_usage_weighted_averaging(
        model: SwitchTransformersForConditionalGeneration,
        batch: Dict[str, torch.Tensor],
        prune_router_non_core: bool = False,
) -> SwitchTransformersForConditionalGeneration:
    grouper = ExpertsGrouperForSwitch(
        config=model.config,
        similarity_fn="cosine",
        similarity_base="router-logits"
    )
    grouper.compute_all_similarities(
        model=model,
        batch=batch,
    )
    grouper.compute_all_usages(
        model=model,
        batch=batch,
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=8,
        merging_encoder_layers=[3, 5, 7, 9, 11],
        merging_decoder_layers=[1, 3, 5, 7, 9, 11]
    )
    model = merge_by_groups_with_usage_frequency_weighting(
        model=model,
        grouper=grouper,
        encoder_merging_layers=[3, 5, 7, 9, 11],
        decoder_merging_layers=[1, 3, 5, 7, 9, 11],
    )

    if prune_router_non_core:
        # Prune the fake experts and router classifier rows
        num_experts = model.config.num_experts
        router_bias = model.config.router_bias
        # Encoder
        for layer_idx in range(3, 12, 2):
            encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
            encoder_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            core_expert_indices = torch.tensor(core_experts[encoder_name], dtype=torch.int).sort().values
            # prune router
            encoder_mlp.router.num_experts = len(core_expert_indices)
            with torch.no_grad():
                pruned_router_cls = encoder_mlp.router.classifier.weight.data[core_expert_indices]
                encoder_mlp.router.classifier = torch.nn.Linear(
                    encoder_mlp.router.classifier.in_features,
                    encoder_mlp.router.num_experts,
                    bias=router_bias
                )
                encoder_mlp.router.classifier.weight.data.copy_(pruned_router_cls)
                if router_bias:
                    pruned_router_bias = encoder_mlp.router.classifier.bias.data[core_expert_indices]
                    encoder_mlp.router.classifier.bias.data.copy_(pruned_router_bias)
            # prune experts
            for expert_idx in range(num_experts):
                if expert_idx not in core_expert_indices:
                    del encoder_mlp.experts[f"expert_{expert_idx}"]
            assert len(encoder_mlp.experts) == encoder_mlp.router.classifier.out_features
        # Decoder
        for layer_idx in range(1, 12, 2):
            decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
            decoder_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            core_expert_indices = torch.tensor(core_experts[decoder_name], dtype=torch.int).sort().values
            # prune router
            decoder_mlp.router.num_experts = len(core_expert_indices)
            with torch.no_grad():
                pruned_router_cls = decoder_mlp.router.classifier.weight.data[core_expert_indices]
                decoder_mlp.router.classifier = torch.nn.Linear(
                    decoder_mlp.router.classifier.in_features,
                    decoder_mlp.router.num_experts,
                    bias=router_bias
                )
                decoder_mlp.router.classifier.weight.data.copy_(pruned_router_cls)
                if router_bias:
                    pruned_router_bias = decoder_mlp.router.classifier.bias.data[core_expert_indices]
                    decoder_mlp.router.classifier.bias.data.copy_(pruned_router_bias)
            # prune experts
            for expert_idx in range(num_experts):
                if expert_idx not in core_expert_indices:
                    del decoder_mlp.experts[f"expert_{expert_idx}"]
            assert len(decoder_mlp.experts) == decoder_mlp.router.classifier.out_features

    return model


def profile_inference_flops(
        checkpoint: str,
        task: str,
        batch_size: Optional[int] = 256,
        losparse_checkpoint: Optional[str] = None,
):
    if "t5" in checkpoint:
        model = T5ForConditionalGeneration.from_pretrained(checkpoint)
    elif "merged" in checkpoint:
        model = load_merged_switch_transformers_from_checkpoint(checkpoint)
        if losparse_checkpoint is not None:
            d_model = model.config.d_model
            d_ff = model.config.d_ff
            low_rank_factor = 32
            parameter_ratio = (d_model + d_ff) * low_rank_factor / (d_model * d_ff)
            model = substitute_switch_transformers_expert_weights(
                model,
                use_svd=False,
                parameter_ratio=parameter_ratio,
                encoder_compression_layers=list(range(3, 12, 2)),
                decoder_compression_layers=list(range(1, 12, 2)),
            )
            model.load_state_dict(torch.load(os.path.join(losparse_checkpoint, "pytorch_model.bin")))
        for layer_idx in range(1, 12, 2):
            model.encoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
            model.decoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
    else:
        model = SwitchTransformersForConditionalGeneration.from_pretrained(checkpoint)
        for layer_idx in range(1, 12, 2):
            model.encoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
            model.decoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
    tokenizer = T5TokenizerFast.from_pretrained(checkpoint)
    dataloader = prepare_dataloader(task, tokenizer, batch_size)
    batch = next(iter(dataloader))
    batch = {k: v.cuda() for k, v in batch.items()}

    model = model.cuda().eval()
    print("Start profiling...")
    model = model.float()
    model = ModelWrapper(model)

    with torch.no_grad():
        flops, macs, params = calculate_flops(
            model=model,
            kwargs=batch,
            print_detailed=True,
            include_backPropagation=False,
            is_sparse=True
        )
    print(f"Task: {task}.")
    print(f"Checkpoint: {checkpoint}.")
    print(f"Batch size: {batch_size}.")
    print(f"Input sequence length: {batch['input_ids'].shape[1]}")
    print(f"Output sequence length: {batch['labels'].shape[1]}")
    print(f"FLOPs: {flops}. MACs: {macs}. Params: {params}.")


if __name__ == '__main__':
    Fire(profile_inference_flops)

mcsmoe/profile-flops.py
UNITES-Lab/MC-SMoE
from typing import Optional, Union

import torch
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from transformers import (
    GPT2TokenizerFast,
)

from mcsmoe.calflops import calculate_flops
from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2se2_to_casual_lm,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForLanguageModeling,
)
from mcsmoe.merging.utils import load_merged_fsgpt_moe_from_checkpoint
from mcsmoe.models import (
    FSGPTForCausalLM,
    FSGPTMoEForCausalLM,
)
from mcsmoe.pruning.losparse import substitute_fsgpt_moe_expert_weights

set_seed(233)


class ModelWrapper(torch.nn.Module):
    def __init__(self, model: Union[FSGPTForCausalLM, FSGPTMoEForCausalLM]):
        super().__init__()
        self.model = model

    def forward(
            self,
            input_ids: torch.LongTensor,
            attention_mask: torch.FloatTensor,
            labels: torch.LongTensor,
    ):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        ).logits


def prepare_dataloader(task: str, tokenizer: GPT2TokenizerFast, batch_size: int):
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"]
    dataset = dataset.map(
        Seq2SeqDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    dataset = dataset.map(
        lambda x: tokenize_seq2se2_to_casual_lm(tokenizer=tokenizer, batch=x),
        num_proc=8,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
        keys_to_ignore=["answer_idx", "choice_idx", "idx"]
    )
    dataloader = DataLoader(
        dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=batch_size,
        num_workers=8
    )
    return dataloader


def profile_inference_flops(
        checkpoint: str,
        task: str,
        batch_size: Optional[int] = 256,
        losparse_checkpoint: Optional[str] = None,
):
    if "dense" in checkpoint:
        model = FSGPTForCausalLM.from_pretrained(
            checkpoint,
            device_map="auto",
        )
    elif "merging" in checkpoint:
        model = load_merged_fsgpt_moe_from_checkpoint(checkpoint)
        if losparse_checkpoint is not None:
            d_model = model.config.d_model
            d_ff = model.config.d_ff
            low_rank_factor = 32
            parameter_ratio = (d_model + d_ff) * low_rank_factor / (d_model * d_ff)
            model = substitute_fsgpt_moe_expert_weights(
                model,
                use_svd=False,
                parameter_ratio=parameter_ratio,
                encoder_compression_layers=list(range(3, 12, 2)),
                decoder_compression_layers=list(range(1, 12, 2)),
            )
            model.load_state_dict(torch.load(os.path.join(losparse_checkpoint, "pytorch_model.bin")))
        for layer_idx in range(1, 12, 2):
            model.encoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
            model.decoder.block[layer_idx].layer[-1].mlp.router.expert_capacity = 2048
    else:
        model = FSGPTMoEForCausalLM.from_pretrained(
            checkpoint,
            device_map="auto",
            moe_eval_capacity_token_fraction=1.0
        )
    tokenizer = GPT2TokenizerFast.from_pretrained(checkpoint)
    dataloader = prepare_dataloader(task, tokenizer, batch_size)
    batch = next(iter(dataloader))
    batch = {k: v.cuda() for k, v in batch.items()}

    print("Start profiling...")
    model = model.bfloat16()
    model = ModelWrapper(model)

    with torch.no_grad(), torch.autocast("cuda"):
        flops, macs, params = calculate_flops(
            model=model,
            kwargs=batch,
            print_detailed=True,
            include_backPropagation=False,
            is_sparse=True
        )
    print(f"Task: {task}.")
    print(f"Checkpoint: {checkpoint}.")
    print(f"Batch size: {batch_size}.")
    print(f"Input sequence length: {batch['input_ids'].shape[1]}")
    print(f"Output sequence length: {batch['labels'].shape[1]}")
    print(f"FLOPs: {flops}. MACs: {macs}. Params: {params}.")


if __name__ == '__main__':
    Fire(profile_inference_flops)

mcsmoe/profile-fsgpt-flops.py
UNITES-Lab/MC-SMoE
import time
from typing import Optional, Dict

import torch
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from transformers import (
    T5TokenizerFast,
    T5ForConditionalGeneration,
    SwitchTransformersForConditionalGeneration
)

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups_with_usage_frequency_weighting
)

set_seed(233)


def prepare_dataloader(task: str, tokenizer: T5TokenizerFast, batch_size: int):
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"]
    dataset = dataset.map(
        Seq2SeqDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=6,
        remove_columns=dataset.column_names
    )
    dataset = dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=6,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt')
    dataloader = DataLoader(
        dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=batch_size,
        num_workers=8
    )
    return dataloader


def merge_switch_by_usage_weighted_averaging(
        model: SwitchTransformersForConditionalGeneration,
        batch: Dict[str, torch.Tensor],
        prune_router_non_core: bool = False,
) -> SwitchTransformersForConditionalGeneration:
    grouper = ExpertsGrouperForSwitch(
        config=model.config,
        similarity_fn="cosine",
        similarity_base="router-logits"
    )
    grouper.compute_all_similarities(
        model=model,
        batch=batch,
    )
    grouper.compute_all_usages(
        model=model,
        batch=batch,
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=8,
        merging_encoder_layers=[3, 5, 7, 9, 11],
        merging_decoder_layers=[1, 3, 5, 7, 9, 11]
    )
    model = merge_by_groups_with_usage_frequency_weighting(
        model=model,
        grouper=grouper,
        encoder_merging_layers=[3, 5, 7, 9, 11],
        decoder_merging_layers=[1, 3, 5, 7, 9, 11],
    )

    if prune_router_non_core:
        # Prune the fake experts and router classifier rows
        num_experts = model.config.num_experts
        router_bias = model.config.router_bias
        # Encoder
        for layer_idx in range(3, 12, 2):
            encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
            encoder_name = f"encoder.block.{layer_idx}.layer.1.mlp"
            core_expert_indices = torch.tensor(core_experts[encoder_name], dtype=torch.int).sort().values
            # prune router
            encoder_mlp.router.num_experts = len(core_expert_indices)
            with torch.no_grad():
                pruned_router_cls = encoder_mlp.router.classifier.weight.data[core_expert_indices]
                encoder_mlp.router.classifier = torch.nn.Linear(
                    encoder_mlp.router.classifier.in_features,
                    encoder_mlp.router.num_experts,
                    bias=router_bias
                )
                encoder_mlp.router.classifier.weight.data.copy_(pruned_router_cls)
                if router_bias:
                    pruned_router_bias = encoder_mlp.router.classifier.bias.data[core_expert_indices]
                    encoder_mlp.router.classifier.bias.data.copy_(pruned_router_bias)
            # prune experts
            for expert_idx in range(num_experts):
                if expert_idx not in core_expert_indices:
                    del encoder_mlp.experts[f"expert_{expert_idx}"]
            assert len(encoder_mlp.experts) == encoder_mlp.router.classifier.out_features
        # Decoder
        for layer_idx in range(1, 12, 2):
            decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
            decoder_name = f"decoder.block.{layer_idx}.layer.2.mlp"
            core_expert_indices = torch.tensor(core_experts[decoder_name], dtype=torch.int).sort().values
            # prune router
            decoder_mlp.router.num_experts = len(core_expert_indices)
            with torch.no_grad():
                pruned_router_cls = decoder_mlp.router.classifier.weight.data[core_expert_indices]
                decoder_mlp.router.classifier = torch.nn.Linear(
                    decoder_mlp.router.classifier.in_features,
                    decoder_mlp.router.num_experts,
                    bias=router_bias
                )
                decoder_mlp.router.classifier.weight.data.copy_(pruned_router_cls)
                if router_bias:
                    pruned_router_bias = decoder_mlp.router.classifier.bias.data[core_expert_indices]
                    decoder_mlp.router.classifier.bias.data.copy_(pruned_router_bias)
            # prune experts
            for expert_idx in range(num_experts):
                if expert_idx not in core_expert_indices:
                    del decoder_mlp.experts[f"expert_{expert_idx}"]
            assert len(decoder_mlp.experts) == decoder_mlp.router.classifier.out_features

    return model


def profile_inference_latency(
        checkpoint: str,
        task: str,
        batch_size: Optional[int] = 256,
        merging_strategy: Optional[str] = None,
):
    if "t5" in checkpoint:
        model = T5ForConditionalGeneration.from_pretrained(checkpoint)
    else:
        model = SwitchTransformersForConditionalGeneration.from_pretrained(checkpoint)
    tokenizer = T5TokenizerFast.from_pretrained(checkpoint)
    dataloader = prepare_dataloader(task, tokenizer, batch_size)
    batch = next(iter(dataloader))
    batch = {k: v.cuda() for k, v in batch.items()}

    if merging_strategy == "redirect":
        model = merge_switch_by_usage_weighted_averaging(model, batch, prune_router_non_core=False)
    elif merging_strategy == "fuse":
        model = merge_switch_by_usage_weighted_averaging(model, batch, prune_router_non_core=True)

    if "t5" not in checkpoint:
        batch["output_router_logits"] = False

    model = model.cuda().eval()
    print("Start profiling...")
    num_runs = 100
    # 1. Test BF16
    model = model.half()
    # warmup
    for i in range(10):
        with torch.no_grad():
            model(**batch)

    # test
    torch.cuda.synchronize()
    bf16_start = time.time()
    for i in range(num_runs):
        with torch.no_grad():
            model(**batch)
    torch.cuda.synchronize()
    bf16_end = time.time()

    # 2. Test FP32
    model = model.float()
    # warmup
    for i in range(10):
        with torch.no_grad():
            model(**batch)

    # test
    torch.cuda.synchronize()
    fp32_start = time.time()
    for i in range(num_runs):
        with torch.no_grad():
            model(**batch)
    torch.cuda.synchronize()
    fp32_end = time.time()

    print(f"********************************************************")
    print(f"****** Latency profile for {checkpoint} on {task} ******")
    if merging_strategy is not None:
        print(f"Merging strategy: {merging_strategy} to 8 experts")
    print(f"Model number of parameters: {model.num_parameters()}")
    print(f"Number of runs: {num_runs}")
    print(f"Batch size: {batch_size}")
    print(f"Input sequence length: {batch['input_ids'].shape[1]}, target sequence length: {batch['labels'].shape[1]}")
    print(f"****** Results ******")
    print(f"FP32 Latency: {(fp32_end - fp32_start) / num_runs * 1000} ms")
    print(f"BF16 Latency: {(bf16_end - bf16_start) / num_runs * 1000} ms")


if __name__ == '__main__':
    Fire(profile_inference_latency)

mcsmoe/profile-latency.py
UNITES-Lab/MC-SMoE
import os
import resource
from typing import Optional, Union, List

import torch
import transformers.utils.logging as hf_logging
from accelerate.utils import set_seed
from datasets import load_dataset
from evaluate import load
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    AutoTokenizer
)

from mcsmoe.data import (
    CasualZeroShotDataPreProcessor,
    tokenize_casual_zero_shot,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForLanguageModeling,
    build_index_for_dataset,
    gather_predictions_references_by_casual_lm_loss
)
from mcsmoe.merging.grouping_fsgpt import (
    ExpertsGrouperForFSGPT,
    prune_non_core_experts_by_groups
)
from mcsmoe.models import FSGPTMoEForCausalLM, FSGPTMoEConfig

set_seed(42)
hf_logging.set_verbosity_info()
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (2048, rlimit[1]))


def _sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


def merge_fsgpt_by_usage_frequency_weighting(
        output_dir: Optional[str] = None,
        checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        num_samples_for_merging: Optional[int] = 512,
        samples_batch_size: Optional[int] = 32,
        eval_batch_size: Optional[int] = 32,
        num_groups: Optional[int] = 16,
        merging_layers: Optional[Union[str, List, int]] = None,
        debug: Optional[bool] = False,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, "prune")
        os.makedirs(output_dir, exist_ok=True)
    merging_layers = _sanitize_merging_layers(merging_layers)

    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    if debug:
        config = FSGPTMoEConfig.from_pretrained(
            checkpoint,
            moe_eval_capacity_token_fraction=1.0,
        )
        model = FSGPTMoEForCausalLM(config).bfloat16().cuda()
        print("[Merging]Debug mode, using fake model")
    else:
        model = FSGPTMoEForCausalLM.from_pretrained(
            checkpoint,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            moe_eval_capacity_token_fraction=1.0,
        )
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
        keys_to_ignore=["answer_idx", "choice_idx", "idx"]
    )
    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    dataset = raw_dataset["train"].select(range(min(num_samples_for_merging, len(raw_dataset["train"]))))
    dataset = build_index_for_dataset(dataset)
    dataset = dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_dataset = dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x, keep_input_only=True),
        num_proc=8,
        batched=True,
        remove_columns=dataset.column_names,
        load_from_cache_file=False
    )
    dataloader = DataLoader(
        tokenized_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=samples_batch_size,
        num_workers=8
    )

    grouper = ExpertsGrouperForFSGPT(
        config=model.config,
        similarity_base="random",
    )
    grouper.compute_all_similarities(
        model=model,
        dataloader=dataloader,
    )
    grouper.compute_all_usages(
        model=model,
        dataloader=dataloader,
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=num_groups,
        merging_layers=merging_layers,
    )
    print(f"[Merging]Core experts: {core_experts}")
    print(f"[Merging]Number of parameters before merging: {model.num_parameters()}")
    model = prune_non_core_experts_by_groups(
        model=model,
        grouper=grouper,
        core_experts=core_experts,
        merging_layers=merging_layers,
    )
    print(f"[Merging]Number of parameters after merging: {model.num_parameters()}")

    # ========= evaluate downstream zero-shot =========
    eval_dataset = raw_dataset["validation"]
    eval_dataset = build_index_for_dataset(eval_dataset)
    eval_dataset = eval_dataset.map(
        CasualZeroShotDataPreProcessor(benchmark=task),
        batched=True,
        num_proc=8,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_casual_zero_shot(tokenizer=tokenizer, batch=x),
        num_proc=8,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        batch_size=eval_batch_size,
        num_workers=4,
        collate_fn=data_collator
    )

    model.eval()
    ids_list = []
    answer_ids_list = []
    choice_ids_list = []
    losses_list = []
    try:
        metric = load(*TASK_MAPPING_DATASET_ARGUMENTS[task])
    except FileNotFoundError:
        print(f"[Evaluation warning] No metric found for task {task}, using accuracy instead.")
        metric = load("accuracy")
    for eval_step, eval_batch in tqdm(enumerate(eval_dataloader),
                                      total=len(eval_dataloader), desc="Evaluating"):
        ids_list += eval_batch.pop("idx")
        answer_ids_list += eval_batch.pop("answer_idx")
        choice_ids_list += eval_batch.pop("choice_idx")
        eval_batch = {k: v.cuda() for k, v in eval_batch.items()}
        with torch.no_grad():
            with torch.autocast("cuda"):
                logits = model(**eval_batch).logits
            # eval probs
            loss_fct = torch.nn.CrossEntropyLoss()
            labels = eval_batch["labels"]
            shift_labels = labels.new_zeros(labels.shape)
            shift_labels[:, :-1] = labels[:, 1:].clone()
            shift_labels[:, -1] = model.config.pad_token_id
            losses = torch.stack([
                loss_fct(logits[i], shift_labels[i].to(logits.device)) for i in range(logits.shape[0])
            ])
            losses_list += losses.tolist()
    predictions_references = gather_predictions_references_by_casual_lm_loss(
        ids_list=ids_list,
        answer_ids_list=answer_ids_list,
        choice_ids_list=choice_ids_list,
        losses_list=losses_list,
    )
    eval_res = metric.compute(predictions=predictions_references["predictions"],
                              references=predictions_references["references"])
    print(f"{task} evaluation result: {eval_res}")


if __name__ == '__main__':
    Fire(merge_fsgpt_by_usage_frequency_weighting)

mcsmoe/prune-fsgpt-non-core-experts.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    prune_non_core_experts_by_groups,
    # merge_by_groups_decoder_only as merge_by_groups,
)
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)
from mcsmoe.utils.training_utils import freeze_switch_routers_for_finetuning

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


@dataclass
class SwitchDistillationConfig:
    max_steps: Optional[int] = 10000
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    mlm_probability: Optional[float] = 0.15
    mean_noise_span_length: Optional[int] = 3
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.1
    hd_lambda: Optional[float] = 0.01
    hd_cos_sim: Optional[bool] = False
    next_rank_for_norm: Optional[int] = 384


def merge_and_distill_downstream_for_recover(
        output_dir: Optional[str] = None,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        merging_strategy: Optional[str] = "drop",
        num_samples_for_merging: Optional[int] = 32,
        similarity_base: Optional[str] = "weight",
        similarity_fn: Optional[str] = "cosine",
        num_groups: Optional[int] = 16,
        globally_group: Optional[bool] = False,
        encoder_merging_layers: Optional[Union[str, List, int]] = None,
        decoder_merging_layers: Optional[Union[str, List, int]] = None,
        exact_fisher: Optional[bool] = True,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-3,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 100,
        num_epochs: Optional[int] = 10,
        max_train_steps: Optional[int] = 20000,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, merging_strategy, similarity_base)
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        student_checkpoint = teacher_checkpoint
    if task in ['squad', 'openbookqs']:
        exact_fisher = False
        print(f"[WARNING]exact_fisher is set to False for {task} task")

    encoder_merging_layers = sanitize_merging_layers(encoder_merging_layers)
    decoder_merging_layers = sanitize_merging_layers(decoder_merging_layers)

    training_config = SwitchDistillationConfig(
        max_steps=max_train_steps,
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        mlm_probability=0.15,
        mean_noise_span_length=3,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = HFSwitch.from_pretrained(student_checkpoint)
    student_model = freeze_switch_routers_for_finetuning(student_model)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=training_config.reg_lambda,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"prune-non-core-{task}-global-{merging_strategy}-{similarity_base}" if globally_group else (
            f"prune-non-core-{task}-uniform-{merging_strategy}-{similarity_base}")
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )
    # random select a subset of training dataset for merging
    dataset_for_merging = tokenized_train_dataset.shuffle(seed=2333).select(range(num_samples_for_merging))

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    dataloader_for_merging = DataLoader(
        dataset_for_merging,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=num_samples_for_merging,
        num_workers=4
    )

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )

    print(f"Pruning non-core experts with {merging_strategy} strategy")
    print(f"Similarities are computed based on {similarity_base} and {similarity_fn}")
    print(f"Number of groups: {num_groups}")
    print(f"Merge encoder layers: {encoder_merging_layers}")
    print(f"Merge decoder layers: {decoder_merging_layers}")

    grouper = ExpertsGrouperForSwitch(
        config=model.student.config,
        similarity_fn=similarity_fn,
        similarity_base=similarity_base
    )
    grouper.compute_all_similarities(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    grouper.compute_all_usages(
        model=model.student,
        batch=next(iter(dataloader_for_merging)),
    )
    core_experts = grouper.group_experts_into_clusters_by_routing_guided_globally(
        average_num_groups=num_groups,
        merging_encoder_layers=encoder_merging_layers,
        merging_decoder_layers=decoder_merging_layers,
    ) if globally_group else grouper.group_experts_into_clusters_by_routing_guided(
        num_groups=num_groups,
    )
    print(grouper.group_state_dict())

    model.student = prune_non_core_experts_by_groups(
        model=model.student,
        grouper=grouper,
        core_experts=core_experts,
        strategy=merging_strategy,
        encoder_merging_layers=encoder_merging_layers,
        decoder_merging_layers=decoder_merging_layers,
    )

    print(f"Number of parameters before pruning non-dorminant experts: {model.teacher.num_parameters()}")
    print(f"Number of parameters after pruning non-dorminant experts: {model.student.num_parameters()}")

    no_weight_decay = ["bias", "LayerNorm", "mat_u", "mat_v"]  # no weight decay for orthogonalization matrices
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                hd_loss = outputs.hd_loss
                task_loss = outputs.task_loss
                norm_loss = outputs.norm_loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_hd_loss": hd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_norm_loss": norm_loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)
                # model.student.train()

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(merge_and_distill_downstream_for_recover)

mcsmoe/prune-non-core-experts-and-kd-downstream.py
UNITES-Lab/MC-SMoE
from .losparse import (
    LoSparsePruner,
    LoSparsePrunerForFSGPT,
    substitute_fsgpt_moe_expert_weights,
    substitute_switch_transformers_expert_weights
)
from .task_specific import (
    TaskSpecificExpertPruner,
    transform_switch_model_for_task_specific_pruning
)
from .task_specific_fsgpt import (
    TaskSpecificExpertPrunerForFSGPT,
    transform_fsgpt_model_for_task_specific_pruning
)


mcsmoe/pruning/__init__.py
UNITES-Lab/MC-SMoE
""" Adapted from LoSparse (arXiv 2306.11222)  """
import math
import random
import re
from dataclasses import dataclass
from typing import Optional, List, Tuple, Callable, Set, Union

import torch
from deepspeed.utils import safe_get_full_grad, safe_get_full_fp32_param
from torch import nn
from tqdm import tqdm
from transformers import (
    SwitchTransformersForConditionalGeneration,
    SwitchTransformersSparseMLP
)

from mcsmoe.models.fsgpt_moe import (
    FSGPTMoEForCausalLM,
    FSGPTMoESparseMLP
)
from mcsmoe.models.sparse import LoSparseLinear

__all__ = [
    'LoSparseConfig',
    'substitute_switch_transformers_expert_weights',
    'substitute_fsgpt_moe_expert_weights',
    'LoSparsePruner',
    'LoSparsePrunerForFSGPT',
]


@dataclass
class LoSparseConfig:
    beta1: Optional[float] = 0.85
    beta2: Optional[int] = 1  # disable uncertainty when beta2 = 1
    deltaT: Optional[int] = 1
    initial_threshold: Optional[float] = 1.0
    final_threshold: Optional[float] = 0.25
    initial_warmup: Optional[int] = 1
    final_warmup: Optional[int] = 4
    warmup_steps: Optional[int] = 100


def low_rank_decomposition(
        weight: torch.Tensor,
        rank_ratio: Optional[float] = 0.1,
        parameter_ratio: Optional[float] = 0.15,
        remove_criteria: Optional[str] = 'max_eigenvalue',
        return_dict: Optional[bool] = False,
):
    """
    Parameters
    ----------
    weight: torch.Tensor
        The matrix to decompose, of shape (H, W)
    rank_ratio: float, optional, default 0.1
        The ratio of the reduced rank to the original rank:
            rank_of_decomposed_matrix / rank_of_input_weight
    parameter_ratio: float, optional, default 0.15
        The ratio of the number of parameters of the decomposed matrix to the original matrix:
            parameter_num_of_decomposed_matrix / (H * W).
        If specify, override rank_ratio
    remove_criteria: str, optional, default 'max_eigenvalue'
        The criteria to remove the small eigenvalues, of ['max_eigenvalue', 'random', 'min_eigenvalue']
    return_dict: bool, optional, default False
        Return a dict if True, else return a tuple (L, R)
    debug: bool, optional, default False
        Print debug information if True
    """
    matrix_dimension = len(weight.size())
    assert matrix_dimension == 2, "Only Support 2D matrix"
    height, width = weight.size()

    # Use SVD to decompose a matrix, default full_matrices is False to save parameters
    dtype = weight.dtype
    weight = weight.float()
    u_, s_, v_ = torch.linalg.svd(weight, full_matrices=False)
    u_ = u_.to(dtype)
    s_ = s_.to(dtype)
    v_ = v_.to(dtype)
    rank = torch.count_nonzero(s_)

    if parameter_ratio is not None:
        reduced_rank = math.ceil(parameter_ratio * (height * width) / (height + width))
    else:
        reduced_rank = math.ceil(rank * rank_ratio)

    if remove_criteria == 'max_eigenvalue':
        l_ = u_ @ (torch.sqrt(torch.diag(s_)[:, 0:reduced_rank]))
        r_ = torch.sqrt(torch.diag(s_)[0:reduced_rank, :]) @ v_
    elif remove_criteria == 'random':
        selected_index = random.choices(range(len(s_)), k=reduced_rank)
        l_ = u_ @ (torch.sqrt(torch.diag(s_)[:, selected_index]))
        r_ = torch.sqrt(torch.diag(s_)[selected_index, :]) @ v_
    elif remove_criteria == 'min_eigenvalue':
        len_s = len(s_)
        l_ = u_ @ (torch.sqrt(torch.diag(s_)[:, len_s - reduced_rank:]))
        r_ = torch.sqrt(torch.diag(s_)[len_s - reduced_rank:, :]) @ v_
    else:
        raise NameError("remove criteria not support")

    if return_dict:
        return {"L": l_, "R": r_, "U": u_, "S": s_, "Vh": v_, 'reduced_rank': reduced_rank}
    else:
        return l_, r_


def prune(module):
    for attr_str in dir(module):
        target_attr = getattr(module, attr_str)
        if type(target_attr) == LoSparseLinear:
            print("====================================================")
            print(attr_str, target_attr)
            target_attr.prune_sparse()
    for name, immediate_child_module in module.named_children():
        prune(immediate_child_module)


def _substitute_single_linear_weight(
        module: nn.Module,
        parameter_ratio: float,
        has_sparse: bool,
        use_svd: bool,
        **kwargs
) -> nn.Module:
    """
    Substitute a single Linear weight with to LinearLoSparse

    Examples
    --------
    >>> linear = nn.Linear(16, 32)
    >>> linear = _substitute_single_linear_weight(linear, parameter_ratio=0.15, has_sparse=True, use_svd=True)
    Reduced Rank: 2 | Num Parameters: 96
    >>> linear
    LinearLoSparse(
      (right): Linear(in_features=16, out_features=2, bias=False)
      (left): Linear(in_features=2, out_features=32, bias=False)
      (sparse): Linear(in_features=16, out_features=32, bias=False)
    )
    """
    has_bias = module.bias is not None

    if use_svd:
        # Decompose a matrix by SVD
        device = module.weight.device
        output = low_rank_decomposition(module.weight.cuda(), parameter_ratio=parameter_ratio,
                                        return_dict=True, **kwargs)
        l_, r_, reduced_rank = output['L'], output['R'], output['reduced_rank']
        s_ = module.weight.cuda() - torch.mm(l_, r_)

        # Create a nn.Module and assign decomposed weights to the parameters
        linear_loras = LoSparseLinear(module.in_features, module.out_features, reduced_rank,
                                      has_bias=has_bias, has_sparse=has_sparse)
        linear_loras.initialize_weight(l_.to(device), r_.to(device), s_.to(device), module.bias)

    else:
        height, width = module.weight.shape
        reduced_rank = math.ceil(parameter_ratio * (height * width) / (height + width))
        l_ = torch.zeros(height, reduced_rank, requires_grad=True)
        r_ = torch.zeros(reduced_rank, width, requires_grad=True)
        s_ = torch.zeros(height, width, requires_grad=True)

        # Create a nn.Module and assign decomposed weights to the parameters
        linear_loras = LoSparseLinear(module.in_features, module.out_features, reduced_rank,
                                      has_bias=has_bias, has_sparse=has_sparse)

        linear_loras.initialize_weight(l_, r_, s_, module.bias)

    return linear_loras


def _get_moe_mlp_real_expert_names(
        mlp: Union[SwitchTransformersSparseMLP, FSGPTMoESparseMLP]
) -> Set[str]:
    names = set()
    for name, _ in mlp.named_parameters():
        # capture "expert_?" in the name
        match = re.search(r"expert_(\d+)", name)
        if match:
            names.add(match.group(0))
    return names


def substitute_switch_transformers_expert_weights(
        model: SwitchTransformersForConditionalGeneration,
        parameter_ratio: Optional[float] = 0.15,
        has_sparse: Optional[bool] = True,
        use_svd: Optional[bool] = True,
        encoder_compression_layers: Optional[List[int]] = None,
        decoder_compression_layers: Optional[List[int]] = None,
        **kwargs
) -> SwitchTransformersForConditionalGeneration:
    """
    Substitute the model's MLP weights with LoSparse

    Parameters
    ----------
    model: SwitchTransformersForConditionalGeneration
        The model to substitute
    parameter_ratio: float, optional, default 0.15
        The ratio of the number of parameters of the decomposed matrix to the original matrix:
            low rank matrix parameter / original matrix parameter
    has_sparse: bool, optional, default True.
        Use LoRaS if True, else use Low Rank only
    use_svd: bool, optional, default True.
        Use SVD if True, else use Low Rank only
    encoder_compression_layers: List[int], optional, default None
        The layers to substitute in the encoder. If None, substitute all layers
    decoder_compression_layers: List[int], optional, default None
        The layers to substitute in the decoder. If None, substitute all layers

    Returns
    -------
    model: SwitchTransformersForConditionalGeneration
        The model with LoSparse weights
    """
    print(f"[LoSparse] Before substituting: {model.num_parameters()}")
    if encoder_compression_layers is not None:
        print(f"[LoSparse] Substituting partial encoder layers: {encoder_compression_layers}")
    if decoder_compression_layers is not None:
        print(f"[LoSparse] Substituting partial decoder layers: {decoder_compression_layers}")
    config = model.config
    num_layers = config.num_layers
    if hasattr(config, 'sparse_encoder_layer_list'):
        sparse_layer_list = config.sparse_encoder_layer_list
    else:
        sparse_layer_list = [i for i in range(1, num_layers, config.encoder_sparse_step)]

    for layer_idx in tqdm(sparse_layer_list, desc="Substituting MLP weights"):
        if encoder_compression_layers is None or layer_idx in encoder_compression_layers:
            encoder_mlp = model.encoder.block[layer_idx].layer[-1].mlp
            for expert_name in _get_moe_mlp_real_expert_names(encoder_mlp):
                encoder_mlp.experts[expert_name].wi = _substitute_single_linear_weight(
                    module=encoder_mlp.experts[expert_name].wi,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )
                encoder_mlp.experts[expert_name].wo = _substitute_single_linear_weight(
                    module=encoder_mlp.experts[expert_name].wo,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )
        if decoder_compression_layers is None or layer_idx in decoder_compression_layers:
            decoder_mlp = model.decoder.block[layer_idx].layer[-1].mlp
            for expert_name in _get_moe_mlp_real_expert_names(decoder_mlp):
                decoder_mlp.experts[expert_name].wi = _substitute_single_linear_weight(
                    module=decoder_mlp.experts[expert_name].wi,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )
                decoder_mlp.experts[expert_name].wo = _substitute_single_linear_weight(
                    module=decoder_mlp.experts[expert_name].wo,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )

    print(f"[LoSparse] After substituting: {model.num_parameters()}")
    return model


def substitute_fsgpt_moe_expert_weights(
        model: FSGPTMoEForCausalLM,
        parameter_ratio: Optional[float] = 0.15,
        has_sparse: Optional[bool] = True,
        use_svd: Optional[bool] = True,
        compression_layers: Optional[List[int]] = None,
        **kwargs
) -> FSGPTMoEForCausalLM:
    print(f"[LoSparse] Before substituting: {model.num_parameters()}")
    if compression_layers is not None:
        print(f"[LoSparse] Substituting partial decoder layers: {compression_layers}")
    config = model.config
    num_layers = config.num_layers
    sparse_layer_list = [i for i in range(1, num_layers, config.sparse_step)]

    for layer_idx in tqdm(sparse_layer_list, desc="Substituting MLP weights"):
        if compression_layers is None or layer_idx in compression_layers:
            ffn = model.decoder.layers[layer_idx].ffn
            for expert_name in _get_moe_mlp_real_expert_names(ffn):
                ffn.experts[expert_name].fc1 = _substitute_single_linear_weight(
                    module=ffn.experts[expert_name].fc1,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )
                ffn.experts[expert_name].fc2 = _substitute_single_linear_weight(
                    module=ffn.experts[expert_name].fc2,
                    parameter_ratio=parameter_ratio,
                    has_sparse=has_sparse,
                    use_svd=use_svd,
                    **kwargs
                )

    print(f"[LoSparse] After substituting: {model.num_parameters()}")
    return model


class LoSparsePruner(object):
    def __init__(
            self,
            config: LoSparseConfig,
            total_step: int,
            mask_param_names: Optional[List[str]] = None,
            pruner_name: Optional[str] = 'PLATON',
            structured_method: Optional[str] = 'mean',
            structured_direction: Optional[str] = 'row'
    ):
        print(f"[LoSparse] config: {config}")
        if mask_param_names is None:
            mask_param_names = ['sparse']
        self.config = config
        self.importance = {}
        self.exp_avg_ipt = {}
        self.exp_avg_unc = {}
        self.mask_param_names = mask_param_names
        self.total_step = total_step
        self.pruner_name = pruner_name
        self.beta1 = self.config.beta1
        self.beta2 = self.config.beta2
        self.deltaT = self.config.deltaT
        self.structured_method = structured_method
        self.structured_direction = structured_direction

    def whether_mask_para(
            self,
            name: str
    ):
        return any(mask_name in name for mask_name in self.mask_param_names)

    def structured_prune(
            self,
            is_dict_mat: torch.Tensor
    ) -> torch.Tensor:
        num_row, num_col = is_dict_mat.shape
        if self.structured_direction == 'row_col':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        elif self.structured_direction == 'row':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        elif self.structured_direction == 'col':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        else:
            raise ValueError("Unimplemented Structured Direction: %s" % self.structured_direction)

    def schedule_threshold_comb(
            self,
            step: int
    ) -> Tuple[float, bool]:
        # Schedule the remaining ratio
        total_step = self.total_step
        initial_threshold = self.config.initial_threshold
        final_threshold = self.config.final_threshold
        initial_warmup = self.config.initial_warmup
        final_warmup = self.config.final_warmup
        warmup_steps = self.config.warmup_steps
        if step <= initial_warmup * warmup_steps:
            threshold = initial_threshold
            mask_ind = False
        elif step > (total_step - final_warmup * warmup_steps):
            threshold = final_threshold
            mask_ind = True
        else:
            spars_warmup_steps = initial_warmup * warmup_steps
            spars_schedule_steps = (final_warmup + initial_warmup) * warmup_steps
            mul_coeff = 1 - (step - spars_warmup_steps) / (total_step - spars_schedule_steps)
            threshold = final_threshold + (initial_threshold - final_threshold) * (mul_coeff ** 3)
            mask_ind = True if step % self.deltaT == 0 else False
        return threshold, mask_ind

    def update_ipt_with_local_window(
            self,
            named_parameters_caller: Callable,
            global_step: int
    ):
        # Calculate the sensitivity and uncertainty
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                if name not in self.exp_avg_ipt:
                    self.exp_avg_ipt[name] = torch.zeros_like(safe_get_full_fp32_param(param))
                    self.importance[name] = torch.zeros_like(safe_get_full_fp32_param(param))
                    if self.beta2 > 0 and self.beta2 != 1:
                        self.exp_avg_unc[name] = torch.zeros_like(safe_get_full_fp32_param(param))
                if self.pruner_name == 'Magnitude':
                    # Calculate the score of magnitude pruning
                    self.importance[name] = safe_get_full_fp32_param(param).abs().detach()
                elif self.pruner_name == 'PLATON':
                    local_step = global_step % self.deltaT
                    update_step = global_step // self.deltaT
                    if local_step == 0:
                        self.exp_avg_ipt[name] = self.beta1 * self.exp_avg_ipt[name] + (
                                1 - self.beta1) * self.importance[name]
                        if 0 < self.beta2 < 1:
                            self.exp_avg_unc[name] = (
                                    self.beta2 * self.exp_avg_unc[name] + (1 - self.beta2) * (
                                    self.importance[name] - self.exp_avg_ipt[name]).abs()
                            )
                        elif self.beta2 == 2.:
                            self.exp_avg_unc[name] = (
                                    (update_step * self.exp_avg_unc[name] + (
                                            self.importance[name] - self.exp_avg_ipt[name]) ** 2) / (
                                            update_step + 1)
                            )
                        self.importance[name] = (
                                safe_get_full_fp32_param(param) * safe_get_full_grad(param)
                        ).abs().detach()
                    else:
                        self.importance[name] = (self.importance[name] * local_step + (
                                safe_get_full_fp32_param(param) * safe_get_full_grad(param)).abs().detach()) / (
                                                        local_step + 1)
                else:
                    raise ValueError("Incorrect Pruner Name.")

    def mask_with_threshold(
            self,
            named_parameters_caller: Callable,
            threshold: float
    ) -> float:
        # Calculate the final importance score
        is_dict = dict()
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                if self.pruner_name == 'Magnitude':
                    is_dict[name] = self.importance[name]
                elif self.pruner_name == 'PLATON':
                    if 0 < self.beta2 < 1:
                        is_dict[name] = self.exp_avg_ipt[name] * self.exp_avg_unc[name]
                    elif self.beta2 == 1.:
                        is_dict[name] = self.exp_avg_ipt[name]
                    elif self.beta2 == 2.:
                        is_dict[name] = self.exp_avg_ipt[name] * self.exp_avg_unc.sqrt()
                    else:
                        # Handling the unaccepted beta2 to default setting
                        is_dict[name] = self.exp_avg_ipt[name] * (
                                self.importance[name] - self.exp_avg_ipt[name]).abs()
                else:
                    raise ValueError("Incorrect Pruner Name.")
                if self.structured_method is not None and len(is_dict[name].shape) == 2:
                    is_dict[name] = self.structured_prune(is_dict[name])
        # Calculate the mask threshold
        all_is = torch.cat([is_dict[n].view(-1) for n in is_dict.keys()])
        mask_threshold = torch.kthvalue(
            all_is.to(torch.float32),
            int(all_is.shape[0] * (1 - threshold))
        )[0].item()
        # Mask weights whose importance lower than threshold
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                param.data.masked_fill_(is_dict[name] < mask_threshold, 0.0)
        return mask_threshold

    def update_and_pruning(
            self,
            named_parameters_caller: Callable,
            global_step: int
    ) -> Tuple[float, float]:
        # Update importance score after optimizer stepping
        self.update_ipt_with_local_window(named_parameters_caller, global_step)
        # Get the remaining ratio
        threshold, mask_ind = self.schedule_threshold_comb(global_step)
        if mask_ind:
            # Mask weights during masking horizon
            mask_threshold = self.mask_with_threshold(named_parameters_caller, threshold)
        else:
            mask_threshold = None
        return threshold, mask_threshold


class LoSparsePrunerForFSGPT(object):
    def __init__(
            self,
            config: LoSparseConfig,
            total_step: int,
            mask_param_names: Optional[List[str]] = None,
            pruner_name: Optional[str] = 'PLATON',
            structured_method: Optional[str] = 'mean',
            structured_direction: Optional[str] = 'row'
    ):
        print(f"[LoSparse] config: {config}")
        if mask_param_names is None:
            mask_param_names = ['sparse']
        self.config = config
        self.importance = {}
        self.exp_avg_ipt = {}
        self.exp_avg_unc = {}
        self.mask_param_names = mask_param_names
        self.total_step = total_step
        self.pruner_name = pruner_name
        self.beta1 = self.config.beta1
        self.beta2 = self.config.beta2
        self.deltaT = self.config.deltaT
        self.structured_method = structured_method
        self.structured_direction = structured_direction

    def whether_mask_para(
            self,
            name: str
    ):
        return any(mask_name in name for mask_name in self.mask_param_names)

    def structured_prune(
            self,
            is_dict_mat: torch.Tensor
    ) -> torch.Tensor:
        num_row, num_col = is_dict_mat.shape
        if self.structured_direction == 'row_col':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        elif self.structured_direction == 'row':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=1, keepdim=True).repeat((1, num_col))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        elif self.structured_direction == 'col':
            if self.structured_method == "mean":
                return torch.mean(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "sum":
                return torch.sum(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "max":
                return torch.max(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            elif self.structured_method == "min":
                return torch.min(is_dict_mat, dim=0, keepdim=True).repeat((num_row, 1))
            else:
                raise ValueError("Unimplemented Structured Method: %s" % self.structured_method)
        else:
            raise ValueError("Unimplemented Structured Direction: %s" % self.structured_direction)

    def schedule_threshold_comb(
            self,
            step: int
    ) -> Tuple[float, bool]:
        # Schedule the remaining ratio
        total_step = self.total_step
        initial_threshold = self.config.initial_threshold
        final_threshold = self.config.final_threshold
        initial_warmup = self.config.initial_warmup
        final_warmup = self.config.final_warmup
        warmup_steps = self.config.warmup_steps
        if step <= initial_warmup * warmup_steps:
            threshold = initial_threshold
            mask_ind = False
        elif step > (total_step - final_warmup * warmup_steps):
            threshold = final_threshold
            mask_ind = True
        else:
            spars_warmup_steps = initial_warmup * warmup_steps
            spars_schedule_steps = (final_warmup + initial_warmup) * warmup_steps
            mul_coeff = 1 - (step - spars_warmup_steps) / (total_step - spars_schedule_steps)
            threshold = final_threshold + (initial_threshold - final_threshold) * (mul_coeff ** 3)
            mask_ind = True if step % self.deltaT == 0 else False
        return threshold, mask_ind

    def update_ipt_with_local_window(
            self,
            named_parameters_caller: Callable,
            global_step: int
    ):
        # Calculate the sensitivity and uncertainty
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                if name not in self.exp_avg_ipt:
                    self.exp_avg_ipt[name] = torch.zeros_like(param)
                    self.importance[name] = torch.zeros_like(param)
                    if self.beta2 > 0 and self.beta2 != 1:
                        self.exp_avg_unc[name] = torch.zeros_like(param)
                if self.pruner_name == 'Magnitude':
                    # Calculate the score of magnitude pruning
                    self.importance[name] = param.abs().detach()
                elif self.pruner_name == 'PLATON':
                    local_step = global_step % self.deltaT
                    update_step = global_step // self.deltaT
                    if local_step == 0:
                        self.exp_avg_ipt[name] = self.beta1 * self.exp_avg_ipt[name] + (
                                1 - self.beta1) * self.importance[name]
                        if 0 < self.beta2 < 1:
                            self.exp_avg_unc[name] = (
                                    self.beta2 * self.exp_avg_unc[name] + (1 - self.beta2) * (
                                    self.importance[name] - self.exp_avg_ipt[name]).abs()
                            )
                        elif self.beta2 == 2.:
                            self.exp_avg_unc[name] = (
                                    (update_step * self.exp_avg_unc[name] + (
                                            self.importance[name] - self.exp_avg_ipt[name]) ** 2) / (
                                            update_step + 1)
                            )
                        self.importance[name] = (
                                param * param.grad
                        ).abs().detach()
                    else:
                        self.importance[name] = (self.importance[name] * local_step + (
                                param * param.grad).abs().detach()) / (local_step + 1)
                else:
                    raise ValueError("Incorrect Pruner Name.")

    def mask_with_threshold(
            self,
            named_parameters_caller: Callable,
            threshold: float
    ) -> float:
        # Calculate the final importance score
        is_dict = dict()
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                if self.pruner_name == 'Magnitude':
                    is_dict[name] = self.importance[name]
                elif self.pruner_name == 'PLATON':
                    if 0 < self.beta2 < 1:
                        is_dict[name] = self.exp_avg_ipt[name] * self.exp_avg_unc[name]
                    elif self.beta2 == 1.:
                        is_dict[name] = self.exp_avg_ipt[name]
                    elif self.beta2 == 2.:
                        is_dict[name] = self.exp_avg_ipt[name] * self.exp_avg_unc.sqrt()
                    else:
                        # Handling the unaccepted beta2 to default setting
                        is_dict[name] = self.exp_avg_ipt[name] * (
                                self.importance[name] - self.exp_avg_ipt[name]).abs()
                else:
                    raise ValueError("Incorrect Pruner Name.")
                if self.structured_method is not None and len(is_dict[name].shape) == 2:
                    is_dict[name] = self.structured_prune(is_dict[name])
        # Calculate the mask threshold
        all_is = torch.cat([is_dict[n].view(-1) for n in is_dict.keys()])
        mask_threshold = torch.kthvalue(
            all_is.to(torch.float32),
            int(all_is.shape[0] * (1 - threshold))
        )[0].item()
        # Mask weights whose importance lower than threshold
        for name, param in named_parameters_caller():
            if self.whether_mask_para(name):
                param.data.masked_fill_(is_dict[name] < mask_threshold, 0.0)
        return mask_threshold

    def update_and_pruning(
            self,
            named_parameters_caller: Callable,
            global_step: int
    ) -> Tuple[float, float]:
        # Update importance score after optimizer stepping
        self.update_ipt_with_local_window(named_parameters_caller, global_step)
        # Get the remaining ratio
        threshold, mask_ind = self.schedule_threshold_comb(global_step)
        if mask_ind:
            # Mask weights during masking horizon
            mask_threshold = self.mask_with_threshold(named_parameters_caller, threshold)
        else:
            mask_threshold = None
        return threshold, mask_threshold

mcsmoe/pruning/losparse.py
UNITES-Lab/MC-SMoE
import re
from collections import OrderedDict
from typing import Optional, Tuple, Dict, List

import torch
from torch.nn import functional as F
from transformers import (
    SwitchTransformersForConditionalGeneration,
    SwitchTransformersConfig,
    SwitchTransformersSparseMLP,
    SwitchTransformersTop1Router
)


class TaskSpecificExpertPruner(object):
    """
    Reproduce "Task-Speciﬁc Expert Pruning for Sparse Mixture-of-Experts" paper
    """

    def __init__(
            self,
            model_config: SwitchTransformersConfig,
            threshold_beta: Optional[float] = 0.1,
            minimum_num_experts: Optional[int] = 8,
            encoder_pruning_layers: Optional[List[int]] = None,
            decoder_pruning_layers: Optional[List[int]] = None,
    ):
        self.num_experts = model_config.num_experts
        self.sparse_layer_indices = list(range(1, model_config.num_layers, model_config.encoder_sparse_step))

        self.experts_alive = None
        self.threshold_beta = threshold_beta
        self.minimum_num_experts = minimum_num_experts

        self._reset_all(encoder_pruning_layers, decoder_pruning_layers)

    def _reset_all(
            self,
            encoder_pruning_layers: Optional[List[int]],
            decoder_pruning_layers: Optional[List[int]],
    ):
        self.experts_alive = OrderedDict()
        for layer_idx in self.sparse_layer_indices:
            if encoder_pruning_layers is None or layer_idx in encoder_pruning_layers:
                encoder_mlp_name = f"encoder.block.{layer_idx}.layer.1.mlp"
                self.experts_alive[encoder_mlp_name] = [True] * self.num_experts
            if decoder_pruning_layers is None or layer_idx in decoder_pruning_layers:
                decoder_mlp_name = f"decoder.block.{layer_idx}.layer.2.mlp"
                self.experts_alive[decoder_mlp_name] = [True] * self.num_experts

    def execute_pruning(
            self,
            model: SwitchTransformersForConditionalGeneration,
            usage_frequency_state_dict: Dict[str, torch.Tensor],
            prune_to_minimum: Optional[bool] = False,
    ) -> SwitchTransformersForConditionalGeneration:
        """
        Execute pruning based on current EMA usage frequency state dict

        Parameters
        ----------
        model: SwitchTransformersForConditionalGeneration
            Model to be pruned.
        usage_frequency_state_dict: Dict[str, torch.Tensor]
            EMA usage frequency state dict.
        prune_to_minimum: Optional[bool]
            Whether to prune to minimum number of experts.

        Returns
        -------
        SwitchTransformersForConditionalGeneration
            Pruned model.

        """
        print("[Task-Specific Pruning]Start!")
        if prune_to_minimum:
            print(f"[Task-Specific Pruning]Prune to minimum number of experts {self.minimum_num_experts}")
        # 1. Update experts alive
        for name in self.experts_alive:
            num_alive = sum(self.experts_alive[name])
            usage_frequency = usage_frequency_state_dict[name]
            # normalize usage frequency in alive experts and set dead experts' usage frequency to 0
            usage_frequency = usage_frequency / usage_frequency[self.experts_alive[name]].sum()
            usage_frequency[~torch.tensor(self.experts_alive[name], dtype=torch.bool)] = 0.0
            if num_alive < self.minimum_num_experts:
                raise ValueError(
                    f"Number of alive experts {num_alive} is less than minimum number of experts {self.minimum_num_experts}")
            if num_alive == self.minimum_num_experts:
                continue
            threshold = self.threshold_beta / num_alive
            if (usage_frequency >= threshold).sum() < self.minimum_num_experts or prune_to_minimum:
                # select top-m experts
                _, top_m_indices = torch.topk(usage_frequency, self.minimum_num_experts)
                self.experts_alive[name] = [True if i in top_m_indices else False for i in range(self.num_experts)]
            else:
                # select experts with usage frequency > threshold
                self.experts_alive[name] = [True if f > threshold else False for f in usage_frequency]
            print(f"Number of alive experts in {name}: {sum(self.experts_alive[name])} / {self.num_experts}")

        # 2. Prune experts
        for name in self.experts_alive:
            layer_idx = int(re.findall(r"\d+", name)[0])
            encoder_or_decoder = name.split(".")[0]
            if encoder_or_decoder == "encoder":
                model.encoder.block[layer_idx].layer[-1].mlp = self._prune_mlp_experts(
                    self.experts_alive[name],
                    model.encoder.block[layer_idx].layer[-1].mlp
                )
            elif encoder_or_decoder == "decoder":
                model.decoder.block[layer_idx].layer[-1].mlp = self._prune_mlp_experts(
                    self.experts_alive[name],
                    model.decoder.block[layer_idx].layer[-1].mlp
                )
            else:
                raise ValueError(f"Unknown encoder_or_decoder {encoder_or_decoder}")

        return model

    @staticmethod
    def _prune_mlp_experts(
            experts_alive: List[bool],
            mlp: SwitchTransformersSparseMLP
    ) -> SwitchTransformersSparseMLP:
        """
        mlp.router should have member `experts_mask` of tensor
        """
        dead_experts_mask = ~torch.tensor(experts_alive, dtype=torch.bool)
        with torch.no_grad():
            mlp.router.experts_mask.data[dead_experts_mask] = -1e6

        return mlp


def transform_switch_model_for_task_specific_pruning(
        model: SwitchTransformersForConditionalGeneration
) -> SwitchTransformersForConditionalGeneration:
    """
    Add expert mask to MLP.
    """
    config = model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.encoder_sparse_step))

    def _custom_forward(self, hidden_states: torch.Tensor) -> Tuple:
        _, router_logits = self._compute_router_probabilities(hidden_states)

        router_logits = router_logits * self.experts_mask
        router_probs = F.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)

        expert_index = torch.argmax(router_probs, dim=-1)
        expert_index = F.one_hot(expert_index, num_classes=self.num_experts)
        # Since experts are pruned, no need to do max-capacity dropping
        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
        return expert_index, router_probs, router_logits

    for layer_idx in sparse_layer_indices:
        model.encoder.block[layer_idx].layer[-1].mlp.router.experts_mask = torch.nn.Parameter(torch.tensor(
            [1.0] * model.encoder.block[layer_idx].layer[-1].mlp.router.num_experts), requires_grad=False)
        model.encoder.block[layer_idx].layer[-1].mlp.router.forward = _custom_forward.__get__(
            model.encoder.block[layer_idx].layer[-1].mlp.router, SwitchTransformersTop1Router
        )
        model.decoder.block[layer_idx].layer[-1].mlp.router.experts_mask = torch.nn.Parameter(torch.tensor(
            [1.0] * model.decoder.block[layer_idx].layer[-1].mlp.router.num_experts), requires_grad=False)
        model.decoder.block[layer_idx].layer[-1].mlp.router.forward = _custom_forward.__get__(
            model.decoder.block[layer_idx].layer[-1].mlp.router, SwitchTransformersTop1Router
        )

    return model

mcsmoe/pruning/task_specific.py
UNITES-Lab/MC-SMoE
import re
from collections import OrderedDict
from typing import Optional, Tuple, Dict, List

import torch

from mcsmoe.models.fsgpt_moe import (
    FSGPTMoEForCausalLM,
    FSGPTMoEConfig,
    FSGPTMoESparseMLP,
    FSGPTMoETop2Router
)


class TaskSpecificExpertPrunerForFSGPT(object):
    """
    Reproduce "Task-Speciﬁc Expert Pruning for Sparse Mixture-of-Experts" paper
    """

    def __init__(
            self,
            model_config: FSGPTMoEConfig,
            threshold_beta: Optional[float] = 0.1,
            minimum_num_experts: Optional[int] = 8,
            pruning_layers: Optional[List[int]] = None,
    ):
        self.num_experts = model_config.num_experts
        self.sparse_layer_indices = list(range(1, model_config.num_layers, model_config.sparse_step))

        self.experts_alive = None
        self.threshold_beta = threshold_beta
        self.minimum_num_experts = minimum_num_experts

        self._reset_all(pruning_layers)

    def _reset_all(
            self,
            pruning_layers: Optional[List[int]],
    ):
        self.experts_alive = OrderedDict()
        for layer_idx in self.sparse_layer_indices:
            if pruning_layers is None or layer_idx in pruning_layers:
                ffn_name = f"decoder.layers.{layer_idx}.ffn"
                self.experts_alive[ffn_name] = [True] * self.num_experts

    def execute_pruning(
            self,
            model: FSGPTMoEForCausalLM,
            usage_frequency_state_dict: Dict[str, torch.Tensor],
            prune_to_minimum: Optional[bool] = False,
    ) -> FSGPTMoEForCausalLM:
        print("[Task-Specific Pruning]Start!")
        if prune_to_minimum:
            print(f"[Task-Specific Pruning]Prune to minimum number of experts {self.minimum_num_experts}")
        # 1. Update experts alive
        for name in self.experts_alive:
            num_alive = sum(self.experts_alive[name])
            usage_frequency = usage_frequency_state_dict[name]
            # normalize usage frequency in alive experts and set dead experts' usage frequency to 0
            usage_frequency = usage_frequency / usage_frequency[self.experts_alive[name]].sum()
            usage_frequency[~torch.tensor(self.experts_alive[name], dtype=torch.bool)] = 0.0
            if num_alive < self.minimum_num_experts:
                raise ValueError(
                    f"Number of alive experts {num_alive} is less than minimum number of experts {self.minimum_num_experts}")
            if num_alive == self.minimum_num_experts:
                continue
            threshold = self.threshold_beta / num_alive
            if (usage_frequency >= threshold).sum() < self.minimum_num_experts or prune_to_minimum:
                # select top-m experts
                _, top_m_indices = torch.topk(usage_frequency, self.minimum_num_experts)
                self.experts_alive[name] = [True if i in top_m_indices else False for i in range(self.num_experts)]
            else:
                # select experts with usage frequency > threshold
                self.experts_alive[name] = [True if f > threshold else False for f in usage_frequency]
            print(f"Number of alive experts in {name}: {sum(self.experts_alive[name])} / {self.num_experts}")

        # 2. Prune experts
        for name in self.experts_alive:
            layer_idx = int(re.findall(r"\d+", name)[0])
            model.decoder.layers[layer_idx].ffn = self._prune_ffn_experts(
                self.experts_alive[name],
                model.decoder.layers[layer_idx].ffn
            )

        return model

    @staticmethod
    def _prune_ffn_experts(
            experts_alive: List[bool],
            ffn: FSGPTMoESparseMLP
    ) -> FSGPTMoESparseMLP:
        """
        mlp.router should have member `experts_mask` of tensor
        """
        dead_experts_mask = ~torch.tensor(experts_alive, dtype=torch.bool)
        with torch.no_grad(), torch.autocast("cuda"):
            ffn.router.experts_mask.data[dead_experts_mask] = -1e9

        return ffn


def transform_fsgpt_model_for_task_specific_pruning(
        model: FSGPTMoEForCausalLM
) -> FSGPTMoEForCausalLM:
    """
    Add expert mask to MLP.
    """
    config = model.config
    sparse_layer_indices = list(range(1, config.num_layers, config.sparse_step))

    def _custom_forward(self, hidden_states: torch.Tensor, padding_mask: Optional[torch.LongTensor] = None) -> Tuple:
        self.input_dtype = hidden_states.dtype
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.reshape((batch_size * sequence_length), hidden_dim)
        hidden_states = hidden_states.to(self.dtype)
        self._cast_classifier()
        router_logits = self.classifier(hidden_states)

        router_logits = router_logits * self.experts_mask

        top_1_mask, router_probs = self.route_tokens(router_logits, self.input_dtype, padding_mask)
        return top_1_mask, router_probs, router_logits

    for layer_idx in sparse_layer_indices:
        device = model.decoder.layers[layer_idx].ffn.router.classifier.weight.device
        model.decoder.layers[layer_idx].ffn.router.experts_mask = torch.nn.Parameter(torch.tensor(
            [1.0] * model.decoder.layers[layer_idx].ffn.router.num_experts, device=device), requires_grad=False)
        model.decoder.layers[layer_idx].ffn.router.forward = _custom_forward.__get__(
            model.decoder.layers[layer_idx].ffn.router, FSGPTMoETop2Router
        )

    return model

mcsmoe/pruning/task_specific_fsgpt.py
UNITES-Lab/MC-SMoE
import math
import os
from dataclasses import dataclass
from typing import Optional, Tuple, List, Union

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5TokenizerFast,
    get_scheduler,
    SwitchTransformersForConditionalGeneration as HFSwitch,
)
from transformers.utils import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.merging import (
    ExpertsGrouperForSwitch,
    merge_by_groups,
    merge_by_groups_within_and_across_models
)
from mcsmoe.models import (
    SwitchTransformersWrapperForDistillation
)

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_warning()


def sanitize_merging_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


@dataclass
class SwitchDistillationConfig:
    max_steps: Optional[int] = 10000
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 10
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1
    mlm_probability: Optional[float] = 0.15
    mean_noise_span_length: Optional[int] = 3
    # === distillation parameters ===
    kd_temperature: Optional[float] = 2.0
    mlm_lambda: Optional[float] = 1.0
    kd_lambda: Optional[float] = 0.2
    reg_lambda: Optional[float] = 0.0
    hd_lambda: Optional[float] = 0.0
    hd_cos_sim: Optional[bool] = False
    next_rank_for_norm: Optional[int] = 384


def random_merge_and_distill_downstream_for_recover(
        output_dir: Optional[str] = None,
        teacher_checkpoint: Optional[str] = None,
        student_checkpoint: Optional[str] = None,
        task: Optional[str] = "sst2",
        num_groups: Optional[int] = 8,
        strategy: Optional[str] = "average",
        encoder_merging_layers: Optional[Union[str, List, int]] = None,
        decoder_merging_layers: Optional[Union[str, List, int]] = None,
        # === training parameters ===
        no_eval_until_epochs: Optional[Union[int, float]] = 0,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 1e-3,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 100,
        num_epochs: Optional[int] = 10,
        max_train_steps: Optional[int] = 20000,
        preprocessing_num_workers: Optional[int] = None,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        # === distillation parameters ===
        kd_temperature: Optional[float] = 2.0,
        mlm_lambda: Optional[float] = 1.0,
        kd_lambda: Optional[float] = 0.2,
        hd_lambda: Optional[float] = 0.0,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    else:
        output_dir = os.path.join(output_dir, "random-merge")
        os.makedirs(output_dir, exist_ok=True)
    if teacher_checkpoint is None:
        raise ValueError("teacher_checkpoint must be specified")
    if student_checkpoint is None:
        student_checkpoint = teacher_checkpoint

    encoder_merging_layers = sanitize_merging_layers(encoder_merging_layers)
    decoder_merging_layers = sanitize_merging_layers(decoder_merging_layers)

    training_config = SwitchDistillationConfig(
        max_steps=max_train_steps,
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        mlm_probability=0.15,
        mean_noise_span_length=3,
        dropout=0.1,
        kd_temperature=kd_temperature,
        mlm_lambda=mlm_lambda,
        kd_lambda=kd_lambda,
        hd_lambda=hd_lambda,
        hd_cos_sim=False,
        reg_lambda=0
    )

    teacher_model = HFSwitch.from_pretrained(teacher_checkpoint)
    student_model = HFSwitch.from_pretrained(student_checkpoint)

    model = SwitchTransformersWrapperForDistillation(
        student=student_model,
        teacher=teacher_model,
        mlm_lambda=training_config.mlm_lambda,
        kd_lambda=training_config.kd_lambda,
        kd_temperature=training_config.kd_temperature,
        hd_lambda=training_config.hd_lambda,
        hd_cos_sim=training_config.hd_cos_sim,
        reg_lambda=training_config.reg_lambda,
    )

    tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        run_name = f"random-merge-{strategy}-{task}-{num_groups}"
        wandb.init(project="mc-smoe",
                   config={**student_model.config.__dict__,
                           **training_config.__dict__,
                           'num_params': student_model.num_parameters()},
                   name=run_name)

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")
    dataset_for_merging = tokenized_train_dataset.select(range(128))

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=4
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=4
    )
    merging_dataloader = DataLoader(
        dataset_for_merging,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=32,
        num_workers=4
    )

    print(f"Random merging")
    print(f"Number of groups: {num_groups}")
    print(f"Merge encoder layers: {encoder_merging_layers}")
    print(f"Merge decoder layers: {decoder_merging_layers}")

    grouper = ExpertsGrouperForSwitch(
        config=model.student.config,
    )
    grouper.group_experts_randomly(num_groups=num_groups)
    print(grouper.group_state_dict())
    if strategy == "average":
        model.student = merge_by_groups(
            model=model.student,
            grouper=grouper,
            strategy="average",
            encoder_merging_layers=encoder_merging_layers,
            decoder_merging_layers=decoder_merging_layers,
            permute=False
        )
    elif strategy == "repair":
        model.student = merge_by_groups(
            model=model.student,
            grouper=grouper,
            strategy="average",
            encoder_merging_layers=encoder_merging_layers,
            decoder_merging_layers=decoder_merging_layers,
            permute=True,
            permute_strategy="activation-matching",
            dataloader=merging_dataloader,
        )
    elif strategy == "zipit":
        model.student = merge_by_groups_within_and_across_models(
            switch_model=model.student,
            grouper=grouper,
            dataloader=merging_dataloader,
            encoder_merging_layers=encoder_merging_layers,
            decoder_merging_layers=decoder_merging_layers,
            dominant_alone=False,
            usage_weighted=False
        )
    elif strategy == "git-rebasin":
        model.student = merge_by_groups(
            model=model.student,
            grouper=grouper,
            strategy="average",
            encoder_merging_layers=encoder_merging_layers,
            decoder_merging_layers=decoder_merging_layers,
            permute=True,
            permute_strategy="weight-matching",
        )

    print(f"Number of parameters before merging: {model.teacher.num_parameters()}")
    print(f"Number of parameters after merging: {model.student.num_parameters()}")

    no_weight_decay = ["bias", "LayerNorm", "mat_u", "mat_v"]  # no weight decay for orthogonalization matrices
    optimizer_grouped_params = [
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if not any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": weight_decay,
        },
        {
            "params": [
                p for n, p in model.student.named_parameters()
                if any(nd in n for nd in no_weight_decay) and p.requires_grad
            ],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_params,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    # ========================= Training ================================
    # Not using `accelerate_run_train` method here

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                kd_loss = outputs.kd_loss
                hd_loss = outputs.hd_loss
                task_loss = outputs.task_loss
                norm_loss = outputs.norm_loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "train_kd_loss": kd_loss.item(),
                           "train_hd_loss": hd_loss.item(),
                           "train_task_loss": task_loss.item(),
                           "train_norm_loss": norm_loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.student.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.student_logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["task_loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["task_loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['task_loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)
                # model.student.train()

            if completed_steps >= max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(random_merge_and_distill_downstream_for_recover)

mcsmoe/random-merge.py
UNITES-Lab/MC-SMoE
import math
import os
import os.path
from dataclasses import dataclass
from typing import Optional, Tuple, Union, List

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import T5TokenizerFast, get_scheduler, SwitchTransformersForConditionalGeneration
from transformers import logging as hf_logging

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)
from mcsmoe.merging import ExpertUsageFrequencyTracker
from mcsmoe.pruning import TaskSpecificExpertPruner, transform_switch_model_for_task_specific_pruning
from mcsmoe.utils.training_utils import freeze_switch_routers_for_finetuning

logger = get_logger(__name__)
logger.setLevel(20)
hf_logging.set_verbosity_error()


def sanitize_pruning_layers(layers: Union[str, List, int]):
    if layers is None:
        layers = list(range(1, 12, 2))
    elif isinstance(layers, str) and len(layers) > 0:
        layers = [int(x) for x in layers.split(",")]
    elif isinstance(layers, str) and len(layers) == 0:
        layers = []
    elif isinstance(layers, int):
        layers = [layers]
    return layers


@dataclass
class SwitchFineTuningConfig:
    task: Optional[str] = "cola"
    num_epochs: Optional[int] = 5
    batch_size: Optional[int] = 32
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 16
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1


def finetune_on_downstream_task_with_task_specific_pruning(
        checkpoint: str,
        output_dir: str,
        task: str = "cola",
        no_eval_until_epochs: Optional[float] = 1,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        num_epochs: Optional[int] = 10,
        inspect_interval_steps: Optional[int] = 4,
        threshold_beta: Optional[float] = 0.1,
        minimum_num_experts: Optional[int] = 8,
        encoder_pruning_layers: Optional[Union[str, List, int]] = None,
        decoder_pruning_layers: Optional[Union[str, List, int]] = None,
        num_pruning_times: Optional[int] = 4,
        num_eval_steps: Optional[int] = None,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 5e-4,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 16,
        save_each_epoch: Optional[bool] = False,
        preprocessing_num_workers: Optional[int] = None,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    if task not in TASK_MAPPING_DATASET_ARGUMENTS:
        raise ValueError(f"task must be one of {list(TASK_MAPPING_DATASET_ARGUMENTS.keys())}")
    if task not in output_dir:
        output_dir = os.path.join(output_dir, task)

    encoder_pruning_layers = sanitize_pruning_layers(encoder_pruning_layers)
    decoder_pruning_layers = sanitize_pruning_layers(decoder_pruning_layers)

    training_config = SwitchFineTuningConfig(
        task=task,
        num_epochs=num_epochs,
        adam_epsilon=1e-8,
        adam_betas=(0.9, 0.98),
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        dropout=0.1,
    )

    model = SwitchTransformersForConditionalGeneration.from_pretrained(checkpoint)
    model = freeze_switch_routers_for_finetuning(model)
    model = transform_switch_model_for_task_specific_pruning(model)

    try:
        tokenizer = T5TokenizerFast.from_pretrained(checkpoint)
    except OSError:
        tokenizer = T5TokenizerFast.from_pretrained("google/switch-base-32")

    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        print(f"Number of parameters in Switch Transformers: {model.num_parameters()}")
        wandb.init(project="mc-smoe",
                   config={**model.config.__dict__, **training_config.__dict__},
                   name=f"switch-task-specific-{task}-{per_device_train_batch_size}-bs-{learning_rate}-lr")

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

    tokenized_train_dataset = train_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=train_dataset.column_names,
        load_from_cache_file=False
    )
    tokenized_eval_dataset = eval_dataset.map(
        lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
        num_proc=preprocessing_num_workers,
        batched=True,
        remove_columns=eval_dataset.column_names,
        load_from_cache_file=False
    )

    # random select a subset of training dataset for computing usage
    dataset_for_usage = tokenized_train_dataset.shuffle(seed=2333).select(range(128))

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    usage_dataloader = DataLoader(
        dataset_for_usage,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=1,
        num_workers=5
    )

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=5
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=5
    )

    # optimizer
    no_decay = ["bias", "layer_norm.weight", "LayerNorm", "layernorm", "layer_norm", "ln"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_parameters,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    pruning_interval_steps = int(max_train_steps / 2 / num_pruning_times)
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )

    usage_tracker = ExpertUsageFrequencyTracker(
        named_parameters_caller=model.named_parameters,
        compute_every_n_steps=inspect_interval_steps,
        beta=0.6,
        device='cuda'
    )

    task_specific_pruner = TaskSpecificExpertPruner(
        model_config=model.config,
        threshold_beta=threshold_beta,
        minimum_num_experts=minimum_num_experts,
        encoder_pruning_layers=encoder_pruning_layers,
        decoder_pruning_layers=decoder_pruning_layers,
    )

    model, optimizer, train_dataloader, eval_dataloader, usage_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, usage_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset['validation']
    )

    num_eval_steps = num_update_steps_per_epoch if num_eval_steps is None else num_eval_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1

                # Task-specific pruning
                usage_tracker.step(
                    model_outputs=outputs,
                    global_step=completed_steps,
                )
                if completed_steps % pruning_interval_steps == 0 and completed_steps > 0:
                    task_specific_pruner.execute_pruning(
                        model=model.module,
                        usage_frequency_state_dict=usage_tracker.usage_frequency_state_dict(),
                        prune_to_minimum=(completed_steps >= max_train_steps // 2)
                    )

            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)

        if save_each_epoch:
            save_dir = os.path.join(output_dir, f"epoch-{epoch + 1}")
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(save_dir,
                                            is_main_process=accelerator.is_local_main_process,
                                            save_function=accelerator.save)
            if accelerator.is_local_main_process:
                tokenizer.save_pretrained(save_dir)

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(finetune_on_downstream_task_with_task_specific_pruning)

mcsmoe/switch-finetune-with-task-specific-pruning.py
UNITES-Lab/MC-SMoE
import math
import os
import os.path
from dataclasses import dataclass
from typing import Optional, Tuple

import torch
import wandb
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import set_seed
from datasets import load_dataset
from fire import Fire
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import (
    T5ForConditionalGeneration,
    T5TokenizerFast,
    get_scheduler,
)

from mcsmoe.data import (
    Seq2SeqDataPreProcessor,
    tokenize_seq2seq,
    TASK_MAPPING_DATASET_ARGUMENTS,
    DataCollatorForSeq2Seq,
    get_evaluate_fn,
    EXTRA_KEYS_FOR_EVAL,
    keep_only_supporting_facts_in_context_for_hotpotqa
)

logger = get_logger(__name__)
logger.setLevel(20)


@dataclass
class T5FineTuningConfig:
    task: Optional[str] = "cola"
    num_epochs: Optional[int] = 5
    batch_size: Optional[int] = 32
    adam_epsilon: Optional[float] = 1e-6
    adam_betas: Optional[Tuple[float, float]] = (0.9, 0.98)
    lr: Optional[float] = 5e-4
    warmup_steps: Optional[int] = 16
    weight_decay: Optional[float] = 0.01
    dropout: Optional[float] = 0.1


def finetune_on_downstream_task(
        output_dir: Optional[str] = None,
        task: Optional[str] = "cola",
        no_eval_until_epochs: Optional[int] = 0,
        per_device_train_batch_size: Optional[int] = 8,
        per_device_eval_batch_size: Optional[int] = 8,
        num_epochs: Optional[int] = 10,
        report: Optional[bool] = True,
        log_steps: Optional[int] = 1,
        weight_decay: Optional[float] = 0.01,
        learning_rate: Optional[float] = 5e-4,
        gradient_accumulation_steps: Optional[int] = 32,
        warmup_steps: Optional[int] = 16,
        save_each_epoch: Optional[bool] = False,
        preprocessing_num_workers: Optional[int] = None,
):
    if output_dir is None:
        raise ValueError("output_dir must be specified")
    if task not in TASK_MAPPING_DATASET_ARGUMENTS:
        raise ValueError(f"task must be one of {list(TASK_MAPPING_DATASET_ARGUMENTS.keys())}")
    if task not in output_dir:
        output_dir = os.path.join(output_dir, task)

    training_config = T5FineTuningConfig(
        task=task,
        num_epochs=num_epochs,
        adam_epsilon=1e-8,
        adam_betas=(0.9, 0.98),
        lr=learning_rate,
        warmup_steps=warmup_steps,
        weight_decay=weight_decay,
        dropout=0.1,
    )

    model = T5ForConditionalGeneration.from_pretrained("t5-base")
    tokenizer = T5TokenizerFast.from_pretrained("t5-base")
    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer,
                                           max_length=tokenizer.model_max_length,
                                           return_tensors='pt',
                                           keys_to_ignore=EXTRA_KEYS_FOR_EVAL)

    accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps)
    set_seed(42)

    if accelerator.is_local_main_process:
        print(f"Number of parameters in T5: {model.num_parameters()}")
        if report:
            wandb.init(project="mc-smoe",
                       config={**model.config.__dict__, **training_config.__dict__},
                       name=f"t5-{task}-{per_device_train_batch_size}-bs-{learning_rate}-lr")

    raw_dataset = load_dataset(*TASK_MAPPING_DATASET_ARGUMENTS[task])

    train_dataset = raw_dataset["train"]
    eval_dataset = raw_dataset["validation"] if task != "mnli" else (
        raw_dataset["validation_matched"], raw_dataset["validation_mismatched"]  # TODO
    )

    with accelerator.main_process_first():
        if task == "hotpotqa":
            train_dataset = train_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
            eval_dataset = eval_dataset.map(
                keep_only_supporting_facts_in_context_for_hotpotqa,
                batched=False,
                num_proc=preprocessing_num_workers
            )
        train_dataset = train_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=train_dataset.column_names
        )
        eval_dataset = eval_dataset.map(
            Seq2SeqDataPreProcessor(benchmark=task, keep_specific_keys=EXTRA_KEYS_FOR_EVAL),
            batched=True,
            num_proc=preprocessing_num_workers,
            remove_columns=eval_dataset.column_names
        )

        tokenized_train_dataset = train_dataset.map(
            lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=False),
            num_proc=preprocessing_num_workers,
            batched=True,
            remove_columns=train_dataset.column_names,
            load_from_cache_file=False
        )
        tokenized_eval_dataset = eval_dataset.map(
            lambda x: tokenize_seq2seq(tokenizer=tokenizer, batch=x, keep_other_keys=True),
            num_proc=preprocessing_num_workers,
            batched=True,
            remove_columns=eval_dataset.column_names,
            load_from_cache_file=False
        )

    logger.info(f"Number of training examples: {len(tokenized_train_dataset)}")
    logger.info(f"Number of validation examples: {len(tokenized_eval_dataset)}")

    train_dataloader = DataLoader(
        tokenized_train_dataset,
        shuffle=True,
        collate_fn=data_collator,
        batch_size=per_device_train_batch_size,
        num_workers=5
    )
    eval_dataloader = DataLoader(
        tokenized_eval_dataset,
        shuffle=False,
        collate_fn=data_collator,
        batch_size=per_device_eval_batch_size,
        num_workers=5
    )

    # optimizer
    no_decay = ["bias", "layer_norm.weight", "LayerNorm", "layernorm", "layer_norm", "ln"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": weight_decay,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]

    optimizer = torch.optim.AdamW(optimizer_grouped_parameters,
                                  lr=learning_rate,
                                  eps=training_config.adam_epsilon,
                                  betas=training_config.adam_betas,
                                  weight_decay=weight_decay)

    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_epochs * num_update_steps_per_epoch
    no_eval_until_steps = no_eval_until_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps
    )
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    evaluate_fn = get_evaluate_fn(
        task=task,
        tokenizer=tokenizer,
        raw_eval_dataset=raw_dataset["validation"]
    )

    num_eval_steps = num_update_steps_per_epoch
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {max_train_steps}")

    progress_bar = tqdm(range(max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = 0

    # Train!
    for epoch in range(num_epochs):
        # model.train()
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                loss = outputs.loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % log_steps == 0 and (
                    accelerator.sync_gradients
            ):
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                wandb.log({"train_loss": loss.item(),
                           "epoch": completed_steps / num_update_steps_per_epoch,
                           "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if (completed_steps % num_eval_steps == 0 and completed_steps >= no_eval_until_steps and
                accelerator.sync_gradients) or (completed_steps <= 1 and step == 0 and no_eval_until_steps == 0
            ):
                model.eval()
                losses = []
                output_labels = []
                output_predictions = []
                output_ids = [] if task in ["squad", "copa", "multirc", "squad_v2", "hotpotqa"] else None
                for eval_step, eval_batch in enumerate(eval_dataloader):
                    extra_keys_eval_batch = {}
                    for key in list(eval_batch.keys()):
                        if key in EXTRA_KEYS_FOR_EVAL:
                            extra_keys_eval_batch[key] = eval_batch.pop(key)
                    with torch.no_grad():
                        outputs = model(**eval_batch)
                    eval_labels = accelerator.gather(eval_batch['labels'])
                    output_labels += torch.cat([
                        eval_labels,
                        torch.ones(eval_labels.shape[0], tokenizer.model_max_length - eval_labels.shape[1],
                                   dtype=eval_labels.dtype,
                                   device=eval_labels.device) * -100
                    ], dim=-1)
                    eval_logits = accelerator.gather(outputs.logits)
                    output_predictions += eval_logits.argmax(dim=-1).tolist()
                    if task in ["squad", "squad_v2", "hotpotqa"]:
                        output_ids += extra_keys_eval_batch["id"]
                    elif task == "copa" or task == "multirc":
                        output_ids += extra_keys_eval_batch["idx"]
                    losses.append(accelerator.gather_for_metrics(outputs["loss"]))
                losses = torch.cat(losses)
                eval_loss = torch.mean(losses)
                output_labels = torch.stack(output_labels, dim=0)
                eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels, ids=output_ids)
                metric_key = list(eval_res.keys())[0]
                eval_res["loss"] = eval_loss.item()

                if eval_res[metric_key] > best_eval:
                    best_eval = eval_res[metric_key]
                    accelerator.wait_for_everyone()
                    wandb.summary["best_" + metric_key] = best_eval
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(os.path.join(output_dir, "best"),
                                                    is_main_process=accelerator.is_local_main_process,
                                                    save_function=accelerator.save)
                    if accelerator.is_local_main_process:
                        tokenizer.save_pretrained(os.path.join(output_dir, "best"))
                        print(f"Best model saved with best evaluation {metric_key}: {eval_res[metric_key]}")
                        print(f"All evaluation results: {eval_res}")

                if accelerator.is_local_main_process:
                    print(f"Step {completed_steps}: eval loss {eval_res['loss']}")
                    eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                    wandb.log(eval_res, step=completed_steps)

                # model.train()

        if save_each_epoch:
            save_dir = os.path.join(output_dir, f"epoch-{epoch + 1}")
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(save_dir,
                                            is_main_process=accelerator.is_local_main_process,
                                            save_function=accelerator.save)
            if accelerator.is_local_main_process:
                tokenizer.save_pretrained(save_dir)

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(output_dir, "latest"),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)

    if accelerator.is_local_main_process:
        tokenizer.save_pretrained(os.path.join(output_dir, "latest"))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


if __name__ == "__main__":
    Fire(finetune_on_downstream_task)

mcsmoe/t5-finetune-on-downstream.py
UNITES-Lab/MC-SMoE
from .constants import (
    FP32_EPS,
    BF16_EPS
)
from .dummy_test import (
    DUMMY_LABELS,
    DUMMY_INPUT_IDS
)
from .sparsity import (
    compute_weight_stable_rank
)

mcsmoe/utils/__init__.py
UNITES-Lab/MC-SMoE
FP32_EPS = 1e-7
BF16_EPS = 1e-3

mcsmoe/utils/constants.py
UNITES-Lab/MC-SMoE
import torch

__all__ = ['DUMMY_INPUT_IDS', 'DUMMY_LABELS']

DUMMY_INPUT_IDS = torch.randint(0, 100, (1, 128))
DUMMY_LABELS = torch.randint(0, 100, (1, 128))

mcsmoe/utils/dummy_test.py
UNITES-Lab/MC-SMoE
import torch

from .constants import FP32_EPS


def compute_weight_stable_rank(weight: torch.Tensor) -> float:
    """
    Compute the stable rank of a weight matrix

    Parameters
    ----------
    weight: torch.Tensor
        The weight matrix

    Returns
    -------
    float
        The stable rank of the weight matrix

    Examples
    --------
    >>> import torch
    >>> compute_weight_stable_rank(torch.randn(100, 100)) > 20
    True
    """
    singular_values = torch.linalg.svdvals(weight)
    return torch.sum(singular_values ** 2).item() / (torch.max(singular_values).item() ** 2 + FP32_EPS)

mcsmoe/utils/sparsity.py
UNITES-Lab/MC-SMoE
"""
The function to train the model
"""
import math
import os
from dataclasses import dataclass
from typing import Optional, Union

import torch
from accelerate.accelerator import Accelerator
from accelerate.data_loader import DataLoaderShard
from accelerate.logging import MultiProcessAdapter
from accelerate.optimizer import AcceleratedOptimizer
from accelerate.scheduler import AcceleratedScheduler
from tqdm import tqdm
from transformers import (
    PreTrainedModel, PreTrainedTokenizerFast, PreTrainedTokenizer, SwitchTransformersForConditionalGeneration
)


@dataclass
class TrainingArguments:
    overrode_max_train_steps: bool
    output_dir: str
    per_device_train_batch_size: int
    num_epochs: Optional[int]
    max_train_steps: Optional[int]
    checkpoint_steps: Optional[int] = None
    gradient_accumulation_steps: Optional[int] = 1
    num_eval_steps: Optional[int] = None
    eval_at_the_beginning: bool = False
    log_steps: Optional[int] = 100
    no_eval_until_epochs: Optional[int] = 0


def accelerate_run_train(
        accelerator: Accelerator,
        model: PreTrainedModel,
        optimizer: AcceleratedOptimizer,
        lr_scheduler: AcceleratedScheduler,
        training_args: TrainingArguments,
        train_dataloader: DataLoaderShard,
        logger: MultiProcessAdapter,
        eval_dataloader: Optional[DataLoaderShard] = None,
        wandb=None,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None,
        evaluate_fn=None,
        evaluate_loss_key="loss",
        print_outputs=False,
):
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / training_args.gradient_accumulation_steps)
    no_eval_until_steps = training_args.no_eval_until_epochs * num_update_steps_per_epoch
    if training_args.num_eval_steps is None:
        #  Eval at the end of each epoch
        training_args.num_eval_steps = num_update_steps_per_epoch
    if training_args.overrode_max_train_steps:
        training_args.max_train_steps = (
                training_args.num_epochs * num_update_steps_per_epoch
        )
    if training_args.checkpoint_steps is None:
        # Default to never saving checkpoints
        training_args.checkpoint_steps = training_args.max_train_steps + 2
    training_args.num_epochs = math.ceil(
        training_args.max_train_steps / num_update_steps_per_epoch
    )
    per_device_train_batch_size = training_args.per_device_train_batch_size
    gradient_accumulation_steps = training_args.gradient_accumulation_steps
    total_batch_size = (
            per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps
    ) if accelerator.num_processes is not None else (
            per_device_train_batch_size * gradient_accumulation_steps
    )

    logger.info("***** Running training *****")
    logger.info(f"  Num Epochs = {training_args.num_epochs}")
    logger.info(f"  Num Samples = {len(train_dataloader)}")
    logger.info(f"  Instantaneous batch size per device = {per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {training_args.max_train_steps}")

    progress_bar = tqdm(range(training_args.max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    best_eval = float('inf')

    # Train!
    for epoch in range(training_args.num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(model):
                outputs = model(**batch)
                if print_outputs and accelerator.is_local_main_process:
                    print(outputs)
                loss = outputs.loss
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.is_local_main_process and completed_steps % training_args.log_steps == 0:
                logger.info(f"epoch {epoch}, step {step}: loss {loss.item()}")
                if wandb is not None:
                    wandb.log({"train_loss": loss.item(),
                               "epoch": completed_steps / num_update_steps_per_epoch,
                               "learning_rate": lr_scheduler.get_lr()}, step=completed_steps)
            if accelerator.sync_gradients:
                progress_bar.update(1)
                completed_steps += 1
            if completed_steps % training_args.checkpoint_steps == 0 and completed_steps > 0:
                accelerator.wait_for_everyone()
                unwrapped_model = accelerator.unwrap_model(model)
                unwrapped_model.save_pretrained(
                    os.path.join(training_args.output_dir, "checkpoint-{}".format(completed_steps)),
                    is_main_process=accelerator.is_local_main_process,
                    save_function=accelerator.save)
            if (completed_steps % training_args.num_eval_steps == 0 and
                completed_steps >= no_eval_until_steps and accelerator.sync_gradients) or (
                    step == 0 and training_args.eval_at_the_beginning
            ):
                if eval_dataloader is not None:
                    model.eval()
                    losses = []
                    output_labels = []
                    output_predictions = []
                    for eval_step, eval_batch in enumerate(eval_dataloader):
                        with torch.no_grad():
                            outputs = model(**eval_batch)
                        if evaluate_fn is not None:
                            eval_labels = accelerator.gather(eval_batch['labels'])
                            output_labels += torch.cat([
                                eval_labels,
                                torch.ones(eval_labels.shape[0],
                                           tokenizer.model_max_length - eval_labels.shape[1],
                                           dtype=eval_labels.dtype,
                                           device=eval_labels.device) * -100
                            ], dim=-1)
                            eval_logits = accelerator.gather(outputs.logits)
                            output_predictions += eval_logits.argmax(dim=-1).tolist()
                        loss = outputs[evaluate_loss_key]
                        losses.append(accelerator.gather_for_metrics(loss))
                    losses = torch.cat(losses)
                    eval_loss = torch.mean(losses)
                    if evaluate_fn is not None:
                        output_labels = torch.stack(output_labels, dim=0)
                        eval_res = evaluate_fn(predictions=output_predictions, labels=output_labels)
                    else:
                        eval_res = dict()
                    metric_key = list(eval_res.keys())[0] if len(eval_res) > 0 else evaluate_loss_key
                    eval_res[evaluate_loss_key] = eval_loss.item()
                    try:
                        perplexity = math.exp(eval_res[evaluate_loss_key])
                    except OverflowError:
                        perplexity = float('inf')
                    eval_res['perplexity'] = perplexity

                    is_language_modeling = (metric_key == evaluate_loss_key)

                    if (is_language_modeling and eval_res[metric_key] < best_eval) or (
                            not is_language_modeling and eval_res[metric_key] > -best_eval
                    ):  # best_eval is inited as inf
                        best_eval = eval_res[metric_key] if is_language_modeling else -eval_res[metric_key]
                        wandb.summary["best_" + metric_key] = eval_res[metric_key]
                        accelerator.wait_for_everyone()
                        unwrapped_model = accelerator.unwrap_model(model)
                        unwrapped_model.save_pretrained(os.path.join(training_args.output_dir, "best"),
                                                        is_main_process=accelerator.is_local_main_process,
                                                        save_function=accelerator.save)
                        if accelerator.is_local_main_process:
                            tokenizer.save_pretrained(os.path.join(training_args.output_dir, "best"))
                            print("Best model saved")

                    if accelerator.is_local_main_process:
                        print(
                            f"epoch {epoch} step {completed_steps}: eval loss {eval_res[evaluate_loss_key]}, perplexity {perplexity}")
                        if wandb is not None:
                            eval_res = {("eval_" + k): v for k, v in eval_res.items()}
                            wandb.log(eval_res, step=completed_steps)

            if completed_steps >= training_args.max_train_steps:
                break

    # Finish Training!
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(os.path.join(training_args.output_dir, "checkpoint-{}".format(completed_steps)),
                                    is_main_process=accelerator.is_local_main_process,
                                    save_function=accelerator.save)
    tokenizer.save_pretrained(os.path.join(training_args.output_dir, "checkpoint-{}".format(completed_steps)))

    if accelerator.is_local_main_process and wandb is not None:
        wandb.finish()


def freeze_switch_routers_for_finetuning(
        model: SwitchTransformersForConditionalGeneration
) -> SwitchTransformersForConditionalGeneration:
    model.router_z_loss_coef = 0
    model.router_aux_loss_coef = 0
    for name, param in model.named_parameters():
        if "router" in name:
            param.requires_grad = False
        else:
            param.requires_grad = True
    return model

mcsmoe/utils/training_utils.py
UNITES-Lab/MC-SMoE