fusion-bench 0.2.13__py3-none-any.whl → 0.2.14__py3-none-any.whl

fusion_bench/method/smile_upscaling/smile_upscaling.py

@@ -13,348 +13,18 @@ from fusion_bench.method import BaseAlgorithm
 from fusion_bench.method.simple_average import simple_average
 from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
 from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.models.smile_moe.linear_from_module import (
+    ExpertNotTrainedError,
+    SmileCompressedLinear,
+    SmileGate,
+    SmileMoELinear,
+)
 from fusion_bench.models.utils import get_attr, set_attr
 from fusion_bench.utils.parameters import print_parameters
 
 log = logging.getLogger(__name__)
 
 
-class ExpertNotTrainedError(Exception):
-    pass
-
-
-def _is_all_zeros(tensor: Tensor | List[Tensor]) -> bool:
-    """
-    Check if a tensor or a list of tensors are all zeros.
-
-    Args:
-        tensor (Tensor | List[Tensor]): A tensor or a list of tensors.
-
-    Returns:
-        bool: True if all elements are zeros, False otherwise.
-    """
-    if isinstance(tensor, Tensor):
-        return torch.allclose(tensor, torch.zeros_like(tensor))
-    else:
-        return all(_is_all_zeros(t) for t in tensor)
-
-
-def _svd(w: Tensor, full_matrices=True) -> Tuple[Tensor, Tensor, Tensor]:
-    """
-    Perform Singular Value Decomposition (SVD) on a tensor.
-
-    Args:
-        w (Tensor): The input tensor.
-        full_matrices (bool): Whether to compute the full-sized U and V matrices.
-
-    Returns:
-        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
-    """
-    u, s, vh = torch.linalg.svd(
-        w, full_matrices=full_matrices, driver="gesvd" if w.is_cuda else None
-    )
-    v = vh.T
-    return u, s, v
-
-
-def svd(
-    w: Tensor, full_matrices=True, accelerator=None
-) -> Tuple[Tensor, Tensor, Tensor]:
-    """
-    Perform SVD on a tensor, optionally using a specified accelerator.
-
-    Args:
-        w (Tensor): The input tensor.
-        full_matrices (bool): Whether to compute the full-sized U and V matrices.
-        accelerator (str): The device to perform the computation on.
-
-    Returns:
-        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
-    """
-    if accelerator is None:
-        return _svd(w, full_matrices=full_matrices)
-    original_device = w.device
-    w = w.to(accelerator)
-    u, s, v = _svd(w)
-    return u.to(original_device), s.to(original_device), v.to(original_device)
-
-
-class SmileGate(nn.Module):
-    def __init__(
-        self,
-        input_features: int,
-        w_diff_list: List[Tensor],
-        k: int,
-        svd_list=None,  # cached `svd_list`, pass it to avoid recomputing
-        upscaling_accelerator=None,
-    ):
-        """
-        Initialize the SmileGate module.
-
-        Args:
-            input_features (int): The number of input features.
-            w_diff_list (List[Tensor]): A list of weight difference tensors.
-            k (int): The number of singular values to keep.
-            svd_list (List[Tuple[Tensor, Tensor, Tensor]]): Cached SVD results.
-            upscaling_accelerator (str): The device to perform the computation on.
-        """
-        super().__init__()
-        self.input_features = input_features
-        self.num_experts = len(w_diff_list)
-        weights = []
-        for i, w_diff in enumerate(w_diff_list):
-            if svd_list is None:
-                u, s, v = svd(w_diff, accelerator=upscaling_accelerator)
-            else:
-                u, s, v = svd_list[i]
-            u = u[:, :k]
-            s = s[:k]
-            v = v[:, :k]
-
-            # weights.append((s * v).T)
-            weights.append(v.T)
-        self.k = s.size(0)  # k is the actual k after truncation
-
-        weights = (
-            torch.stack(weights, dim=0)
-            .reshape(self.num_experts * self.k, -1)
-            .contiguous()
-        )
-        self.weights = nn.Parameter(
-            weights
-        )  # weights should be a tensor of shape (num_experts * k, n)
-
-    def forward(self, x: Tensor):
-        """
-        Forward pass of the SmileGate module.
-
-        Args:
-            x (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The routing weights.
-        """
-        batch_size = x.size(0)
-        if self.num_experts == 1:
-            return torch.ones(batch_size, 1, device=x.device, dtype=x.dtype)
-
-        routing_weights = F.linear(x, self.weights).view(
-            batch_size, self.num_experts, self.k
-        )
-        routing_weights = routing_weights.norm(p=2, dim=2)
-        return routing_weights
-
-
-class SmileCompressedLinear(nn.Module):
-    def __init__(self, model: nn.Linear, k: int, svd_cache=None):
-        """
-        Initialize the SmileCompressedLinear module.
-
-        Args:
-            model (nn.Linear): The linear model to compress.
-            k (int): The number of singular values to keep.
-            svd_cache (Tuple[Tensor, Tensor, Tensor]): Cached SVD results.
-        """
-        super().__init__()
-        if svd_cache is None:
-            u, s, v = svd(model.weight)
-        else:
-            u, s, v = svd_cache
-        if k > 0:
-            u = u[:, :k]
-            s = s[:k]
-            v = v[:, :k]
-
-        self.u = nn.Parameter(u)
-        self.svh = nn.Parameter((s * v).T)
-
-        if model.bias is not None:
-            self.bias = nn.Parameter(model.bias.data, requires_grad=True)
-        else:
-            self.register_parameter("bias", None)
-
-    def forward(self, x):
-        """
-        Forward pass of the SmileCompressedLinear module.
-
-        Args:
-            x (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The output tensor.
-        """
-        x = F.linear(x, self.svh)
-        x = F.linear(x, self.u, self.bias)
-        return x
-
-
-class SmileMoELinear(nn.Module):
-    @torch.no_grad()
-    def __init__(
-        self,
-        pretrained_model: nn.Linear,
-        finetuned_models: List[nn.Linear],
-        gate_k: int,
-        k: int,
-        top_k: int = 1,
-        full_matrices=True,
-        upscaling_accelerator=None,
-        routing_use_diff=True,
-    ):
-        """
-        Initialize the SmileMoELinear module.
-
-        Args:
-            pretrained_model (nn.Linear): The pretrained linear model.
-            finetuned_models (List[nn.Linear]): A list of fine-tuned linear models.
-            gate_k (int): The number of singular values to keep for the gate.
-            k (int): The number of singular values to keep for the experts.
-            top_k (int): The number of top experts to select.
-            full_matrices (bool): Whether to compute the full-sized U and V matrices.
-            upscaling_accelerator (str): The device to perform the computation on.
-            routing_use_diff (bool): Whether to use weight differences for routing.
-        """
-        super().__init__()
-        self.num_experts = len(finetuned_models)
-        self.top_k = top_k
-        self.k = k
-        self.gate_k = gate_k
-        self.in_features = pretrained_model.in_features
-        self.out_features = pretrained_model.out_features
-
-        w_diff_list = [m.weight - pretrained_model.weight for m in finetuned_models]
-        if _is_all_zeros(w_diff_list):
-            # All fine-tuned models are identical to the pretrained model
-            raise ExpertNotTrainedError()
-
-        if routing_use_diff or k > 0:
-            svd_cache_list = [
-                svd(w, full_matrices=full_matrices, accelerator=upscaling_accelerator)
-                for w in w_diff_list
-            ]  # the svd cache list to avoid recomputing
-
-        # construct the gate network
-        if routing_use_diff:
-            self.gate = SmileGate(
-                input_features=self.in_features,
-                w_diff_list=w_diff_list,
-                k=gate_k,
-                svd_list=svd_cache_list,
-                upscaling_accelerator=upscaling_accelerator,
-            )
-        else:
-            self.gate = SmileGate(
-                input_features=self.in_features,
-                w_diff_list=[m.weight for m in finetuned_models],
-                k=gate_k,
-                svd_list=None,
-                upscaling_accelerator=upscaling_accelerator,
-            )
-
-        # construct experts
-        for m, w_diff in zip(finetuned_models, w_diff_list):
-            m.weight.data = w_diff
-        if k > 0:
-            experts = [
-                SmileCompressedLinear(m, k, svd_cache=svd_cache)
-                for m, svd_cache in zip(finetuned_models, svd_cache_list)
-            ]
-        else:
-            # if k is not set (<0), we use the full fine-tuned model
-            experts = finetuned_models
-        self.experts = nn.ModuleList(experts)
-
-        if pretrained_model.bias is not None:
-            for m in experts:
-                m.bias.data = m.bias.data - pretrained_model.bias
-        # assign the pretrained model (the shared part)
-        self.pretrained_model = pretrained_model
-
-    def forward(self, hidden_states: Tensor):
-        """
-        Forward pass of the SmileMoELinear module.
-
-        Args:
-            hidden_states (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The output tensor.
-        """
-        pretrained_out = self.pretrained_model(hidden_states)
-
-        input_shape = hidden_states.size()
-        hidden_states = hidden_states.view(-1, self.in_features)
-
-        router_logits = self.gate(hidden_states)
-        routing_weights = F.softmax(router_logits, dim=1)
-        # sample the expert according to the routing weights
-        routing_weights, selected_experts = torch.topk(
-            routing_weights, self.top_k, dim=-1
-        )
-        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
-
-        final_hidden_states = torch.zeros(
-            (hidden_states.size(0), self.out_features),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be sollicitated
-        expert_mask = torch.nn.functional.one_hot(
-            selected_experts, num_classes=self.num_experts
-        ).permute(2, 1, 0)
-
-        # Loop over all available experts in the model and perform the computation on each expert
-        for expert_idx in range(self.num_experts):
-            expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx])
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, self.in_features)
-            if current_state.numel() == 0:
-                continue
-            current_hidden_states = (
-                expert_layer(current_state) * routing_weights[top_x, idx, None]
-            )
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(
-                0, top_x, current_hidden_states.to(hidden_states.dtype)
-            )
-        final_hidden_states = final_hidden_states.reshape(
-            *input_shape[:-1], self.out_features
-        )
-        final_hidden_states = pretrained_out + final_hidden_states
-        return final_hidden_states
-
-    @property
-    def weight(self):
-        """
-        Mimic linear layer. Bacause in some cases, user might indicate the device (or dtype of parameters) of the linear layer using `linear_layer.weight.device`
-        """
-        return self.pretrained_model.weight
-
-    @property
-    def bias(self):
-        return self.pretrained_model.bias
-
-    def __repr__(self):
-        return (
-            f"SingularMoELinear("
-            f"in_features={self.pretrained_model.in_features}, "
-            f"out_features={self.pretrained_model.out_features}, "
-            f"num_experts={self.num_experts}, "
-            f"top_k={self.top_k}, "
-            f"gate_k={self.gate_k}, "
-            f"k={self.k}"
-            f")"
-        )
-
-
 class SmileUpscalingAlgorithm(
     SimpleProfilerMixin,
     BaseAlgorithm,

fusion_bench/models/modeling_smile_mistral/modeling_smile_mistral.py

@@ -27,6 +27,8 @@ from transformers.models.mistral.modeling_mistral import (
     MistralRotaryEmbedding,
 )
 
+from fusion_bench.models.smile_moe.linear_from_hf_config import SmileLinear
+
 from .configuration_smile_mistral import SmileMistralConfig
 
 logger = logging.getLogger(__name__)
@@ -80,209 +82,6 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
-class SmileGate(nn.Module):
-    __constants__ = ["in_features", "num_experts", "k"]
-    in_features: int
-    num_experts: int
-    k: int
-    weight: Tensor
-
-    def __init__(
-        self,
-        in_features: int,
-        num_experts: int,
-        k: int,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.input_features = in_features
-        self.num_experts = num_experts
-        self.k = k
-
-        self.weight = nn.Parameter(
-            torch.empty(num_experts * k, in_features, **factory_kwargs)
-        )
-
-    def forward(self, x: Tensor):
-        batch_size = x.size(0)
-        if self.num_experts == 1:
-            return torch.ones(batch_size, 1, device=x.device, dtype=x.dtype)
-
-        routing_weights = F.linear(x, self.weight).view(
-            batch_size, self.num_experts, self.k
-        )
-        routing_weights = routing_weights.norm(p=2, dim=2)
-        return routing_weights
-
-
-class SmileLinearExpert(nn.Module):
-    __constants__ = ["in_features", "out_features", "k"]
-    in_features: int
-    out_features: int
-    k: int
-
-    def __init__(
-        self,
-        in_features,
-        out_features,
-        k: int,
-        bias: bool,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.k = k
-
-        self.u = nn.Parameter(torch.empty(out_features, k, **factory_kwargs))
-        self.svh = nn.Parameter(torch.empty(k, in_features, **factory_kwargs))
-
-        if bias:
-            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
-        else:
-            self.register_parameter("bias", None)
-
-    def forward(self, x):
-        x = F.linear(x, self.svh)
-        x = F.linear(x, self.u, self.bias)
-        return x
-
-
-class SmileLinear(nn.Module):
-    @torch.no_grad()
-    def __init__(
-        self,
-        config: SmileMistralConfig,
-        in_features,
-        out_features,
-        bias: bool,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.num_local_experts = config.num_local_experts
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.rank_of_expert = config.rank_of_expert
-        self.rank_of_router = config.rank_of_router
-        self.in_features = in_features
-        self.out_features = out_features
-
-        # construct the gate network
-        self.gate = SmileGate(
-            in_features=in_features,
-            num_experts=self.num_local_experts,
-            k=self.rank_of_router,
-            **factory_kwargs,
-        )
-
-        # the shared linear
-        self.shared_linear = nn.Linear(
-            in_features, out_features, bias=bias, **factory_kwargs
-        )
-
-        # construct experts
-        if self.rank_of_expert > 0:
-            self.experts = nn.ModuleList(
-                [
-                    SmileLinearExpert(
-                        in_features=in_features,
-                        out_features=out_features,
-                        bias=bias,
-                        k=self.rank_of_expert,
-                        **factory_kwargs,
-                    )
-                    for _ in range(self.num_local_experts)
-                ]
-            )
-        else:
-            self.experts = nn.ModuleList(
-                [
-                    nn.Linear(in_features, out_features, bias=bias, **factory_kwargs)
-                    for _ in range(self.num_local_experts)
-                ]
-            )
-
-    def forward(self, hidden_states: Tensor):
-        pretrained_out = self.shared_linear(hidden_states)
-
-        input_shape = hidden_states.size()
-        hidden_states = hidden_states.view(-1, self.in_features)
-
-        router_logits = self.gate(hidden_states)
-        routing_weights = F.softmax(router_logits, dim=1)
-        # sample the expert according to the routing weights
-        routing_weights, selected_experts = torch.topk(
-            routing_weights, self.num_experts_per_tok, dim=-1
-        )
-        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
-
-        final_hidden_states = torch.zeros(
-            (hidden_states.size(0), self.out_features),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be sollicitated
-        expert_mask = torch.nn.functional.one_hot(
-            selected_experts, num_classes=self.num_local_experts
-        ).permute(2, 1, 0)
-
-        # Loop over all available experts in the model and perform the computation on each expert
-        for expert_idx in range(self.num_local_experts):
-            expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx])
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, self.in_features)
-            if current_state.numel() == 0:
-                continue
-            current_hidden_states = (
-                expert_layer(current_state) * routing_weights[top_x, idx, None]
-            )
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(
-                0, top_x, current_hidden_states.to(hidden_states.dtype)
-            )
-        final_hidden_states = final_hidden_states.reshape(
-            *input_shape[:-1], self.out_features
-        )
-        final_hidden_states = pretrained_out + final_hidden_states
-        return final_hidden_states
-
-    @property
-    def weight(self):
-        """
-        Mimic linear layer. Bacause in some cases, user might indicate the device (or dtype of parameters) of the linear layer using `linear_layer.weight.device`
-        """
-        return self.shared_linear.weight
-
-    @property
-    def bias(self):
-        return self.shared_linear.bias
-
-    def __repr__(self):
-        return (
-            f"SingularMoELinear("
-            f"in_features={self.shared_linear.in_features}, "
-            f"out_features={self.shared_linear.out_features}, "
-            f"num_local_experts={self.num_local_experts}, "
-            f"num_experts_per_tok={self.num_experts_per_tok}, "
-            f"rank_of_router={self.rank_of_router}, "
-            f"rank_of_expert={self.rank_of_expert}"
-            f")"
-        )
-
-
 class SmileMistralAttention(nn.Module):
     """
     Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer

fusion_bench/models/modeling_smile_qwen2/__init__.py

@@ -0,0 +1,8 @@
+from . import register
+from .configuration_smile_qwen2 import SmileQwen2Config
+from .modeling_smile_qwen2 import (
+    SmileQwen2ForCausalLM,
+    SmileQwen2ForQuestionAnswering,
+    SmileQwen2ForSequenceClassification,
+    SmileQwen2Model,
+)

fusion_bench/models/modeling_smile_qwen2/configuration_smile_qwen2.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+class SmileQwen2Config(Qwen2Config):
+    model_type = "smile_qwen2"
+
+    def __init__(
+        self,
+        num_experts_per_tok: int = 1,
+        rank_of_router: int = None,
+        rank_of_expert: int = None,
+        num_local_experts: int = None,
+        **kwargs,
+    ):
+        self.num_experts_per_tok = num_experts_per_tok
+        self.rank_of_router = rank_of_router
+        self.rank_of_expert = rank_of_expert
+        self.num_local_experts = num_local_experts
+
+        super().__init__(**kwargs)