### Install compressed-tensors from Source

Source: https://github.com/vllm-project/compressed-tensors/blob/main/README.md

Clone the repository and install the compressed-tensors library in editable mode.

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

--------------------------------

### Install compressed-tensors from PyPI

Source: https://github.com/vllm-project/compressed-tensors/blob/main/README.md

Install the stable release of the compressed-tensors library using pip.

```bash
pip install compressed-tensors
```

--------------------------------

### Distributed Inference Setup

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Sets up distributed inference using torchrun, initializing NCCL and loading the model with offloading. Ensure the script is run with `torchrun`.

```python
# run with: torchrun --nproc-per-node=4 script.py
from transformers import AutoModelForCausalLM
from compressed_tensors.offload import init_dist, load_offloaded_model

init_dist()  # initialize NCCL

with load_offloaded_model():
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.1-70B",
        device_map="auto",
        torch_dtype="bfloat16",
    )
```

--------------------------------

### Get Quantization Scheme Directly

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Retrieve the quantization scheme for a specific preset name and target layers. This example demonstrates getting the scheme for 'W4A16' on Linear layers and printing its weight bits and group size.

```python
scheme = preset_name_to_scheme("W4A16", targets=["Linear"])
print(f"Weight bits: {scheme.weights.num_bits}")  # Output: Weight bits: 4
print(f"Group size: {scheme.weights.group_size}")  # Output: Group size: 128
```

--------------------------------

### Install compressed-tensors Nightly Release

Source: https://github.com/vllm-project/compressed-tensors/blob/main/README.md

Install the pre-release (nightly) version of the compressed-tensors library using pip.

```bash
pip install --pre compressed-tensors
```

--------------------------------

### Save a Compressed Model with PTQ

Source: https://github.com/vllm-project/compressed-tensors/blob/main/README.md

Example demonstrating post-training quantization (PTQ) and saving a quantized model using compressed-tensors. Requires PyTorch, Hugging Face Transformers, and Datasets libraries.

```python
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda:0", torch_dtype="auto")

config = QuantizationConfig.parse_file("./examples/bit_packing/int4_config.json")
config.quantization_status = QuantizationStatus.CALIBRATION
apply_quantization_config(model, config)

dataset = load_dataset("ptb_text_only")["train"]
tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_function(examples):
    return tokenizer(examples["sentence"], padding=False, truncation=True, max_length=1024)

tokenized_dataset = dataset.map(tokenize_function, batched=True)
data_loader = DataLoader(tokenized_dataset, batch_size=1, collate_fn=DefaultDataCollator())

with torch.no_grad():
    for idx, sample in tqdm(enumerate(data_loader), desc="Running calibration"):
        sample = {key: value.to(device) for key,value in sample.items()}
        _ = model(**sample)

        if idx >= 512:
            break

model.apply(freeze_module_quantization)
model.apply(compress_quantized_weights)

output_dir = "./ex_llama1.1b_w4a16_packed_quantize"
compressor = ModelCompressor.from_pretrained_model(model)
compressor.compress_model(model)
model.save_pretrained(output_dir)
```

--------------------------------

### Use Preset Quantization Scheme in Config

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Configure quantization by specifying preset names within the QuantizationConfig. This example shows applying 'FP8_DYNAMIC' to all Linear layers.

```python
config = QuantizationConfig(
    config_groups={
        "FP8_DYNAMIC": ["Linear"],  # Apply FP8_DYNAMIC to all Linear layers
    }
)
```

--------------------------------

### Quantize Tensor to INT8

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Perform tensor quantization using the `quantize` function. This example demonstrates quantizing a float tensor to INT8 with specified quantization arguments, including number of bits, type, strategy, and symmetry.

```python
import torch
from compressed_tensors.quantization.lifecycle import quantize, dequantize, fake_quantize
from compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy

# Create sample tensor
x = torch.randn(4, 128, dtype=torch.float32, device="cuda")

# Define quantization args
args = QuantizationArgs(
    num_bits=8,
    type=QuantizationType.INT,
    strategy=QuantizationStrategy.TENSOR,
    symmetric=True,
)

# Calculate scale and zero point
scale = x.abs().max() / 127.0
zero_point = torch.zeros(1, device=x.device)

# Quantize tensor to int8
x_quantized = quantize(
    x=x,
    scale=scale,
    zero_point=zero_point,
    args=args,
    dtype=torch.int8,
)
print(f"Quantized dtype: {x_quantized.dtype}")  # torch.int8
print(f"Quantized range: [{x_quantized.min()}, {x_quantized.max()}]")
```

--------------------------------

### Get Execution Device for Module

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Returns the device that input tensors should be moved to before running the module's forward pass. For offloaded modules, returns `module._parameters.onload_device`. For non-offloaded modules, returns the device of the first parameter/buffer.

```python
get_execution_device(module, default=None)
```

--------------------------------

### Get Offloaded Device for Module

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Returns the device where the module's weights reside when not in use (between forward passes). For offloaded modules, returns `module._parameters.offload_device`. For non-offloaded modules, returns the device of the first parameter/buffer.

```python
get_offloaded_device(module, default=None)
```

