Search
✴️

Mixed Precision & QLORA & Gradient Checkpointing

Category
BlogPost
Venue
Backbone
Text
PPT

Mixed Precision

1.
FP32 weight를 FP16 copy weight 만들어 줌
a.
copy weight가 forward, backward에 사용
2.
FP16 copy weight를 통해 forward 진행
3.
Forward로 계산된 FP16 prediction을 FP32로 캐스팅 진행
4.
FP32 prediction을 이용해 FP32 loss 계산하고, scaling factor 곱함
5.
scaled FP32 loss를 FP16으로 캐스팅 진행
6.
scaled FP16 loss를 통해 backpropagation 진행후 gradient 계산
7.
FP16 gradient를 FP32로 캐스팅 하고, 다시 scaling factor로 나눔
8.
FP32 gradient를 통해 FP32 weight update

QLORA

QLoRA의 storage data type은 4-bit / computation data type은 BFloat16
→ QLoRA weight가 사용될 때, BFloat16으로 dequantize를 하고 16-bit에서 행렬 연산 진행
4-bit NormalFloat(NF) Quantization
NF data type은 이론적으로 가장 이상적인 data type이며 Quantile Quantization 기법에 사용
Quantile Quantization: data를 동일한 크기의 quantile로 분류하는 방법. Empirical cumulative distribution function을 통해 tensor의 quantile을 추정하는 식으로 작동.
(각 quantization bin에는 동일한 input tensor의 개수가 존재)
Limitation: quantile 추정이 expensive → 빠른 추정을 하기위해 SRAM quantile같은 approximation을 통해 추정하는 방법이 제안 되었지만, outliers에 대해 large quantization error가 존재해서 불안정
Input tensor들이 고정된 quantization constant의 분포로부터 나오면 앞선 문제들 해결 가능 → input tensors have the same quantiles (동일한 분포를 갖는다)
Pretrained neural network weights follow zero-centered normal distribtuion with standard deivation σ\sigma
σ\sigma scaling을 통해 모든 weight를 single fixed distribution으로 변형 가능 → data type과 weights에 대응하는 quantiles들이 동일한 range [-1,1] 갖게끔 정규화 적용
zero-mean normal-dist를 위한 data type을 구하는 Process
1.
estimate 2k+12^k+1 quantiles of a N(0,1)N(0,1) distribtuion to obtain a k-bit quantile quantization data type for norm-dist
2.
Take this data type and normalize its values into [-1,1] range
3.
quantize an input weight tensor by normalizing it into [-1,1] through absolute maximum rescaling
weight range와 data type range를 맞춘후, data type의 qiq_i 값들에 대한 quantize 진행 (2k2^k)
위에 (3) 과정이 ‘σ\sigma scaling을 통해 모든 weight를 single fixed distribution으로 변형’하는 역할 수행
Double Quantization
추가적인 memory 사용량을 아끼기 위해 소개된 Quantization constant를 quantizing하는 방법
Paged Optimzier
CPU와 GPU간의 연동을 통해 GPU memory가 고갈되면, CPU RAM과 disk에서 memory paging하는 것과 비슷하게 동작. 자동으로 CPU RAM으로 backup해서 process가 끊임 없이 연산을 계속할 수 있게 함. (optimzier update step이 필요할 때, CPU → GPU 이동)
QLoRA
One storage data type (4-bit NF)
One computation data type (16-bit BF)
forward/backward process를 진행할 때, data type에 대해 dequantize를 수행
⇒ Model을 NF4 형태로 불러옴 > Computation시에 BF16으로 dequantization ‘MODEL Parameter Frozen’ + LORA Weight Update > Training Done.

Gradient Checkpointing

1.
중간 결과 저장 (Layer의 INPUT/OUTPUT 결과값 저장): 지정한 Layer의 중간결과만 forward pass때 저장함. weight output layer부터 ouput layer까지만 forward를 호출해서 gradient재 계산이 가능함.
2.
그래디언트 재계산(model.backward()호출시에 input, output 결과값 이용해서 gradient값 재계산): 역전파 시, 체크포인팅이 적용된 레이어의 그래디언트는 저장된 중간 결과를 사용하여 재계산. 이는 forward 함수의 일부를 다시 호출하여 이루어질 수 있음.
즉, checkpoint부터 계산하고자 하는 node까지 forward pass 후에 다시 gradient 계산.
3.
메모리 효율성: 이 방법은 중간 그래디언트를 모두 저장하는 대신 필요할 때만 계산하여 메모리 사용을 크게 줄일 수 있는 장점이 존재. 그 결과 더 크고 복잡한 모델을 제한된 메모리 리소스로 학습할 수 있게 됨.
4.
계산 오버헤드: 하지만 forward 의 추가 호출은 당연히 계산 오버헤드를 발생시킬 수 있음. 중간 결과를 다시 계산해야 하기 때문에, 전체적인 학습 시간이 늘어날 수 있다는 당연한 결과를 초래함.