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

fusion_bench/method/pruning/sparsegpt_utils/sparsegpt.py

@@ -0,0 +1,128 @@
+import math
+import time
+
+import torch
+import torch.nn as nn
+import transformers
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+
+
+## SparseGPT: https://github.com/IST-DASLab/sparsegpt/tree/f5c25005a61f96a0933ca2f95705a963585aafaa
+class SparseGPT:
+
+    def __init__(self, layer):
+        self.layer = layer
+        self.dev = self.layer.weight.device
+        W = layer.weight.data.clone()
+        if isinstance(self.layer, nn.Conv2d):
+            W = W.flatten(1)
+        if isinstance(self.layer, transformers.Conv1D):
+            W = W.t()
+        self.rows = W.shape[0]
+        self.columns = W.shape[1]
+        self.H = torch.zeros((self.columns, self.columns), device=self.dev)
+        self.nsamples = 0
+
+    def add_batch(self, inp, out):
+        if len(inp.shape) == 2:
+            inp = inp.unsqueeze(0)
+        tmp = inp.shape[0]
+        if isinstance(self.layer, nn.Linear) or isinstance(
+            self.layer, transformers.Conv1D
+        ):
+            if len(inp.shape) == 3:
+                inp = inp.reshape((-1, inp.shape[-1]))
+            inp = inp.t()
+        self.H *= self.nsamples / (self.nsamples + tmp)
+        self.nsamples += tmp
+        inp = math.sqrt(2 / self.nsamples) * inp.float()
+        self.H += inp.matmul(inp.t())
+
+    def fasterprune(self, sparsity, prune_n=0, prune_m=0, blocksize=128, percdamp=0.01):
+        W = self.layer.weight.data.clone()
+        if isinstance(self.layer, nn.Conv2d):
+            W = W.flatten(1)
+        if isinstance(self.layer, transformers.Conv1D):
+            W = W.t()
+        W = W.float()
+
+        tick = time.time()
+
+        H = self.H
+        del self.H
+        dead = torch.diag(H) == 0
+        H[dead, dead] = 1
+        W[:, dead] = 0
+
+        Losses = torch.zeros(self.rows, device=self.dev)
+
+        damp = percdamp * torch.mean(torch.diag(H))
+        diag = torch.arange(self.columns, device=self.dev)
+        H[diag, diag] += damp
+        H = torch.linalg.cholesky(H)
+        H = torch.cholesky_inverse(H)
+        H = torch.linalg.cholesky(H, upper=True)
+        Hinv = H
+
+        mask = None
+
+        for i1 in range(0, self.columns, blocksize):
+            i2 = min(i1 + blocksize, self.columns)
+            count = i2 - i1
+
+            W1 = W[:, i1:i2].clone()
+            Q1 = torch.zeros_like(W1)
+            Err1 = torch.zeros_like(W1)
+            Losses1 = torch.zeros_like(W1)
+            Hinv1 = Hinv[i1:i2, i1:i2]
+
+            if prune_n == 0:
+                if mask is not None:
+                    mask1 = mask[:, i1:i2]
+                else:
+                    tmp = W1**2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
+                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
+                    mask1 = tmp <= thresh
+            else:
+                mask1 = torch.zeros_like(W1) == 1
+
+            for i in range(count):
+                w = W1[:, i]
+                d = Hinv1[i, i]
+
+                if prune_n != 0 and i % prune_m == 0:
+                    tmp = (
+                        W1[:, i : (i + prune_m)] ** 2
+                        / (torch.diag(Hinv1)[i : (i + prune_m)].reshape((1, -1))) ** 2
+                    )
+                    mask1.scatter_(
+                        1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True
+                    )
+
+                q = w.clone()
+                q[mask1[:, i]] = 0
+
+                Q1[:, i] = q
+                Losses1[:, i] = (w - q) ** 2 / d**2
+
+                err1 = (w - q) / d
+                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
+                Err1[:, i] = err1
+
+            W[:, i1:i2] = Q1
+            Losses += torch.sum(Losses1, 1) / 2
+
+            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])
+
+        torch.cuda.synchronize()
+        if isinstance(self.layer, transformers.Conv1D):
+            W = W.t()
+        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(
+            self.layer.weight.data.dtype
+        )
+
+    def free(self):
+        self.H = None
+        torch.cuda.empty_cache()

fusion_bench/method/pruning/wanda_utils/data.py

@@ -1,13 +1,15 @@
 # Code adapted from https://github.com/IST-DASLab/sparsegpt/blob/master/datautils.py
 
+import os
 import random
 from typing import List, Optional, Tuple, cast  # noqa: F401
 
-from datasets import load_dataset
 from torch import Tensor
 from tqdm.auto import tqdm
 from transformers import PreTrainedTokenizer
 
+from datasets import load_dataset
+
 
 # Wrapper for tokenized input IDs
 class TokenizerWrapper:
@@ -61,6 +63,7 @@ def get_c4(
     seqlen: int,
     tokenizer,
     data_path: str = "allenai/c4",
+    cache_dir: str = ".cache/allenai--c4",
 ) -> Tuple[List[Tuple[Tensor, Tensor]], TokenizerWrapper]:
     """
     Load and process the c4 dataset.
@@ -76,19 +79,35 @@ def get_c4(
         tuple (Tuple[List[Tuple[Tensor, Tensor]], TokenizerWrapper]): Tuple containing the training samples and the validation dataset.
     """
     # Load train and validation datasets
-    traindata = load_dataset(
-        data_path,
-        # "allenai--c4", # https://github.com/huggingface/datasets/issues/6559
-        data_files={"train": "en/c4-train.00000-of-01024.json.gz"},
-        split="train",
-    )
-    valdata = load_dataset(
-        data_path,
-        # "allenai--c4",
-        data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
-        split="validation",
-    )
-
+    if os.path.exists(f"{cache_dir}/en/c4-train.00000-of-01024.json.gz"):
+        traindata = load_dataset(
+            "json",
+            data_files={"train": f"{cache_dir}/en/c4-train.00000-of-01024.json.gz"},
+            split="train",
+        )
+    else:
+        traindata = load_dataset(
+            data_path,
+            # "allenai--c4", # https://github.com/huggingface/datasets/issues/6559
+            data_files={"train": "en/c4-train.00000-of-01024.json.gz"},
+            split="train",
+        )
+
+    if os.path.exists(f"{cache_dir}/en/c4-validation.00000-of-00008.json.gz"):
+        valdata = load_dataset(
+            "json",
+            data_files={
+                "validation": f"{cache_dir}/en/c4-validation.00000-of-00008.json.gz",
+            },
+            split="validation",
+        )
+    else:
+        valdata = load_dataset(
+            data_path,
+            # "allenai--c4",
+            data_files={"validation": "en/c4-validation.00000-of-00008.json.gz"},
+            split="validation",
+        )
     # Generate samples from training set
     if seed is not None:
         random.seed(seed)

fusion_bench/method/randes/__init__.py

@@ -0,0 +1,15 @@
+R"""
+RanDeS: Randomized Delta Superposition
+
+Implementation of "RanDeS: Randomized Delta Superposition for Multi-Model Compression"
+paper link: http://arxiv.org/abs/2505.11204
+
+Modified from https://github.com/Zhou-Hangyu/randes
+"""
+
+from .base_algorithm import SuperposedAlgorithmBase
+from .modelsoup import SuperposedModelSoupAlgorithm
+from .task_arithmetic import (
+    SuperposedTaskArithmeticAlgorithm,
+    SuperposedTaskArithmeticLoRAAlgorithm,
+)