Use model compression pathways

[kylesayrs]

Purpose

• Use in-memory model compression pathway in order to reduce memory requirements when saving models
• These changes along with postprocessing changes move users towards a pattern where they are aware of the status of the model (frozen/compressed) and call save_pretrained manually

Prerequisites

• [Tests] Use proper offloading utils in test_compress_tensor_utils #1449

Changes

• Modify save_pretrained_wrapper to use compress_model(model) rather than compress(state_dict)
• Modify save_pretrained_wrapper so that the state dict is only retrieved if not skipping compression stats
• Modify save_pretrained_wrapper to save dictionary and python files, even if there is no explicit compressor
• Modify save_checkpoint (used by training) to decompress after the checkpoint is saved

Example/Testing Changes

As far as I can tell, below lists all of the instances where a model undergoes saving (no immediately followed by script exit)

File Path	Solution
examples/trl_mixin/ex_trl_constant.py test_oneshot_and_finetune.py	Decompress in between stages
examples/quantization_2of4_sparse_w4a16/llama7b_sparse_w4a16.py test_oneshot_and_finetune_with_tokenizer.py	Do not save in between stages to avoid compressed state
test_oneshot_then_finetune.py	No work is required, as model is decompressed upon loading from disk
test_compress_tensor_utils.py	Fix test to use dispatch_model (which is actually used by transformers) rather than cpu_offload

Testing

State Dict	In Memory

oneshot_save.py

import torch
from transformers import AutoModelForCausalLM
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from pttp import TensorProfiler

#MODEL_ID = "DeepSeek-V3_local_bf16"
MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"

with TensorProfiler() as prof:
    prof.mark_event("Load model")
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)

    prof.mark_event("Oneshot")
    oneshot(
        model=model,
        recipe=QuantizationModifier(targets="Linear", scheme="W4A16"),
        trust_remote_code_model=True,
    )

    prof.mark_event("Save model")
    model.save_pretrained("sav_testing", save_compressed=True, skip_compression_stats=True)

prof.save_memory_timeline("save_timeline.png")

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[brian-dellabetta]

exciting!