[llm-compressor] Args Dataclasses: 평평한 Kwargs를 세 개의 구조로 분리하기

들어가며

oneshot()는 40개 이상의 평평한 키워드 인자를 받는다. 하지만 그 인자들은 서로 다른 컴포넌트가 필요로 한다. model, precision은 HuggingFace from_pretrained에 전달되어야 하고, dataset, num_calibration_samples는 캘리브레이션 데이터로더가 쓰고, recipe, stage는 세션 초기화가 쓴다. 이 분리를 담당하는 것이 src/llmcompressor/args/ 아래의 세 dataclass다.

• ModelArguments — 모델 로딩/저장
• DatasetArguments — 캘리브레이션 데이터와 파이프라인
• RecipeArguments — 레시피 경로, 스테이지, 오버라이드 변수

사용자 입장에서는 이 dataclass들을 직접 import할 일이 거의 없다. 그러나 llm-compressor 내부에서는 거의 모든 함수가 이 세 객체를 받아 동작한다. 이 글은 각 dataclass의 필드 구성과, 평평한 kwargs가 어떻게 세 조각으로 분리되는지를 분석한다.

공식 문서

• transformers HfArgumentParser

핵심 구조/코드 분석

ModelArguments: 모델 로딩과 저장 파라미터

src/llmcompressor/args/model_arguments.py는 HuggingFace transformers의 HfArgumentParser가 기대하는 형식의 dataclass다.

@dataclass
class ModelArguments:
    model: str = field(                               # HF 허브 ID 또는 로컬 경로 (필수)
        metadata={"help": "A pretrained model ..."},
    )
    config_name: str | None = field(default=None)     # 별도 config 경로 (대부분 None)
    tokenizer: str | None = field(default=None)       # 별도 토크나이저 경로
    processor: str | None = field(default=None)       # 멀티모달 프로세서 경로
    use_auth_token: bool = field(default=False)       # HF 사설 모델 접근용 토큰 사용
    precision: str = field(default="auto")            # 가중치 dtype ("auto"/"fp16"/"bf16"/"fp32")
    tie_word_embeddings: bool = field(default=True)   # 입출력 임베딩 공유 유지 (False = 강제 untie)
    trust_remote_code_model: bool = field(default=False)  # 원격 커스텀 모델 코드 실행 허용
    save_compressed: bool = field(default=True)       # 저장 시 compressed-tensors 포맷 사용
    model_revision: str = field(default="main")       # HF 허브 브랜치/태그/커밋

필드	기본값	영향
model	필수	AutoModelForCausalLM.from_pretrained 의 첫 인자
precision	"auto"	torch_dtype 로 변환되어 로딩 시 캐스팅
tie_word_embeddings	True	False 시 lm_head.weight 를 embeddings와 분리해 별도 양자화 가능
save_compressed	True	save_pretrained 호출 시 compressed-tensors 직렬화 경로 선택
trust_remote_code_model	False	HF 허브의 modeling_*.py 실행 허용. 보안상 명시적 opt-in 필요

특히 tie_word_embeddings는 llm-compressor 특유의 고려사항이다. 많은 LLM이 입력 임베딩 행렬과 lm_head의 가중치를 공유(tie)하는데, 양자화를 적용할 때는 lm_head만 FP16으로 두고 나머지를 저비트로 양자화하는 것이 정확도 손실을 줄이는 일반적 전략이다. 이 때 공유된 상태로는 양자화가 불가능하므로, llm-compressor는 필요 시 이 둘을 분리(untie)한다.

RecipeArguments: 레시피와 세션 변수

src/llmcompressor/args/recipe_arguments.py는 가장 단순하다. 레시피 경로와 몇 개의 실행 옵션만 담는다.

@dataclass
class RecipeArguments:
    recipe: str | None = field(default=None)          # 레시피 YAML 경로 또는 Modifier 리스트
    recipe_args: list[str] | None = field(default=None)  # "key=value" 형식 오버라이드 목록
    clear_sparse_session: bool | None = field(default=False)  # 실행 간 Lifecycle 상태 리셋
    stage: str | None = field(default=None)           # 멀티스테이지 레시피에서 실행할 스테이지 이름

recipe_args는 특별하다. 레시피 YAML 안에 num_bits: $NUM_BITS처럼 Jinja 변수를 쓴 경우, 이 리스트의 값이 치환된다. CLI에서 --recipe_args "NUM_BITS=4,GROUP_SIZE=128"처럼 넘기면 parse_args가 이를 딕셔너리로 변환해 저장한다.

DatasetArguments: 캘리브레이션과 파이프라인

src/llmcompressor/args/dataset_arguments.py는 세 dataclass 중 가장 크고, 상속 계층을 활용한다. DVCDatasetArguments → CustomDatasetArguments → DatasetArguments 순서로 상속된다.

@dataclass
class DVCDatasetArguments:
    dvc_data_repository: str | None = field(default=None)  # DVC 레포 경로 (데이터 버저닝용)


@dataclass
class CustomDatasetArguments(DVCDatasetArguments):
    dataset_path: str | None = field(default=None)    # JSON/CSV/DVC 파일 경로
    text_column: str = field(default="text")          # 토크나이저에 넣을 텍스트 컬럼 키
    preprocessing_func: None | str | Callable = field(default=None)  # 채팅 템플릿 적용 함수
    batch_size: int = field(default=1)                # 캘리브레이션 배치 크기
    data_collator: str | Callable = field(default="truncation")  # 배치 수집 방식


@dataclass
class DatasetArguments(CustomDatasetArguments):
    dataset: str | Dataset | DatasetDict | DataLoader | None = field(default=None)
    num_calibration_samples: int = field(default=512) # 캘리브레이션 샘플 수
    max_seq_length: int = field(default=384)          # 토크나이즈 최대 길이
    shuffle_calibration_samples: bool = field(default=True)
    pad_to_max_length: bool = field(default=True)
    pipeline: str | None = field(default="independent")   # basic/sequential/data_free/independent
    sequential_targets: list[str] | None = field(default=None)  # sequential 대상 레이어 타입
    sequential_offload_device: str = field(default="cpu")       # 활성화 오프로드 디바이스
    quantization_aware_calibration: bool = field(default=True)
    sequential_prefetch: bool = field(default=False)   # 다음 배치 미리 로딩 (백그라운드 스레드)
    moe_calibrate_all_experts: bool = field(default=True)  # MoE 모든 전문가에 토큰 공급
    tracing_ignore: list[str] = field(default_factory=list)    # 트레이싱에서 건너뛸 함수 이름

필드	의미
num_calibration_samples	GPTQ/AWQ 권장치는 128~512. 적으면 통계 부족, 많으면 시간 낭비
max_seq_length	각 샘플의 토큰 수 상한. 128~2048 사이가 일반적
pipeline	Pipeline Registry가 최종 결정할 파이프라인의 힌트. 사용자 명시값이 Modifier 호환성을 만족하지 않으면 자동 승격
sequential_targets	sequential 파이프라인이 "레이어 단위"로 쪼갤 때, 어느 모듈 타입을 레이어 경계로 볼지. 보통 LlamaDecoderLayer 같은 이름
sequential_offload_device	중간 활성화를 CPU/다른 GPU로 내릴 디바이스. 단일 GPU에서 70B급 모델을 양자화할 때 중요
moe_calibrate_all_experts	MoE 모델 캘리브레이션 시 모든 expert에 토큰을 강제 공급. False면 라우팅된 expert만 통계를 봄

DatasetArguments는 상속 덕분에 "커스텀 데이터셋만 쓴다면 CustomDatasetArguments만 import해서 써도 된다"는 유연성을 제공한다. 실제로는 oneshot()의 parse_args가 항상 최하위 DatasetArguments를 사용한다.

parse_args: 평평한 Kwargs 분리기

세 dataclass를 kwargs에서 떼어내는 dispatcher는 src/llmcompressor/args/utils.py의 parse_args다.

def parse_args(
    **kwargs,                     # oneshot() 에서 넘어온 모든 키워드 인자
) -> tuple[
    ModelArguments,
    DatasetArguments,
    RecipeArguments | None,
    str | None,                  # output_dir (세 dataclass 중 어디에도 속하지 않음)
]:
    output_dir = kwargs.pop("output_dir", None)

    # HuggingFace 의 다중 dataclass 파서
    parser_args = (ModelArguments, DatasetArguments, RecipeArguments)
    parser = HfArgumentParser(parser_args)
    parsed_args = parser.parse_dict(kwargs)

    model_args, dataset_args, recipe_args = parsed_args

    # recipe_args 는 "key=value" 문자열 리스트 → 딕셔너리로 변환
    if recipe_args.recipe_args is not None:
        if not isinstance(recipe_args.recipe_args, dict):
            arg_dict = {}
            for recipe_arg in recipe_args.recipe_args:
                key, value = recipe_arg.split("=")
                arg_dict[key] = value
            recipe_args.recipe_args = arg_dict

    # deprecation 경고
    if dataset_args.remove_columns is not None:
        logger.warning(
            "`remove_columns` argument is depreciated. ...",
            DeprecationWarning,
        )

    # processor 를 tokenizer 로 자동 할당
    resolve_processor_from_model_args(model_args)

    return model_args, dataset_args, recipe_args, output_dir

