Removed recipeStage, recipeModifer, stageModifier classes

[shanjiaz]

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:

• Add validations for recipe class.
• Simplify recipe.py structure

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Hello @shanjiaz, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello team, gemini-code-assist here to provide a summary of this pull request. This PR, authored by shanjiaz, is a work in progress aimed at simplifying the internal representation of recipes and modifiers. The core change involves removing the recipeStage, recipeModifier, and stageModifier classes and updating the codebase to use a flatter structure where the Recipe class directly holds a list of Modifier instances.

Highlights

• Recipe Structure Simplification: The hierarchical structure involving recipeStage, recipeModifier, and stageModifier has been removed. The Recipe class now directly manages a list of Modifier instances.
• Codebase Adaptation: Various parts of the codebase, including the lifecycle, session, pipelines, and transformers, have been updated to reflect the removal of the old stage/modifier classes and interact directly with the Recipe's list of Modifier objects.
• Recipe Parsing and Serialization: The Recipe class includes a new model validator (parse_from_dict) to maintain backward compatibility for loading recipes defined in the old staged YAML format. The serialization logic (yaml, dict) has been updated to handle the new internal structure, although the YAML output format attempts to retain some grouping structure.
• Recipe Combination Logic Removed: The simplify_combine_recipes method and related logic for combining multiple recipes have been removed from the Recipe class and the update_and_save_recipe utility.

Changelog

Click here to see the changelog

• src/llmcompressor/core/lifecycle.py
  • Updated imports and type hints to use Modifier instead of StageModifiers.
  • Changed the modifiers field type from List[StageModifiers] to List[Modifier].
• src/llmcompressor/core/session.py
  • Removed the get_modifiers method, which previously aggregated modifiers from StageModifiers.
• src/llmcompressor/entrypoints/oneshot.py
  • Updated access to modifiers from session.get_modifiers() to session.lifecycle.modifiers.
• src/llmcompressor/modifiers/init.py
  • Removed StageModifiers from the __all__ list.
• src/llmcompressor/pipelines/independent/pipeline.py
  • Removed import of StageModifiers.
  • Updated the __call__ method to iterate directly over session.lifecycle.modifiers instead of creating temporary StageModifiers instances.
• src/llmcompressor/pipelines/layer_sequential/pipeline.py
  • Updated access to modifiers from session.get_modifiers() to session.lifecycle.modifiers.
• src/llmcompressor/pipelines/registry.py
  • Added a blank line in from_modifiers.
• src/llmcompressor/pipelines/sequential/helpers.py
  • Removed a blank line in get_targets_from_modifiers.
• src/llmcompressor/pipelines/sequential/pipeline.py
  • Updated access to modifiers from session.get_modifiers() to session.lifecycle.modifiers.
• src/llmcompressor/recipe/init.py
  • Removed imports and __all__ entries for RecipeBase, RecipeModifier, and RecipeStage.
• src/llmcompressor/recipe/recipe.py
  • Removed imports for RecipeBase, RecipeModifier, RecipeStage. Added imports for BaseModel, ConfigDict, and ModifierFactory.
  • Changed Recipe to inherit from BaseModel.
  • Added explicit fields (version, args, stage, modifiers) and model_config.
  • Updated from_modifiers to directly assign modifiers and stage.
  • Added parse_from_dict model validator to handle parsing the old staged recipe format.
  • Updated simplify_recipe to filter modifiers based on target_stage.
  • Updated simplify_combine_recipes to extend the modifiers list and update args.
  • Updated create_modifier to instantiate Modifier objects from the stored list.
  • Removed remap_stages and extract_dict_stages methods.
  • Updated dict and yaml methods to use the new get_yaml_serializable_dict helper.
  • Removed the _get_yaml_dict method.
  • Renamed get_yaml_serializable_stage_dict to get_yaml_serializable_dict and updated its logic for the new structure.
  • Added a traceback.print_stack() call in the yaml method (likely for debugging).
• src/llmcompressor/transformers/compression/helpers.py
  • Added import for Modifier.
  • Updated infer_sparsity_structure_from_stage_modifiers to accept List[Modifier] and iterate directly.
• src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py
  • Simplified update_and_save_recipe to save the current recipe, potentially renaming the file if an existing recipe is found.
  • Removed the logic for combining recipes using simplify_combine_recipes.
• tests/llmcompressor/helpers.py
  • Removed a test recipe string that used multiple stages.
• tests/llmcompressor/recipe/test_recipe.py
  • Updated tests to reflect the removal of StageModifiers and the new structure.
  • Added create_modifier() calls in test_serialization.
• tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py
  • Commented out the shutil.rmtree call in tearDown.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

Stages removed, stages gone,
A flatter recipe carries on.
Modifiers now, a simple list,
Complexity, gently dismissed.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request focuses on refactoring the recipe handling mechanism by removing RecipeStage, RecipeModifier, and StageModifier classes. This simplification is a significant step towards a more streamlined recipe structure. The changes are generally well-implemented and align with the PR's objectives.

I've identified a few areas that need attention, primarily a debugging artifact, a potential unintended behavior change, and a docstring update. Given this is a Work In Progress (WIP), these points should help in finalizing the changes.

Overall, good progress on simplifying the recipe system!

Summary of Findings

• Debugging Artifact: A traceback.print_stack() call was found in src/llmcompressor/recipe/recipe.py, which should be removed.
• Docstring Example Update Needed: The docstring example for Recipe.create_modifier() in src/llmcompressor/recipe/recipe.py needs to be updated to reflect the new return type and structure after refactoring.
• Potential Behavior Change in Recipe Saving: The update_and_save_recipe function in src/llmcompressor/transformers/sparsification/compressed_tensors_utils.py no longer merges existing recipes with the current session's recipe. Instead, it saves the new recipe separately if an old one exists. This behavior change should be confirmed.
• Test Cleanup: The tearDown method in tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py has its directory cleanup logic commented out. This should be restored for proper test hygiene.
• Whitespace Changes: Minor whitespace changes (addition/removal of blank lines) were observed in several files. These are stylistic and do not affect functionality. (Not commented inline due to severity settings)
• Recipe Modifier State Management: In src/llmcompressor/recipe/recipe.py, the self.modifiers attribute can be either List[Dict[str, Any]] or List[Modifier]. The dict() and yaml() methods rely on get_yaml_serializable_dict which expects List[Modifier]. This implies create_modifier() must be called before serialization if the recipe was loaded from a file/string. This implicit dependency could be documented or the design made more robust. (Not commented inline due to severity settings)

Merge Readiness

This pull request is currently marked as Work In Progress (WIP), and there are a few critical and medium severity issues identified that should be addressed before considering it for merging. Specifically, the debugging traceback call needs to be removed, the docstring example updated, the recipe saving behavior confirmed, and test cleanup restored. Once these points are resolved, the PR will be in a much better state. As a reviewer, I am not authorized to approve pull requests, but I recommend addressing these findings.

[brian-dellabetta]

This is looking great!
🤩

Would just make sure Dipika is good with the changes, and we'll probably want to kick off a weekly/nightly test before merging when it's all ready

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[brian-dellabetta]

looks good but got a breakpoint() in code. also can you remove [WIP] from PR header if it's ready?

[brian-dellabetta]

I think this is looking good!

[kylesayrs]

Looks like all references to modifiers are caught, good job

grep -r '\.modifier' src/ tests/ --include="*.py" | grep -v 'from'

[kylesayrs]

Really nice job. My big question is whether we can fold Recipe.create_modifier into the Recipe construction step, since it seems like recipe creation is always followed by create_modifier. This would simplify the Recipe.modifier type hint too.

Other than that, just a few nits and questions. Awesome job!

[kylesayrs]

Did you mean to pass a recipe_stage argument to Recipe here?

Also, why use model_validate rather than create_instance?

[shanjiaz]

This is inside the create_instance, if I call create_instance here it'll end up in infinite recursive calls?

[kylesayrs]

I mean that the recipe which you have created does not know which stage it is

Maybe you want something like this

recipe_dict = _load_json_or_yaml_string(path_or_modifiers)
recipe_dict = filter_dict(obj, target_stage=target_stage)
Recipe.create_recipe(modifiers=recipe_dict, stage=target_stage)

[kylesayrs]

Same here, did you mean to pass stage to reipce creation and why not create using create_instance?

[shanjiaz]

same as above.

[kylesayrs]

If this isn't being used then we should just remove it

[kylesayrs]

nit: remove

[kylesayrs]

Why is this test removed?

[shanjiaz]

This is the only test that assumes recipe should work with multiple stages. I could rewrite the test, at the moment I thought removing the second stage is faster.

[kylesayrs]

Do you know what this assertion test is really doing? Do you think we can remove it?

[shanjiaz]

Yeah, this line assert len(modifiers := stage_modifiers[0].modifiers) == 1 is just unwrapping the stage_modifier structure, and since we don't have stage_modifier anymore, it's safe to remove.