Support for Decompressing Models from HF Hub

[Satrat]

The original compression PR neglected to test loading a compressed model from the hub, and was only working with local compression loads. This PR adds support for loading a compressed model from the hub with SparseAutoModelForCausalLM.from_pretrained.

Test

Loads the compressed model from the hub, checks the decompressed weights match the original

from sparseml.transformers import SparseAutoModelForCausalLM, SparseAutoTokenizer
import torch

ORIG_MODEL_PATH = "neuralmagic/TinyLlama-1.1B-Chat-v1.0-pruned2.4"
COMPRESSED_MODEL_PATH = "mgoin/TinyLlama-1.1B-Chat-v1.0-pruned2.4-compressed"

# Compress and export the model
model = SparseAutoModelForCausalLM.from_pretrained(ORIG_MODEL_PATH, device_map="auto", torch_dtype="auto")
tokenizer = SparseAutoTokenizer.from_pretrained(ORIG_MODEL_PATH)

model_compressed = SparseAutoModelForCausalLM.from_pretrained(COMPRESSED_MODEL_PATH, device_map="auto", torch_dtype="auto")

og_state_dict = model.state_dict()
reconstructed_state_dict = model_compressed.state_dict()
assert(len(og_state_dict) == len(reconstructed_state_dict))
for key in og_state_dict.keys():
    dense_tensor = og_state_dict[key]
    reconstructed_tensor = reconstructed_state_dict[key]
    assert torch.equal(dense_tensor.cpu(), reconstructed_tensor.cpu())