Enable Sparse compression

[rahul-tuli]

This PR makes compressors composable. After this change, both sparse and quantization compression can be applied to a sparse quantized model.

Notable Changes:

• Reworked the logic for inferring sparse compressor format:
  • Previously, if a model was quantized, the sparse compressor was always set to the Identity (dense) compressor. Now, an appropriate sparse compressor is selected (currently only sparse_bitmask or sparse_24).
  • A special condition for Marlin-style compression was added (Sparse compressor should be set to Dense if marlin_24 compressor is used).
  • Updated the SPARSITY_THRESHOLD to 50% as discussed offline

Dependencies:

• Add sparsity structure enum neuralmagic/compressed-tensors#197
• neuralmagic/compressed-tensors#182
• vllm-project/llm-compressor#821

Test Script

from transformers import AutoTokenizer

from llmcompressor.modifiers.quantization import QuantizationModifier
from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot

MODEL_ID = "nm-testing/llama2.c-stories110M-pruned50-compressed-tensors"

# Load model
model = SparseAutoModelForCausalLM.from_pretrained(
    MODEL_ID, device_map="auto", torch_dtype="auto"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

recipe = """
quant_stage:
    quant_modifiers:
        QuantizationModifier:
            ignore: ["lm_head"]
            config_groups:
                group_0:
                    weights:
                        num_bits: 4
                        type: int
                        strategy: tensor
                        dynamic: false
                        symmetric: true
                    input_activations:
                        num_bits: 4
                        type: int
                        strategy: tensor
                        dynamic: true
                        symmetric: true
                    targets: ["Linear"]
    pruning_modifiers:
        ConstantPruningModifier:
            targets: [
                're:.*q_proj.weight',
                're:.*k_proj.weight', 
                're:.*v_proj.weight',
                're:.*o_proj.weight',
                're:.*gate_proj.weight',
                're:.*up_proj.weight',
                're:.*down_proj.weight',
            ]
            start: 0
"""

# Apply quantization.
oneshot(model=model, recipe=recipe)

# Confirm generations of the quantized model look sane.
print("========== SAMPLE GENERATION ==============")
input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))
print("==========================================")

# Save to disk in compressed-tensors format.
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8"
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

[kylesayrs]

These changes highlight how tied sparsity and quantization are (quantization format now depends on sparsity structure, and sparsity config now depends on quantization format).

I think this is the proper way to do it (first infer format, then use format to generate configs), but I hope in the future that we can structure the classes in a way such that the flow is clearer and users have an easier time understanding and using the tools

[dsikka]

We should add a couple of e2e tests to run sample generation in llm-compressor.
We can then eventually expand on them to also run in vllm (once we've integrated).

[kylesayrs]

To make this more generic, why not pass in the quantization config? We will for sure have different compression formats which affect sparsity in the future

[dsikka]

Yeah we should make this generic as marlin will likely not be the only case

[kylesayrs]

Overall structure is good (having quantization config be dependent on sparsity structure is better than it being dependent on a sparsity config)

I definitely recommend genericising the is_marlin a little bit by passing the compression format directly so it's easier to support more formats in the future

[dsikka]

Just summarizing offline discussions:

Although marlin_24 should no longer be the default, we still need to make sure we have a pathway to enable it, such as through the addition of a "format" argument in the recipe.