[Performance] Add memory compression and decompression pathways

[kylesayrs]

Purpose

• Reduce memory requirements when compressing a model

Memory Visualization

Compression Memory Improvement

Format	State dict compression	Model compression
Quantized
Sparse
Sparse24 Stacked

Model Compression and Decompression

Format	Model Compression + Decompression
Quantized
Sparse
Sparse24 Stacked

Demonstration Script

import torch
from pttp import TensorProfiler
from transformers import AutoModelForCausalLM, AutoConfig
from compressed_tensors.compressors import ModelCompressor

for name, model_stub, comp_stub in [
    (
        "quantized-only",
        "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-uncompressed",
        "nm-testing/llama2.c-stories42M-gsm8k-quantized-only-compressed",
    ),
    (
        "sparse-only",
        "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-uncompressed",
        "nm-testing/llama2.c-stories42M-gsm8k-sparse-only-compressed",
    ),
    (
        "stacked",
        "nm-testing/llama2.c-stories42M-gsm8k-stacked-uncompressed",
        "nm-testing/llama2.c-stories42M-gsm8k-stacked-compressed",
    )
]:
    from transformers.utils.quantization_config import CompressedTensorsConfig

    with TensorProfiler() as prof:
        prof.mark_event("Start load")
        config = AutoConfig.from_pretrained(model_stub)
        config.tie_word_embeddings = False
        model = AutoModelForCausalLM.from_pretrained(model_stub, torch_dtype=torch.float32, device_map="cuda:0", config=config)
        compressor = ModelCompressor.from_pretrained(comp_stub)

        prof.mark_event("Start compress")
        compressor.compress_model(model)

        prof.mark_event("Start decompress")
        compressor.decompress_model(model)

    prof.save_memory_timeline(f"cdc_{name}.png")

Prerequisites

• Enable module state_dict compression, simplify compression logic #302
• [Accelerate] allow get_execution_device to be used when initializing a model #303

Changes

• Implement compress_model and decompress_model, which both act on a model in memory rather than a state dict or model on disk
  • Because compress_model compresses each module independently, implement show_progress on compress methods to squelch tqdm prints for each module compression
  • Implement decompress_from_state_dict for sparsity compressors
  • Extend get_nested_mappings_from_state_dict to support returning unmatched params, similar to get_nested_weight_mappings
    • I personally dislike this usage, as I think it leads to multiple sources of truth as to which modules should be compressed. IMO, a module should be (de)compressed if and only if it is listed in the config. This function is used to create another source of truth, which is that a module should be compressed if and only if it has the relevant compression params. Currently, we use both, which means taking the intersection.
• Misc
  • Fix bug on decompress_from_state_dict where scheme was gotten instead of weight args
  • Clarify name where variable weight_name was referring to a module path, not a weight name
  • Change sparse24 decompressor behavior where the weight was being moved to an arbitrary cuda device in fp8 cases. This violates the assumption that all ops are performed on the cpu
  • Remove remove_suffix util which can be replaced with str.removesuffix as of python3.9+ (which is the minimum we support, double check with @dsikka @rahul-tuli
  • Use get_execution_device when initing params for CompressedLinear

Testing

• Added test_compress_model which tests that memory compression is equivalent to dict compression
• Added test_decompress_model which tests that hfquantizer decompression (from disk) is equivalent to decompression from memory

[dsikka]

Is this ready for review? Still draft

[rahul-tuli]

Could we add a test to compress a model with sparsity+quantization?

[rahul-tuli]

LGTM pending conflict, good work!

[rahul-tuli]

Looks good pending verification that sparse only models can be compressed using these changes!

[dsikka]

Why do we need to use CompressedLinear for compression? What about if we’re compressing something that isn’t a linear layer?

[dsikka]

Looks good pending verification that sparse only models can be compressed using these changes!

Can you share sparse + fp8 models recipes where we have non-uniform sparsity and/or quantization cases?
@rahul-tuli

cc @kylesayrs

[brian-dellabetta]

It's beautiful Kyle 🥇 . Love the detailed summary and charts showing the improvement

[brian-dellabetta]

LGTM!

[rahul-tuli]

LGTM!