SparseML Compression Pt 1: saving w/compression configs

[Satrat]

This initial PR for safetensors compression sets up the CompressionConfig and ModelCompressor registries and implements bitmask compression on save. See the corresponding internal docs PR for design details

To be implemented in follow-up PRs

• bitmask decompression
• saving/loading/inferring sparsity config from model
• SparseAutoModel load/save interface

Example

from sparseml.transformers import SparseAutoModelForCausalLM
from sparseml.transformers.compression import BitmaskConfig, BitmaskCompressor
from safetensors import safe_open
import os

MODEL_PATH = "zoo:llama2-7b-gsm8k_llama2_pretrain-pruned50.oneshot"
OUTPUT_PATH = "./test_compress_output"

model = SparseAutoModelForCausalLM.from_pretrained(MODEL_PATH)

sparsity_config = BitmaskConfig()
compressor = BitmaskCompressor(config=sparsity_config)

model_state_dict = model.state_dict()
sparse_state_dict = compressor.compress(model_state_dict)


model.save_pretrained(OUTPUT_PATH, safe_serialization=True, state_dict=sparse_state_dict)

safetensors_path = os.path.join(OUTPUT_PATH, "model-00001-of-00002.safetensors")
with safe_open(safetensors_path, framework="pt", device=0) as f:
    test_name = "model.layers.4.self_attn.k_proj.weight"
    bitmask = f.get_tensor(test_name + ".bitmask")
    shape = f.get_tensor(test_name + ".shape")
    values = f.get_tensor(test_name + ".compressed")
    row_offsets = f.get_tensor(test_name + ".row_offsets")
    print(f"bitmask: {bitmask}")
    print(f"shape: {shape}")
    print(f"values: {values}")
    print(f"row offsets: {row_offsets}")

[dbogunowicz]

Couple of nitpicks. All in all great job, definitely well done on a bite-sized scope of this PR

[mgoin]

Nice job Sara, this looks good to me for saving. Just some comments on state dict conflicts and tensor devices

[dbogunowicz]

🎅