Merge pull request #108 from microsoft/sparse_masks

Sparse masks
microsoft · Sep 19, 2024 · 0163d2f · 0163d2f
2 parents 36b612b + 0bc1662
commit 0163d2f
Show file tree

Hide file tree

Showing 12 changed files with 2,716 additions and 3 deletions.
diff --git a/mttl/arguments.py b/mttl/arguments.py
@@ -404,6 +404,8 @@ class TrainingArgs(DataArgs):
     save_if_loaded_from_ckpt: bool = True
     dataset_type: str = None
 
+    profile: bool = False  # if 'True' will profile the model training
+
     @property
     def dataset_config(self):
         if self.dataset_type is not None:

diff --git a/mttl/models/lightning/expert_module.py b/mttl/models/lightning/expert_module.py
@@ -139,7 +139,9 @@ def __init__(self, model_object=None, **kwargs):
     def on_save_checkpoint(self, ckpt):
         super().on_save_checkpoint(ckpt)
 
-        ckpt["expert_info"] = self.as_expert(self.training_config).expert_info.asdict()
+        ckpt["expert_info"] = self.model.as_expert(
+            self.training_config
+        ).expert_info.asdict()
 
 
 class MultiExpertModule(LightningTrainingMixin, LightningEfficientCheckpoint):
@@ -172,6 +174,9 @@ def __init__(
     def add_expert_instance(self, *args, **kwargs):
         return self.model.add_expert_instance(*args, **kwargs)
 
+    def add_experts_from_library(self, *args, **kwargs):
+        return self.model.add_experts_from_library(*args, **kwargs)
+
     def on_save_checkpoint(self, ckpt):
         super().on_save_checkpoint(ckpt)
 

diff --git a/mttl/models/modifiers/__init__.py b/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