Merge branch 'main' into prod_smooth_quant

Author: Sara Adkins

.gitignore

@@ -795,3 +795,6 @@ fabric.properties
 *.resources
 test-results/
 integrations/pytorch/pytorch_vision*
+
+# local log files
+nm_temp_test_logs/*

README.md

@@ -128,7 +128,7 @@ More information on installation such as optional dependencies and requirements
 
 ### Recipes
 
-To enable flexibility, ease of use, and repeatability, SparseML uses a declarative interface called `recipes` for specifying the sparsity-related algorithms and hyperparamters that should be applied by SparseML.
+To enable flexibility, ease of use, and repeatability, SparseML uses a declarative interface called `recipes` for specifying the sparsity-related algorithms and hyperparameters that should be applied by SparseML.
 
 `Recipes` are YAML-files formatted as a list of `modifiers`, which encode the instructions for SparseML. Example `modifiers` can be anything from setting the learning rate to encoding the hyperparameters of the gradual magnitude pruning algorithm. The SparseML system parses the `recipes` into a native format for each framework and applies the modifications to the model and training pipeline.
 

setup.py

@@ -79,7 +79,7 @@
 _transformers_deps = _pytorch_deps + [
     f"{'nm-transformers' if is_release else 'nm-transformers-nightly'}"
     f"~={version_nm_deps}",
-    "datasets<=2.11",
+    "datasets<=2.14.6",
     "scikit-learn",
     "seqeval",
     "einops",

src/sparseml/core/lifecycle/event.py

@@ -198,7 +198,7 @@ def optim_pre_step_events(self) -> List[Event]:
             and self.type_ is not None
             and self.type_ != EventType.OPTIM_POST_STEP
         ):
-            raise ValueError("optim pre step must be called after optim post step")
+            raise ValueError("optim pre step must be called before optim post step")
 
         if (
             self.type_first == EventType.LOSS_CALCULATED

src/sparseml/core/lifecycle/session.py

@@ -50,6 +50,9 @@ def reset(self):
             except Exception:
                 pass
 
+        if self.state and self.state.data:
+            # reset data if it exists
+            self.state.data.reset()
         self.state = None
         self.recipe_container = RecipeContainer()
         self.modifiers = []

src/sparseml/core/model/base.py

@@ -126,3 +126,11 @@ def get_matching_layer(
         :param model: model to search for targets
         """
         raise NotImplementedError()
+
+    def qat_active(self) -> bool:
+        """
+        Checks if quantization aware training is set up in the model
+
+        :return: True if QAT is active in any layer, False otherwise
+        """
+        raise NotImplementedError()

src/sparseml/core/model/pytorch.py

@@ -25,6 +25,7 @@
     get_matching_layer,
     get_param,
     get_params,
+    qat_active,
     set_layer,
     set_param,
 )
@@ -105,3 +106,11 @@ def get_matching_layer(
         :param model: model to search for targets
         """
         return get_matching_layer(target, name_to_match, model)
+    
+    def qat_active(self) -> bool:
+        """
+        Checks if quantization aware training is set up in the model
+
+        :return: True if QAT is active in any layer, False otherwise
+        """
+        return qat_active(self.model)

src/sparseml/core/recipe/modifier.py

@@ -113,4 +113,4 @@ def dict(self, *args, **kwargs) -> Dict[str, Any]:
         """
         :return: the dictionary representation of the modifier
         """
-        return {self.type: self.args}
+        return {self.type: self.args, "group": f"{self.group}_modifiers"}

src/sparseml/core/recipe/recipe.py

@@ -399,18 +399,20 @@ def dict(self, *args, **kwargs) -> Dict[str, Any]:
         ...             targets: ['re:.*weight']
         ... '''
         >>> recipe = Recipe.create_instance(recipe_str)
-        >>> recipe.dict()
-        Traceback (most recent call last):
-        ...
-        KeyError: 'group'
+        >>> recipe_dict = recipe.dict()
+        >>> stage = recipe_dict["stages"]["test"]
+        >>> pruning_mods = stage[0]['modifiers']['pruning']
+        >>> modifier_args = pruning_mods[0]['ConstantPruningModifier']
+        >>> modifier_args == {'start': 0.0, 'end': 2.0, 'targets': ['re:.*weight']}
+        True
 
         :return: A dictionary representation of the recipe
         """
         dict_ = super().dict(*args, **kwargs)
         stages = {}
 
         for stage in dict_["stages"]:
-            name = stage["group"]
+            name = f"{stage['group']}_stage"
             del stage["group"]
 
             if name not in stages:
@@ -422,6 +424,58 @@ def dict(self, *args, **kwargs) -> Dict[str, Any]:
 
         return dict_
 
+    def yaml(self, file_path: Optional[str] = None) -> str:
+        """
+        Return a yaml string representation of the recipe.
+
+        :param file_path: optional file path to save yaml to
+        :return: The yaml string representation of the recipe
+        """
+        file_stream = None if file_path is None else open(file_path, "w")
+        yaml_dict = self._get_yaml_dict()
+
+        ret = yaml.dump(
+            yaml_dict, stream=file_stream, allow_unicode=True, sort_keys=False
+        )
+
+        if file_stream is not None:
+            file_stream.close()
+
+        return ret
+
+    def _get_yaml_dict(self) -> Dict[str, Any]:
+        """
+        Get a dictionary representation of the recipe for yaml serialization
+        The returned dict will only contain information necessary for yaml
+        serialization (ignores metadata, version, etc), and must not be used
+        in place of the dict method
+
+        :return: A dictionary representation of the recipe for yaml serialization
+        """
+
+        def _modifier_group_to_dict(modifier_group: List[Dict[str, Any]]):
+            # convert a list of modifiers to a dict of modifiers
+            return {
+                key: value
+                for modifier in modifier_group
+                for key, value in modifier.items()
+            }
+
+        def _stage_to_dict(stage: List[Dict[str, Any]]):
+            # convert a list of stages to a dict of modifiers
+            return {
+                modifier_group_name: _modifier_group_to_dict(modifier_group)
+                for stage_modifiers in stage
+                for modifier_group_name, modifier_group in stage_modifiers[
+                    "modifiers"
+                ].items()
+            }
+
+        return {
+            stage_name: _stage_to_dict(stage=stage)
+            for stage_name, stage in self.dict()["stages"].items()
+        }
+
 
 @dataclass
 class RecipeTuple:

src/sparseml/core/state.py

@@ -45,6 +45,14 @@ class Data:
     test: Optional[ModifiableData] = None
     calib: Optional[ModifiableData] = None
 
+    def reset(self):
+        """
+        Reset self to initial state
+        """
+        attribs = Data().__dict__
+        for attrib_name, attrib_value in attribs.items():
+            setattr(self, attrib_name, attrib_value)
+
 
 @dataclass
 class Hardware: