[llm-compressor] Gemma4 MoE 모델 양자화를 위한 llm-compressor 지원 추가 분석

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

PR 링크: vllm-project/llm-compressor#2565 상태: Merged | 변경: +None / -None

들어가며

최근 대규모 언어 모델(LLM)의 발전은 Mixture-of-Experts (MoE) 아키텍처를 통해 더욱 가속화되고 있습니다. MoE는 모델의 크기를 크게 늘리지 않으면서도 파라미터 효율성을 높여 성능을 향상시키는 강력한 방법론입니다. 하지만 MoE 모델은 일반적인 Dense 모델과 구조가 달라, 양자화(Quantization)와 같은 최적화 기법을 적용하기 위해서는 특별한 고려가 필요합니다.

본 기술 블로그 글에서는 vllm-project/llm-compressor 레포지토리에 제출된 PR(#123, 가상 PR 번호)을 분석합니다. 이 PR은 Gemma4 MoE 모델에 대한 양자화 지원을 추가하여, MoE 모델의 전문가(expert) 단위 최적화를 가능하게 하는 중요한 개선을 담고 있습니다. 특히, FP8 Dynamic 및 NVFP4 (W4A4) 양자화 예제를 포함하여 실제 적용 사례를 제시합니다.

이 글에서는 해당 PR의 코드 변경사항을 상세히 분석하고, 각 변경이 왜 좋은 최적화인지, 그리고 MoE 모델 최적화에 있어 어떤 일반적인 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

이번 PR은 주로 Gemma4 MoE 모델을 llm-compressor 프레임워크에서 지원하기 위한 핵심적인 수정과 새로운 예제 파일 추가에 집중하고 있습니다. 주요 변경 사항은 다음과 같습니다.

1. `src/llmcompressor/modeling/gemma4.py` - Gemma4 MoE 모델 지원 모듈

이 파일은 Gemma4 MoE 모델의 전문가(expert)들을 llm-compressor의 양자화 파이프라인에서 올바르게 처리할 수 있도록 하는 핵심 로직을 담고 있습니다.

핵심 변경:

SequentialGemma4TextExperts 클래스 추가: 이 클래스는 Hugging Face의 Gemma4TextExperts 모듈을 래핑하여, MoE 모델의 전문가 가중치를 양자화 및 최적화가 가능한 형태로 '언패킹(unpacking)'하는 역할을 합니다. MoE 모델은 종종 여러 전문가의 가중치를 하나의 텐서로 묶어 저장하는데, 이를 개별 전문가 MLP 모듈로 분리해야 각 전문가 레이어를 타겟팅하여 양자화를 적용할 수 있습니다.
```
@MoECalibrationModule.register("Gemma4TextExperts")
class SequentialGemma4TextExperts(MoECalibrationModule):
    # ... (클래스 정의)
    def __init__(
        self,
        original: Gemma4TextExperts,
        config: Gemma4Config,
        calibrate_all_experts: bool = True,
    ):
        super().__init__()
        # ...
        expert_list = Gemma4TextExpertsList(config.text_config, original)
        for i, expert in enumerate(expert_list):
            self.add_module(str(i), expert)
    # ...
```
__init__ 메서드에서 Gemma4TextExpertsList를 사용하여 원본 Gemma4TextExperts 객체로부터 개별 전문가 MLP 모듈을 생성하고, 이를 self.add_module(str(i), expert)를 통해 클래스의 자식 모듈로 등록합니다. 이렇게 하면 QuantizationModifier에서 targets="Linear"와 같이 지정했을 때, 각 전문가의 Linear 레이어들이 개별적으로 타겟팅될 수 있습니다.

Gemma4TextExpertsList 클래스 추가: 이 클래스는 Gemma4TextExperts의 3D 전문가 가중치 텐서를 실제 transformers.models.gemma4.modeling_gemma4.Gemma4TextMLP 모듈의 리스트로 변환합니다. 이는 Hugging Face의 Gemma4TextExperts가 gate_up_proj와 down_proj를 3차원 텐서로 관리하는 방식 때문에 필요합니다.

class Gemma4TextExpertsList(torch.nn.ModuleList):
    # ...
    def __init__(self, config: Gemma4TextConfig, original: Gemma4TextExperts):
        from transformers.models.gemma4.modeling_gemma4 import Gemma4TextMLP

        # ...
        super().__init__([
            Gemma4TextMLP(config, layer_idx=0) for _ in range(self.num_experts)
        ])

        gate_up_data = original.gate_up_proj.data  # [num_experts, 2*inter, hidden]
        down_data = original.down_proj.data  # [num_experts, hidden, inter]

        for i in range(self.num_experts):
            gate_up = gate_up_data[i]  # [2*intermediate, hidden]
            down = down_data[i]  # [hidden, intermediate]

            # ... (가중치 복사 및 contiguous)
            self[i].gate_proj.weight.data = (
                gate_up[:intermediate_size, :].clone().contiguous()
            )
            self[i].up_proj.weight.data = (
                gate_up[intermediate_size:, :].clone().contiguous()
            )
            self[i].down_proj.weight.data = down.clone().contiguous()

리뷰어 kylesayrs가 지적했듯이, clone().contiguous()는 clone()만으로도 충분할 수 있습니다. 하지만 명시적으로 contiguous를 보장하는 것은 잠재적인 메모리 접근 문제를 방지하는 데 도움이 될 수 있습니다.

2. `examples/quantization_w4a4_fp4/gemma4_example.py` - NVFP4 (W4A4) 양자화 예제

이 파일은 Gemma4 MoE 모델에 대해 4비트 가중치 및 4비트 활성화(NVFP4) 양자화를 수행하는 방법을 보여주는 예제 스크립트입니다.

핵심 변경:

모델 및 프로세서 로드:

MODEL_ID = "google/gemma-4-26B-A4B-it"
model = Gemma4ForConditionalGeneration.from_pretrained(MODEL_ID, dtype="auto")
processor = AutoProcessor.from_pretrained(MODEL_ID)

QuantizationModifier 설정:
```
recipe = QuantizationModifier(
    targets="Linear",
    scheme="NVFP4",
    ignore=[
        "lm_head",
        "re:.*embed.*",
        "re:.*router",
        "re:.*vision_tower.*",
    ],
)
```
scheme="NVFP4"를 사용하여 4비트 양자화를 지정합니다. ignore 리스트는 양자화에서 제외할 레이어들을 명시합니다. re:.*router는 MoE 모델의 라우터 레이어가 양자화 대상에서 제외됨을 의미합니다. 이는 라우터는 일반적으로 작은 크기이고, 양자화 시 성능 저하가 클 수 있기 때문입니다.
데이터셋 로딩 및 전처리:
```
DATASET_ID = "neuralmagic/calibration"
NUM_CALIBRATION_SAMPLES = 20
MAX_SEQUENCE_LENGTH = 8192
# ... (preprocess_function, data_collator 정의)
ds = load_dataset(DATASET_ID, name="LLM", split=f"train[:{NUM_CALIBRATION_SAMPLES}]")
ds = ds.map(preprocess_function, batched=False, remove_columns=ds.column_names)
```
양자화 후 성능을 보정하기 위한 calibration 데이터셋을 로드하고, 모델에 맞게 전처리하는 함수(preprocess_function)와 데이터 콜레이터(data_collator)를 정의합니다. 리뷰어 kylesayrs는 data_collator에서 torch.tensor(value, dtype=model.dtype).squeeze(0)로 변경하는 것을 제안했습니다. 이는 모델의 원래 데이터 타입(model.dtype)을 사용하여 불필요한 타입 변환을 방지하려는 의도로 보입니다.
oneshot 함수 호출 및 모델 저장:
```
oneshot(
    model=model,
    recipe=recipe,
    dataset=ds,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    data_collator=data_collator,
)
# ... (model.save_pretrained)
```
llmcompressor.oneshot 함수를 호출하여 실제 양자화 과정을 수행합니다. calibration 데이터셋과 관련 설정을 전달하여 최적의 양자화 파라미터를 찾습니다. 이후 save_pretrained를 통해 양자화된 모델을 저장합니다.

3. `examples/quantization_w8a8_fp8/gemma4_example.py` - FP8 Dynamic 양자화 예제

이 파일은 Gemma4 MoE 모델에 대해 FP8 Dynamic 양자화를 수행하는 예제입니다.

핵심 변경:

QuantizationModifier 설정:

recipe = QuantizationModifier(
    targets="Linear",
    scheme="FP8_DYNAMIC",
    ignore=[
        "lm_head",
        "re:.*embed.*",
        "re:.*router",
        "re:.*vision_tower.*",
    ],
)

scheme="FP8_DYNAMIC"을 사용하여 FP8 양자화를 지정합니다. 이는 가중치는 FP8으로, 활성화는 동적으로 FP8으로 양자화하는 방식입니다.

oneshot 함수 호출 및 모델 저장:
```
oneshot(model=model, recipe=recipe)
# ... (model.save_pretrained)
```
이 예제에서는 calibration 데이터셋을 명시적으로 사용하지 않고 oneshot을 호출합니다. 이는 FP8_DYNAMIC 스킴이 calibration 없이도 적용 가능하거나, 혹은 기본 설정을 사용함을 의미할 수 있습니다. (주: 실제 llm-compressor의 동작 방식에 따라 다를 수 있습니다.)

리뷰어 dsikka와 kylesayrs가 /raid/engine/dsikka/와 같은 절대 경로를 제거하고 상대 경로를 사용하도록 수정을 요청했습니다. 이는 예제 코드가 특정 환경에 종속되지 않도록 하기 위함입니다.

왜 이게 좋은가?

이 PR은 여러 가지 측면에서 좋은 최적화 및 개선을 보여줍니다.

MoE 모델 지원 확대: llm-compressor는 원래 Dense 모델에 강점을 보였으나, 이 PR을 통해 Gemma4와 같은 최신 MoE 아키텍처에 대한 지원을 공식적으로 추가했습니다. 이는 MoE 모델의 활용도를 높이고, 더 많은 사용자가 MoE 모델을 효율적으로 최적화할 수 있게 합니다.
전문가 단위 최적화 가능: MoE 모델의 핵심은 전문가(expert)입니다. 이 PR은 전문가 가중치를 개별 Linear 모듈로 분리함으로써, 각 전문가에 대해 독립적인 양자화 설정을 적용하거나, 특정 전문가만 타겟팅하는 등 세밀한 최적화가 가능해집니다. 이는 모델의 성능 저하를 최소화하면서 압축률을 극대화하는 데 중요합니다.
다양한 양자화 스킴 지원: FP8 Dynamic 및 NVFP4 (W4A4)와 같은 최신 양자화 스킴을 예제로 제공함으로써, 사용자는 다양한 비트 수와 양자화 기법을 쉽게 시도해볼 수 있습니다. 특히 NVFP4는 4비트 양자화로 상당한 모델 크기 감소와 추론 속도 향상을 기대할 수 있습니다.
실제 적용 예제 제공: 단순히 기능 추가에 그치지 않고, 실제 Hugging Face 모델(google/gemma-4-26B-A4B-it)을 사용한 예제 스크립트를 제공하여 사용자들이 쉽게 따라 하고 자신의 모델에 적용해볼 수 있도록 합니다. 이는 라이브러리의 실용성을 크게 높입니다.
코드 품질 및 유지보수성 향상: 리뷰 댓글을 통해 제기된 경로 문제, clone().contiguous() 중복 제거 등은 코드의 견고성과 유지보수성을 향상시키는 데 기여합니다. Gemma4TextExpertsList와 같이 명확한 역할을 가진 클래스를 분리하는 것은 코드의 가독성을 높입니다.

일반적 교훈:

아키텍처별 특화 지원의 중요성: MoE와 같이 특수한 아키텍처는 일반적인 Dense 모델과는 다른 방식으로 최적화해야 합니다. 이러한 아키텍처의 내부 구조를 이해하고, 이를 프레임워크가 처리할 수 있는 형태로 변환하는 것이 중요합니다.
실용적인 예제의 가치: 아무리 강력한 기능이라도 사용자가 쉽게 이해하고 적용할 수 있는 예제가 없다면 그 가치가 반감됩니다. 실제 모델과 데이터셋을 사용한 예제는 라이브러리 채택률을 높이는 데 결정적인 역할을 합니다.
리뷰 프로세스의 중요성: 코드 리뷰는 단순히 버그를 찾는 것을 넘어, 코드 품질, 성능, 유지보수성 등 다양한 측면에서 개선점을 발견하고 공유하는 중요한 과정입니다. 리뷰어들의 피드백을 적극적으로 반영하는 것이 프로젝트의 완성도를 높입니다.

결론

이번 PR은 llm-compressor가 Gemma4 MoE 모델에 대한 양자화 지원을 성공적으로 추가함으로써, MoE 모델 최적화 분야에 중요한 기여를 했습니다. 전문가 단위의 세밀한 최적화와 다양한 양자화 스킴 지원, 그리고 실제 적용 예제 제공은 이 PR의 가치를 더욱 높입니다. 앞으로도 llm-compressor가 다양한 최신 LLM 아키텍처에 대한 최적화 지원을 확대해나가기를 기대합니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

PR Analysis 의 다른글

이전글 [sglang] SGLang NIXL 이기종 TP 환경에서 디스어그리게이션 KV 캐시 전송 버그 수정 및 성능 개선
현재글 : [llm-compressor] Gemma4 MoE 모델 양자화를 위한 llm-compressor 지원 추가 분석
다음글 [sglang] sglang 성능 최적화: torch.compile 퓨전 복원을 통한 TopK 후처리 개선

댓글

관련 포스트

PR Analysis 의 다른글