QAT Scheduling and Granularity
마지막 수정:
QAT는 그냥 learning rate를 낮춰서 fine-tune하는 작업이 아니다.
학습 중에 quantization noise를 일부러 넣기 때문에 schedule이 중요하다.
너무 갑자기 낮은 bit로 밀어 넣으면
-> loss가 튀고 회복이 어려움
scale이 계속 흔들리면
-> 모델이 적응할 target이 바뀜
좋은 QAT는 모델이 낮은 precision 환경에 천천히 적응하도록 만든다.
Prepare
- Bits
- FP -> INT8 fake
- LR
- warmup starts
- Observer
- collect ranges
Stabilize
- Bits
- INT8
- LR
- warmup to base LR
- Observer
- scales move
Lower bits
- Bits
- INT6
- LR
- cosine decay
- Observer
- freeze when stable
Target
- Bits
- INT4
- LR
- small updates
- Observer
- fixed grid
Learning rate warmup으로 충격을 줄인다
QAT를 시작하면 forward에 fake quantization이 들어간다. 같은 weight라도 모델이 보는 값이 달라진다.
before QAT:
w_fp 그대로 사용
during QAT:
w_fp -> fake quantized value 사용
이 변화는 loss spike를 만들 수 있다. 그래서 초반에는 learning rate를 작게 시작해서 점진적으로 올린다.
warmup:
small LR -> base LR
Warmup은 quantization noise가 들어간 새로운 training regime에 모델이 적응할 시간을 준다.
Observer freeze는 calibration과 training의 경계다
Observer는 activation이나 weight range를 관찰해서 scale을 계산한다.
QAT 초반에는 observer가 계속 range를 업데이트할 수 있다.
training batch
-> observe range
-> update scale
-> fake quantize
하지만 scale이 끝까지 계속 바뀌면 모델 입장에서는 목표가 흔들린다.
오늘의 grid
내일의 grid
다음 epoch의 grid
그래서 어느 시점 이후에는 observer를 freeze한다.
observer freeze:
scale / zero-point 고정
-> 모델은 고정된 quantization grid에 적응
실무적으로는 scale 변화가 충분히 작아졌을 때 freeze한다. 예를 들어 relative scale change가 1% 아래로 안정되면 freeze 후보가 된다.
Progressive quantization은 bit-width를 천천히 낮춘다
목표가 INT4라면 처음부터 INT4로 시작하지 않을 수 있다.
INT8 warmup
-> INT6 intermediate
-> INT4 target
이런 progressive schedule은 모델이 coarse grid에 단계적으로 적응하게 한다.
처음부터 INT4:
큰 quantization shock
INT8 -> INT6 -> INT4:
작은 shock 여러 번
특히 INT4 이하에서는 progressive quantization이 안정성에 도움이 된다.
BatchNorm은 먼저 fold한다
CNN에서는 BatchNorm folding 순서가 중요하다.
BatchNorm을 나중에 fold하면 weight distribution이 바뀐다. 그러면 QAT 중 관찰한 scale과 export 때의 weight range가 어긋날 수 있다.
좋은 순서:
fold BatchNorm
-> QAT
-> export
나쁜 순서:
QAT
-> fold BatchNorm
-> scale mismatch
Transformer에는 BatchNorm이 보통 없지만, 이 원리는 여전히 중요하다. Export 전에 graph transformation이 weight나 activation range를 바꾸면 QAT가 학습한 grid와 달라질 수 있다.
Granularity는 gradient survival에도 영향을 준다
Per-tensor fake quantization은 tensor 전체가 scale 하나를 공유한다.
global max가 scale을 결정
-> 작은 channel은 integer range를 조금만 사용
-> boundary나 zero 근처에 값이 몰림
이 경우 일부 channel은 STE gradient가 잘 살아남지 못할 수 있다.
Per-channel fake quantization은 output channel마다 scale을 둔다.
각 channel이 자기 range에 맞는 scale 사용
-> 더 많은 weight가 유효 grid 안에서 움직임
-> gradient가 더 잘 살아남음
그래서 QAT에서도 granularity는 accuracy만이 아니라 training dynamics 문제다.
Transformer에서는 selective QAT를 먼저 생각한다
Transformer 전체를 QAT하는 것은 비싸다.
먼저 sublayer sensitivity를 보는 것이 좋다.
attention projections
FFN up projection
FFN down projection
embedding / lm head
특정 sublayer만 민감하다면 전체 QAT보다 selective QAT나 mixed precision이 더 싸다.
민감한 layer:
QAT or higher precision
덜 민감한 layer:
PTQ 유지
QAT의 목표는 모든 곳에 훈련 비용을 쓰는 것이 아니라, PTQ가 감당하지 못한 quantization error를 필요한 만큼만 줄이는 것이다.
판단 규칙
INT8 target:
대부분 PTQ 먼저, QAT는 드묾
INT4 target:
PTQ 실패 진단 후 QAT 후보
below INT4:
QAT, specialized methods, architecture support 검토
scale이 흔들림:
observer freeze 타이밍 확인
loss가 transition마다 크게 튐:
progressive quantization / LR schedule 조정
QAT는 “더 오래 학습하면 해결된다”가 아니다. Schedule, observer, granularity가 맞아야 짧은 fine-tuning으로 quantized regime에 적응할 수 있다.
확인
- Observer freeze는 왜 필요한가?
- INT4 QAT에서 progressive quantization이 도움이 되는 이유는 무엇인가?
- QAT에서 per-channel granularity가 gradient survival에 영향을 주는 이유는 무엇인가?