Refactor Initialized Model Export

[Satrat]

• Fixing bug where exporting a model initialized directly from the HuggingFace hub failed
• Moved config updates for kv-cache injection to the export function, and add support for kv-cache injection export to instantiated models
• Automatically infer transformers integration based on model properties

Example

from sparseml.transformers import SparseAutoModelForCausalLM, SparseAutoTokenizer
from sparseml.transformers import oneshot
from sparseml.modifiers import SparseGPTModifier
from sparseml import export

ORIG_MODEL_PATH = "roneneldan/TinyStories-1M"
OUTPUT_DIR = "./test_output_here"

model = SparseAutoModelForCausalLM.from_pretrained(ORIG_MODEL_PATH, device_map="auto")
tokenizer = SparseAutoTokenizer.from_pretrained(ORIG_MODEL_PATH)

oneshot(
    model=model,
    tokenizer=tokenizer,
    dataset="open_platypus",
    recipe=SparseGPTModifier(sparsity=0.5),
    output_dir="./test_export_output",
    num_calibration_samples=512,
)

export(
    model=model,
    tokenizer=tokenizer,
    target_path=OUTPUT_DIR
)

[dbogunowicz]

Left a few comments, all in all, great changes!

[dbogunowicz]

just curious, why is this change needed now?

[Satrat]

Honestly, I am not quite sure. I would occasionally get an error about deleting a file that didn't exist when exporting and this fixed it

[dbogunowicz]

why config_class is a deciding attribute here?

[Satrat]

This is really just to tell if its a PreTrainedModel, but I didn't want to have to add the transformers dependency to this part of the repo