DOCSEARCHbyAlgolia
Submit Docs

compressed-tensors

repository·main·Indexed 43 days ago

https://github.com/vllm-project/compressed-tensors

A library extending the safetensors format to provide a unified and extensible way to store and manage compressed tensor data. It supports various quantization schemes (GPTQ, AWQ, SmoothQuant, MXFP4), sparsity patterns, and transforms for LLM deployment, integrating with vLLM, Hugging Face, PyTorch, and SGLang. It includes an offload module for transparently moving weights between CPU, GPU, and disk to handle models larger than available VRAM.

Tokens
15.7K
Snippets
50
Records
67
Agent score
53%

What's inside compressed-tensors

  1. What is compressed-tensors?

    main

    The compressed-tensors library is an extension of the safetensors format designed to provide a unified, extensible way to store and manage compressed tensor data. It addresses the fragmentation in the model compression landscape by providing a single format that supports various quantization schemes (GPTQ, AWQ, SmoothQuant, etc.), sparsity patterns (unstructured and 2:4), and transforms (Hadamard, etc.).

    Key capabilities include:

    • Unified Checkpoint Format: Consistent representation for diverse compression schemes.
    • Flexible Quantization: Supports activation quantization, mixed precision, low/arbitrary-bit, and KV cache quantization.
    • Sparsity & Transforms: Handles semi-structured sparsity and rotation-based quantization.
    • Model Offloading: Transparent CPU/disk/distributed offloading for large models.
    • Integration: Works with Hugging Face, PyTorch, vLLM, and SGLang.

  2. How the Offload Module works

    main

    The compressed_tensors.offload module enables running models larger than available GPU memory by replacing standard _parameters and _buffers with an OffloadCache.

    Core Workflow:

    1. Module Wrapping: The offload_module function wraps a module and replaces its parameters with an OffloadCache implementation (e.g., CPUCache, DiskCache, or DistributedCPUCache).
    2. Forward Hook: The module's .forward() method is wrapped to automatically handle the movement of inputs and weights.
    3. Onload/Offload Lifecycle: When a layer is needed for computation, the weights are moved from the offload device (e.g., CPU) to the execution device (e.g., GPU). After computation, they can be moved back to save memory.
    4. Caching: The OffloadCache implementations optimize access. For example, a CPUCache ensures that while the first access triggers a CPU→GPU copy, subsequent reads hit the cache to minimize latency.

  3. How OffloadCache implementations work

    main

    The OffloadCache is an abstract MutableMapping used to manage tensors that are not resident on the primary device. Different implementations serve different hardware constraints:

    • CPUCache: Offloads tensors to CPU RAM. Use this when GPU VRAM is insufficient but CPU RAM is available. Onloading is a standard .to(device) call.
    • DeviceCache: Offloads tensors to a CUDA device. Use this for consistent management when a module is already resident on a device.
    • DiskCache: Offloads tensors to disk as safetensors files. In-memory, offloaded tensors are represented as meta tensors. Use this when neither CPU nor GPU memory can hold the weights. Requires weights to be in safetensors format.
    • Distributed Variants:
      • DistributedCPUCache: Uses POSIX shared memory (/dev/shm) so all ranks share the same physical CPU memory, avoiding redundant copies.
      • DistributedDeviceCache: Replicates tensors across all ranks on their local GPUs.
      • DistributedDiskCache: All ranks share the same disk files.
    class OffloadCache(MutableMapping, ABC):
    onload_device: torch.device | str
    offload_device: torch.device | Literal["disk"]

  4. How the `compressed_tensors` Offload Module works

    main

    The compressed_tensors.offload module enables inference for models larger than available GPU memory by transparently moving weights between CPU, GPU, and disk.

    It achieves this through two primary mechanisms:

    1. OffloadCache Substitution: The module replaces a torch.nn.Module's _parameters and _buffers dictionaries with OffloadCache instances. This OffloadCache is a MutableMapping that intercepts access:
      • On write (__setitem__): The tensor is immediately moved to its storage device (e.g., CPU or disk).
      • On read (__getitem__): The tensor is automatically onloaded to the execution device before being returned.
    2. Forward Wrapper: The offload_module function wraps the module's forward method. This wrapper automatically moves all input tensors to the specified onload_device before execution, ensuring activations and weights reside on the same device without requiring manual device management in the model code.

  5. Manage the Onload/Offload Lifecycle

    main

    The offloading lifecycle follows these steps during a forward pass:

    1. Input tensors are received.
    2. The forward wrapper moves inputs to the onload_device.
    3. PyTorch accesses parameters via _parameters[name], triggering OffloadCache.__getitem__ which calls onload() to move weights to the device.
    4. Computation runs on the onload_device.
    5. After the forward pass, onloaded weights are released (garbage collected) to free memory, unless disable_offloading is active.

  6. Install compressed-tensors

    main

    You can install the library via PyPI for stable or nightly builds, or build directly from the source repository.

    Stable Release

    pip install compressed-tensors

    Nightly Release

    pip install --pre compressed-tensors

    From Source

    git clone https://github.com/vllm-project/compressed-tensors
