Removed recipeStage, recipeModifer, stageModifier classes (#1514)

SUMMARY:
Removed recipeStage, recipeModifer, stageModifier classes and
dependencies.


TEST PLAN:
Tested locally

DETAILS & CONCERNS:
1. `recipe.py` is where most changes happened. 
- Rewrote the parsing logic in `parse_from_dict`. 
- `Recipe` class now is basically just a list of `Modifiers`, all other
changes are downstream dependencies.
- All the serilization helpers and deserialization helpers now live in
`utils`. Had to rewrite serialization logic.
- Stage is stored as a string by `self.stage`.
2. Had to update a test in `test_recipe`. We support creating a recipe
instance by reading a multi-stage recipe but no longer support
serializing multi-stage recipes directly. The current structure assumes
one recipe could only have one stage.
3. Had to rewrite `update_and_save_recipe` function since now each
recipe only has one stage and therefore can't handle merging multi-stage
recipes. The new setup:
- `infer_recipe_from_model_path` will find existing recipe path. 
- `yaml` method now takes in an optional argument `existing_recipe_path`
and merge the existing recipe with the current recipe. Since there would
only be at most one existing recipe at a time, we should be safe.

TODO:
- [x]  Add validations for recipe class. 
- [x] Simplify `recipe.py` structure

---------

Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>
Co-authored-by: Kyle Sayers <kylesayrs@gmail.com>

Author: shanjiaz

src/llmcompressor/core/lifecycle.py

@@ -12,7 +12,6 @@
 
 from llmcompressor.core.events import Event, EventType
 from llmcompressor.core.state import State
-from llmcompressor.modifiers import StageModifiers
 from llmcompressor.recipe import Recipe, RecipeArgsInput, RecipeInput, RecipeStageInput
 
 __all__ = ["CompressionLifecycle"]
@@ -33,7 +32,6 @@ class CompressionLifecycle:
 
     state: State = field(default_factory=State)
     recipe: Recipe = field(default_factory=Recipe)
-    modifiers: List[StageModifiers] = field(default_factory=list)
 
     initialized_: bool = False
     finalized: bool = False
@@ -60,7 +58,7 @@ def reset(self):
         """
         logger.debug("Resetting compression lifecycle")
 
-        for mod in self.modifiers:
+        for mod in self.recipe.modifiers:
             if not mod.initialized or mod.finalized:
                 continue
             try:
@@ -92,21 +90,26 @@ def initialize(
             return
 
         logger.debug("Initializing compression lifecycle")
-        self.recipe = Recipe.simplify_recipe(
-            recipe=recipe, target_stage=recipe_stage, override_args=recipe_args
-        )
-        self.modifiers = self.recipe.create_modifier()
+        if not recipe:
+            self.recipe = Recipe()
+        else:
+            self.recipe = Recipe.create_instance(
+                path_or_modifiers=recipe, target_stage=recipe_stage
+            )
+            if recipe_args:
+                self.recipe.args = {**recipe_args}
 
         mod_data = []
-        for mod in self.modifiers:
+        for mod in self.recipe.modifiers:
             data = mod.initialize(state=self.state, **kwargs)
             logger.debug("Initialized modifier: {}", mod)
             if data is not None:
                 mod_data.append(data)
 
         self.initialized_ = True
         logger.info(
-            "Compression lifecycle initialized for {} modifiers", len(self.modifiers)
+            "Compression lifecycle initialized for {} modifiers",
+            len(self.recipe.modifiers),
         )
 
         return mod_data
@@ -130,7 +133,7 @@ def finalize(self, **kwargs) -> List[Any]:
 
         logger.debug("Finalizing compression lifecycle")
         mod_data = []
-        for mod in self.modifiers:
+        for mod in self.recipe.modifiers:
             data = mod.finalize(state=self.state, **kwargs)
             logger.debug("Finalized modifier: {}", mod)
             if data is not None:
@@ -139,7 +142,8 @@ def finalize(self, **kwargs) -> List[Any]:
         self.finalized = True
 
         logger.info(
-            "Compression lifecycle finalized for {} modifiers", len(self.modifiers)
+            "Compression lifecycle finalized for {} modifiers",
+            len(self.recipe.modifiers),
         )
 
         return mod_data
@@ -196,7 +200,7 @@ def event(
 
         event = Event(type_=event_type)
         mod_data = []
-        for mod in self.modifiers:
+        for mod in self.recipe.modifiers:
             data = mod.update_event(state=self.state, event=event, **kwargs)
             logger.debug("Updated event with modifier: {}", mod)
             if data is not None:

src/llmcompressor/core/session.py

@@ -220,17 +220,6 @@ def get_serialized_recipe(self) -> Optional[str]:
 
         logger.warning("Recipe not found in session - it may have been reset")
 
-    def get_modifiers(self):
-        """
-        Get all modifiers across all stages
-        """
-        stage_modifiers = self.lifecycle.modifiers
-        return [
-            modifier
-            for stage_modifier in stage_modifiers
-            for modifier in stage_modifier.modifiers
-        ]  # noqa: E127
-
     def _log_model_info(self):
         # Log model level logs if cadence reached
         current_index = self._lifecycle.global_step

src/llmcompressor/entrypoints/oneshot.py

@@ -186,9 +186,8 @@ def apply_recipe_modifiers(
             recipe_args=self.recipe_args.recipe_args,
             calib_data=calibration_dataloader,
         )
-
         user_pipeline = self.dataset_args.pipeline
-        modifiers = session.get_modifiers()
+        modifiers = session.lifecycle.recipe.modifiers
         pipeline = CalibrationPipeline.from_modifiers(modifiers, user=user_pipeline)
         pipeline(self.model, calibration_dataloader, self.dataset_args)
 

src/llmcompressor/modifiers/__init__.py

@@ -1,11 +1,9 @@
 from .factory import ModifierFactory
 from .interface import ModifierInterface
 from .modifier import Modifier
-from .stage import StageModifiers
 
 __all__ = [
     "ModifierFactory",
     "ModifierInterface",
     "Modifier",
-    "StageModifiers",
 ]

src/llmcompressor/modifiers/stage.py

@@ -5,7 +5,6 @@
 from torch.utils.data.dataloader import DataLoader
 
 from llmcompressor.core import active_session
-from llmcompressor.modifiers.stage import StageModifiers
 from llmcompressor.pipelines.registry import CalibrationPipeline
 from llmcompressor.utils.helpers import patch_attr
 
@@ -34,18 +33,15 @@ def __call__(
         _logger = logger.patch(lambda r: r.update(function="IndependentPipeline"))
 
         session = active_session()
-        modifiers = session.get_modifiers()
-        with patch_attr(session.lifecycle, "modifiers", None):
-            for index, modifier in enumerate(modifiers):
-                mod_type = str(type(modifier).__name__)
-                session.lifecycle.modifiers = [
-                    StageModifiers(modifiers=[modifier], group=mod_type, index=index)
-                ]
-
+        modifiers = session.lifecycle.recipe.modifiers
+        with patch_attr(session.lifecycle.recipe, "modifiers", None):
+            for modifier in modifiers:
+                mod_type = type(modifier).__name__
+                session.lifecycle.recipe.modifiers = [modifier]
                 pipeline = CalibrationPipeline.from_modifiers([modifier])
                 pipeline_name = pipeline.__class__.__name__
                 _logger.info(f"Inferred `{pipeline_name}` for `{mod_type}`")
 
                 pipeline(model, dataloader, dataset_args)
 
-            # restore modifiers on exit so model can be compressed based on recipe
+                # restore modifiers on exit so model can be compressed based on recipe

src/llmcompressor/pipelines/independent/pipeline.py

@@ -63,7 +63,7 @@ def __call__(
         model_device = get_execution_device(model)
 
         # find layers
-        modifiers = session.get_modifiers()
+        modifiers = session.lifecycle.recipe.modifiers
         sequential_targets = get_sequential_targets(modifiers, model, dataset_args)
         layers = match_modules(model, sequential_targets)
 

src/llmcompressor/pipelines/layer_sequential/pipeline.py

@@ -57,8 +57,9 @@ def __call__(
         model_device = get_execution_device(model)
 
         # prepare to trace subgraphs
-        modifiers = session.get_modifiers()
+        modifiers = session.lifecycle.recipe.modifiers
         sequential_targets = get_sequential_targets(modifiers, model, dataset_args)
+
         ignore = dataset_args.tracing_ignore
 
         # trace subgraphs

src/llmcompressor/pipelines/sequential/pipeline.py

@@ -1,17 +1,11 @@
-from .base import RecipeBase
 from .metadata import DatasetMetaData, LayerMetaData, ModelMetaData, ParamMetaData
-from .modifier import RecipeModifier
 from .recipe import Recipe, RecipeArgsInput, RecipeInput, RecipeStageInput
-from .stage import RecipeStage
 
 __all__ = [
     "DatasetMetaData",
     "ParamMetaData",
     "LayerMetaData",
     "ModelMetaData",
-    "RecipeBase",
-    "RecipeModifier",
-    "RecipeStage",
     "Recipe",
     "RecipeInput",
     "RecipeStageInput",