Sparse masks

mttl/callbacks.py

@@ -24,6 +24,50 @@
 DEBUG = False
 
 
+class UpdateSparseMask(pl.Callback):
+    def __init__(self, update_interval=5, dm=None, save_mask_dir=None):
+        super().__init__()
+        self.update_interval = update_interval
+        self.update_counter = 0
+
+        #
+        self.dm = dm
+        self.save_mask_dir = save_mask_dir
+
+    def update_mask(self, pl_module, batch):
+
+        from mttl.models.modifiers.sparse_mask import (
+            make_sparse_model_during_training,
+            save_mask,
+        )
+
+        # make_sparse_model(pl_module, self.dm, keep_ratio=self.keep_ratio)
+        make_sparse_model_during_training(pl_module, batch)
+        # save mask
+        # f_name = f"{self.save_mask_dir}/mask"
+        # save_mask(pl_module, f_name)
+
+    def on_train_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        """
+        only updates the mask on epoch=0
+        """
+        if trainer.current_epoch == 0:
+            self.update_counter += 1
+            if (
+                self.update_counter % self.update_interval == 0
+                or self.update_counter == 1  # to set mask at the beginning
+            ):
+                # Update mask
+                self.update_mask(pl_module, batch)
+                self.update_counter = 0  # Reset counter for next interval
+
+
 class LiveCheckpointCallback(pl.Callback):
     """A better model checkpoint callback, that works in synchrony with LiveLogMixin."""
 

mttl/config.py

@@ -404,6 +404,12 @@ class TrainingArgs(DataArgs):
     pipeline_eval_tasks: str = None
     save_if_loaded_from_ckpt: bool = True
     dataset_type: str = None
+
+
+    keep_ratio: float = 0.05
+    BLOCK_SIZE: int = 16 # used for block-sparsity, decides the size of the block
+    sps_type: str = "block_sparse" # block_sparse,regular_sparse
+    sps_impl: str = "sp_add+sp_mm" #sp_add+sp_mm, scatter+filte
 
     @property
     def dataset_config(self):

mttl/models/modifiers/__init__.py

@@ -6,3 +6,4 @@
 import mttl.models.modifiers.lora  # noqa: F401
 import mttl.models.modifiers.mlp  # noqa: F401
 import mttl.models.modifiers.prompt_tuning  # noqa: F401
+import mttl.models.modifiers.sparse_mask  # noqa: F401