[llm-compressor] State & ModelLayer: 압축 상태 저장소
들어가며
Lifecycle이 "상태 머신"이라면 State는 "상태 저장소"다. 모델, 옵티마이저, 데이터셋, 하드웨어 정보를 한 객체에 모아두고 각 Modifier가 state=... 인자로 받아 쓴다. src/llmcompressor/core/state.py에 정의된 네 개의 dataclass(Data, Hardware, State, ModifiedState)를 분석한다.
핵심 구조/코드 분석
Data와 Hardware 서브 상태
@dataclass
class Data:
train: Any | None = None # 훈련 데이터 (QAT/Sparse Finetune 시)
val: Any | None = None # 검증 데이터
test: Any | None = None # 테스트 데이터
calib: Any | None = None # 캘리브레이션 데이터 (oneshot PTQ)
@dataclass
class Hardware:
device: str | None = None # 현재 디바이스 (cuda:0, cpu)
devices: list[str] | None = None # 모든 디바이스 목록
rank: int | None = None # 분산 학습 rank
world_size: int | None = None # 전체 노드 수
local_rank: int | None = None # 로컬 GPU rank
local_world_size: int | None = None # 로컬 GPU 개수
distributed: bool | None = None # 분산 학습 여부
distributed_strategy: str | None = None # DDP/FSDP/Deepspeed 등
Data는 네 종류의 데이터 중 어느 것이 세팅되었는지로 시나리오를 암시한다. calib만 세팅되면 oneshot PTQ, train/val까지 세팅되면 QAT라는 식이다. Hardware는 분산 훈련용 정보를 전부 포함하지만, oneshot PTQ에서는 device만 주로 쓰인다.
State 최상위 컨테이너
@dataclass
class State:
model: Any = None # 압축 대상 모델
teacher_model: Any = None # 지식 증류용 교사 (옵션)
optimizer: Any = None # 옵티마이저 (훈련 시)
optim_wrapped: bool = None # 옵티마이저 콜백 부착 여부
loss: Any = None # 현재 손실값
batch_data: Any = None # 현재 배치 데이터
data: Data = field(default_factory=Data) # 데이터셋 집합
hardware: Hardware = field(default_factory=Hardware) # 하드웨어 정보
loss_masks: list[torch.Tensor] | None = None # 손실 마스크 리스트
current_batch_idx: int = -1 # 현재 배치 인덱스
sequential_prefetch: bool = False # sequential pipeline 프리페치 여부
| 필드 | 용도 |
|---|---|
model |
nn.Module 인스턴스. Modifier가 파라미터 직접 접근 |
teacher_model |
증류 시 교사 모델. 기본 PTQ에서는 None |
optimizer |
파인튜닝 시 torch.optim.Optimizer |
batch_data |
파이프라인이 넘긴 현재 배치 |
data.calib |
DataLoader 인스턴스, 캘리브레이션 순회용 |
hardware.device |
현재 디바이스, Modifier가 텐서 생성 시 참조 |
loss_masks |
캘리브레이션 시 패딩 토큰 마스킹용 |
current_batch_idx |
몇 번째 배치인지. GPTQ 같은 누적 통계 Modifier가 배치별 가중치를 주기 위해 사용 |
sequential_prefetch |
백그라운드 배치 프리페치 플래그 |
compression_ready: 간단한 가드
@property
def compression_ready(self) -> bool:
ready = self.model is not None and self.optimizer is not None
return ready
옵티마이저가 None이면 False를 반환한다. 이는 훈련 시나리오 체크용이다. 순수 oneshot PTQ는 옵티마이저가 없으므로 compression_ready가 False지만, 이 속성 자체가 실패 조건이 아니라 단순 정보성 플래그이기 때문에 PTQ에서는 참조되지 않는다.
update: 인자 기반 필드 주입
def update(
self,
model: Any = None,
teacher_model: Any = None,
optimizer: Any = None,
attach_optim_callbacks: bool = True, # 옵티마이저 콜백 부착 여부
train_data: Any = None,
val_data: Any = None,
test_data: Any = None,
calib_data: Any = None,
copy_data: bool = True, # 데이터 복사 여부 (True=deepcopy)
start: float = None, # 훈련 시작 에폭
steps_per_epoch: int = None,
batches_per_step: int = None,
**kwargs,
) -> dict:
if model is not None:
self.model = model
if teacher_model is not None:
self.teacher_model = teacher_model
if optimizer is not None:
self.optim_wrapped = attach_optim_callbacks
self.optimizer = optimizer
if train_data is not None:
self.data.train = train_data if not copy_data else deepcopy(train_data)
if val_data is not None:
self.data.val = val_data if not copy_data else deepcopy(val_data)
if test_data is not None:
self.data.test = test_data if not copy_data else deepcopy(test_data)
if calib_data is not None:
self.data.calib = calib_data if not copy_data else deepcopy(calib_data)
if "device" in kwargs:
self.hardware.device = kwargs["device"]
return kwargs
눈에 띄는 것은 copy_data 플래그다. 기본값이 True라 데이터셋이 deepcopy된다. 이는 "사용자가 밖에서 DataLoader를 계속 쓰고 있더라도, llm-compressor가 내부에서 그 상태를 건드리지 않는다"를 보장한다. 큰 DataLoader는 복사 비용이 있지만, PyTorch DataLoader는 상태가 작아 부담이 크지 않다. 사용자가 의도적으로 복사를 끄려면 copy_data=False를 넘기면 된다.
또 update 메서드는 모든 필드를 조건부로 업데이트한다. 즉 None을 넘기면 해당 필드는 건드리지 않는다. 이는 부분 업데이트를 가능하게 한다. 예를 들어 두 번째 스테이지에서 모델은 그대로 두고 캘리브레이션 데이터만 교체하는 것이 가능하다.
ModifiedState: 실행 결과 스냅샷
@dataclass
class ModifiedState:
model: Any | None = None # 수정된 모델 참조
optimizer: Any | None = None # 수정된 옵티마이저
loss: Any | None = None # 수정된 손실
modifier_data: list[dict[str, Any]] | None = None # Modifier 별 반환 데이터
def __init__(self, model, optimizer, loss, modifier_data):
self.model = model
self.optimizer = optimizer
self.loss = loss
self.modifier_data = modifier_data
이 dataclass는 CompressionSession.initialize / finalize / event의 반환값이다. 실제로 모델은 인플레이스로 수정되므로 이 반환값은 "방금 일어난 변경의 스냅샷"이다. 사용자가 반환값을 받지 않아도 모델은 이미 수정되어 있다. modifier_data는 각 Modifier가 실행 중 반환한 디버그/상태 정보를 모은 리스트로, 로깅이나 커스텀 후처리에 활용된다.
커스텀 __init__이 dataclass 자동 생성 __init__을 덮어쓰는데, 실제 코드는 같은 동작을 한다. 이는 과거 호환성이거나, 향후 초기화 로직을 추가하기 위한 자리로 보인다.
core/model_layer.py: 레이어 참조 헬퍼
src/llmcompressor/core/model_layer.py는 매우 작은 파일로, 모델 레이어를 가리키는 경량 참조를 정의한다. 이는 Lifecycle이나 Modifier가 "이 레이어 이름으로 이 모듈을 가리킨다"를 표현할 때 쓰는 간단한 래퍼다. 실제 코드는 ModelParameterizedLayer 같은 매우 작은 dataclass로 구성되어 있으며, 모델 내 모듈의 name과 module 참조를 함께 들고 있다.
왜 이 설계인가
1. 단일 State 컨테이너. 모든 Modifier가 state= 인자 하나로 필요한 정보를 모두 받는다. 개별 파라미터로 인자를 10개씩 받는 것보다 훨씬 깨끗하다. 새 필드를 추가해도 Modifier 시그니처가 변하지 않는다.
2. Data와 Hardware 분리. State의 최상위에 평평하게 두지 않고 서브 dataclass로 나눈 이유는 의미적 그룹핑이다. state.data.calib은 state.calib_data보다 읽기 쉽고, state.hardware.device는 state.hw_device보다 의도가 명확하다.
3. copy_data=True 기본값. 기본값으로 깊은 복사를 하는 것은 "외부 상태를 건드리지 않는다"는 보수적 방어다. 성능이 우선인 사용자는 명시적으로 꺼야 한다. 안전한 기본값이다.
4. update의 부분 업데이트. 모든 필드가 조건부로 업데이트되어 "하나만 바꾸는" 시나리오가 자연스럽다. 두 번째 initialize 호출에서 모델 재설정 없이 새 데이터만 주입할 수 있다.
5. ModifiedState 별도 타입. State를 그대로 반환해도 되지만, "반환된 시점의 스냅샷"임을 명시하기 위해 별도 타입으로 분리한다. 타입 힌트만 봐도 "이것은 read-only 결과값"임을 알 수 있다.
마무리
State는 llm-compressor 내부 모든 컴포넌트가 공유하는 "공용 메모장"이다. 데이터 플로우의 중심 자료구조로서, 얇지만 모든 Modifier와 Pipeline이 의존한다. 다음 글에서는 배치 단위 훅을 구성하는 Events 시스템을 분석한다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] Lifecycle: Modifier 초기화-이벤트-종료 상태 머신
- 현재글 : [llm-compressor] State & ModelLayer: 압축 상태 저장소
- 다음글 [llm-compressor] Events: 배치 라이프사이클 훅과 에폭 계산 로직
댓글