--------------------------------

### Temporarily unwrap offload forward wrapper

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `unwrap_offload_forward` as a context manager to temporarily remove the offload forward wrapper. This allows modification of the underlying forward function, for example, by other libraries like PEFT or quantization hooks. The offload wrapper is re-applied upon exiting the context.

```python
with unwrap_offload_forward(module):
    module.forward = my_patched_forward  # patching the real forward
# on exit: offload wrapper is applied around my_patched_forward
```

--------------------------------

### Create QuantizationConfig with Preset Schemes

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Instantiate a QuantizationConfig using preset schemes. Presets simplify configuration by mapping common names to target layers.

```python
# Or use preset schemes
config_preset = QuantizationConfig(
    config_groups={
        "W8A8": ["Linear"],  # Preset name mapped to target layers
    },
    quantization_status=QuantizationStatus.INITIALIZED
)

# Set to calibration mode
config.quantization_status = QuantizationStatus.CALIBRATION
```

--------------------------------

### Get Available CUDA Device Memory

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `get_device_memory` to get a mapping of available CUDA devices to their total memory in bytes. Returns an empty dictionary if no CUDA devices are found.

```python
get_device_memory()
```

--------------------------------

### Apply Quantization Configuration to Model

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Load a model and apply a custom quantization configuration. This prepares the model for calibration by attaching quantization schemes to specified target modules.

```python
import torch
from transformers import AutoModelForCausalLM
from compressed_tensors.quantization import (
    QuantizationConfig,
    QuantizationStatus,
    apply_quantization_config,
)

# Load model
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="cuda:0",
    torch_dtype="auto"
)

# Create quantization config
config = QuantizationConfig(
    config_groups={
        "group_1": {
            "weights": {
                "num_bits": 8,
                "type": "int",
                "symmetric": True,
                "strategy": "channel"
            },
            "input_activations": {
                "num_bits": 8,
                "type": "int",
                "symmetric": True,
                "strategy": "token",
                "dynamic": True
            },
            "targets": ["Linear"]
        }
    },
    quantization_status=QuantizationStatus.CALIBRATION
)

# Apply quantization config to model
apply_quantization_config(model, config)

# Model is now ready for calibration
# Each target module now has quantization_scheme and quantization_status attributes
for name, module in model.named_modules():
    if hasattr(module, "quantization_scheme"):
        print(f"{name}: quantized={module.quantization_scheme is not None}")
```

--------------------------------

### Reference Preset Quantization Schemes

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Demonstrates how to reference preset quantization schemes by name within configuration groups. The library offers common presets like W8A16 for 8-bit weight-only quantization.

```python
from compressed_tensors.quantization import QuantizationConfig, preset_name_to_scheme

# Available presets:
# - W8A16: 8-bit weight-only quantization
```

--------------------------------

### Get Execution Device

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Retrieves the current execution device for a given module. This helps in understanding where a module's parameters are currently residing (e.g., CPU, GPU).

```python
import torch
from compressed_tensors.offload import (
    update_offload_parameter,
    get_execution_device,
    disable_offloading,
    disable_onloading,
    dispatch_model,
)
from transformers import AutoModelForCausalLM

# Load model with device mapping
model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
    device_map="auto",  # Automatically distributes across devices
    torch_dtype=torch.float16,
)

# Update a module's parameter (works with offloaded modules)
for name, module in model.named_modules():
    if hasattr(module, "weight_scale"):
        new_scale = torch.ones_like(module.weight_scale)
        update_offload_parameter(module, "weight_scale", new_scale)

# Get the execution device for a module
for name, module in model.named_modules():
    if hasattr(module, "weight"):
        device = get_execution_device(module)
        print(f"{name}: executes on {device}")
        break
```

--------------------------------

### Get Device Map of a Dispatched Model

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `get_device_map` to retrieve the `DeviceMap` of a dispatched model, detailing the `(onload_device, offload_device)` for each module. This is useful for inspecting or serializing the dispatch configuration.

```python
device_map = get_device_map(model)
# -> {"model.layers.0": (cuda:0, cpu), "model.layers.1": (cuda:0, disk), ...}
```

--------------------------------

### apply_quantization_config

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Initializes a model for quantization by attaching quantization schemes to target modules.

```APIDOC
## apply_quantization_config

### Description
Initializes a model for quantization by attaching quantization schemes to target modules. It prepares the model for calibration or inference with quantized parameters.

### Request Body
- **model** (torch.nn.Module) - Required - The model to be quantized.
- **config** (QuantizationConfig) - Required - The configuration defining quantization groups and status.
```

--------------------------------

### Define and Apply Quantization Configuration

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define a quantization configuration with specified groups and status, then apply it to a model. This is useful for setting up quantization parameters before calibration.

```python
config = QuantizationConfig(
    config_groups={
        "group_1": {
            "weights": {
                "num_bits": 8,
                "type": "int",
                "symmetric": True,
                "strategy": "tensor"
            },
            "input_activations": {
                "num_bits": 8,
                "type": "int",
                "symmetric": True,
                "strategy": "tensor"
            },
            "targets": ["Linear"]
        }
    },
    quantization_status=QuantizationStatus.CALIBRATION
)

apply_quantization_config(model, config)
```

--------------------------------

### Manually Dispatch Model Weights with a Device Map

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Manually specifies the device for each module using a device map. The `offload_dir` parameter is required when using 'disk' as a target device.

```python
from compressed_tensors.offload import dispatch_with_map

device_map = {
    "model.embed_tokens": ("cuda:0", "cuda:0"),
    "model.layers.0":     ("cuda:0", "cpu"),
    "model.layers.1":     ("cuda:0", "disk"),
    # ...
}
dispatch_with_map(model, device_map, offload_dir="/tmp/weights")
```

--------------------------------

### dispatch_model

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

The primary entry point for production deployment, providing automatic, memory-aware, multi-GPU dispatch.

```APIDOC
## dispatch_model

### Description
Dispatches a model across devices with explicit memory constraints.

### Parameters
- **device_memory** (dict) - Optional - Mapping of device to available bytes. Defaults to querying all CUDA devices.
- **extra_memory** (int) - Optional - Bytes to reserve for activations. If None, estimated from model config.
- **no_split_modules** (list) - Optional - Names of module classes that cannot be split across devices.

### Request Example
model = dispatch_model(model, device_memory={torch.device("cuda:0"): 16e9})
```

--------------------------------

### Prepare Calibration Data Loader

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Load and tokenize calibration data using Hugging Face's `datasets` and `transformers` libraries. A `DataLoader` is created for iterating through the tokenized dataset with a specified batch size and collator.

```python
tokenizer = AutoTokenizer.from_pretrained(model_name)
dataset = load_dataset("garage-bAInd/Open-Platypus", split=f"train[:{num_calibration_samples}]")

def tokenize(examples):
    return tokenizer(examples["output"], padding=False, truncation=True, max_length=512)

tokenized = dataset.map(tokenize, batched=True)
loader = DataLoader(tokenized, batch_size=1, collate_fn=DefaultDataCollator())
```

--------------------------------

### Define FP8 Dynamic Token-wise Activation Quantization Args

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define QuantizationArgs for FP8 dynamic token-wise activation quantization. This configuration is suitable for activation compression.

```python
# FP8 dynamic token-wise activation quantization
fp8_activations = QuantizationArgs(
    num_bits=8,
    type=QuantizationType.FLOAT,
    strategy=QuantizationStrategy.TOKEN,
    symmetric=True,
    dynamic=True,
)
```

--------------------------------

### Calculate Quantization Parameters

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Compute scale and zero point for symmetric or asymmetric quantization strategies.

```python
import torch
from compressed_tensors.quantization.utils import calculate_qparams
from compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy

# Sample tensor to quantize
tensor = torch.randn(512, 1024, device="cuda")

# Define quantization args
args = QuantizationArgs(
    num_bits=8,
    type=QuantizationType.INT,
    strategy=QuantizationStrategy.TENSOR,
    symmetric=True,
)

# Calculate min/max values
min_val, max_val = torch.aminmax(tensor)

# Calculate scale and zero point
scale, zero_point = calculate_qparams(
    min_vals=min_val,
    max_vals=max_val,
    quantization_args=args,
)
print(f"Scale: {scale.item():.6f}")
print(f"Zero point: {zero_point.item()}")

# For asymmetric quantization
args_asym = QuantizationArgs(
    num_bits=4,
    type=QuantizationType.INT,
    strategy=QuantizationStrategy.GROUP,
    group_size=128,
    symmetric=False,
)

# Reshape for group-wise calculation
reshaped = tensor.reshape(-1, args_asym.group_size)
min_vals = reshaped.min(dim=-1).values
max_vals = reshaped.max(dim=-1).values

scale_group, zp_group = calculate_qparams(
    min_vals=min_vals,
    max_vals=max_vals,
    quantization_args=args_asym,
)
print(f"Group scales shape: {scale_group.shape}")  # Per-group scales
```

--------------------------------

### Compress Model and Save

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Use `ModelCompressor` to compress the model based on the defined quantization configuration, then save the compressed model and update its configuration. This is the final step in the quantization workflow.

```python
compressor = ModelCompressor(quantization_config=config)
compressor.compress_model(model)
model.save_pretrained(output_dir)
compressor.update_config(output_dir)

print(f"Quantized model saved to {output_dir}")
```

--------------------------------

### Calibration Hook for Quantization Parameters

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

A forward hook to compute and update quantization parameters (scales and zero points) for weights and input activations during calibration. Ensure the module has a 'quantization_scheme' attribute.

```python
def calibration_hook(module, input, output):
    scheme = getattr(module, "quantization_scheme", None)
    if not scheme:
        return

    # Update weight quantization parameters
    if scheme.weights:
        min_val, max_val = torch.aminmax(module.weight.data)
        scale, zp = calculate_qparams(min_val, max_val, scheme.weights)
        update_offload_parameter(module, "weight_scale", scale)
        update_offload_parameter(module, "weight_zero_point", zp)

    # Update activation quantization parameters
    if scheme.input_activations and len(input) > 0:
        min_val, max_val = torch.aminmax(input[0])
        scale, zp = calculate_qparams(min_val, max_val, scheme.input_activations)
        update_offload_parameter(module, "input_scale", scale)
        update_offload_parameter(module, "input_zero_point", zp)
```

