Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.

Conversation

rahul-tuli
Copy link
Member

@rahul-tuli rahul-tuli commented Apr 19, 2024

This PR updates sparseml pydantic support to v2

Steps:

  • use pydantic's official bump pydantic tool for pre-liminary migration
  • update setup.py
  • fix model configs and typing
  • update recipe serialization, also fixes a bug with original serialization

Test methodology:

  • make test targets=torch passes locally
  • sanity check with a dummy recipe containing multiple stages and multiple groups within each stage
  • test script attached with description

Sanity Script:

# local/feature/main.py


from sparseml.core import Recipe
import pprint as pp

import yaml

pp = pp.PrettyPrinter(indent=4)

recipe = """

first_oneshot_stage:
    pruning_modifiers:
        ConstantPruningModifier:
            start_epoch: 0
            end_epoch: 5
            targets: __ALL_PRUNABLE__
        
        MagnitudePruningModifier:
            start_epoch: 5
            end_epoch: 10
            init_sparsity: 0.1
            final_sparsity: 0.5
            targets: __ALL_PRUNABLE__
    
    quantization_modifiers:
        QuantizationModifier:
            start_epoch: 10
            end_epoch: 15
            bits: 8
            targets: __ALL_PRUNABLE__

second_oneshot_stage:
    pruning_modifiers:
        ConstantPruningModifier:
            start_epoch: 15
            end_epoch: 20
            targets: __ALL_PRUNABLE__
        
        MagnitudePruningModifier:
            start_epoch: 20
            end_epoch: 25
            init_sparsity: 0.1
            final_sparsity: 0.5
            targets: __ALL_PRUNABLE__
    
    quantization_modifiers:
        QuantizationModifier:
            start_epoch: 25
            end_epoch: 30
            bits: 8
            targets: __ALL_PRUNABLE__
"""

recipe = Recipe.create_instance(recipe)

yaml_str = recipe.yaml()
# print(yaml_str)
# print("++"* 20)

recipe_2 = Recipe.create_instance(yaml_str)
recipe_2_yaml = recipe_2.yaml()
# print(recipe_2_yaml)


assert yaml_str.strip() == recipe_2_yaml.strip()

Note: tests on GHA will fail until equivalent sparsezoo diff neuralmagic/sparsezoo#483 lands

Remove stale/misleading doctest
@rahul-tuli rahul-tuli changed the title [WIP] Update pydantic Update pydantic support to v2 Apr 23, 2024
@rahul-tuli rahul-tuli self-assigned this Apr 23, 2024
@bfineran bfineran merged commit 82e204a into main Apr 23, 2024
@bfineran bfineran deleted the update-pydantic branch April 23, 2024 20:15
@mgoin mgoin mentioned this pull request May 7, 2024
8 tasks
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants