SparseML Compression Pt 2: Load compressed weights

[Satrat]

This PR implements ModelCompressor.decompress(), which will decompress the weights in the safetensors file one by one. Also includes a bunch of helper functions for reading safetensors files and dealing with the compressed format. See the corresponding internal docs PR for design details

Note: #2177 needs to be merged first

To be implemented in follow-up PR

• inferring sparsity config from model
• SparseAutoModel save/load interface

Example

Sample code for compressing a model with 50% sparsity(See PR #2177), then reloading the compressed weights as a dense model

from sparseml.transformers import SparseAutoModelForCausalLM
from sparseml.transformers.compression import BitmaskConfig, BitmaskCompressor
from sparseml.utils.pytorch.utils import measure_cuda_memory
from tqdm import tqdm
import torch

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

torch.cuda.set_device(0)
with measure_cuda_memory() as m:
    model = SparseAutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="cuda:0")
print(f"Load dense model peak GPU {m.overall_peak_memory / float(2**30):.4f} GB")

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

# compresses the model using Bitmask compression
with measure_cuda_memory() as m:
    model_state_dict = model.state_dict()
    sparse_state_dict = compressor.compress(model_state_dict)

    # save the compressed model
    model.save_pretrained(
        OUTPUT_PATH, 
        safe_serialization=True, 
        state_dict=sparse_state_dict
    )

print(f"Save compressed model peak GPU {m.overall_peak_memory / float(2**30):.4f} GB")

# use the dense state dict to reload the model
torch.cuda.set_device(1)
with measure_cuda_memory() as m:
    model_again = SparseAutoModelForCausalLM.from_pretrained(
        OUTPUT_PATH, 
        device_map="cuda:1"
    )

    #returns iterator
    dense_state_dict = compressor.decompress(OUTPUT_PATH)
    for name, data in tqdm(dense_state_dict, desc="Decompressing model"):
        BitmaskCompressor.replace_layer(name, data, model_again)

print(f"Load compressed model peak GPU {m.overall_peak_memory / float(2**30):.4f} GB")

Load dense model peak GPU 25.2276 GB
Compressing model: 100%|████████████████████████████████████████████████████████████████████████████████████████| 291/291 [01:28<00:00, 3.29it/s]
Save compressed model peak GPU 25.2276 GB
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 6.27it/s]
Decompressing model: 291it [01:11, 4.08it/s]
Load compressed model peak GPU 25.7159 GB

The base branch was changed.

[mgoin]

LGTM thanks, just one line I think was missed

[mgoin]

thanks!