--------------------------------

### Load Model with Offloading using Context Manager

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `load_offloaded_model` as a context manager to patch `from_pretrained` for `transformers` models. It enables device offloading on rank 0 and loading on the meta device for other ranks, followed by conversion to `compressed_tensors` offloading.

```python
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
    )
```

--------------------------------

### Model Loading and Dispatch

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

APIs for loading models with offloading and dispatching them across devices.

```APIDOC
## POST /offload/load_offloaded_model

### Description
Loads a model with compressed-tensors offloading enabled. This is typically used as a context manager.

### Method
POST

### Endpoint
/offload/load_offloaded_model

### Parameters
None

### Request Example
```python
from transformers import AutoModelForCausalLM
from compressed_tensors.offload import load_offloaded_model

with load_offloaded_model():
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.1-70B",
        device_map="auto", 
        torch_dtype="bfloat16",
    )
```

### Response
#### Success Response (200)
- **model** (object) - The loaded model with offloading enabled.

#### Response Example
```json
{
  "model": "loaded_model_object"
}
```

## POST /offload/dispatch_model

### Description
Dispatches a model across available GPUs and offloads the remainder to CPU. This function automatically determines the optimal distribution.

### Method
POST

### Endpoint
/offload/dispatch_model

### Parameters
#### Request Body
- **model** (object) - Required - The model to dispatch (typically loaded on CPU or meta device).

### Request Example
```python
from compressed_tensors.offload import dispatch_model

model = ...  # load on CPU or meta device
model = dispatch_model(model)
```

### Response
#### Success Response (200)
- **dispatched_model** (object) - The model dispatched across available devices.

#### Response Example
```json
{
  "dispatched_model": "dispatched_model_object"
}
```
```

--------------------------------

### Define Block-wise FP8 Quantization Args

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define QuantizationArgs for block-wise FP8 quantization, similar to DeepSeek's approach. This specifies block structure for quantization.

```python
# Block-wise FP8 (DeepSeek-style) quantization
fp8_block_weights = QuantizationArgs(
    num_bits=8,
    type=QuantizationType.FLOAT,
    strategy=QuantizationStrategy.BLOCK,
    symmetric=True,
    dynamic=False,
    block_structure=[128, 128],  # 128x128 blocks
)
```

--------------------------------

### Initialize Distributed Process Group in Script

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Initializes the distributed process group using NCCL and environment variables. Sets the CUDA device to the local rank. Requires TORCHELASTIC_RUN_ID environment variable.

```python
# At the top of my_script.py:
from compressed_tensors.offload import init_dist
init_dist()
```

--------------------------------

### Load QuantizationConfig from JSON

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Load a QuantizationConfig object from a JSON string. This is useful for defining custom quantization configurations.

```python
from compressed_tensors.quantization import QuantizationConfig, QuantizationStatus

# Load quantization config from JSON file
config = QuantizationConfig.model_validate_json ('''
{
    "quant_method": "compressed-tensors",
    "format": "naive-quantized",
    "config_groups": {
        "group_1": {
            "weights": {
                "num_bits": 8,
                "type": "int",
                "symmetric": true,
                "strategy": "tensor"
            },
            "input_activations": {
                "num_bits": 8,
                "type": "int",
                "symmetric": true,
                "strategy": "tensor"
            },
            "targets": ["Linear"]
        }
    }
}
''')
```

--------------------------------

### Dispatch Model for Multi-GPU

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Dispatches a model across available GPUs and offloads the remainder to CPU. Load the model on CPU or meta device before dispatching.

```python
from compressed_tensors.offload import dispatch_model

model = ...  # load on CPU or meta device
model = dispatch_model(model)  # automatically fills available GPUs, offloads remainder to CPU
```

--------------------------------

### Run Model Calibration

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Iterate through the calibration data loader and run the model in no_grad mode to collect statistics for quantization. The loop breaks after processing the specified number of calibration samples.

```python
with torch.no_grad():
    for idx, batch in enumerate(tqdm(loader, desc="Calibrating")):
        batch = {k: v.to(device) for k, v in batch.items()}
        model(**batch)
        if idx >= num_calibration_samples:
            break
```

--------------------------------

### Compress and Save a Quantized Model

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Orchestrate model compression using ModelCompressor. This involves loading a model, applying quantization, compressing weights in-place, saving the compressed model, and updating its configuration.

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from compressed_tensors.compressors import ModelCompressor
from compressed_tensors.quantization import (
    QuantizationConfig,
    QuantizationStatus,
    apply_quantization_config,
)

# Load and prepare model
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="cuda:0", torch_dtype="auto"
)

# Apply quantization config
config = QuantizationConfig(
    config_groups={"W8A8": ["Linear"]},
    quantization_status=QuantizationStatus.FROZEN
)
apply_quantization_config(model, config)

# Create compressor from quantized model
compressor = ModelCompressor.from_pretrained_model(model)

# Compress model in-place (converts weights to compressed format)
compressor.compress_model(model)
# Model weights are now in compressed format (e.g., int8)

# Save compressed model
output_dir = "./compressed_model"
model.save_pretrained(output_dir)
compressor.update_config(output_dir)  # Updates config.json with quantization info

# Later: decompress model for inference
compressor.decompress_model(model)
# Model weights are now back to float format with quantization applied
```

--------------------------------

### Inspect Model Dispatch Configuration

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Retrieves the current device map and execution/offload devices for model modules. Useful for debugging and understanding the current state of weight distribution.

```python
from compressed_tensors.offload import get_device_map, get_execution_device, get_offloaded_device

device_map = get_device_map(model)
exec_dev = get_execution_device(model.layers[0])   # cuda:0
offload_dev = get_offloaded_device(model.layers[0]) # cpu
```

--------------------------------

### Load Model with Offloading

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Loads a model using compressed-tensors offloading. Use the `load_offloaded_model` context manager. The `device_map` can be set to 'auto' or 'auto_offload'.

```python
from transformers import AutoModelForCausalLM
from compressed_tensors.offload import load_offloaded_model

with load_offloaded_model():
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.1-70B",
        device_map="auto",      # or "auto_offload" to restrict to cpu/disk only
        torch_dtype="bfloat16",
    )
# model is now using compressed-tensors offloading
```

--------------------------------

### Initialize Distributed Process Group

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Initializes the distributed process group using NCCL and environment variables. Sets the CUDA device to the local rank. Requires TORCHELASTIC_RUN_ID environment variable.

```bash
torchrun --nproc-per-node=4 my_script.py
```

--------------------------------

### Manage Offloading with Context Managers

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use context managers to control whether tensors are cached in memory or accessed directly from the offload device.

```python
# Without: each access triggers a CPU→GPU copy
for _ in range(3):
    result = module.weight  # 3 separate copies

# With: first access copies, subsequent reads hit cache
with disable_offloading():
    for _ in range(3):
        result = module.weight  # 1 copy, 2 cache hits
```

```python
# Inspect the raw offloaded tensor (e.g., CPU tensor) without triggering a copy
with disable_onloading():
    cpu_tensor = module.weight   # returns the CPU tensor directly
    module.weight = new_tensor   # assigns without triggering offload
```

--------------------------------

### Dispatch Model with Explicit Device Map

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `dispatch_with_map` to dispatch a model according to a detailed `DeviceMap`, which specifies the `(onload_device, offload_device)` for each module. An optional `offload_dir` can be provided for disk offloading.

```python
device_map = {
    "model.embed_tokens":  ("cuda:0", "cuda:0"),
    "model.layers.0":      ("cuda:0", "cpu"),
    "model.layers.1":      ("cuda:0", "disk"),
}
dispatch_with_map(model, device_map, offload_dir="/tmp/offload")
```

--------------------------------

### Save and Load Offloaded Model with Accelerate

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Converts a model to the Hugging Face accelerate format for saving, then converts it back for continued inference. Use `to_accelerate` before saving and `from_accelerate` after loading.

```python
from compressed_tensors.offload import to_accelerate, from_accelerate

# Convert to accelerate for saving
to_accelerate(model)
model.save_pretrained("./output")

# Convert back for continued inference
from_accelerate(model)
```

--------------------------------

### Offload Module Registration

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Dynamically adds submodules to an offloaded model, ensuring consistent offloading configuration.

```APIDOC
## POST /register_offload_module

### Description
Registers a new submodule on a parent module, propagating offloading if the parent is offloaded.

### Method
POST

### Endpoint
/register_offload_module

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **base** (object) - Required - The parent module to which the new submodule will be attached.
- **name** (string) - Required - The name of the new submodule.
- **module** (object) - Required - The submodule to register.

### Request Example
```json
{
  "base": "model",
  "name": "new_layer",
  "module": "new_layer_instance"
}
```

### Response
#### Success Response (200)
- **status** (string) - Indicates successful registration.

#### Response Example
```json
{
  "status": "registered"
}
```
```

--------------------------------

### QuantizationArgs API

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Defines the specific parameters for weight, input, or output quantization, including bit-width, strategy, and type.

```APIDOC
## QuantizationArgs

### Description
Specifies the technical details for quantization, such as bit-width, strategy (tensor, channel, group, block, token), and data type.

### Parameters
- **num_bits** (int) - Required - Number of bits for quantization.
- **type** (string) - Required - Data type (INT or FLOAT).
- **strategy** (string) - Required - Quantization strategy (CHANNEL, GROUP, TOKEN, BLOCK, TENSOR).
- **symmetric** (boolean) - Optional - Whether to use symmetric quantization.
- **dynamic** (boolean) - Optional - Whether to use dynamic quantization.
- **group_size** (int) - Optional - Size for group-wise quantization.
- **block_structure** (list) - Optional - Dimensions for block-wise quantization.
```

--------------------------------

### Distributed Utilities

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Utilities for managing distributed PyTorch environments.

```APIDOC
## GET /dist_utils/is_distributed

