Support for targets and ignore in Sparsity Compressors

[rahul-tuli]

This PR introduces support for using targets and ignore in sparsity compressors. It has been tested against the llm-compressor repository at commit a47137d8 (on main).

Changes Made

• Cleaned up several utilities and added corresponding tests.
• Updated the BaseSparsity.compress(...) methods to accept a new compression_targets argument.
• Enhanced the ModelCompressor to directly populate the compression_targets argument.

Verification

The functionality was verified using the following script:

Verification Script

"""
Usage: python verification.py
Tested against llm-compressor commit a47137d8
"""


from transformers import AutoTokenizer, AutoModelForCausalLM
from llmcompressor.transformers.compression.sparsity_config import SparsityConfigMetadata
from llmcompressor.transformers import oneshot
from safetensors import safe_open

MODEL_ID = "nm-testing/llama2.c-stories42M-pruned2.4"

def check_first_layer(save_dir, check_compressed=True):
    with safe_open(f"{save_dir}/model.safetensors", framework="pt", device=0) as f:
        layer_0_keys = [key for key in f.keys() if "model.layers.0" in key]
        if check_compressed:
            assert any("compressed" in key for key in layer_0_keys), "First layer is not compressed as expected."
        else:
            assert not any("compressed" in key for key in layer_0_keys), "First layer is compressed unexpectedly."

def main():
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

    # Apply oneshot to wrap save_pretrained 
    oneshot(model=model)

    # Compress and save the model
    sparsity_config = SparsityConfigMetadata.from_pretrained(model, compress=True)
    save_dir_compressed = f"{MODEL_ID.split('/')[1]}-2of4-compressed"
    model.save_pretrained(save_dir_compressed, sparsity_config=sparsity_config)
    tokenizer.save_pretrained(save_dir_compressed)

    # Verify the first layer is compressed
    check_first_layer(save_dir_compressed, check_compressed=True)

    # Ignore the first layer and save the model again
    sparsity_config.ignore.append("re:model.layers.0.*")
    save_dir_ignored = f"{MODEL_ID.split('/')[1]}-2of4-ignored-first-layer"
    model.save_pretrained(save_dir_ignored, sparsity_config=sparsity_config)
    tokenizer.save_pretrained(save_dir_ignored)

    # Verify the first layer is not compressed
    check_first_layer(save_dir_ignored, check_compressed=False)

if __name__ == "__main__":
    main()

The script passes successfully without any assertions.

Script Output

2024-11-27T10:18:45.295223+0000 | one_shot | INFO - *** One Shot ***
2024-11-27T10:18:45.295382+0000 | initialize | INFO - Compression lifecycle initialized for 0 modifiers
2024-11-27T10:18:45.295428+0000 | finalize | INFO - Compression lifecycle finalized for 0 modifiers
Calculating model sparsity: 100%|███████████████████████████████████████████████████████████████████████████████████████████| 75/75 [00:00<00:00, 1068.21it/s]
Checking whether model follows 2:4 sparsity structure: 100%|████████████████████████████████████████████████████████████████| 57/57 [00:00<00:00, 1572.75it/s]
Calculating model sparsity: 100%|███████████████████████████████████████████████████████████████████████████████████████████| 75/75 [00:00<00:00, 1562.54it/s]
Compressing model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████| 75/75 [00:00<00:00, 1205.22it/s]
2024-11-27T10:18:46.694477+0000 | get_serialized_recipe | WARNING - Recipe not found in session - it may have been reset
Calculating model sparsity: 100%|███████████████████████████████████████████████████████████████████████████████████████████| 75/75 [00:00<00:00, 1892.06it/s]
Compressing model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████| 75/75 [00:00<00:00, 1701.64it/s]
2024-11-27T10:18:47.582677+0000 | get_serialized_recipe | WARNING - Recipe not found in session - it may have been reset

[kylesayrs]

I'm not sure how/if this is related to #822 (it's listed as a dependency)

1. Doesn't this list of targets need to be accounted for during decompression?
2. Don't these changes throw away any weights which are not targeted for sparse compression?

[rahul-tuli]

I'm not sure how/if this is related to #822 (it's listed as a dependency)

1. Doesn't this list of targets need to be accounted for during decompression?
2. Don't these changes throw away any weights which are not targeted for sparse compression?

Point 1: Decompression takes care of that using COMPRESSION_PARAM_NAMES
Point 2: Fixed

It is listed as a dependency for #822 because without this we cannot enable sparse compression + quantization compression. These changes are needed for #822 to work fine.

[kylesayrs]

LGTM!

[dsikka]

nit: use a more descriptive variable name for other_params

[rahul-tuli]

Done!

[dsikka]

Can you also update this docstring to indicate what the output is expected to look like when return_other_params is True?

[rahul-tuli]

Updated!

[dsikka]

In what cases would we not have a match? Why do we need the other_params dictionary if the output from decompress is used to replace weights?

compressed-tensors/src/compressed_tensors/compressors/model_compressors/model_compressor.py

Line 306 in 525ef3a

self._replace_weights(dense_gen, model)

[rahul-tuli]

decompress has to return uncompressed_params from the safetensors file as well, because those also must be populated in the original potentially empty model. think of cases where we also have quantization in the mix, things like weight_scale must also be populated

We do not have a match for all the uncompressed_params, again things like scales

[horheynm]

len(".weight") can be hardcoded to the actual length. Should also add that it comes from len(".weight")