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

fusion_bench/method/expert_sparsity/utils/calibration_data.py

@@ -0,0 +1,153 @@
+"""
+This module contains the code for loading the calibration data.
+
+- C4
+- Math
+"""
+
+import itertools
+import logging
+import os
+
+import torch
+import transformers
+from datasets import load_dataset
+from transformers import PreTrainedTokenizer, default_data_collator
+from transformers.testing_utils import CaptureLogger
+from huggingface_hub import hf_hub_download
+
+logger = logging.getLogger(__name__)
+
+
+DATASETS = {
+    # C4: Please download first part of the C4 training data `c4-train.00000-of-01024.json` from [allenai/c4](https://huggingface.co/datasets/allenai/c4/blob/main/en/c4-train.00000-of-01024.json.gz).
+    "c4": lambda: load_dataset(
+        "json",
+        data_files={
+            "train": hf_hub_download(
+                "allenai/c4",
+                filename="en/c4-train.00000-of-01024.json.gz",
+                repo_type="dataset",
+            )
+        },
+    ),
+    # MATH: You can use our pre-built calibration set in `./data/math_pretrain_style.json`. To reproduce our construction, please download the training set of [MATH](https://github.com/hendrycks/math) and use our [script](data/math_calib_construction.py).
+    # NOTE: I have uploaded the math_pretrain_style.json to my huggingface repo:
+    # https://huggingface.co/datasets/tanganke/math_pretrain_style/tree/main.
+    "math": lambda: load_dataset(
+        "json",
+        data_files={
+            "train": hf_hub_download(
+                "tanganke/math_pretrain_style",
+                filename="math_pretrain_style.json",
+                repo_type="dataset",
+            )
+        },
+    ),
+}
+
+
+def build_calib_loader(
+    dataset: str,
+    tokenizer: PreTrainedTokenizer,
+    max_block_size: int,
+    n_blocks_for_stat: int,
+    batch_size: int,
+    num_workers: int,
+    seed: int = 42,
+):
+    # dataset can be a string or a dataset object.
+    # If it is a string, it can be the name of the dataset in DATASETS or the path to the dataset (a json file).
+    if isinstance(dataset, str):
+        if dataset in DATASETS:
+            all_set = DATASETS[dataset]()
+        else:
+            assert os.path.exists(dataset), f"Dataset {dataset} not found."
+            all_set = load_dataset("json", data_files={"train": dataset})
+    else:
+        assert dataset is not None, "Dataset is not provided."
+        all_set = dataset
+
+    block_size = tokenizer.model_max_length
+    if block_size > max_block_size:
+        logger.info(
+            "The chosen tokenizer supports a `model_max_length` that is longer than the default `max_block_size` value"
+            f" of {max_block_size}. If you would like to use a longer `block_size` up to `tokenizer.model_max_length` you can"
+            " override this default with `--max_block_size xxx`."
+        )
+        block_size = max_block_size
+
+    if n_blocks_for_stat > 0:  # Random choose `n_blocks_for_stat` blocks
+        calib_set = (
+            all_set["train"]
+            .shuffle(seed=seed)
+            .select(range(min(n_blocks_for_stat * 16, len(all_set["train"]))))
+        )
+    else:  # Use the whole set
+        logger.warning("n_blocks_for_stat <= 0, using the whole dataset.")
+        calib_set = all_set["train"].shuffle(seed=seed)
+
+    logger.info(f"Calibration dataset: {calib_set}")
+    text_column_name = (
+        "text" if "text" in calib_set.features else list(calib_set.features)[0]
+    )
+
+    tok_logger = transformers.utils.logging.get_logger(
+        "transformers.tokenization_utils_base"
+    )
+
+    def tokenize_function(examples):
+        with CaptureLogger(tok_logger) as cl:
+            output = tokenizer(examples[text_column_name])
+        # clm input could be much much longer than block_size
+        if "Token indices sequence length is longer than the" in cl.out:
+            tok_logger.warning(
+                "^^^^^^^^^^^^^^^^ Please ignore the warning above - this long input will be chunked into smaller bits"
+                " before being passed to the model."
+            )
+        return output
+
+    tokenized_calib_set = calib_set.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=list(calib_set.features),
+    )
+
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {
+            k: list(itertools.chain(*examples[k])) for k in examples.keys()
+        }
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    lm_calib_set = tokenized_calib_set.map(
+        group_texts,
+        batched=True,
+    )
+
+    if n_blocks_for_stat > 0:
+        assert len(lm_calib_set) > n_blocks_for_stat
+        lm_calib_set = lm_calib_set.select(range(n_blocks_for_stat))
+
+    calib_loader = torch.utils.data.DataLoader(
+        lm_calib_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=False,
+        shuffle=False,
+        collate_fn=default_data_collator,
+    )
+
+    return calib_loader

