# Copyright (C) 2025 Arcee AI # SPDX-License-Identifier: LGPL-3.0-only from typing import Any, Dict, List, Optional import torch import torch.nn.functional as F from typing_extensions import override from mergekit.architecture import WeightInfo from mergekit.common import ModelReference from mergekit.graph import Task from mergekit.merge_methods.base import ( ConfigParameterDef, MergeMethod, MergeTensorInput, ) from mergekit.merge_methods.rectify_embed import rectify_embed_sizes class DynamicThresholdFusion: def approximate_quantiles(self, tensor, q): # Flatten the tensor flat_tensor = tensor.view(-1) # If tensor is too large, sample it if flat_tensor.numel() > 1e6: flat_tensor = flat_tensor[torch.randperm(flat_tensor.numel())[:1000000]] # Sort the (possibly sampled) tensor sorted_tensor, _ = torch.sort(flat_tensor) # Compute quantile indices quantile_indices = (q * (sorted_tensor.numel() - 1)).long() # Return quantiles return sorted_tensor[quantile_indices] def calculate_dynamic_threshold(self, importance_scores, tukey_fence=1.5): # Approximate median and quantiles median = self.approximate_quantiles(importance_scores, torch.tensor([0.5]))[0] q1, q3 = self.approximate_quantiles( importance_scores, torch.tensor([0.25, 0.75]) ) # Calculate IQR iqr = q3 - q1 # Set threshold as median + tukey_fence * IQR dynamic_threshold = median + tukey_fence * iqr return dynamic_threshold def compute_fusion_mask(self, importance_scores, tukey_fence=1.5): threshold = self.calculate_dynamic_threshold(importance_scores, tukey_fence) fusion_mask = (importance_scores >= threshold).float() return fusion_mask, threshold class MultiFusionMergeTask(Task[torch.Tensor]): gather_tensors: MergeTensorInput base_model: ModelReference weight_info: WeightInfo importance_metric: str = "delta_mag" tukey_fence: float = 1.5 def uses_accelerator(self) -> bool: return True def arguments(self) -> Dict[str, Task]: return {"tensors": self.gather_tensors} def execute(self, tensors: Dict[ModelReference, torch.Tensor]) -> torch.Tensor: if len(tensors) == 1: return list(tensors.values())[0] elif len(tensors) != 2: raise RuntimeError("MutliFusion merge expects exactly two models") elif self.base_model not in tensors: raise RuntimeError("Base model not in input tensors") [a, b] = list(tensors.items()) if a[0] != self.base_model: [a, b] = [b, a] prepped_tensors = [a[1], b[1]] rectify_embed_sizes(self.weight_info, prepped_tensors) importance_scores = self._compute_importance( prepped_tensors[1], prepped_tensors[0] ) dynamic_threshold_fusion = DynamicThresholdFusion() fusion_mask, _threshold = dynamic_threshold_fusion.compute_fusion_mask( importance_scores, tukey_fence=self.tukey_fence ) delta = prepped_tensors[1] - prepped_tensors[0] masked_delta = delta * fusion_mask fused = prepped_tensors[0] + masked_delta return fused def _compute_importance( self, params: torch.Tensor, base_params: torch.Tensor, eps: float = 1e-8 ) -> torch.Tensor: if self.importance_metric == "kl_div": return self._compute_kl_div_importance(params, base_params, eps) elif self.importance_metric == "delta_mag": return self._compute_delta_mag_importance(params, base_params) elif self.importance_metric == "cosine_sim": return self._compute_cosine_sim_importance(params, base_params) elif self.importance_metric == "fisher_grad": return self._compute_fisher_grad_importance(params, base_params) else: raise ValueError(f"Unknown importance metric: {self.importance_metric}") def _compute_kl_div_importance( self, params: torch.Tensor, base_params: torch.Tensor, eps: float = 1e-8 ) -> torch.Tensor: diff = (params - base_params).abs() p = F.softmax(params, dim=-1) + eps q = F.softmax(base_params, dim=-1) + eps kl_div = torch.sum(p * torch.log(p / q), dim=-1) return diff * kl_div.unsqueeze(-1) def _compute_delta_mag_importance( self, params: torch.Tensor, base_params: torch.Tensor ) -> torch.Tensor: # Magnitude of delta - used by TIES/DARE/DELLA delta = params - base_params return delta.abs() def _compute_cosine_sim_importance( self, params: torch.Tensor, base_params: torch.Tensor ) -> torch.Tensor: # Cosine similarity based - inspired by Model Stock delta = params - base_params delta_flat = delta.view(-1) base_flat = base_params.view(-1) # Compute cosine similarity between delta and base dot_product = torch.dot(delta_flat, base_flat) norm_delta = torch.norm(delta_flat) norm_base = torch.norm(base_flat) # Avoid division by zero if norm_delta == 0 or norm_base == 0: return torch.zeros_like(delta) cosine_sim = dot_product / (norm_delta * norm_base) # Convert similarity to importance (higher similarity = more important) importance = delta.abs() * (1 + cosine_sim.abs()) return importance.view_as(delta) def _compute_fisher_grad_importance( self, params: torch.Tensor, base_params: torch.Tensor ) -> torch.Tensor: # Fisher/gradient-based importance - inspired by Karcher/Fisher information # Since we don't have access to gradients/data, we use a proxy based on # the magnitude and variance of the delta delta = params - base_params # Compute variance along the last dimension as a proxy for Fisher information if delta.dim() > 1: variance = torch.var(delta, dim=-1, keepdim=True) else: variance = delta.var().unsqueeze(0) ## # Importance combines magnitude and variance ## importance = delta.abs() * (variance + 1e-8) ## return importance # Use variance directly as importance (SCE-style) rather than variance * magnitude importance = variance + 1e-8 return importance class MultiFusionMerge(MergeMethod): def name(self) -> str: return "multi_fusion" @override def pretty_name(self) -> Optional[str]: return "Multi Fusion" @override def reference_url(self) -> Optional[str]: return "https://huggingface.co/Naphula/Slimaki-Tavern-24B-v1.3" def parameters(self) -> List[ConfigParameterDef]: return [ ConfigParameterDef( name="importance_metric", required=False, default_value="delta_mag", ), ConfigParameterDef( name="tukey_fence", required=False, default_value=1.5, ) ] def make_task( self, output_weight: WeightInfo, tensors: MergeTensorInput, base_model: Optional[ModelReference], parameters: Dict[str, Any], **kwargs, ) -> Task[torch.Tensor]: return MultiFusionMergeTask( gather_tensors=tensors, weight_info=output_weight, base_model=base_model, importance_metric=parameters["importance_metric"], tukey_fence=parameters["tukey_fence"] )