본문으로 건너뛰기

[SGLang] AutoRound: 자동 라운딩 최적화 양자화

들어가며

양자화에서 라운딩(반올림)은 정확도 손실의 주요 원인이다. Round-to-Nearest(RTN)는 가장 가까운 양자화 레벨로 반올림하지만, 이것이 항상 최적은 아니다. AutoRound는 라운딩 방향을 학습하여 출력 오차를 최소화하는 기법이다. SGLang은 python/sglang/srt/layers/quantization/auto_round.py에서 AutoRound 양자화 모델의 추론을 지원한다.

구조도

AutoRoundConfig
├── weight_bits: {2, 3, 4, 8}
├── group_size: int
├── sym: bool (대칭 양자화 여부)
├── packing_format: "auto_round:auto_gptq" | "auto_round:auto_awq"
├── backend: "auto" | "gptq" | "awq" | "marlin" | ...
├── block_name_to_quantize: Optional[List[str]]
├── extra_config: Optional[Dict]  (레이어별 설정)
└── data_type: "int"

get_quant_method()
├── packing_format에 "gptq" → apply_gptq_quant_layer()
└── packing_format에 "awq"  → apply_awq_quant_layer()
    ├── Marlin 지원? → Marlin 커널 사용
    └── 미지원? → 기본 AWQ/GPTQ 커널 사용

핵심 코드 분석

1. AutoRoundConfig: 유연한 설정 체계

AutoRound는 2/3/4/8비트를 지원하며, 패킹 포맷에 따라 GPTQ 또는 AWQ 백엔드를 선택한다.

class AutoRoundConfig(QuantizationConfig):
    SUPPORTED_BITS = {2, 3, 4, 8}
    SUPPORTED_DTYPES = {"int"}
    SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq"}
    SUPPORTED_BACKENDS = {
        "auto", "gptq", "gptq:marlin",
        "awq", "awq:marlin", "marlin"
    }

    def __init__(self, weight_bits, group_size, sym=True,
                 packing_format="auto_round:auto_gptq",
                 backend="auto", ...):
        self.weight_bits = weight_bits
        self.group_size = group_size
        self.sym = sym
        self.pack_factor = Fraction(32, weight_bits)

pack_factorFraction을 사용하여 3비트처럼 32로 나누어 떨어지지 않는 경우에도 정확한 패킹 비율을 계산한다.

2. 레이어별 개별 설정

extra_config를 통해 레이어마다 다른 양자화 설정을 적용할 수 있다.

def get_layer_config(self, layer, layer_name: str):
    def get_config(name: str, quantized: bool = True):
        if not self.extra_config:
            return (
                self.weight_bits if quantized else 16,
                self.group_size if quantized else -1,
                self.sym if quantized else True,
            )
        # 정확한 이름 매칭
        if name in self.extra_config:
            cfg = self.extra_config[name]
            return (
                cfg.get("bits", self.weight_bits if quantized else 16),
                cfg.get("group_size", self.group_size if quantized else -1),
                cfg.get("sym", self.sym if quantized else True),
            )

이 설계는 혼합 정밀도(mixed precision)를 지원한다. 예를 들어 attention 레이어는 4비트, MLP는 8비트로 양자화할 수 있다.

3. 정규식 기반 레이어 매칭

레이어 이름에 정규식 패턴을 사용할 수 있다.

REGEX_SPECIAL_CHARS = set(r"*+?^$()[]{}|\\")
for pattern, cfg in self.extra_config.items():
    if not isinstance(pattern, str) or not any(
        c in REGEX_SPECIAL_CHARS for c in pattern
    ):
        continue
    try:
        if re.fullmatch(pattern, name):
            return (
                cfg.get("bits", self.weight_bits),
                cfg.get("group_size", self.group_size),
                cfg.get("sym", self.sym),
            )
    except re.error:
        continue

정규식 특수 문자가 포함된 키만 패턴으로 취급하여, 일반 레이어 이름과의 충돌을 방지한다.

4. FusedMoE와 Fused QKV 처리

