Sparse MoE

mttl/models/modifiers/sm_updater.py

@@ -11,9 +11,9 @@
 
 from mttl.logging import logger
 from mttl.models.modifiers.base import Modifier
-from mttl.models.modifiers.sm_config import SparseMaskConfig
-from mttl.models.modifiers.sparse_utils.sparse_linear import MaskedLinear, SparseLinear
-from mttl.models.modifiers.sparse_utils.utils import (
+from mttl.models.modifiers.sparse_mask_config import SparseMaskConfig
+from mttl.models.modifiers.sparsity.sparse_linear import MaskedLinear, SparseLinear
+from mttl.models.modifiers.sparsity.sparse_utils.utils import (
     get_2d_indices_from_csr_matrix,
     get_top_k_sparcity,
     scipy_csr_to_torch_csr,

mttl/models/modifiers/sparse_mask.py

@@ -3,28 +3,19 @@
 from dataclasses import dataclass
 from typing import Union
 
-import numpy as np
-import torch
-from scipy.sparse import csr_matrix
 from torch import nn
 from triton.ops.blocksparse.matmul import dsd_lut, sdd_lut
 
 from mttl.logging import logger
 from mttl.models.modifiers.base import Modifier, ModifierConfig
-from mttl.models.modifiers.sm_config import SparseMaskConfig
-from mttl.models.modifiers.sm_updater import MaskUpdater
-from mttl.models.modifiers.sparse_utils.sparse_linear import (
+from mttl.models.modifiers.sparse_mask_config import SparseMaskConfig
+from mttl.models.modifiers.sparsity.mask_updater import MaskUpdater
+from mttl.models.modifiers.sparsity.sparse_linear import (
     MaskedLinear,
     ScatteredSparseLinearModule,
     SparseLinear,
     SparseLinearConfig,
 )
-from mttl.models.modifiers.sparse_utils.utils import (
-    get_2d_indices_from_csr_matrix,
-    get_top_k_sparcity,
-    scipy_csr_to_torch_csr,
-    torch_csr_to_scipy_csr,
-)
 
 
 class SparseMaskAdapter(Modifier):
@@ -55,7 +46,7 @@ def __init__(
 
     def forward(self, input):
         if self.maks_update_mode and self.training:
-            return self.mask_updater(self.sparse_layer, input)
+            return self.mask_updater(input)
         return self.sparse_layer(input)
 
     def prepare_for_mask_update(self):

mttl/models/modifiers/sm_config.py → mttl/models/modifiers/sparse_mask_config.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from mttl.models.modifiers.sparse_utils.sparse_linear import SparseLinearConfig
+from mttl.models.modifiers.sparsity.sparse_linear import SparseLinearConfig
 
 
 @dataclass