핵심 트릭은 HfArgumentParser.parse_dict이다. 이 메서드는 주어진 kwargs 딕셔너리를 스캔해, 각 dataclass가 선언한 필드 이름과 매칭되는 것만 해당 dataclass로 라우팅한다. 세 dataclass의 필드 이름이 서로 겹치지 않기 때문에 충돌 없이 분리된다.

두 번째 중요한 부분은 recipe_args.recipe_args의 이중 변환이다. CLI에서 넘겨진 ["NUM_BITS=4", "GROUP_SIZE=128"] 문자열 리스트가 {"NUM_BITS": "4", "GROUP_SIZE": "128"} 딕셔너리로 바뀐다. 레시피 YAML에서 Jinja 템플릿 변수 치환 시 이 딕셔너리가 사용된다.

마지막 resolve_processor_from_model_args는 "사용자가 processor를 명시하지 않았지만 모델이 멀티모달이라면 tokenizer 경로를 processor로 넣는다" 같은 편의 로직이다. 이 덕분에 사용자는 LLM과 VLM을 구분 없이 동일 API로 사용할 수 있다.

왜 이 설계인가

1. Dataclass 세 개로 관심사를 분리한다. ModelArguments를 받는 함수는 데이터셋 옵션을 건드릴 수 없고, 반대도 마찬가지다. 타입 힌트만 봐도 함수의 책임이 드러난다. 리팩터링 시에도 한 dataclass의 필드를 추가해도 다른 dataclass와 충돌할 수 없다.

2. HuggingFace HfArgumentParser를 재사용한다. HuggingFace의 transformers.TrainingArguments와 같은 방식으로 파싱되므로, 이미 transformers에 익숙한 사용자는 바로 적응할 수 있다. 또 HfArgumentParser는 CLI argparse 인자 생성과 JSON 파싱까지 지원하므로, --model meta-llama/Llama-3-8B 같은 커맨드라인 진입점을 거의 공짜로 얻는다.

3. DatasetArguments 상속 계층. 커스텀 데이터셋 시나리오는 HuggingFace 데이터셋보다 필드가 적다. 상속 계층으로 이를 표현하면, CustomDatasetArguments만 사용하는 유닛 테스트에서는 HF 전용 필드를 무시할 수 있다. 상속은 OOP 이론적으로는 과한 결합이지만, dataclass에서는 "필드 집합의 부분집합 관계"로 해석되므로 실용적이다.

4. recipe_args 이중 타입. 이 필드는 입력 시 list[str]로 들어오지만 parse_args를 거치면 dict[str, str]로 변한다. 정적 타입 관점에서는 안티패턴이지만, 사용자가 CLI와 Python 양쪽에서 같은 키를 쓰려면 이 변환이 필요하다. 한쪽 전용 타입으로 통일하면 다른 쪽 사용성이 떨어진다.

5. output_dir만 분리. 세 dataclass 중 어느 곳에도 포함되지 않는 output_dir을 parse_args가 직접 pop해서 튜플의 네 번째 요소로 반환한다. 이는 output_dir이 "모델을 설명하는 속성"도 "데이터셋 속성"도 "레시피 속성"도 아닌 "이번 실행의 부작용 경로"이기 때문이다. 별도 dataclass를 만들기엔 너무 단일한 값이어서 최상위 변수로 둔다.

마무리

이 세 dataclass는 llm-compressor 내부의 "공식 API"다. 진입점만이 평평한 kwargs를 다루고, 그 아래 모든 함수는 이 구조화된 객체를 받아 동작한다. 다음 글에서는 RecipeArguments.recipe가 가리키는 Recipe DSL을 해부한다.

참고 자료

• 공식 문서: HuggingFace HfArgumentParser
• 소스 코드:
  • src/llmcompressor/args/model_arguments.py
  • src/llmcompressor/args/dataset_arguments.py
  • src/llmcompressor/args/recipe_arguments.py
  • src/llmcompressor/args/utils.py
• 관련 포스트:
  • Oneshot 진입점
  • Recipe DSL
  • Recipe Metadata
  • CompressionSession
  • Dataset Calibration