[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_factor에 Fraction을 사용하여 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비트 |
설계 근거
- 패킹 포맷 재활용: AutoRound는 자체 포맷 대신 AWQ/GPTQ의 패킹 포맷을 재활용한다. 이를 통해 기존 커널(Marlin, Triton)을 그대로 사용할 수 있다.
- 혼합 정밀도:
extra_config로 레이어별 비트 수를 다르게 설정하면, 정확도에 민감한 레이어는 8비트, 나머지는 4비트로 양자화하여 전체 품질을 향상시킨다. - Fraction 패킹: 3비트처럼 비표준 비트 수도
Fraction(32, 3)으로 정확한 패킹 비율을 처리한다. - 하위 호환: GPTQ/AWQ 모델과 동일한 추론 경로를 공유하므로, 커널 최적화의 이점을 그대로 받는다.
관련 포스트
참고
- SGLang 소스:
python/sglang/srt/layers/quantization/auto_round.py - AutoRound 논문: AutoRound: Optimize Weight Rounding via Signed Gradient Descent
- AutoRound 라이브러리: intel/auto-round
관련 포스트
SGLang 의 다른글
- 이전글 [SGLang] BitsAndBytes: QLoRA와 NF4 동적 양자화
- 현재글 : [SGLang] AutoRound: 자동 라운딩 최적화 양자화
- 다음글 [SGLang] Compressed Tensors: 통합 양자화 프레임워크
댓글