### Description
Returns `True` if `torch.distributed` is available and has been initialized.

### Method
GET

### Endpoint
/dist_utils/is_distributed

### Parameters
None

### Response
#### Success Response (200)
- **is_distributed** (bool) - `True` if distributed is initialized, `False` otherwise.

#### Response Example
```json
{
  "is_distributed": true
}
```

## GET /dist_utils/is_rank0

### Description
Returns `True` if not in a distributed context, or if the current process is the source rank (typically rank 0).

### Method
GET

### Endpoint
/dist_utils/is_rank0

### Parameters
None

### Response
#### Success Response (200)
- **is_rank0** (bool) - `True` if the current process is rank 0 or not in a distributed context.

#### Response Example
```json
{
  "is_rank0": true
}
```

## POST /dist_utils/init_dist

### Description
Initializes the distributed process group using NCCL and environment variables set by `torchrun`. Sets the CUDA device to the local rank. Requires the `TORCHELASTIC_RUN_ID` environment variable.

### Method
POST

### Endpoint
/dist_utils/init_dist

### Parameters
None

### Request Example
```bash
torchrun --nproc-per-node=4 my_script.py
```

### Response
#### Success Response (200)
- **status** (string) - Indicates successful initialization.

#### Response Example
```json
{
  "status": "initialized"
}
```

## POST /dist_utils/as_broadcastable

### Description
Returns a view of the tensor compatible with `dist.broadcast`. Works around an NCCL limitation: FP8 dtypes cannot be broadcast on pre-Hopper hardware. FP8 tensors are reinterpreted as `uint8` for the broadcast and share the same underlying storage.

### Method
POST

### Endpoint
/dist_utils/as_broadcastable

### Parameters
#### Request Body
- **tensor** (torch.Tensor) - Required - The tensor to make broadcastable.

### Request Example
```json
{
  "tensor": "your_fp8_tensor_data"
}
```

### Response
#### Success Response (200)
- **broadcastable_tensor** (torch.Tensor) - A view of the tensor compatible with `dist.broadcast`.

#### Response Example
```json
{
  "broadcastable_tensor": "broadcastable_tensor_data"
}
```
```

--------------------------------

### Dispatch Model with Explicit Device Memory Constraints

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `dispatch_model` for production deployment with automatic, memory-aware, multi-GPU dispatch. Specify `device_memory` to set available bytes per device.

```python
model = dispatch_model(model, device_memory={torch.device("cuda:0"): 16e9})
```

--------------------------------

### Set Onload Device for Model Execution

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `set_onload_device` to move all modules in a model to a specified `onload_device` without altering weight storage locations. This is simpler than `dispatch_model` when weights are already correctly placed.

```python
# Move all execution to cuda:0, keeping offloads unchanged
model = set_onload_device(model, onload_device="cuda:0")
```

--------------------------------

### Offloading Control API

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Context managers and utilities to control offloading and onloading behavior globally.

```APIDOC
## disable_offloading()

### Description
Context manager that disables offloading globally. Onloaded tensors are cached to prevent repeated device-to-device copies.

## disable_onloading()

### Description
Context manager that disables onloading globally. Parameter reads return raw offloaded tensors; assignments do not trigger offloading.

## update_offload_parameter(module, name, data)

### Description
Updates an offloaded parameter or buffer in-place.

### Parameters
- **module** (object) - Required - The module containing the parameter.
- **name** (str) - Required - The name of the parameter.
- **data** (tensor) - Required - The new data to update.
```

--------------------------------

### offload_module

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Attaches offloading to a single torch.nn.Module by replacing parameters and buffers with cache instances and wrapping the forward method.

```APIDOC
## POST /offload_module

### Description
Replaces module parameters and buffers with cache instances and wraps the forward function to handle device movement.

### Parameters
#### Request Body
- **module** (torch.nn.Module) - Required - The module to offload.
- **onload_device** (str) - Required - The device where the module will execute (e.g., "cuda:0").
- **offload_device** (str) - Required - The device where weights are stored (e.g., "cpu").

### Request Example
{
  "module": "layer",
  "onload_device": "cuda:0",
  "offload_device": "cpu"
}
```

--------------------------------

### Offload a single linear layer to CPU

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Use `offload_module` to attach offloading to a specific `torch.nn.Module`, specifying the device for loading and offloading. This is useful for fine-grained control over module offloading.

```python
offload_module(layer, onload_device="cuda:0", offload_device="cpu")
```

--------------------------------

### Convert Model Checkpoints

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Convert checkpoints from external formats to compressed-tensors format.

```python
from compressed_tensors.entrypoints.convert import (
    convert_checkpoint,
    ModelOptNvfp4Converter,
)
from compressed_tensors.quantization import QuantizationArgs, QuantizationType
```

--------------------------------

### Convert NVIDIA ModelOpt NVFP4 to compressed-tensors

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Converts a model checkpoint from NVIDIA ModelOpt NVFP4 format to the compressed-tensors format. Specify target layers using regex and configure KV cache quantization.

```python
MODEL_ID = "nvidia/Qwen3-32B-NVFP4"
SAVE_DIR = "./converted_model"

