[Cherry-Pick] Fix GHA transformers errors (#2176)

* [Export API Refactor][Fix] Varying sample and model input names on correctness validation (#2131)

* initial commit

* Apply suggestions from code review

* quality

* solving the actual problem

* Fix GHA transformer errors (#2175)

* initial commit

* add test_helpers

* revert

* fix lm_head edge case

* Remove leftover print

---------

Co-authored-by: Sara Adkins <sara@neuralmagic.com>

---------

Co-authored-by: Sara Adkins <sara@neuralmagic.com>

Author: dbogunowicz

src/sparseml/modifiers/utils/layer_compressor.py

@@ -122,7 +122,12 @@ def revert_layer_wrappers(self):
         Reverts wrapped root modules back to their original structure
         """
         for name, module_wrapper in self.modules.items():
-            set_layer(name, module_wrapper.layer, self.layer)
+            full_name = self._get_full_submodule_name(name)
+            if len(name) == 0:  # special case if layer has no children (i.e. lm_head)
+                with summon_full_params_context(self.model):
+                    set_layer(full_name, module_wrapper.layer, self.model)
+            else:
+                set_layer(name, module_wrapper.layer, self.layer)
             module_wrapper.free()
         self.modules = None
 

tests/sparseml/transformers/utils/test_helpers.py

@@ -11,15 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import os
-from collections import OrderedDict
+import shutil
 
 import pytest
 import torch
-import transformers
+from transformers import AutoConfig, AutoModelForCausalLM
 
-from huggingface_hub import snapshot_download
+from accelerate import init_empty_weights
 from sparseml.transformers.utils.helpers import (
     create_fake_dataloader,
     infer_recipe_from_model_path,
@@ -32,84 +31,49 @@
 
 
 @pytest.fixture()
-def generative_model_path(tmp_path):
-    return snapshot_download("roneneldan/TinyStories-1M", local_dir=tmp_path)
+def generative_model():
+    return "roneneldan/TinyStories-1M"
 
 
 @pytest.fixture()
-def model_path(tmp_path):
-    return Model(
-        "zoo:mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block_quantized",
-        tmp_path,
-    ).training.path
+def bert_model():
+    return "zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/pruned95_obs_quant-none"  # noqa E501
 
 
 @pytest.fixture()
 def sequence_length():
-    return 384
+    return 320
 
 
-@pytest.fixture()
-def dummy_inputs():
-    input_ids = torch.zeros((1, 32), dtype=torch.int64)
-    attention_mask = torch.ones((1, 32), dtype=torch.int64)
+def test_create_fake_dataloader(generative_model, sequence_length):
+    config = AutoConfig.from_pretrained(generative_model)
+    tokenizer = initialize_tokenizer(
+        generative_model, sequence_length=sequence_length, task="text-generation"
+    )
+    with init_empty_weights():
+        model = AutoModelForCausalLM.from_config(config)
 
-    return OrderedDict(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
+    expected_input_names = ["input_ids", "attention_mask"]
+    num_samples = 2
+    data_loader, input_names = create_fake_dataloader(
+        model=model,
+        tokenizer=tokenizer,
+        num_samples=num_samples,
     )
 
+    assert input_names == expected_input_names
+    for i, sample in enumerate(data_loader):
+        assert sample["input_ids"].shape == torch.Size([1, sequence_length])
+        assert sample["attention_mask"].shape == torch.Size([1, sequence_length])
+        assert set(sample.keys()) == set(expected_input_names)
+    assert i == num_samples - 1
 
-@pytest.mark.parametrize(
-    "stub",
-    [
-        "zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/pruned95_obs_quant-none",  # noqa E501
-    ],
-)
-def test_is_transformer_model(tmp_path, stub):
-    zoo_model = Model(stub, tmp_path)
+
+def test_is_transformer_model(tmp_path, bert_model):
+    zoo_model = Model(bert_model, tmp_path)
     source_path = zoo_model.training.path
     assert is_transformer_model(source_path)
-
-
-@pytest.mark.parametrize(
-    "stub",
-    [
-        "zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/pruned95_obs_quant-none",  # noqa E501
-    ],
-)
-def test_save_zoo_directory(stub, tmp_path_factory):
-    path_to_training_outputs = tmp_path_factory.mktemp("outputs")
-    save_dir = tmp_path_factory.mktemp("save_dir")
-
-    zoo_model = Model(stub, path_to_training_outputs)
-    zoo_model.download()
-
-    save_zoo_directory(
-        output_dir=save_dir,
-        training_outputs_dir=path_to_training_outputs,
-    )
-    new_zoo_model = Model(str(save_dir))
-    assert new_zoo_model.validate(minimal_validation=True, validate_onnxruntime=False)
-
-
-@pytest.mark.parametrize(
-    "model_path, recipe_found",
-    [
-        ("roneneldan/TinyStories-1M", False),
-        ("mgoin/all-MiniLM-L6-v2-quant-ds", True),
-        (
-            "zoo:mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block_quantized",  # noqa E501
-            True,
-        ),
-    ],
-)
-def test_infer_recipe_from_model_path(model_path, recipe_found):
-    recipe = infer_recipe_from_model_path(model_path)
-    if recipe_found:
-        assert isinstance(recipe, str)
-        return
-    assert recipe is None
+    shutil.rmtree(tmp_path)
 
 
 def test_infer_recipe_from_local_model_path(tmp_path):
@@ -124,6 +88,16 @@ def test_infer_recipe_from_local_model_path(tmp_path):
     assert recipe == recipe_path.as_posix()
 
 
+@pytest.fixture(autouse=True)
+def model_path_and_recipe_path(tmp_path):
+    model_path = tmp_path / "model.onnx"
+    recipe_path = tmp_path / "recipe.yaml"
+    recipe_path.touch()
+    model_path.touch()
+
+    return model_path, recipe_path
+
+
 @pytest.mark.parametrize(
     "model_path",
     [
@@ -140,16 +114,6 @@ def test_resolve_recipe_file(model_path, model_path_and_recipe_path):
     )
 
 
-@pytest.fixture()
-def model_path_and_recipe_path(tmp_path):
-    model_path = tmp_path / "model.onnx"
-    recipe_path = tmp_path / "recipe.yaml"
-    recipe_path.touch()
-    model_path.touch()
-
-    return model_path, recipe_path
-
-
 def test_resolve_recipe_file_from_local_path(model_path_and_recipe_path):
     model_path, recipe_path = model_path_and_recipe_path
     assert recipe_path.as_posix() == resolve_recipe_file(
@@ -165,24 +129,40 @@ def test_resolve_recipe_file_from_local_path(model_path_and_recipe_path):
     )
 
 
-def test_create_fake_dataloader(generative_model_path, sequence_length):
-    expected_input_names = ["input_ids", "attention_mask"]
-    sequence_length = 32
-    num_samples = 2
+@pytest.mark.parametrize(
+    "model, recipe_found",
+    [
+        ("roneneldan/TinyStories-1M", False),
+        ("mgoin/all-MiniLM-L6-v2-quant-ds", True),
+        (
+            "zoo:mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block_quantized",  # noqa E501
+            True,
+        ),
+    ],
+)
+def test_infer_recipe_from_model_path(model, recipe_found):
+    recipe = infer_recipe_from_model_path(model)
+    if recipe_found:
+        assert isinstance(recipe, str)
+        return
+    assert recipe is None
 
-    model = transformers.AutoModelForCausalLM.from_pretrained(generative_model_path)
-    tokenizer = initialize_tokenizer(
-        generative_model_path, sequence_length=sequence_length, task="text-generation"
-    )
-    data_loader, input_names = create_fake_dataloader(
-        model=model,
-        tokenizer=tokenizer,
-        num_samples=num_samples,
-    )
 
-    assert input_names == expected_input_names
-    for i, sample in enumerate(data_loader):
-        assert sample["input_ids"].shape == torch.Size([1, sequence_length])
-        assert sample["attention_mask"].shape == torch.Size([1, sequence_length])
-        assert set(sample.keys()) == set(expected_input_names)
-    assert i == num_samples - 1
+@pytest.mark.parametrize(
+    "stub",
+    [
+        "zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/pruned95_obs_quant-none",  # noqa E501
+    ],
+)
+def test_save_zoo_directory(tmp_path, stub):
+    path_to_training_outputs = Model(stub).path
+    save_dir = tmp_path
+
+    save_zoo_directory(
+        output_dir=save_dir,
+        training_outputs_dir=path_to_training_outputs,
+    )
+    zoo_model = Model(str(save_dir))
+    assert zoo_model.validate(minimal_validation=True, validate_onnxruntime=False)
+    shutil.rmtree(path_to_training_outputs)
+    shutil.rmtree(save_dir)