[llm-compressor] Modeling Overrides: DeepSeek/Llama4/Qwen3 등 모델별 패치
들어가며
대부분의 LLM은 표준 Transformer 아키텍처를 따라 llm-compressor로 문제없이 압축된다. 하지만 일부 모델은 특수 구조나 최적화 때문에 기본 파이프라인이 실패한다. 대표적으로 MoE(Mixture of Experts) 모델이다. MoE 라우팅은 expert마다 다른 토큰을 받으므로, 일부 expert는 캘리브레이션 중 토큰을 전혀 받지 못해 관측 통계가 비게 된다. llm-compressor는 src/llmcompressor/modeling/ 디렉토리에 모델별 override를 둬서 이런 문제를 해결한다.
핵심 구조/코드 분석
디렉토리 구성
src/llmcompressor/modeling/
├── __init__.py
├── moe_context.py # MoE 캘리브레이션 공통 컨텍스트
├── offset_norm.py # RMSNorm 오프셋 보정
├── fuse.py # Fused weight 처리
├── deepseek_v3.py # DeepSeek-V3
├── deepseekv32/ # DeepSeek-V3.2 (서브디렉토리)
├── llama4.py # Llama-4
├── gemma4.py # Gemma-4
├── glm4_moe.py # GLM-4 MoE
├── glm_moe_dsa.py # GLM MoE with dense shared attention
├── gpt_oss.py # GPT-OSS
├── granite4.py # Granite-4
├── qwen3_moe.py # Qwen-3 MoE
├── qwen3_5_moe.py # Qwen-3.5 MoE
├── qwen3_next_moe.py # Qwen-3 Next MoE
├── qwen3_vl_moe.py # Qwen-3 VL MoE
├── afmoe.py # AFMoE
└── patch/
└── qwen3_omni_patch.py # Qwen-3 Omni 특수 패치
15개 이상의 모델별 파일이 있다. 각 파일은 모델의 특정 부분을 monkey-patch하거나 override한다.
moe_context.py: MoE 캘리브레이션 컨텍스트
가장 중요한 공통 유틸리티다. Oneshot 진입점 글에서 잠깐 본 moe_calibration_context가 여기에 정의되어 있다.
@contextmanager
def moe_calibration_context(model, calibrate_all_experts: bool = True):
"""
Context manager that temporarily modifies MoE routing during calibration.
When `calibrate_all_experts=True`, all experts receive all tokens (instead
of only their routed tokens). This ensures every expert's observer has
sufficient statistics for quantization.
"""
original_forwards = {}
# 1) 모델에서 MoE 블록 찾기
moe_modules = _find_moe_modules(model)
# 2) 각 MoE 모듈의 forward 를 "전수 라우팅" 버전으로 교체
for mod in moe_modules:
original_forwards[mod] = mod.forward
def _all_experts_forward(self, hidden_states, *args, **kwargs):
"""
Modified forward: pass hidden_states through EVERY expert,
regardless of routing.
"""
batch_size, seq_len, hidden_size = hidden_states.shape
# 각 expert 에 대해 모든 토큰 통과
for expert in self.experts:
_ = expert(hidden_states) # 통계만 수집, 결과 무시
# 원본 routing 도 실행해 올바른 출력 반환
return original_forwards[self](hidden_states, *args, **kwargs)
mod.forward = _all_experts_forward.__get__(mod)
try:
yield
finally:
# 3) 원본 forward 복원
for mod, orig in original_forwards.items():
mod.forward = orig
핵심: calibrate_all_experts=True일 때, 각 expert는 라우팅 결정과 무관하게 모든 토큰을 받는다. 이는 관측자가 모든 expert의 가중치 분포를 정확히 파악하도록 보장한다. 단점은 캘리브레이션 시간이 늘어나는 것이지만, 양자화 정확도 향상이 훨씬 크다.
deepseek_v3.py: DeepSeek-V3 특수 처리
DeepSeek-V3는 **Multi-head Latent Attention (MLA)**이라는 특수 attention을 쓴다. 표준 q/k/v projection 대신 "latent projection"을 거치는 구조다. 기본 AWQ/SmoothQuant 매핑이 이를 인식하지 못하면 잘못된 변환이 일어난다.
# deepseek_v3.py
def patch_for_calibration(model):
"""Apply DeepSeek-V3 specific patches"""
# 1) MLA 의 latent projection 을 q/k/v 로 분해 (개념적)
# 2) Shared expert 처리 — 그냥 Linear 로 취급 가능하도록 flag 조정
# 3) FFN 구조가 SwiGLU 변형이므로 activation smoothing 매핑 업데이트
for layer in model.model.layers:
if hasattr(layer, "self_attn"):
_patch_mla(layer.self_attn)
if hasattr(layer, "mlp") and hasattr(layer.mlp, "shared_experts"):
_mark_shared_expert_as_linear(layer.mlp)
각 모델별 파일은 해당 모델의 특수 구조를 llm-compressor의 일반 파이프라인과 호환되게 만든다.
fuse.py: Fused weight 처리
어떤 모델은 q/k/v를 하나의 큰 Linear로 fused해 저장한다(qkv_proj). 이런 경우 일반적인 per-projection 매핑이 작동하지 않는다. fuse.py는 이를 분해하거나, 또는 fused 상태 그대로 처리 가능하게 어댑터를 제공한다.
def unfuse_qkv(module: torch.nn.Linear, num_heads: int, head_dim: int):
"""Split fused qkv_proj into separate q_proj, k_proj, v_proj modules"""
out_features = module.out_features
assert out_features == 3 * num_heads * head_dim
qkv_weight = module.weight.data
q_w, k_w, v_w = qkv_weight.split(num_heads * head_dim, dim=0)
q_proj = torch.nn.Linear(module.in_features, num_heads * head_dim, bias=False)
k_proj = torch.nn.Linear(module.in_features, num_heads * head_dim, bias=False)
v_proj = torch.nn.Linear(module.in_features, num_heads * head_dim, bias=False)
q_proj.weight.data = q_w
k_proj.weight.data = k_w
v_proj.weight.data = v_w
return q_proj, k_proj, v_proj
이 분해는 캘리브레이션 시 일시적으로 적용되고, 저장 전에 다시 fused로 합쳐질 수 있다.
offset_norm.py: RMSNorm 오프셋
RMSNorm에는 trainable weight가 있고 일부 변형에는 bias 오프셋도 있다. 양자화 시 이들을 어떻게 처리할지가 문제다. norm_calibration_context는 RMSNorm을 표준 형태로 정규화해 다른 Modifier들이 일관되게 다룰 수 있게 한다.
패치 등록 메커니즘
모델별 패치는 __init__.py에서 모델 클래스 이름 기반으로 자동 디스패치된다.
# modeling/__init__.py
_MODEL_PATCHES = {
"DeepseekV3ForCausalLM": "deepseek_v3",
"DeepseekV32ForCausalLM": "deepseekv32",
"Llama4ForCausalLM": "llama4",
"Gemma4ForCausalLM": "gemma4",
"Qwen3MoeForCausalLM": "qwen3_moe",
...
}
def apply_model_patches(model):
"""Apply model-specific patches based on model class name"""
model_class = type(model).__name__
if model_class in _MODEL_PATCHES:
module_name = _MODEL_PATCHES[model_class]
patch_module = importlib.import_module(f"llmcompressor.modeling.{module_name}")
if hasattr(patch_module, "patch_for_calibration"):
patch_module.patch_for_calibration(model)
사용자는 이 과정을 신경 쓸 필요가 없다. Oneshot 진입점이 자동으로 apply_model_patches를 호출한다.
왜 이 설계인가
1. 모델별 파일 분리. 각 특수 아키텍처를 자체 파일에 두면 서로 간섭하지 않는다. DeepSeek의 패치가 Llama 처리를 건드릴 일이 없다. 새 모델 추가 시 새 파일만 만들면 된다.
2. MoE의 전수 라우팅. calibrate_all_experts=True가 MoE 양자화의 핵심 트릭이다. 이 컨텍스트 없이는 MoE 모델의 대부분 expert가 "통계 부족" 상태가 되어 양자화가 실패한다.
3. 컨텍스트 매니저 기반 패치. 모든 모델 override가 @contextmanager로 감싸져 있어 "캘리브레이션 중에만 적용, 이후 자동 복원"이 보장된다. 모델 원본이 영구적으로 수정될 위험이 없다.
4. 자동 디스패치. 모델 클래스 이름 기반 dispatcher가 있어 사용자가 특정 패치를 수동으로 활성화할 필요가 없다. "이 모델은 llm-compressor가 지원하는가?"만 확인하면 된다.
5. Fuse/unfuse 분리. Fused qkv 같은 구조는 캘리브레이션 중 분해했다가 저장 시 다시 fused로 합치는 방식으로 다룰 수 있다. 이 변환 로직이 fuse.py에 격리되어 있어 다른 Modifier가 단일 Linear로만 동작해도 된다.
마무리
Modeling Overrides는 llm-compressor가 "새 LLM 아키텍처가 나와도 빠르게 지원 가능한" 유연성의 핵심이다. 다음 글은 캘리브레이션 데이터 로더인 Dataset Calibration을 본다.
참고 자료
관련 포스트
llm-compressor 의 다른글
- 이전글 [llm-compressor] Compression Save: compressed-tensors 체크포인트 저장
- 현재글 : [llm-compressor] Modeling Overrides: DeepSeek/Llama4/Qwen3 등 모델별 패치
- 다음글 [llm-compressor] Dataset Calibration: c4/wikitext/ultrachat 로더
댓글