convert_checkpoint(
    model_stub=MODEL_ID,
    save_directory=SAVE_DIR,
    converter=ModelOptNvfp4Converter(
        # Target specific layers using regex patterns
        targets=[
            "re:.*mlp.*\\.(gate_up|gate|up|down)_proj$",
            "re:.*self_attn.*\\.(q|k|v|o)_proj$",
        ],
        # Include KV cache quantization config
        kv_cache_scheme=QuantizationArgs(
            num_bits=8,
            dynamic=False,
            type=QuantizationType.FLOAT,
        ),
    ),
    max_workers=8,  # Parallel processing
)
```

--------------------------------

### Register Calibration Hooks on Model

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Register the defined calibration hook to all modules in the model using `model.apply`. This prepares the model to collect quantization statistics during the calibration run.

```python
model.apply(lambda m: m.register_forward_hook(calibration_hook))
```

--------------------------------

### CPUCache

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Implementation of OffloadCache that offloads tensors to CPU RAM. Suitable when GPU VRAM is insufficient but CPU RAM is available.

```APIDOC
## CPUCache

### Description
Offloads tensors to CPU RAM. Onloading involves a standard `.to(device)` call from CPU to the configured `onload_device` (typically a CUDA device).

- **offload**: Moves the tensor to `torch.device("cpu")`.
- **onload**: Moves the tensor from CPU to the `onload_device`.
- **update_offload**: Performs an in-place copy (`offloaded.copy_(data)`) on the CPU tensor.

### Use Case
Use when GPU VRAM is insufficient and CPU RAM is available. This is the most common offload strategy.

### Distributed Variant
Synchronizes data across ranks during `offload`.
```

--------------------------------

### Load Model for Quantization

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Loads a pre-trained model for further processing, such as quantization. The model is loaded onto the specified device with automatic data type handling.

```python
import torch
from pathlib import Path
from tqdm import tqdm
from torch.utils.data import DataLoader
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, DefaultDataCollator
from compressed_tensors.compressors import ModelCompressor
from compressed_tensors.quantization import (
    QuantizationConfig,
    QuantizationStatus,
    apply_quantization_config,
)
from compressed_tensors.quantization.utils import calculate_qparams
from compressed_tensors.offload import update_offload_parameter

# Configuration
model_name = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
output_dir = "./quantized_model"
num_calibration_samples = 512
device = "cuda:0" if torch.cuda.is_available() else "cpu"

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

--------------------------------

### Device Inspection API

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Utilities to inspect the execution and offload devices of model modules.

```APIDOC
## get_execution_device(module, default=None)

### Description
Returns the device that input tensors should be moved to before running the module's forward pass.

## get_offloaded_device(module, default=None)

### Description
Returns the device where the module's weights reside when not in use.
```

--------------------------------

### Compress Quantized Weights

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Convert floating-point weights to a compressed representation after calibration.

```python
import torch
from compressed_tensors.quantization.lifecycle import compress_quantized_weights
from compressed_tensors.quantization import (
    QuantizationConfig,
    QuantizationStatus,
    apply_quantization_config,
)
from transformers import AutoModelForCausalLM

# Load and quantize model
model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
    device_map="cuda:0",
    torch_dtype=torch.float16,
)

# Apply quantization config
config = QuantizationConfig(
    config_groups={"W4A16": ["Linear"]},
    quantization_status=QuantizationStatus.FROZEN,
)
apply_quantization_config(model, config)

# After calibration, compress weights in-place
# This converts float weights to int4 representation
model.apply(compress_quantized_weights)

# Check compression result
for name, module in model.named_modules():
    if hasattr(module, "weight") and hasattr(module, "quantization_status"):
        if module.quantization_status == QuantizationStatus.COMPRESSED:
            print(f"{name}: weight dtype={module.weight.dtype}")
            # Weights are now int8 (packed int4)
```

--------------------------------

### Convert Model to Accelerate Offloading

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Converts a compressed_tensors-offloaded model back to accelerate offloading. This is necessary before calling `model.save_pretrained()`. For disk-offloaded modules, creates an `OffloadedWeightsLoader`. Sets `model.hf_device_map` as expected by `transformers`.

```python
to_accelerate(model)
model.save_pretrained("./output")
```

--------------------------------

### Define INT8 Channel-wise Weight Quantization Args

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define QuantizationArgs for INT8 channel-wise weight quantization. This configuration specifies bit precision, type, strategy, and symmetry.

```python
from compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy

# INT8 channel-wise weight quantization
int8_weights = QuantizationArgs(
    num_bits=8,
    type=QuantizationType.INT,
    strategy=QuantizationStrategy.CHANNEL,
    symmetric=True,
    dynamic=False,
)
```

--------------------------------

### DiskCache

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Implementation of OffloadCache that offloads tensors to disk as safetensors files. Suitable for models that exceed both GPU VRAM and CPU RAM.