Fused 레이어는 하위 레이어들의 설정이 일관되어야 한다.

# FusedMoE 처리
if self.extra_config and "fusedmoe" in layer.__class__.__name__.lower():
    moe_configs = [
        get_config(name, quantized)
        for name in self.extra_config
        if name.startswith(layer_name)
    ]
    if moe_configs:
        if len(set(moe_configs)) == 1:
            return moe_configs[0]
        raise ValueError(
            f"Fused MoE layer '{layer_name}' requires "
            f"consistent quant config for all sub-layers"
        )

# Fused QKV 처리
for fusion_key, sub_keys in self.packed_modules_mapping.items():
    if fusion_key in layer_name:
        sub_configs = [get_config(name, quantized) for name in sub_names]
        if len(set(sub_configs)) == 1:
            return sub_configs[0]
        raise ValueError(...)

QKV가 하나로 합쳐진(fused) 경우, Q/K/V의 양자화 설정이 모두 동일해야 한다.

5. 양자화 메서드 디스패치

패킹 포맷에 따라 AWQ 또는 GPTQ 경로로 분기한다.

def get_quant_method(self, layer, prefix):
    if "gptq" in self.packing_format or "gptq" in self.backend:
        return self.apply_gptq_quant_layer(layer, prefix)
    if "awq" in self.packing_format or "awq" in self.backend:
        return self.apply_awq_quant_layer(layer, prefix)

6. AWQ 경로: Marlin 자동 선택

AWQ 경로에서도 Marlin 커널 호환 여부를 자동으로 판단한다.

def apply_awq_quant_layer(self, layer, prefix, backend="auto"):
    weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
    if not self.check_quantized(weight_bits):
        return UnquantizedLinearMethod()

    if backend == "auto" or "marlin" in backend:
        use_marlin = (weight_bits in AWQ_TYPE_MAP) and \
            check_marlin_supported(AWQ_TYPE_MAP[weight_bits],
                                   group_size, not sym)

16비트(check_quantized 실패)인 레이어는 양자화를 건너뛴다. 이는 extra_config로 특정 레이어를 비양자화 상태로 유지할 때 사용된다.

7. GPTQ 경로: NPU 분기

GPTQ 경로는 NPU 환경에서 별도의 Ascend 메서드를 사용한다.

def apply_gptq_quant_layer(self, layer, prefix, backend="auto"):
    if _is_npu:
        quant_args = GPTQConfig(
            weight_bits=weight_bits, group_size=group_size,
            lm_head_quantized=False, desc_act=False, dynamic={},
        )
        quant_args.sym = sym
        if isinstance(layer, FusedMoE):
            return GPTQMoEAscendMethod(quant_args)
        if isinstance(layer, (LinearBase, ParallelLMHead)):
            return GPTQLinearAscendMethod(quant_args)

AutoRound vs RTN vs GPTQ

항목 RTN GPTQ AutoRound
라운딩 방식 최근접 Hessian 기반 SGD 최적화
보정 데이터 불필요 필요 (128-256 샘플) 필요 (소량)
양자화 시간 매우 빠름 느림 중간
정확도 (4비트) 낮음 높음 GPTQ와 동등 이상
혼합 정밀도 미지원 미지원 지원 (extra_config)
비트 수 고정 4/8비트 2/3/4/8비트

설계 근거

  1. 패킹 포맷 재활용: AutoRound는 자체 포맷 대신 AWQ/GPTQ의 패킹 포맷을 재활용한다. 이를 통해 기존 커널(Marlin, Triton)을 그대로 사용할 수 있다.
  2. 혼합 정밀도: extra_config로 레이어별 비트 수를 다르게 설정하면, 정확도에 민감한 레이어는 8비트, 나머지는 4비트로 양자화하여 전체 품질을 향상시킨다.
  3. Fraction 패킹: 3비트처럼 비표준 비트 수도 Fraction(32, 3)으로 정확한 패킹 비율을 처리한다.
  4. 하위 호환: GPTQ/AWQ 모델과 동일한 추론 경로를 공유하므로, 커널 최적화의 이점을 그대로 받는다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글