Implementation:Mit han lab Llm awq Get calib dataset

Overview

get_calib_dataset is a concrete tool for loading and preparing calibration datasets for AWQ quantization, provided by the llm-awq library.

Source Location

• Repository: llm-awq (https://github.com/mit-han-lab/llm-awq)
• File: awq/utils/calib_data.py
• Lines: 5-32

Signature

def get_calib_dataset(data="pileval", tokenizer=None, n_samples=512, block_size=512):

Import

from awq.utils.calib_data import get_calib_dataset

I/O Contract

Inputs

Parameter	Type	Required	Default	Description
data	str	No	"pileval"	Name of the calibration dataset to load
tokenizer	PreTrainedTokenizer	Yes	None	HuggingFace tokenizer for encoding text samples
n_samples	int	No	512	Number of text samples to collect before concatenation
block_size	int	No	512	Token length of each output calibration block

Output

• List[torch.Tensor] -- A list of tensors, each with shape [1, block_size], representing contiguous blocks of tokenized calibration text.

Implementation Details

The function performs the following steps:

1. Load dataset: Loads the mit-han-lab/pile-val-backup dataset from HuggingFace (validation split) when data="pileval".
2. Shuffle: Shuffles the dataset with a fixed seed (42) for reproducibility.
3. Filter and collect: Iterates through the dataset, encoding each text sample with the tokenizer. Samples whose encoded length exceeds 512 tokens or that produce empty tensors are skipped. Collection stops after n_samples valid samples.
4. Concatenate: All collected sample tensors are concatenated along the sequence dimension (dim=1) into a single long tensor.
5. Split: The concatenated tensor is split into non-overlapping blocks of size block_size, and returned as a list.

Source Code

def get_calib_dataset(data="pileval", tokenizer=None, n_samples=512, block_size=512):
    if data == "pileval":
        dataset = load_dataset("mit-han-lab/pile-val-backup", split="validation")
    else:
        raise NotImplementedError
    dataset = dataset.shuffle(seed=42)
    samples = []
    n_run = 0
    for data in dataset:
        line = data["text"]
        line = line.strip()
        line_encoded = tokenizer.encode(line)
        if len(line_encoded) > 512:
            continue
        sample = torch.tensor([line_encoded])
        if sample.numel() == 0:
            continue
        samples.append(sample)
        n_run += 1
        if n_run == n_samples:
            break
    cat_samples = torch.cat(samples, dim=1)
    n_split = cat_samples.shape[1] // block_size
    return [cat_samples[:, i * block_size : (i + 1) * block_size] for i in range(n_split)]

Usage Example

from transformers import AutoTokenizer
from awq.utils.calib_data import get_calib_dataset

# Load tokenizer for the target model
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_fast=False)

# Load calibration dataset with default settings
calib_data = get_calib_dataset(
    data="pileval",
    tokenizer=tokenizer,
    n_samples=512,
    block_size=512,
)

print(f"Number of calibration blocks: {len(calib_data)}")
print(f"Each block shape: {calib_data[0].shape}")
# Output example:
# Number of calibration blocks: ~200
# Each block shape: torch.Size([1, 512])

Related Pages

• Principle:Mit_han_lab_Llm_awq_Calibration_Data_Preparation
• Environment:Mit_han_lab_Llm_awq_Python_Runtime_Environment

Knowledge Sources

• Repo|llm-awq|https://github.com/mit-han-lab/llm-awq

Domains

• NLP
• Quantization