fusion_bench/method/gossip/flan_t5_layer_wise_gossip.py

@@ -32,7 +32,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .layer_wise_gossip import ModelScheduler

fusion_bench/method/knots/knots_utils.py

@@ -0,0 +1,23 @@
+import torch
+
+
+def subspace_alignment(
+    delta_weights: list[torch.Tensor],
+    svd_dtype: torch.dtype | None = torch.float64,
+    eps: float = 1e-4,
+):
+    """
+    Reference: Model merging with SVD to tie the Knots. http://arxiv.org/abs/2410.19735
+    """
+    if svd_dtype is None:
+        svd_dtype = delta_weights[0].dtype
+    original_dtype = delta_weights[0].dtype
+    output_dim, input_dim = delta_weights[0].size()
+    concat_task_vector = torch.cat(delta_weights, dim=1)
+    U, S, Vh = torch.linalg.svd(concat_task_vector.to(svd_dtype), full_matrices=False)
+    # Keep only supported basis components
+    U = U[:, S > eps].to(original_dtype)
+    Vh = Vh[S > eps].to(original_dtype)
+    S = S[S > eps].to(original_dtype)
+    Vhs = torch.split(Vh, input_dim, dim=1)
+    return U, S, Vhs

fusion_bench/method/pwe_moe/module.py

@@ -13,14 +13,9 @@ import torch.func
 from torch import Tensor, nn
 from torch.nn import functional as F
 
-log = logging.getLogger(__name__)
-
+from fusion_bench.utils import join_list
 
-def join_list(list_of_list: List[List]):
-    ans = []
-    for item in list_of_list:
-        ans.extend(item)
-    return ans
+log = logging.getLogger(__name__)
 
 
 class PWEMoEGate(nn.Module):

fusion_bench/method/simple_average.py

