[Forkless SparseML Transformers] [Feature Branch] Setting Up The modification module

tests/sparseml/transformers/sparsification/modification/conftest.py

@@ -12,16 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import deepcopy
+
 import pytest
 from transformers import AutoConfig, AutoModel
 
 from accelerate import init_empty_weights
-from sparseml.transformers import (
-    SparseAutoConfig,
-    SparseAutoModel,
-    SparseAutoModelForCausalLM,
-)
-from sparsezoo import Model
+from sparseml.transformers import SparseAutoConfig, SparseAutoModelForCausalLM
+from sparseml.transformers.sparsification.modification import modify_model
 
 
 @pytest.fixture
@@ -55,36 +53,6 @@ def llama_zoo_model():
     return model
 
 
-@pytest.fixture
-def distilbert_zoo_model(tmp_path):
-    stub = "zoo:distilbert-squad_wikipedia_bookcorpus-pruned80.4block_quantized"
-    model_path = Model(stub, tmp_path).training.path
-    model = SparseAutoModel.question_answering_from_pretrained(
-        model_path, model_type="model"
-    )
-    return model
-
-
-@pytest.fixture
-def mobilebert_zoo_model(tmp_path):
-    stub = "zoo:mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block_quantized"
-    model_path = Model(stub, tmp_path).training.path
-    model = SparseAutoModel.question_answering_from_pretrained(
-        model_path, model_type="model"
-    )
-    return model
-
-
-@pytest.fixture
-def bert_zoo_model(tmp_path):
-    stub = "zoo:bert-base-squad_wikipedia_bookcorpus-pruned95.obs_quantized"
-    model_path = Model(stub, tmp_path).training.path
-    model = SparseAutoModel.question_answering_from_pretrained(
-        model_path, model_type="model"
-    )
-    return model
-
-
 @pytest.fixture
 def bert_model():
     config = AutoConfig.from_pretrained("bert-base-uncased")
@@ -133,3 +101,63 @@ def opt_model():
     with init_empty_weights():
         model = AutoModel.from_config(config)
     return model
+
+
+@pytest.fixture
+def helpers():
+    return Helpers
+
+
+class Helpers:
+    @staticmethod
+    def check_model_modified(
+        original_model_, module_to_replace, func_to_validate_replacement
+    ):
+        num_attn_blocks = original_model_.config.num_hidden_layers
+
+        original_model = deepcopy(original_model_)
+        modified_model = modify_model(original_model_)
+
+        modified_modules_original_model = [
+            module
+            for module in original_model.modules()
+            if func_to_validate_replacement(module)
+            and isinstance(module, module_to_replace)
+        ]
+
+        modified_modules_modified_model = [
+            module
+            for module in modified_model.modules()
+            if func_to_validate_replacement(module)
+            and isinstance(module, module_to_replace)
+        ]
+
+        original_modules_original_model = [
+            module
+            for module in original_model.modules()
+            if not func_to_validate_replacement(module)
+            and isinstance(module, module_to_replace)
+        ]
+
+        original_modules_modified_model = [
+            module
+            for module in modified_model.modules()
+            if not func_to_validate_replacement(module)
+            and isinstance(module, module_to_replace)
+        ]
+
+        # make sure that the original model has no modified modules
+        # and that the modified model has no original modules
+        assert (
+            len(modified_modules_original_model)
+            == len(original_modules_modified_model)
+            == 0
+        )
+        # make sure that the original model has N original modules
+        # and that the modified model has N modified modules
+        # where N is the number of transformer's attention blocks
+        assert (
+            len(modified_modules_modified_model)
+            == len(original_modules_original_model)
+            == num_attn_blocks
+        )

tests/sparseml/transformers/sparsification/modification/test_modify_model.py

@@ -33,7 +33,6 @@ def __init__(self):
 
 
 def test_modify_model_without_actual_modification(model):
-
     is_modified = copy(model.modified)
     model = modify_model(model)
     assert model.modified == is_modified == False  # noqa E712

tests/sparseml/transformers/sparsification/modification/test_modifying_bert.py

@@ -12,70 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from copy import deepcopy
 
-from sparseml.transformers.sparsification.modification import modify_model
+from transformers.models.bert.modeling_bert import BertSelfAttention
 
 
-def test_modifying_bert(bert_model):
+def test_modifying_bert(bert_model, helpers):
     from sparseml.transformers.sparsification.modification.modifying_bert import (  # noqa F401
         modify,
     )
 