cd compressed-tensors
pip install -e .
    pip install compressed-tensors

  7. Load an offloaded model with load_offloaded_model

    main

    The recommended way to load offloaded models from pretrained checkpoints. This context manager patches from_pretrained on transformers model classes to handle distributed loading and weight sharing across ranks.

    Special device_map value: Use device_map="auto_offload" to restrict accelerate to only use CPU and disk, ensuring GPU VRAM is fully available for activations during inference.

    Parameters:

    • model_class: The specific class to patch (e.g., AutoModelForCausalLM). Recommended for efficiency.
    • extra_cpu_mem: Bytes to reserve in CPU RAM for non-weight operations (default 5 GB).
    from transformers import AutoModelForCausalLM

with load_offloaded_model():
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3-8B",
        device_map="auto_offload",  # CT-specific: only cpu/disk, no GPU
    )

  8. Use disable_offloading and disable_onloading context managers

    main

    You can control how the offload module behaves using two global context managers:

    • disable_offloading(): Prevents repeated CPU→GPU copies. Onloaded tensors are cached in memory; subsequent reads become cache hits. Use this when processing multiple tokens or batches to avoid redundant device movement.
    • disable_onloading(): Returns the raw offloaded tensors directly without triggering device movement. Use this for debugging or inspecting the raw storage (e.g., seeing the CPU tensor directly).
    # With disable_offloading: 1 copy, 2 cache hits
with disable_offloading():
    for _ in range(3):
        result = module.weight

# With disable_onloading: returns the CPU tensor directly
with disable_onloading():
    cpu_tensor = module.weight

  9. Compress a model to MXFP4

    main

    This guide demonstrates how to load a model (e.g., Llama 3 8B), apply round-to-nearest (RTN) MXFP4 weight quantization, compute scales directly from weights (no calibration data required), and save the compressed model.

    Workflow:

    1. Load Model: Use AutoModelForCausalLM.
    2. Configure Quantization: Use QuantizationConfig to define the scheme (e.g., MXFP4) and target layers (e.g., Linear), while specifying layers to ignore (e.g., lm_head).
    3. Apply Config: Call apply_quantization_config(model, config) to attach a QuantizationScheme to targeted modules.
    4. Compute Scales: Iterate through modules, extract weights, and use calculate_qparams to find scales based on min/max values (optionally using group_size). Update the module using update_offload_parameter.
    5. Compress and Save: Use ModelCompressor.from_pretrained_model(model) to initialize a compressor, call compressor.compress_model(model), save the model, and finally call compressor.update_config(output_dir) to update the config.json with the necessary quantization metadata.
    model_name = "meta-llama/Meta-Llama-3-8B"
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Load the model
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map=device, torch_dtype="auto"
)

# Set-up the quantization config.
config = QuantizationConfig(
    config_groups={"MXFP4": ["Linear"]},
    ignore=["lm_head"],
)
apply_quantization_config(model, config)

# Compute weight scales using round-to-nearest quantization
for name, module in model.named_modules():
    scheme = getattr(module, "quantization_scheme", None)
    if scheme is None or scheme.weights is None:
        continue

    weight = module.weight.data
    args = scheme.weights
    group_size = args.group_size

    if group_size is not None and group_size > 0:
        reshaped = weight.unflatten(-1, (math.ceil(weight.shape[-1] / group_size), group_size))
        min_val = reshaped.amin(dim=-1)
        max_val = reshaped.amax(dim=-1)
    else:
        min_val, max_val = torch.aminmax(weight)

    scale, _ = calculate_qparams(min_val, max_val, args)
    update_offload_parameter(module, "weight_scale", scale)

output_dir = "./Meta-Llama-3-8B-MXFP4"
compressor = ModelCompressor.from_pretrained_model(model)
compressor.compress_model(model)
model.save_pretrained(output_dir)
compressor.update_config(output_dir)

  10. Convert between compressed_tensors and accelerate offloading

    main

    Because the transformers ecosystem expects HuggingFace accelerate hooks, you may need to convert models:

    • to_accelerate(model): Converts a compressed_tensors-offloaded model back to accelerate offloading. Required before calling model.save_pretrained().
    • from_accelerate(model): Converts a model from accelerate offloading to compressed_tensors offloading. This is called automatically by load_offloaded_model.
    # Before saving:
to_accelerate(model)
model.save_pretrained("./output")

# After saving, convert back for inference:
from_accelerate(model)

  11. Compress models using compressed-tensors primitives

    main

    You can use the compressed-tensors primitives directly to compress models. This approach involves:

    1. Applying quantization configurations.
    2. Computing scales using PyTorch hooks or manual calculation.
    3. Saving the resulting models in the compressed-tensors format.

    Note: The direct primitive approach in these examples uses basic round-to-nearest quantization and does not include advanced calibration algorithms.

  12. Initialize distributed environment with init_dist

    main

    To use distributed cache variants (like DistributedCPUCache), initialize the distributed process group using init_dist(). This uses NCCL and environment variables set by torchrun. It requires the TORCHELASTIC_RUN_ID environment variable.

    Usage:

    1. Run with torchrun: torchrun --nproc-per-node=4 my_script.py
    2. Call init_dist() at the top of your script.
    # At the top of my_script.py:
from compressed_tensors.offload import init_dist
init_dist()