@@ -39,12 +39,13 @@ def simple_average(
         >>> import torch.nn as nn
         >>> model1 = nn.Linear(10, 10)
         >>> model2 = nn.Linear(10, 10)
-        >>> averaged_model = simple_averageing([model1, model2])
+        >>> averaged_model = simple_average([model1, model2])
 
         >>> state_dict1 = model1.state_dict()
         >>> state_dict2 = model2.state_dict()
-        >>> averaged_state_dict = simple_averageing([state_dict1, state_dict2])
+        >>> averaged_state_dict = simple_average([state_dict1, state_dict2])
     """
+    assert len(modules) > 0, "modules must be a non-empty list"
     if isinstance(modules[0], nn.Module):
         if base_module is None:
             new_module = deepcopy(modules[0])

fusion_bench/method/task_singular_vector/TSVM.py

@@ -1,5 +1,31 @@
-"""
-Example:
+R"""
+# Task Singular Vector Merging (TSVM) Algorithm Implementation
+
+This module implements the Task Singular Vector Merging algorithm for combining multiple fine-tuned models
+into a single unified model.
+
+## Example Usage:
+
+Merge 8 CLIP-ViT-B/32 models with TSVM:
+
+```bash
+fusion_bench \
+    method=task_singular_vector/TaskSingularVectorMerging \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+```
+
+Merge 8 CLIP-ViT-B/32 models with TSVM and return individual transformed models:
+
+```bash
+fusion_bench \
+    method=task_singular_vector/TaskSingularVectorMerging \
+    method.return_single_task_models=true \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+```
+
+Merge 20 CLIP-VIT-B/32 models with TSVM:
 
 ```bash
 fusion_bench \
@@ -7,14 +33,22 @@ fusion_bench \
     modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TALL20_model_only \
     taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TALL20
 ```
+
+## References:
+
+- Gargiulo, et al. Task Singular Vectors: Reducing Task Interference in Model Merging. 
+    https://arxiv.org/abs/2412.00081
+- See `docs/algorithms/task_singular_vector.md` for more details.
 """
 
+from copy import deepcopy
 from typing import Iterable, List, Optional, Union
 
 import torch
 from omegaconf import ListConfig
 from torch import Tensor, nn
 
+import fusion_bench as fb
 from fusion_bench import BaseAlgorithm
 from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.utils import timeit_context
@@ -35,49 +69,228 @@ from .utils import (
 
 
 class TaskSingularVectorMerging(BaseAlgorithm, LightningFabricMixin):
+    """
+    Task Singular Vector Merging (TSVM) Algorithm
+
+    This class implements a model merging technique that leverages Singular Value
+    Decomposition (SVD) to identify and combine the most important directions in the task vector
+    space. The algorithm is particularly effective for merging multiple models fine-tuned on
+    different tasks while preserving their essential capabilities.
+
+    Key Concepts:
+    - Task Vector: The difference between a fine-tuned model and its pretrained base model,
+      representing the knowledge gained during fine-tuning for a specific task.
+    - Singular Value Decomposition: A matrix factorization technique used to find the principal
+      components (most important directions) in the space of task vectors.
+    - Model Merging: The process of combining multiple models into a single unified model that
+      retains capabilities from all constituent models.
+
+    Algorithm Steps:
+    1. Extract task vectors from all fine-tuned models by subtracting the pretrained model
+    2. Apply SVD to the matrix of task vectors to find principal components
+    3. Reconstruct task vectors using only the most significant singular vectors
+    4. Merge the reconstructed task vectors (either individually scaled or as a sum)
+    5. Add the final merged task vector to the pretrained model to create the unified model
+
+    see `docs/algorithms/task_singular_vector.md` for comprehensive algorithmic details.
+    """
 
     def __init__(
         self,
-        alpha: Union[float, Iterable[float]] = None,
-        remove_keys: Optional[List[str]] = None,
+        alpha: Optional[Union[float, Iterable[float]]] = None,
+        exclude_keys: Optional[List[str]] = None,
+        return_single_task_models: bool = False,
         **kwargs,
     ):
+        """
+        Initialize the Task Singular Vector Merging algorithm.
+
+        Args:
+            alpha (Union[float, Iterable[float]], optional): Scaling factor(s) for task vectors.
+                This parameter controls the strength of the task-specific adaptations in the final model.
+
+                - If a single float: Applied to the final merged task vector after SVD reconstruction.
+                  This uniformly scales the entire merged adaptation.
+
+                - If an iterable of floats: Applied to individual task vectors before SVD and merging.
+                  Must have the same length as the number of models in the modelpool.
+                  This allows for task-specific weighting (e.g., giving more importance to certain tasks).
+
+                - If None: No scaling is applied (equivalent to alpha=1.0).
+
+                Example: alpha=[0.8, 1.2, 0.5] would apply different weights to three different task vectors.
+
+            exclude_keys (Optional[List[str]], optional): List of parameter names to exclude from TSVM.
+                These parameters will not participate in the SVD computation and merging process.
+                Useful for excluding certain layers (e.g., task-specific heads, normalization layers)
+                that should not be merged across tasks. Defaults to an empty list.
+
+                Example: exclude_keys=['classifier.weight', 'classifier.bias'] to skip classification heads.
+
+            return_single_task_models (bool, optional): Whether to return individual transformed models.
+
+                - If True: Returns a dictionary containing both individual models with their transformed
+                  task vectors applied AND the final merged model. The dictionary has the structure:
+
+                  >>> {'model_name_1': transformed_model_1, ..., 'merged': final_merged_model}
+
+                - If False: Returns only the final merged model.
+
+                This is useful for analysis or when you need access to intermediate results.
+                Defaults to False.
+
+            **kwargs: Additional arguments passed to the parent BaseAlgorithm class.
+
+        Note:
+            The choice between single alpha vs. list of alphas affects the merging strategy:
+            - Single alpha: SVD is applied first, then the result is scaled
+            - List of alphas: Individual task vectors are scaled first, then SVD is applied
+        """
         self.alpha = alpha
-        self.remove_keys = remove_keys if remove_keys is not None else []
+        self.exclude_keys = exclude_keys if exclude_keys is not None else []
+        self.return_single_task_models = return_single_task_models
         super().__init__(**kwargs)
 
-    def run(self, modelpool):
-        # Load the pre-trained model and the fine-tuned models
+    def load_pretrained_model_and_task_vectors(self, modelpool: fb.BaseModelPool):
+        """
+        Load the pretrained base model and compute task vectors from all fine-tuned models.
+
+        This method performs the initial step of the TSVM algorithm by:
+        1. Loading the original pretrained model (before any task-specific fine-tuning)
+        2. For each fine-tuned model in the pool:
+           - Load the fine-tuned model
+           - Compute the task vector (fine-tuned params - pretrained params)
+           - Optionally apply individual scaling if alpha is provided as a list
+
+        Task vectors represent the knowledge gained during fine-tuning and are the core
+        data structure that TSVM operates on.
+
+        Args:
+            modelpool (fb.BaseModelPool): Pool containing the pretrained model and all
+                fine-tuned models to be merged.
+
+        Returns:
+            tuple: A tuple containing:
+                - pretrained_model: The original pretrained model (torch.nn.Module)
+                - task_vectors: List of task vectors (List[StateDictType]), where each
+                  task vector is a state dictionary representing the parameter differences
+                  for one specific task
+        """
+        # Load the original pretrained model that serves as the base for all fine-tuned variants
         pretrained_model = modelpool.load_pretrained_model()
-        finetuned_models = list(modelpool.models())
 
-        ptm_check = pretrained_model.state_dict()
-        ft_checks = [model.state_dict() for model in finetuned_models]
-        check_parameterNamesMatch(ft_checks + [ptm_check])
+        # Initialize list to store computed task vectors
+        task_vectors = []
+
+        # Process each fine-tuned model in the modelpool
+        for model_idx, model_name in enumerate(modelpool.model_names):
+            # Load the current fine-tuned model
+            finetuned_model = modelpool.load_model(model_name)
+
+            # Compute task vector: difference between fine-tuned and pretrained parameters
+            # This captures the task-specific adaptations learned during fine-tuning
+            task_vector = state_dict_sub(
+                finetuned_model.state_dict(), pretrained_model.state_dict()
+            )
+            task_vectors.append(task_vector)
 
-        with timeit_context("Flattening out Checkpoints"):
-            task_vectors = [state_dict_sub(check, ptm_check) for check in ft_checks]
-            if isinstance(self.alpha, Iterable):
+            # Apply individual scaling to task vectors if alpha is provided as a list
+            # This allows for task-specific weighting before the SVD computation
+            if self.alpha is not None and isinstance(self.alpha, Iterable):
+                # Ensure the number of alpha values matches the number of models
                 assert len(self.alpha) == len(
-                    task_vectors
-                ), "Alpha and task vectors must have the same length"
-                task_vectors = [
-                    state_dict_mul(state_dict=tv, scalar=alpha)
-                    for alpha, tv in zip(self.alpha, task_vectors)
-                ]
+                    modelpool.model_names
+                ), f"Alpha list length ({len(self.alpha)}) must match number of models ({len(modelpool.model_names)})"
 
+                # Scale the current task vector by its corresponding alpha value
+                task_vectors[-1] = state_dict_mul(
+                    state_dict=task_vectors[-1], scalar=self.alpha[model_idx]
+                )
+
+        return pretrained_model, task_vectors
+
+    def run(self, modelpool: fb.BaseModelPool):
+        """
+        Execute the complete Task Singular Vector Merging algorithm.
+
+        This is the main entry point that orchestrates the entire TSVM process:
+
+        The algorithm leverages the mathematical insight that task vectors often lie in a
+        lower-dimensional subspace, and SVD helps identify the most important directions
+        in this subspace while filtering out noise and interference.
+
+        Args:
+            modelpool (fb.BaseModelPool): Pool of models to merge, including:
+                - The pretrained base model
+                - Multiple fine-tuned models (one per task)
+                All models must have compatible architectures.
+
+        Returns:
+            Union[torch.nn.Module, Dict[str, torch.nn.Module]]:
+                - If return_single_task_models=False: Returns the merged model
+                - If return_single_task_models=True: Returns a dictionary with:
+                  * Individual transformed models keyed by their original names
+                  * Final merged model under the key 'merged'
+
+        Raises:
+            AssertionError: If alpha list length doesn't match the number of models
+        """
+        # Determine the compute device for SVD operations (GPU if available for faster computation)
+        accelerator = self.fabric.device
+
+        # Phase 1: Load pretrained model and compute task vectors from all fine-tuned models
+        pretrained_model, task_vectors = self.load_pretrained_model_and_task_vectors(
+            modelpool
+        )
+
+        # Phase 2: Apply SVD-based merging to the task vectors
+        # This is the core of the TSVM algorithm where:
+        # - Task vectors are organized into a matrix
+        # - SVD finds the principal components (most important directions)
+        # - Task vectors are reconstructed using only the most significant components
+        # - The reconstructed vectors are merged (summed) to create a unified task vector
         new_merged_tv = TSVM_utils.compute_and_sum_svd_mem_reduction(
             task_vectors,
-            exclude_keys=self.remove_keys,
-            accelerator=self.fabric.device,
+            exclude_keys=self.exclude_keys,  # Skip certain parameters from SVD
+            accelerator=accelerator,  # Use GPU if available
+            return_single_task_models=self.return_single_task_models,
         )
 
-        # If alpha is a float, we need to scale the new merged task vector by alpha
-        if self.alpha is not None and isinstance(self.alpha, float):
+        # Handle the case where individual transformed task vectors are also returned
+        if self.return_single_task_models:
+            new_merged_tv, single_task_models = new_merged_tv
+
+        # Phase 3: Apply global scaling to the merged task vector (if alpha is a single value)
+        # This is different from individual scaling applied earlier - here we scale the
+        # final merged result, which affects the overall strength of all merged adaptations
+        if self.alpha is not None and isinstance(self.alpha, (float, int)):
             print(f"Scaling new merged task vector by alpha: {self.alpha}")
             new_merged_tv = state_dict_mul(state_dict=new_merged_tv, scalar=self.alpha)
 
+        # Phase 4: Prepare individual transformed models if requested
+        if self.return_single_task_models:
+            models = {}
+            # Create individual models by adding each transformed task vector to the pretrained base
+            for model_idx, model_name in enumerate(modelpool.model_names):
+                # Create a deep copy to avoid modifying the original pretrained model
+                model = deepcopy(pretrained_model)
+                # Apply the transformed task vector to get the individual model
+                model.load_state_dict(
+                    state_dict_add(model.state_dict(), single_task_models[model_idx])
+                )
+                models[model_name] = model
+
+        # Phase 5: Create the final merged model by adding the merged task vector to pretrained model
+        # This produces a single model that combines capabilities from all input models
         pretrained_model.load_state_dict(
             state_dict_add(new_merged_tv, pretrained_model.state_dict())
         )
-        return pretrained_model
+
+        # Phase 6: Return results based on the requested output format
+        if self.return_single_task_models:
+            # Include the final merged model in the dictionary of results
+            models["merged"] = pretrained_model
+            return models
+        else:
+            # Return only the merged model
+            return pretrained_model

fusion_bench/method/task_singular_vector/utils/TSVM_utils.py

@@ -1,7 +1,9 @@
+import collections
 import math
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import torch
+from torch import Tensor, nn
 
 from fusion_bench.utils.type import StateDictType
 
@@ -314,7 +316,8 @@ def compute_and_sum_svd_mem_reduction(
     task_vectors: List[StateDictType],
     exclude_keys: Optional[List[str]] = None,
     accelerator: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
-) -> StateDictType:
+    return_single_task_models: bool = False,
+):
     """
     Computes the Singular Value Decomposition (SVD) for each vector in the task_vectors,
     reduces the dimensionality of the vectors based on the sv_reduction factor, and concatenate
@@ -326,23 +329,25 @@ def compute_and_sum_svd_mem_reduction(
                             dictionary of vectors.
         exclude_keys (list): A list of keys to exclude from the TSVM.
         accelerator (torch.device): The device to use for the computation.
+        return_single_task_models (bool): Whether to return the single task models after the TSVM.
 
     Returns:
         dict: A dictionary containing the new vectors after SVD computation and merging.
     """
     if exclude_keys is None:
         exclude_keys = []
