[llm-compressor] Recipe Metadata: 직렬화 헬퍼와 모델 메타데이터 구조
들어가며
Recipe DSL 글에서 살펴본 대로, Recipe 객체는 YAML/JSON 문자열을 Modifier 인스턴스 리스트로 변환한다. 그런데 그 반대 방향, 즉 "메모리에 있는 Modifier 리스트를 다시 YAML로 저장"하는 것도 필요하다. 실험을 재현 가능하게 기록하려면 사용된 레시피를 체크포인트와 함께 저장해야 하고, 여러 스테이지로 구성된 레시피를 점진적으로 쌓아가려면 두 YAML을 병합할 수 있어야 한다. 이 글은 src/llmcompressor/recipe/utils.py의 직렬화/병합 헬퍼와, src/llmcompressor/recipe/metadata.py의 메타데이터 Pydantic 모델을 분석한다.
공식 문서
핵심 구조/코드 분석
_load_json_or_yaml_string: 포맷 자동 감지
문자열을 받아 JSON인지 YAML인지 자동으로 판별한다.
def _load_json_or_yaml_string(content: str) -> Dict[str, Any]:
# JSON 먼저 시도, 실패하면 YAML
try:
ret = json.loads(content)
except json.JSONDecodeError:
try:
ret = yaml.safe_load(content)
except yaml.YAMLError as err:
raise ValueError(f"Could not parse recipe from string {content}") from err
# 루트가 dict 이어야 함 — 스칼라/리스트는 레시피가 아님
if not isinstance(ret, dict):
raise ValueError(
f"Could not parse recipe from string {content}. If you meant load from "
"a file, please make sure that the specified file path exists"
)
return ret
yaml.safe_load는 JSON을 부분집합으로 받아들이기 때문에 YAML만 써도 충분할 것 같지만, JSON을 먼저 시도하는 이유는 더 빠르기 때문이다. 특히 yaml.safe_load는 순수 파이썬 구현일 경우 수 밀리초 단위로 느리다. JSON 파서는 C 구현이라 거의 즉시 반환한다.
마지막의 dict 체크는 사용자 실수 방지다. 사용자가 레시피 "파일 경로"를 recipe=에 넣었는데 그 경로가 존재하지 않으면 llm-compressor는 이를 "문자열 레시피"로 간주해 파싱을 시도한다. 경로 문자열은 당연히 dict가 아니라 스칼라로 파싱되므로, 이 체크가 친절한 에러 메시지를 만든다.
_parse_recipe_from_md: 마크다운 Front Matter 추출
마크다운 파일을 레시피로 쓸 수 있는 이유다.
def _parse_recipe_from_md(file_path, yaml_str):
# `---` 로 둘러싸인 YAML front matter 블록을 정규식으로 추출
yaml_delim = r"(?:---|\+\+\+)"
yaml_body = r"(.*?)"
re_pattern = r"^\s*" + yaml_delim + yaml_body + yaml_delim
regex = re.compile(re_pattern, re.S | re.M)
result = regex.search(yaml_str)
if result:
yaml_str = result.group(1)
else:
raise RuntimeError(
"Could not extract YAML front matter from recipe card: {}".format(file_path)
)
return yaml_str
---로 감싼 YAML front matter를 추출한다. Jekyll/Hugo 블로그 포스트의 YAML 헤더와 같은 관례다. +++도 허용하는 것은 Hugo 사용자를 위한 배려다. re.S | re.M은 "다중 라인에서 .이 개행을 포함하도록" 만드는 플래그 조합이다.
get_yaml_serializable_dict: Modifier 리스트 → 레시피 딕셔너리
Recipe.yaml() 호출의 실제 변환기다.
def get_yaml_serializable_dict(modifiers: List[Modifier], stage: str) -> Dict[str, Any]:
stage_dict = {}
stage_name = stage + "_stage" # 스테이지 이름에 `_stage` 접미 추가
stage_dict[stage_name] = {}
for modifier in modifiers:
# 각 Modifier 의 group 속성이 있으면 사용, 없으면 스테이지 이름 재활용
group = getattr(modifier, "group", stage) or stage
group_name = f"{group}_modifiers" # 그룹 이름에 `_modifiers` 접미 추가
modifier_type = modifier.__class__.__name__ # 예: "GPTQModifier"
# Pydantic model_dump() 로 dict 추출, None/내부 필드/group 제외
args = {
k: v
for k, v in modifier.model_dump().items()
if v is not None and not k.endswith("_") and k != "group"
}
# 같은 그룹이 없으면 초기화
if group_name not in stage_dict[stage_name]:
stage_dict[stage_name][group_name] = {}
stage_dict[stage_name][group_name][modifier_type] = args
return stage_dict
출력 구조는 {stage_name: {group_name: {modifier_class_name: args}}} 네 단계 중첩이다. 이는 정확히 Recipe DSL 글에서 본 YAML 형식과 일치한다. _load_json_or_yaml_string → from_dict 역순으로 돌아가는 구조다.
특별히 주목할 점 두 가지.
- None 값 필터링. Modifier의 Pydantic 필드 중 값이
None인 것은 직렬화에서 제외한다. 기본값으로 돌아가므로 굳이 YAML에 적을 필요가 없다. 결과 YAML이 훨씬 깔끔해진다. k.endswith("_")필터링. 언더스코어로 끝나는 필드는 "내부 필드"로 간주되어 제외된다. Modifier 구현체에서 내부 상태를 이 컨벤션으로 표시하면 자동으로 직렬화 대상에서 빠진다.
filter_dict: 특정 스테이지만 선택
def filter_dict(obj: dict, target_stage: Optional[str] = None) -> dict:
if not target_stage:
return obj
return {k: v for k, v in obj.items() if k.startswith(target_stage)}
멀티스테이지 레시피에서 한 스테이지만 실행할 때 사용된다. Recipe.create_instance(..., target_stage="quant_stage")로 호출하면 레시피의 다른 스테이지들은 파싱 단계에서 떨어져 나간다. 이는 from_dict가 빈 딕셔너리를 받아도 동작하도록 설계된 덕에 가능하다.
append_recipe_dict: 두 레시피 딕셔너리 병합
가장 미묘한 헬퍼다. 두 레시피를 합칠 때 스테이지 키 충돌을 어떻게 처리할지가 문제다.
def append_recipe_dict(d1: dict, d2: dict) -> dict:
"""
If both have the same stage key (e.g. 'test_stage'), the result will contain:
'test_stage_0', 'test_stage_1', etc.
"""
result = dict(d1)
for key, val in d2.items():
if key not in result:
result[key] = val
else:
# 스테이지 키 충돌 — 양쪽에 번호 접미 추가
base_key = re.sub(r"_\d+$", "", key)
# 원본이 아직 번호가 안 붙었으면 0번부터 시작
if key == base_key:
result[f"{base_key}_0"] = result.pop(key)
result[f"{base_key}_1"] = val
else:
# 이미 번호가 붙은 키라면 다음 빈 번호 찾기
i = 1
while f"{base_key}_{i}" in result:
i += 1
result[f"{base_key}_{i}"] = val
return result
| 상황 | 결과 |
|---|---|
| 키가 충돌하지 않음 | 그대로 병합 |
test_stage가 양쪽에 있음 |
test_stage_0, test_stage_1로 변경 |
test_stage_0이 한쪽에만 있고 test_stage가 다른 쪽에 있음 |
다음 빈 번호 찾아서 test_stage_N 생성 |
이 동작 덕에 사용자는 "기존 pruning 스테이지가 있는 파일에 같은 이름의 스테이지를 새로 추가해도 덮어쓰이지 않는다"는 보장을 얻는다. 실험 로그의 안전성을 위한 설계다.
metadata.py: 모델/데이터셋 메타데이터 Pydantic 모델
src/llmcompressor/recipe/metadata.py는 레시피 실행 시 참조되는 메타데이터 구조를 Pydantic 모델로 정의한다.
class DatasetMetaData(BaseModel):
name: str = None # 데이터셋 이름 (예: "c4")
version: str = None # 데이터셋 버전
hash: str = None # 데이터셋 해시 (재현성 체크)
shape: list[int] = Field(default_factory=list)
num_classes: int = None # 분류 태스크일 경우
num_train_samples: int = None
num_val_samples: int = None
num_test_samples: int = None
class ParamMetaData(BaseModel):
name: str = None # 파라미터 이름 (예: "model.layers.0.self_attn.q_proj.weight")
shape: list[int] = None # 텐서 shape
weight_hash: str = None # 가중치 해시 (압축 전후 비교용)
class LayerMetaData(BaseModel):
name: str = None # 레이어 이름
type: str = None # 레이어 타입 (예: "Linear")
index: int = None # 레이어 인덱스
attributes: dict[str, Any] = None # 커스텀 속성
input_shapes: list[list[int]] = None # 입력 텐서 shape 리스트
output_shapes: list[list[int]] = None # 출력 텐서 shape 리스트
params: dict[str, ParamMetaData] = None # 이 레이어가 소유한 파라미터들
class ModelMetaData(BaseModel):
architecture: str = None # 모델 아키텍처 이름 (예: "LlamaForCausalLM")
sub_architecture: str = None # 서브 아키텍처 (예: "decoder-only")
input_shapes: list[list[int]] = None
output_shapes: list[list[int]] = None
layers: list[LayerMetaData] = Field(default_factory=list)
layer_prefix: str | None = None # 모델 레이어 접두 (예: "model.layers")
이 모델들은 "실행 당시의 스냅샷"을 기록하는 용도다. 예를 들어 압축 후 체크포인트에 ModelMetaData를 함께 저장해두면, 나중에 "이 체크포인트는 어떤 데이터셋으로 어떤 레이어들에 어떤 해시를 가진 가중치에 적용되었는가"를 검증할 수 있다.
현재 코드베이스에서는 이 메타데이터가 전면적으로 쓰이지는 않지만, 향후 재현성 기능(예: Weights & Biases 연동, 체크포인트 lineage 추적)을 위한 스켈레톤으로 유지되고 있다. Pydantic 모델이기 때문에 JSON 직렬화와 Swagger 스펙 생성이 거의 공짜다.
왜 이 설계인가
1. JSON 우선 파서 순서. JSON이 YAML의 부분집합이지만, JSON을 먼저 시도하면 파싱 성능이 크게 좋아진다. JSON을 사용하는 프로그램적 입력(예: 다른 툴이 생성한 레시피)은 몇 마이크로초만에 처리된다.
2. 마크다운 레시피 지원. 실험 기록과 실행 코드를 한 파일에 두려는 연구자 워크플로우에 맞춘 설계다. Jupyter 노트북과 비슷한 철학이지만, 파일 하나만 있으면 된다는 점이 다르다.
3. None 값과 내부 필드 자동 필터링. 직렬화 시 "사용자가 명시하지 않은 기본값"과 "내부 상태"를 자동으로 제외해, 레시피 YAML이 사람이 읽기 좋은 최소 형태로 유지된다. 이는 Pydantic의 기본 model_dump()만 쓰면 얻을 수 없는 llm-compressor 특유의 컨벤션이다.
4. 충돌 시 번호 접미. 병합 시 덮어쓰지 않고 번호 접미를 추가하는 정책은 데이터 손실을 원천 차단한다. 사용자가 실수로 같은 이름의 스테이지를 두 번 병합해도 양쪽 데이터가 다 보존된다.
5. 메타데이터 Pydantic 모델은 확장 지점. 현재 전면 사용되지는 않지만, 향후 "어느 체크포인트가 어느 데이터로 어떻게 만들어졌는가"를 추적할 준비가 되어 있다. 지금 정의해두면 나중에 기능 추가 시 마이그레이션 부담이 없다.
마무리
Recipe의 메타데이터 계층은 "사람이 쓴 YAML을 객체로 읽고, 객체를 다시 YAML로 쓰고, 두 YAML을 안전하게 병합하고, 재현성을 위해 메타데이터를 남긴다"는 네 가지 기본기를 갖추고 있다. 이 계층 위에 CompressionSession이 얹혀져 Modifier를 실제로 실행한다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] Recipe DSL: YAML로 Modifier를 조합하는 선언적 언어
- 현재글 : [llm-compressor] Recipe Metadata: 직렬화 헬퍼와 모델 메타데이터 구조
- 다음글 [llm-compressor] CompressionSession: 전역 싱글톤 세션과 Lifecycle 래퍼
댓글