[llm-compressor] Compression Save: compressed-tensors 체크포인트 저장

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

들어가며

llm-compressor가 만드는 결과물은 compressed-tensors 포맷의 HuggingFace 체크포인트다. 이 포맷은 vLLM/SGLang 추론 엔진이 직접 로딩할 수 있어야 한다. 압축된 가중치, 스케일, 제로포인트, 그룹 인덱스를 safetensors 파일에 저장하고, config.json에 quantization 메타데이터를 추가하는 등의 직렬화 로직을 src/llmcompressor/transformers/compression/가 담당한다.

핵심 구조/코드 분석

`transformers/compression/` 구성

transformers/compression/
├── __init__.py
├── compressed_tensors_utils.py    # save 시 config/weights 처리
└── helpers.py                     # 보조 유틸

핵심은 compressed_tensors_utils.py다. 여기에 save_pretrained_wrapper 같은 함수가 있어 HF save_pretrained를 감싼다.

`save_pretrained_wrapper`: 압축 저장 진입점

def save_pretrained_wrapper(
    model,
    save_directory: str,
    save_compressed: bool = True,
    **kwargs,
):
    """
    Wrapper around model.save_pretrained() that handles compressed-tensors
    serialization when save_compressed=True.
    """
    if not save_compressed:
        return model.save_pretrained(save_directory, **kwargs)

    # 1) compressed-tensors 의 ModelCompressor 생성
    from compressed_tensors.compressors import ModelCompressor
    compressor = ModelCompressor.from_pretrained_model(model)

    # 2) 가중치를 압축된 형태로 변환
    compressed_state_dict = compressor.compress(model)

    # 3) 모델 저장 (일반 state_dict 대신 compressed 버전 전달)
    model.save_pretrained(
        save_directory,
        state_dict=compressed_state_dict,
        safe_serialization=True,   # safetensors 강제
        **kwargs,
    )

    # 4) config.json 에 quantization_config 섹션 추가
    update_config_json(save_directory, compressor)

    # 5) 레시피도 저장 (재현성)
    if hasattr(model, "_recipe"):
        recipe_yaml = model._recipe.yaml()
        with open(Path(save_directory) / "recipe.yaml", "w") as f:
            f.write(recipe_yaml)

단계	작업
1	`ModelCompressor` 인스턴스 생성 (compressed-tensors 라이브러리)
2	가중치를 INT4/INT8/FP8 형태로 패킹
3	safetensors에 저장
4	config.json에 quantization_config 섹션 추가
5	recipe.yaml 저장 (선택적)

`compressed-tensors`의 역할

실제 비트 패킹은 compressed-tensors 라이브러리가 담당한다. llm-compressor는 이 라이브러리를 호출해:

INT4 가중치 8개를 하나의 INT32에 패킹
FP8 스케일을 FP32 텐서로 저장
Zero point를 필요하면 저장 (symmetric이면 생략)
G_idx 재배열 정보 저장 (GPTQ actorder 시)

각 quantization scheme에 맞는 compressor가 compressed_tensors.compressors에 등록되어 있다.

`update_config_json`: 메타데이터 주입

def update_config_json(save_directory: str, compressor: ModelCompressor):
    """Add quantization_config to config.json"""
    config_path = Path(save_directory) / "config.json"
    with open(config_path) as f:
        config = json.load(f)

    # compressor 가 생성한 quantization_config 추가
    config["quantization_config"] = compressor.get_quantization_config_dict()

    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)

이 quantization_config 섹션이 vLLM/SGLang이 체크포인트를 로딩할 때 참조하는 핵심 메타데이터다. 예시:

{
  "_name_or_path": "meta-llama/Llama-3-8B",
  "architectures": ["LlamaForCausalLM"],
  "hidden_size": 4096,
  ...
  "quantization_config": {
    "format": "pack-quantized",
    "quantization_status": "compressed",
    "config_groups": {
      "group_0": {
        "targets": ["Linear"],
        "weights": {
          "num_bits": 4,
          "type": "int",
          "symmetric": true,
          "strategy": "group",
          "group_size": 128
        }
      }
    },
    "ignore": ["lm_head"]
  }
}

vLLM이 이 config를 보고 "이 모델은 W4A16 compressed-tensors로 양자화되었다. 128 단위 group. lm_head는 제외"를 파악하고 적절한 kernel (Marlin, ExLlama 등)을 호출한다.

`helpers.py`의 보조 함수들

helpers.py에는 특수 케이스 처리 함수들이 있다.

def untie_word_embeddings_if_needed(model):
    """
    If lm_head and embed_tokens share weights but quantization plans differ,
    untie them so each can have its own scale.
    """
    lm_head = model.get_output_embeddings()
    embed = model.get_input_embeddings()

    if lm_head.weight.data_ptr() == embed.weight.data_ptr():
        # 실제 메모리 공유 → 분리
        lm_head.weight = torch.nn.Parameter(lm_head.weight.data.clone())


def strip_internal_attributes(model):
    """
    Remove calibration metadata that shouldn't be in checkpoint
    (e.g., _imatrix_importance, _pruning_mask, observer instances).
    """
    for module in model.modules():
        for attr in [
            "_imatrix_importance",
            "_pruning_mask",
            "weight_observer",
            "input_observer",
            "_imatrix_hook",
        ]:
            if hasattr(module, attr):
                delattr(module, attr)

**untie_word_embeddings_if_needed**는 중요하다. LLaMA·Qwen 같은 모델은 lm_head.weight와 embed_tokens.weight가 같은 메모리를 가리킨다. 양자화 시 lm_head는 FP16으로 두고 embed_tokens는 INT4로 하려면 분리가 필요하다. 이 함수가 .clone()으로 메모리를 분리한다.

**strip_internal_attributes**는 캘리브레이션 중 모듈에 붙은 내부 속성을 제거해 체크포인트가 깨끗하게 유지되도록 한다. _imatrix_importance, _pruning_mask 같은 것들이 대상이다.

저장 후 모델 재로딩

vLLM/SGLang이 이 체크포인트를 로딩할 때는 다음 과정을 거친다.

config.json의 quantization_config 읽기
compressed-tensors 라이브러리의 decompressor 호출
각 Linear에 적절한 QuantizationLinearMethod (GPTQ, AWQ, FP8 등) 부착
safetensors에서 패킹된 가중치 로딩
추론 준비 완료

llm-compressor는 이 흐름을 염두에 두고 "vLLM이 바로 로딩할 수 있는" 형태로 저장한다.

왜 이 설계인가

1. compressed-tensors 라이브러리 재사용. 실제 비트 패킹과 역패킹 로직은 그 라이브러리에 있다. llm-compressor는 "무엇을 압축할지"만 결정하고, "어떻게 압축할지"는 라이브러리에 위임한다. 관심사 분리.

2. save_pretrained 호환. save_pretrained_wrapper가 HF의 save_pretrained를 감싸기 때문에 사용자 경험은 표준 HF와 같다. model.save_pretrained("./out") 한 줄로 끝난다.

3. quantization_config JSON 섹션. vLLM/SGLang은 이 섹션 하나만 파싱하면 모든 양자화 정보를 알 수 있다. llmcompressor와 추론 엔진의 계약이 이 JSON 스키마로 정의된다.

4. Recipe 저장 옵션. recipe.yaml을 함께 저장하면 "이 체크포인트가 어떻게 만들어졌는가"를 추적할 수 있다. 실험 재현성과 감사(audit)에 도움이 된다.

5. 내부 속성 정리. strip_internal_attributes로 캘리브레이션 메타데이터를 지워 체크포인트가 유출되지 않게 한다. 보안과 크기 최적화 양쪽에 도움.

마무리

Compression Save는 llm-compressor와 추론 엔진을 잇는 마지막 다리다. 이 다음은 모델별 override인 Modeling Overrides를 본다.

참고 자료

라이브러리: compressed-tensors
소스 코드:
- src/llmcompressor/transformers/compression/compressed_tensors_utils.py
- src/llmcompressor/transformers/compression/helpers.py
관련 포스트:

llm-compressor 의 다른글

이전글 [llm-compressor] Transformers Tracing: 모델 그래프 추적과 부분 forward
현재글 : [llm-compressor] Compression Save: compressed-tensors 체크포인트 저장
다음글 [llm-compressor] Modeling Overrides: DeepSeek/Llama4/Qwen3 등 모델별 패치

댓글

관련 포스트

llm-compressor 의 다른글