[llm-compressor] Dataset Calibration: c4/wikitext/ultrachat 로더

2026년 4월 13일수정: 2026년 4월 13일

들어가며

양자화 캘리브레이션에는 "LLM이 실제로 만나는 분포"를 대표할 수 있는 데이터가 필요하다. llm-compressor는 공통 데이터셋을 위한 내장 로더를 제공한다. c4, wikitext, ultrachat, openplatypus 같은 것들이다. src/llmcompressor/datasets/가 진입점이고, src/llmcompressor/transformers/data/가 구체적인 데이터셋별 로더를 둔다.

핵심 구조/코드 분석

`datasets/utils.py`의 진입점

def get_calibration_dataloader(
    dataset_args: DatasetArguments,
    processor,
) -> DataLoader:
    """
    Main entry point: given DatasetArguments, return a PyTorch DataLoader
    ready for calibration.
    """
    # 1) 이미 DataLoader 면 그대로 반환
    if isinstance(dataset_args.dataset, DataLoader):
        return dataset_args.dataset

    # 2) 데이터셋 이름 기반 로더 디스패치
    if isinstance(dataset_args.dataset, str):
        dataset = _load_named_dataset(dataset_args.dataset, dataset_args)
    elif isinstance(dataset_args.dataset, (Dataset, DatasetDict)):
        dataset = dataset_args.dataset
    elif dataset_args.dataset_path:
        dataset = _load_custom_dataset(dataset_args)
    else:
        # 3) 데이터셋 없음 — data-free 파이프라인용
        return None

    # 4) 토크나이즈 및 전처리
    dataset = _tokenize_and_preprocess(dataset, processor, dataset_args)

    # 5) 샘플 수 제한과 shuffle
    dataset = _limit_and_shuffle(dataset, dataset_args)

    # 6) DataLoader 생성
    dataloader = DataLoader(
        dataset,
        batch_size=dataset_args.batch_size,
        num_workers=dataset_args.dataloader_num_workers,
        collate_fn=_get_collator(dataset_args),
    )

    return dataloader

단계별 동작:

이미 DataLoader 형태로 받았으면 그대로 사용
문자열이면 이름 기반 로더 호출 (c4, wikitext 등)
커스텀 경로면 dataset_path로 로딩
토크나이즈 + 전처리
num_calibration_samples만큼 자르고 셔플
PyTorch DataLoader로 래핑

`transformers/data/` 내장 데이터셋

transformers/data/
├── base.py                   # 베이스 클래스
├── c4.py                      # C4
├── wikitext.py                # WikiText
├── ultrachat_200k.py          # UltraChat
├── open_platypus.py           # OpenPlatypus
├── cnn_dailymail.py           # CNN/DailyMail
├── gsm8k.py                   # GSM8K
├── evolcodealpaca.py          # EvolCodeAlpaca
├── flickr_30k.py              # Flickr30k (멀티모달)
├── peoples_speech.py          # PeoplesSpeech (오디오)
└── custom.py                  # 사용자 CSV/JSON

각 파일은 "데이터셋 이름 → HuggingFace datasets.load_dataset 호출 + 전처리"를 담당한다.

`base.py`의 공통 베이스

class TextGenerationDataset:
    """Base class for calibration datasets"""

    def __init__(self, dataset_args, processor):
        self.dataset_args = dataset_args
        self.processor = processor

    def get_raw_dataset(self) -> Dataset:
        """Subclass implements: load raw dataset from HF hub"""
        raise NotImplementedError()

    def preprocess(self, sample: dict) -> dict:
        """
        Default preprocessing: tokenize `text_column` up to max_seq_length.
        Subclasses can override for chat templates, format conversion, etc.
        """
        text = sample[self.dataset_args.text_column]
        tokens = self.processor(
            text,
            truncation=True,
            max_length=self.dataset_args.max_seq_length,
            padding=(
                "max_length"
                if self.dataset_args.pad_to_max_length
                else False
            ),
            return_tensors="pt",
        )
        return {k: v.squeeze(0) for k, v in tokens.items()}

    def __call__(self) -> Dataset:
        dataset = self.get_raw_dataset()
        dataset = dataset.map(self.preprocess)
        return dataset

템플릿 메서드 패턴이다. 자식 클래스는 get_raw_dataset만 구현하면 되고, 토크나이즈는 부모가 처리한다. 채팅 데이터셋은 preprocess를 오버라이드해 채팅 템플릿을 적용한다.

`ultrachat_200k.py` 예시: 채팅 포맷

class UltraChat200k(TextGenerationDataset):
    """UltraChat 200k — chat-formatted multi-turn conversations"""

    def get_raw_dataset(self) -> Dataset:
        return load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")

    def preprocess(self, sample: dict) -> dict:
        # 대화 포맷: [{"role": "user", "content": ...}, {"role": "assistant", ...}, ...]
        messages = sample["messages"]

        # 채팅 템플릿 적용 (모델별로 다름)
        text = self.processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
        )

        # 일반 토크나이즈
        tokens = self.processor(
            text,
            truncation=True,
            max_length=self.dataset_args.max_seq_length,
            padding="max_length" if self.dataset_args.pad_to_max_length else False,
        )
        return tokens