-    num_attn_blocks = bert_model.config.num_hidden_layers
-
-    # keep the original model for comparison
-    bert_ = deepcopy(bert_model)
-    bert = modify_model(bert_model)
-
-    # check how many modified "BertSelfAttention" modules are in the original
-    # model (should be 0, as the model is not modified yet)
-    modified_modules_original_model = [
-        module
-        for module in bert_.modules()
-        if _is_bert_attention_modified(module)
-        and module.__class__.__name__ == "BertSelfAttention"
-    ]
-    # check how many modified "BertSelfAttention" modules are
-    # in the modified model (should be num_attn_blocks, as the
-    # model is modified, and has num_attn_blocks attention blocks)
-    modified_modules_modified_model = [
-        module
-        for module in bert.modules()
-        if _is_bert_attention_modified(module)
-        and module.__class__.__name__ == "BertSelfAttention"
-    ]
-    # check how many original "BertSelfAttention"
-    # modules are in the original
-    # model (should be num_attn_blocks, as the model is
-    # not modified yet, and has num_attn_blocks attention blocks)
-    original_modules_original_model = [
-        module
-        for module in bert_.modules()
-        if not _is_bert_attention_modified(module)
-        and module.__class__.__name__ == "BertSelfAttention"
-    ]
-    # check how many original "BertSelfAttention"
-    # modules are in the modified
-    # model (should be 0, as the model is
-    # modified, and should not contain any original
-    # "BertSelfAttention" modules)
-    original_modules_modified_model = [
-        module
-        for module in bert.modules()
-        if not _is_bert_attention_modified(module)
-        and module.__class__.__name__ == "BertSelfAttention"
-    ]
-
-    assert (
-        len(modified_modules_original_model)
-        == len(original_modules_modified_model)
-        == 0
-    )
-    assert (
-        len(modified_modules_modified_model)
-        == len(original_modules_original_model)
-        == num_attn_blocks
+    helpers.check_model_modified(
+        bert_model,
+        module_to_replace=BertSelfAttention,
+        func_to_validate_replacement=_is_bert_attention_modified,
     )
 
 

tests/sparseml/transformers/sparsification/modification/test_modifying_distillbert.py

@@ -12,105 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from copy import deepcopy
+from transformers.models.distilbert.modeling_distilbert import MultiHeadSelfAttention
 
-import pytest
-from transformers import AutoConfig, AutoModel
 
-from accelerate import init_empty_weights
-from sparseml.pytorch.optim.manager import ScheduledModifierManager
-from sparseml.transformers.sparsification.modification import modify_model
-
-
-@pytest.fixture
-def distilbert_recipe():
-    return """version: 1.1.0
-stage_test:
-  stage_test_modifiers:
-      - !QuantizationModifier
-          exclude_module_types: ['QATMatMul']"""
-
-
-@pytest.fixture
-def distilbert_model():
-    config = AutoConfig.from_pretrained("distilbert/distilbert-base-uncased")
-    with init_empty_weights():
-        model = AutoModel.from_config(config)
-    return model
-
-
-def test_modifying_distilbert(distilbert_model):
+def test_modifying_distilbert(distilbert_model, helpers):
     from sparseml.transformers.sparsification.modification.modifying_distilbert import (  # noqa F401
         modify,
     )
 
-    num_attn_blocks = distilbert_model.config.num_hidden_layers
-
-    # keep the original model for comparison
-    distilbert_ = deepcopy(distilbert_model)
-    distilbert = modify_model(distilbert_model)
-
-    # check how many modified "MultiHeadSelfAttention" modules are in the original
-    # model (should be 0, as the model is not modified yet)
-    modified_modules_original_model = [
-        module
-        for module in distilbert_.modules()
-        if _is_distilbert_attention_modified(module)
-        and module.__class__.__name__ == "MultiHeadSelfAttention"
-    ]
-    # check how many modified "MultiHeadSelfAttention" modules are
-    # in the modified model (should be num_attn_blocks, as the
-    # model is modified, and has num_attn_blocks attention blocks)
-    modified_modules_modified_model = [
-        module
-        for module in distilbert.modules()
-        if _is_distilbert_attention_modified(module)
-        and module.__class__.__name__ == "MultiHeadSelfAttention"
-    ]
-    # check how many original "MultiHeadSelfAttention"
-    # modules are in the original
-    # model (should be num_attn_blocks, as the model is
-    # not modified yet, and has num_attn_blocks attention blocks)
-    original_modules_original_model = [
-        module
-        for module in distilbert_.modules()
-        if not _is_distilbert_attention_modified(module)
-        and module.__class__.__name__ == "MultiHeadSelfAttention"
-    ]
-    # check how many original "MultiHeadSelfAttention"
-    # modules are in the modified
-    # model (should be 0, as the model is
-    # modified, and should not contain any original
-    # "MultiHeadSelfAttention" modules)
-    original_modules_modified_model = [
-        module
-        for module in distilbert.modules()
-        if not _is_distilbert_attention_modified(module)
-        and module.__class__.__name__ == "MultiHeadSelfAttention"
-    ]
-
-    assert (
-        len(modified_modules_original_model)
-        == len(original_modules_modified_model)
-        == 0
-    )
-    assert (
-        len(modified_modules_modified_model)
-        == len(original_modules_original_model)
-        == num_attn_blocks
-    )
-
-
-def test_apply_recipe(distilbert_recipe, distilbert_zoo_model):
-    from sparseml.transformers.sparsification.modification.modifying_distilbert import (  # noqa F401
-        modify,
+    helpers.check_model_modified(
+        distilbert_model,
+        module_to_replace=MultiHeadSelfAttention,
+        func_to_validate_replacement=_is_distilbert_attention_modified,
     )
 
-    manager = ScheduledModifierManager.from_yaml(distilbert_recipe)
-    distilbert_zoo_model.train()
-    manager.apply_structure(distilbert_zoo_model)
-    assert True
-
 
 def _is_distilbert_attention_modified(module):
     # only the modified "MultiHeadSelfAttention" modules have the