```APIDOC
## DiskCache

### Description
Offloads tensors to disk as `safetensors` files. In-memory, offloaded tensors are represented as **meta tensors** (tensors with no data, only shape/dtype/stride). A separate `index` dict maps meta tensor identity to its disk location.

- **offload**: Writes the tensor to a `.safetensors` file in `offload_dir` and returns a meta tensor.
- **onload**: Reads the file at the stored path using `safe_open` and reconstructs the tensor.
- **update_offload**: Overwrites the existing file with new data.
- **`__delitem__`**: Removes the tensor's entry from the index and deletes the corresponding file if it was created by this cache (i.e., not an original model file).

### Use Case
Use when neither CPU RAM nor GPU VRAM can hold all weights. Disk offload is the slowest but allows running arbitrarily large models.

### Note on Format
`DiskCache` requires weights to be stored in `safetensors` format. When loading from non-safetensors formats, weights are onloaded and then re-saved in `safetensors` format.

### Distributed Variant
Synchronizes data across ranks during `offload`.
```

--------------------------------

### Apply KV Cache Quantization Configuration

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Applies a quantization configuration to a model, specifically enabling KV cache quantization. This reduces memory usage by quantizing key and value tensors in attention layers.

```python
from compressed_tensors.quantization import (
    QuantizationConfig,
    QuantizationArgs,
    QuantizationType,
    QuantizationStrategy,
    apply_quantization_config,
)
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T",
    device_map="cuda:0",
    torch_dtype="auto",
)

# Configure KV cache quantization (FP8)
config = QuantizationConfig(
    config_groups={
        "weights": {
            "weights": {
                "num_bits": 8,
                "type": "int",
                "symmetric": True,
                "strategy": "channel",
            },
            "targets": ["Linear"],
        }
    },
    kv_cache_scheme=QuantizationArgs(
        num_bits=8,
        type=QuantizationType.FLOAT,
        strategy=QuantizationStrategy.TENSOR,
        symmetric=True,
        dynamic=False,
    ),
)

# Apply config with KV cache quantization
apply_quantization_config(model, config)

# KV cache quantization is applied to self_attn modules
for name, module in model.named_modules():
    if "self_attn" in name and hasattr(module, "quantization_scheme"):
        scheme = module.quantization_scheme
        if scheme.input_activations:
            print(f"{name}: KV cache quantized to {scheme.input_activations.num_bits} bits")
```

--------------------------------

### Define INT4 Group Quantization Args

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define QuantizationArgs for INT4 group quantization with a specified group size. This is useful for more aggressive compression.

```python
# INT4 group quantization with group size 128
int4_weights = QuantizationArgs(
    num_bits=4,
    type=QuantizationType.INT,
    strategy=QuantizationStrategy.GROUP,
    group_size=128,
    symmetric=True,
    dynamic=False,
)
```

--------------------------------

### dispatch_model

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Automatically distributes a model across available CUDA devices with CPU offload fallback.

```APIDOC
## POST /dispatch_model

### Description
Analyzes model size and device memory to automatically distribute modules across GPUs, offloading remaining parts to CPU.

### Parameters
#### Request Body
- **model** (torch.nn.Module) - Required - The model to dispatch.
- **device_memory** (dict) - Optional - Memory constraints for devices.
- **extra_memory** (int) - Optional - Memory reserved for activations/KV cache.
- **no_split_modules** (list) - Optional - List of modules that should not be split.
```

--------------------------------

### Model Conversion API

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Functions to convert models between accelerate and compressed_tensors offloading formats.

```APIDOC
## from_accelerate(model)

### Description
Converts a model from accelerate offloading to compressed_tensors offloading. Called automatically by load_offloaded_model.

### Parameters
- **model** (object) - Required - The model to convert.

### Response
- **device_map** (dict) - The device map.
- **offload_dir** (str) - The optional disk offload directory.

## to_accelerate(model)

### Description
Converts a compressed_tensors-offloaded model back to accelerate offloading. Necessary before calling model.save_pretrained().

### Parameters
- **model** (object) - Required - The model to convert.

### Response
- **hf_device_map** (dict) - The device map expected by transformers.
```

--------------------------------

### dispatch_with_map

Source: https://github.com/vllm-project/compressed-tensors/blob/main/src/compressed_tensors/offload/README.md

Dispatches a model according to an explicit DeviceMap.

```APIDOC
## dispatch_with_map

### Description
Dispatches a model based on a provided dictionary mapping module names to (onload_device, offload_device) tuples.

### Parameters
- **model** (object) - Required - The model to dispatch.
- **device_map** (dict) - Required - Mapping of module names to device tuples.
- **offload_dir** (str) - Optional - Directory for offloading.
```

--------------------------------

### Define Quantization Scheme with Regex Targets

Source: https://context7.com/vllm-project/compressed-tensors/llms.txt

Define a QuantizationScheme using regular expressions to target specific layers, such as query and key projection layers in self-attention modules.

```python
# Use regex patterns to target specific layers
attention_scheme = QuantizationScheme(
    targets=["re:.*self_attn.*\.q_proj$", "re:.*self_attn.*\.k_proj$"],
    weights=QuantizationArgs(
        num_bits=4,
        type=QuantizationType.INT,
        strategy=QuantizationStrategy.GROUP,
        group_size=128,
        symmetric=True,
    ),
)
```