apply_chat_template는 HuggingFace transformers의 기능으로, 모델별 채팅 템플릿을 자동으로 적용한다. LLaMA는 <|begin_of_text|>... 형식, Qwen은 <|im_start|>... 형식 등이 자동 처리된다.

`c4.py`: 가장 일반적인 선택

class C4(TextGenerationDataset):
    """C4 (Common Crawl) — general web text"""

    def get_raw_dataset(self) -> Dataset:
        return load_dataset(
            "allenai/c4",
            "en",
            split="train",
            streaming=self.dataset_args.streaming,
        )
    # preprocess 는 base 의 기본 구현 사용

C4는 거의 모든 캘리브레이션에 안전한 선택이다. 다양한 도메인의 웹 텍스트라 어떤 모델이든 "평균적인 분포"를 볼 수 있다. 128~512 샘플이면 대부분의 양자화에 충분하다.

`custom.py`: JSON/CSV 파일

사용자가 직접 만든 데이터셋을 쓰려면:

class CustomDataset(TextGenerationDataset):
    def get_raw_dataset(self) -> Dataset:
        path = self.dataset_args.dataset_path
        if path.endswith(".json"):
            return load_dataset("json", data_files=path)["train"]
        elif path.endswith(".csv"):
            return load_dataset("csv", data_files=path)["train"]
        else:
            raise ValueError(f"Unsupported file format: {path}")

사용자는 oneshot(dataset_path="./my_data.json", text_column="content")처럼 호출하면 된다.

`flickr_30k.py` 예시: 멀티모달

class Flickr30k(TextGenerationDataset):
    """Flickr30k — image-caption pairs for VLM calibration"""

    def get_raw_dataset(self) -> Dataset:
        return load_dataset("nlphuji/flickr30k", split="test")

    def preprocess(self, sample: dict) -> dict:
        # 이미지 + 텍스트 동시 처리
        image = sample["image"]       # PIL.Image
        caption = sample["caption"][0]  # 첫 번째 캡션 사용

        # VLM processor 는 이미지+텍스트를 함께 받아 토큰화
        inputs = self.processor(
            text=caption,
            images=image,
            return_tensors="pt",
            padding="max_length",
            max_length=self.dataset_args.max_seq_length,
        )
        return {k: v.squeeze(0) for k, v in inputs.items()}

VLM 캘리브레이션은 processor가 Tokenizer 대신 ProcessorMixin이어야 한다. 이미지와 텍스트를 함께 처리하는 멀티모달 모델용이다.

왜 이 설계인가

1. 이름 기반 디스패치. 사용자가 dataset="c4"만 쓰면 자동으로 적절한 로더가 호출된다. 데이터셋 설정을 매번 손으로 쓰지 않아도 된다.

2. 템플릿 메서드 패턴. TextGenerationDataset 베이스가 토크나이즈 로직을 공유하고, 자식이 get_raw_dataset과 (필요 시) preprocess만 오버라이드한다. 새 데이터셋 추가가 쉽다.

3. 채팅 템플릿 자동 적용. UltraChat 같은 채팅 데이터는 모델마다 다른 포맷을 쓴다. apply_chat_template을 쓰면 사용자가 이를 신경 쓸 필요가 없다.

4. 멀티모달 지원. Flickr30k, PeoplesSpeech 같은 데이터셋으로 VLM과 ASR 모델도 캘리브레이션 가능하다. 같은 로더 인터페이스로 통합된다.

5. 커스텀 경로 지원. dataset_path 인자로 CSV/JSON 파일을 직접 쓸 수 있어 도메인 특화 데이터셋(의료, 법률 등)에도 대응한다. text_column으로 어느 컬럼이 텍스트인지 지정한다.

마무리

Dataset Calibration은 llm-compressor의 "데이터 주입구"다. 표준 공개 데이터셋부터 커스텀 CSV까지 일관된 인터페이스로 지원한다. 다음 글부터는 부록 섹션으로 PyTorch Utils를 본다.

참고 자료

소스 코드:
- src/llmcompressor/datasets/utils.py
- src/llmcompressor/transformers/data/
관련 포스트:

llm-compressor 의 다른글

이전글 [llm-compressor] Modeling Overrides: DeepSeek/Llama4/Qwen3 등 모델별 패치
현재글 : [llm-compressor] Dataset Calibration: c4/wikitext/ultrachat 로더
다음글 [llm-compressor] PyTorch Utils: 희소화 통계와 모듈 헬퍼

댓글

관련 포스트

llm-compressor 의 다른글