Quantization Compressor Support

[Satrat]

Requires this compressed-tensors branch: neuralmagic/compressed-tensors#45

• Adds support for saving compressed quantized models within SparseAutoModel saving. Compression type can be passed in via quantization_format or inferred from the model itself
• Simplified a lot of the save/load logic by moving it to helper classes in compressed-tensors

Examples

Very little UX change, similar to sparsity we just pass save_compressed=True to enable compression. By default, we save weights in the fake_quant format is save_compressed isn't set.

from sparseml.transformers import SparseAutoModelForCausalLM, oneshot

recipe="tests/sparseml/transformers/compression/recipes/new_quant_full.yaml"
model_stub = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
dataset = "open_platypus"
max_seq_length = 512
num_calibration_samples = 512
output_dir = "./test_updated_llama1.1b_quant_compressed"

model = SparseAutoModelForCausalLM.from_pretrained(model_stub, device_map="cuda:0")
oneshot(
    model=model,
    dataset=dataset,
    overwrite_output_dir=True,
    output_dir=output_dir,
    max_seq_length=max_seq_length,
    num_calibration_samples=num_calibration_samples,
    recipe=recipe,
    pad_to_max_length=False,
    save_compressed=True
)

Reloading a fake_quant model and then compressing it:

from sparseml.transformers import SparseAutoModelForCausalLM

output_dir_fake = "./test_updated_llama1.1b_quant"
output_dir_compressed = "./test_updated_llama1.1b_quant_compressed"

model_reloaded = SparseAutoModelForCausalLM.from_pretrained(output_dir_fake)
model_reloaded.save_pretrained(output_dir_compressed, save_compressed=True)

You can also specify a quantization compression format by name. Right now we only have support for unpacked int quantization, but as we add additional compression formats for quantization this becomes more relevant

from sparseml.transformers import SparseAutoModelForCausalLM

output_dir_fake = "./test_updated_llama1.1b_quant"
output_dir_compressed = "./test_updated_llama1.1b_quant_compressed"

model_reloaded = SparseAutoModelForCausalLM.from_pretrained(output_dir_fake)
model_reloaded.save_pretrained(output_dir_compressed, quantization_format="int_quantized")

[dbogunowicz]

🥇

The base branch was changed.

[rahul-tuli]

LGTM after proposed fix