-    sv_reduction = 1 / len(task_vectors)
+    num_tasks = len(task_vectors)
+    sv_reduction = 1 / num_tasks
 
-    new_vector = {}
+    new_vector: Dict[str, Tensor] = {}
+    if return_single_task_models:
+        single_task_models = [{} for _ in range(num_tasks)]
     for key in task_vectors[0]:
         original_device = task_vectors[0][key].device
         original_dtype = task_vectors[0][key].dtype
 
-        new_vector[key] = {}
         for i, task_vector in enumerate(task_vectors):
-            vec = task_vector[key].to(accelerator)
-
+            vec = task_vector[key].to(device=accelerator, non_blocking=True)
             if len(task_vector[key].shape) == 2 and key not in exclude_keys:
                 # at current, the SVD is not supported for half precision, so we need to convert to float32
                 if not (
@@ -350,13 +355,14 @@ def compute_and_sum_svd_mem_reduction(
                 ):
                     vec = vec.to(dtype=torch.float32)
 
-                u, s, v = torch.linalg.svd(vec, full_matrices=False)
+                # vec = u @ torch.diag(s) @ vh
+                u, s, vh = torch.linalg.svd(vec, full_matrices=False)
 
                 if i == 0:
                     print(f"Computed SVD for {key}...")
                     sum_u = torch.zeros_like(u, device=accelerator)
                     sum_s = torch.zeros_like(s, device=accelerator)
-                    sum_v = torch.zeros_like(v, device=accelerator)
+                    sum_vh = torch.zeros_like(vh, device=accelerator)
                 reduced_index_s = int(s.shape[0] * sv_reduction)
 
                 # select only the first reduced_index_s columns of u and place them
@@ -367,10 +373,9 @@ def compute_and_sum_svd_mem_reduction(
                     :reduced_index_s
                 ]
                 # select only the first reduced_index_s rows of v and place them
-                sum_v[i * reduced_index_s : (i + 1) * reduced_index_s, :] = v[
+                sum_vh[i * reduced_index_s : (i + 1) * reduced_index_s, :] = vh[
                     :reduced_index_s, :
                 ]
-
             else:
                 # if the vector is not a 2D tensor or is in exclude_keys, compute the mean
                 if i == 0:
@@ -379,22 +384,49 @@ def compute_and_sum_svd_mem_reduction(
                     new_vector[key] += (vec - new_vector[key]) / (i + 1)
 
         if len(task_vector[key].shape) == 2 and key not in exclude_keys:
-            u_u, s_u, v_u = torch.linalg.svd(sum_u, full_matrices=False)
-            u_v, s_v, v_v = torch.linalg.svd(sum_v, full_matrices=False)
+            u_u, s_u, vh_u = torch.linalg.svd(sum_u, full_matrices=False)
+            u_vh, s_vh, vh_vh = torch.linalg.svd(sum_vh, full_matrices=False)
 
             new_vector[key] = torch.linalg.multi_dot(
                 (
                     u_u,
-                    v_u,
+                    vh_u,
                     torch.diag(sum_s),
-                    u_v,
-                    v_v,
+                    u_vh,
+                    vh_vh,
                 )
             )
-        new_vector[key] = new_vector[key].to(
-            device=original_device, dtype=original_dtype, non_blocking=True
-        )
-    return new_vector
+            new_vector[key] = new_vector[key].to(
+                device=original_device, dtype=original_dtype, non_blocking=True
+            )
+            if return_single_task_models:
+                reduced_index_s = int(sum_s.shape[0] * sv_reduction)
+                new_u = u_u @ vh_u
+                new_vh = u_vh @ vh_vh
+                for i in range(num_tasks):
+                    single_task_models[i][key] = torch.linalg.multi_dot(
+                        (
+                            new_u[:, i * reduced_index_s : (i + 1) * reduced_index_s],
+                            torch.diag(
+                                sum_s[i * reduced_index_s : (i + 1) * reduced_index_s]
+                            ),
+                            new_vh[i * reduced_index_s : (i + 1) * reduced_index_s, :],
+                        )
+                    ).to(
+                        device=original_device, dtype=original_dtype, non_blocking=True
+                    )
+        else:
+            new_vector[key] = new_vector[key].to(
+                device=original_device, dtype=original_dtype, non_blocking=True
+            )
+            if return_single_task_models:
+                for i in range(num_tasks):
+                    single_task_models[i][key] = new_vector[key].clone()
+
+    if not return_single_task_models:
+        return new_vector
+    else:
+        return new_vector, single_task_models
 
 
 ###############

fusion_bench/method/task_singular_vector/utils/__init__.py

@@ -5,3 +5,4 @@ from fusion_bench.method.ties_merging.ties_merging_utils import (
 from fusion_bench.utils import state_dict_to_vector, vector_to_state_dict
 
 from . import TSVC_utils, TSVM_utils
+from .task_singular_interference import compute_task_singular_interference