[Export Refactor][Transformers] Enable loading SparseModels

[dbogunowicz]

Feature Description

The main functionality introduced by this PR enables loading and applying the recipe in the new format to the text-generation models.

The notable changes are:

src/sparseml/export/export.py -> allowing the user to provide recipe explicitely to the exporter. If the recipe is not provided, but the recipe.yaml is located side-by-side with the model, there will be an attempt to load it automatically

src/sparseml/pytorch/model_load/helpers.py -> moving the log_model_load out of the transformers.SparseModel, otherwise we would get nasty import dependencies. Plus, it just make sense for it to be more general function.

src/sparseml/transformers/integration_helper_functions.py -> simplifying the create_model function and adapting it to the new changes

src/sparseml/transformers/utils/helpers.py -> just adding an extra transformer-specific wrapper around the general apply_recipe_structure_to_model method. Just in case we need some transformer-specific functionalities.

src/sparseml/transformers/utils/initializers.py -> refactoring initialize_model, so now it becomes a comprehensive, main entry point for loading a SparseModel.

src/sparseml/transformers/utils/load_task_model.py -> look above, this function is an extension of initialize_model (now initialize_sparse_model)

Testing

Tests in tests/sparseml/export/transformers/test_generative_transformers.py demonstrate the current capabilities of the new export module in the context of LLMs.

[Satrat]

nit: helper fn?

[bfineran]

+1, this should ideally be resolved at the registry mixin level (all lower case, remove special characters and whitespace)

[dbogunowicz]

Please take a look at: neuralmagic/sparsezoo#404
It's been sitting for a while in the "purgatory", but has already been reviewed once by Sara.
@Satrat , @bfineran if you approved this sparsezoo feature, we could remove the need for unifiying the integration string (RegistryMixin will take care of it). For the task itself, I have also a good solution, will implement in a second.

[Satrat]

lets clean up this logic to something like this:

if default_recipe_exists and recipe_provided:
     warn that we're overwriting the model's input recipe, probably a bad idea
      use recipe_provided
elif recipe_provided:
    warn that this recipe should have been previously applied to the model
    use recipe_provided
elif default_recipe:
    just the info log needed
    use default_recipe
else: # neither recipe exists
    warn that no recipe was applied
     carry on

[Satrat]

It should be safe to change this to the automodel, I'm just not sure if these torch.nn..... = skip lines can be deleted or not

[dbogunowicz]

@alexm-nm can we have your buy-in here?

[Satrat]

Could we make eval vs train an argument? It depends on if we're doing training(train) or oneshot/export(eval)

[dbogunowicz]

Can we just set this appropriately in the code paths themselves? I feel like adding a boolean argument would be a bit weird. It's just more verbose and informative to call model.eval() or model.train() explicitly in the code.

[Satrat]

I think we can remove this and just use the SparseAutoModel?

[Satrat]

Was there a reason not to replace this with the SparseAutoModel? Other than the OPT edge case, I think we only need to add in precision support

[Satrat]

Oh one other issue I see: the new text_generation.py v2 finetuning script uses the pretrained_from_distil function, that needs to be updated

[bfineran]

could you update the PR description to include an example of initialize_sparse_model

[bfineran]

+1, this should ideally be resolved at the registry mixin level (all lower case, remove special characters and whitespace)

[bfineran]

should be debug for any non critical logs on our regular path

[bfineran]

let's include a usage example here

[bfineran]

LGTM - is the intention to hand off initialize_sparse_model to MLR instead of the previously discussed SparseAutoModel?

[dbogunowicz]

Yes

[Satrat]

LGTM, had one nitpick on the obcq script. Should we update the OBCQ and finetuning scripts to use the new initialize_sparse_model funcion? That can also be a later PR though

[Satrat]

Now that we're using SparseAutoModel for everything this warning message doesn't make much sense. The only thing we lose from not using a supported model is perplexity evaluation, and thats more of a debugging feature anyways. Can we change the message to this:

f"A supported model type({SUPPORTED_MODELS}) could not be "
f"parsed from model_path={model_path}. Perplexity evaluation will only work for supported models"

[Satrat]

I think we can get rid of the if statement altogether, and just condition on the model for setting forward_fn

[Satrat]

Nice, the logging messages are really clear, and good to have this in a helper function