관심 있는 NLP 논문을 읽어보고 간단히 정리했습니다.
혹시 부족하거나 잘못된 내용이 있다면 댓글 부탁드립니다 🙇♂️
(코드 구현에 관한 내용은 마지막에 다루고 있습니다!!)
[University of Washington]
- 기학습된 모델을 4-bit로 quantize한 뒤 Low Rank Adapters(LoRA)를 학습하는 방식
- QLoRA로 학습된 model family, Guanaco를 공개.
- ChatGPT의 99.3% 성능을 발휘할 수 있는 65B 모델을 single GPU에서 24시간 동안 fine-tuning
- 세 개의 tenchiques: (a) 4-bit NormalFloat (NF4), (b) Double Quantization, (c) Paged Optimizers
출처 : https://arxiv.org/abs/2305.14314
1. Introduction
LLM이 부상함에 따라 이를 fine-tuning하는 비용이 너무 크다는 문제가 끊임없이 제기되었고, 이를 해결하기 위한 대표적인 PEFT(Parameter Efficient Fine-Tuning)으로 LoRA가 아주 많이 사용되고 있습니다.
개념 자체가 받아들이기에 크게 어렵지 않은 내용이기도 하고, 구현도 잘 되어 있어서 많은 이들을 구해준 학습 기법으로 자리를 잡았죠.
이와 동시에 모델의 가중치를 더 작은 자료형으로 표현하여 사용하는 메모리의 양을 최소화하면서도 기존에 준하는 성능을 내도록 유도하는 Quantization 기법도 끊임없이 발전하고 있습니다.
물론 자료형을 직접 건드리는 것 외에도 transformer 아키텍쳐의 attention의 연산 방식을 개선/수정함으로써 연산 효율성을 높이는 방식도 꽤나 많구요.
본 논문에서는 두 방식, LoRA와 Quantizaion을 결합한 QLoRA를 제시합니다.
자세한 방식에 대해서는 후술하겠지만, 가중치를 4-bit로 불러오고 실제 계산은 16-bit로 수행하게 됩니다.
또한 천 개 이상의 모델에 대해 다양한 instruction tuning datasets을 fine-tuning했다고 합니다.
그리고 MMLU(Massive Multitask Language Understanding)와 같이 대표적인 벤치마크에서 좋은 결과를 낸 모델이 다른 곳(Vicuna chatbot benchmark)에서도 그럴 것임을 보장하지는 않는다는 연구 결과를 제시합니다.
2. Background
QLoRA를 이해하기 위해 필요한 선수 지식에 대해 설명하는 내용입니다.
Block-wise k-bit Quantization, Low-rank Adapters, Memory Requirement of Parameter-Efficient Finetuning에 대해 순서대로 설명합니다.
2.1. Block-wise k-bit Quantization
정보를 담고 있는 데이터의 자료형을 더 적은 메모리 공간을 차지하는 자료형으로 converting하는 기법을 Quantization이라고 부릅니다.
논문에서는 32-bit float 자료형을 8-bit integer로 변형하는 예시를 다루고 있습니다.
이때 더 작은 자료형의 범위 전체가 사용될 수 있도록 해당 자료형의 최댓값을 이용한 scaling을 수행하는 것으로 이해할 수 있습니다.
구체적인 수식은 아래와 같습니다.
$$\mathrm{X}^{\mathrm{Int8}}=\mathrm{round}\left ( \frac{127}{\mathrm{absmax}\left( \mathrm{X^{\mathrm{FP32}}} \right )}\mathrm{X^{\mathrm{FP32}}}\right )=\mathrm{round}\left( c^{\mathrm{FP32}}\cdot X^{\mathrm{FP32}}\right )$$
이때 사용되는 $c$를 quantization constant 또는 quantization scale이라고 부릅니다.
식을 조금 다른 관점에서 해석해본다면 $\mathrm{X^{\mathrm{FP32}}}$를 $\mathrm{absmax} \mathrm{X^{\mathrm{FP32}}} $로 나누는 과정을 먼저 생각해볼 수 있습니다.
$\mathrm{X^{\mathrm{FP32}}}$는 여러 개의 값이 담긴 리스트라고 생각을 해본다면, 그중의 최댓값으로 나눠줌으로써 가장 큰 값을 1로 만들어주는 것이죠.
결국 $\mathrm{X^{\mathrm{FP32}}}$에 담긴 모든 값들은 -1 이상 1 이하의 값을 갖게 됩니다.
이를 축소하고자 하는 자료형의 절댓값의 최댓값으로 곱해줍니다.
Int8 자료형의 경우 $2^8 = 256$개 만큼의 숫자를 표현할 수 있으므로 -127부터 127까지의 값을 가질 수 있습니다.
따라서 -1 이상 1 이하의 값으로 scaling한 리스트에 127을 곱해줌으로써 -127부터 127 사이의 값을 갖게 되는 것이죠.
하지만 이런 과정을 거치게 되면 절댓값이 최댓값이였던 값을 제외하고는 정수값임을 보장할 수 없기 때문에 round 함수를 적용해준 것입니다.
quantization constant는 결국 변경하고자 하는 자료형의 절댓값의 최댓값을, 현재 자료형의 절댓값의 최댓값으로 나눈 값이 됩니다.
이렇게 압축된 숫자는 이후 원복되어야 하는데, 이때 quantization constant를 거꾸로 나눠줌으로써 자료형을 복구할 수 있습니다.
이런 dequantization을 식으로 표현하면 아래와 같습니다.
$$\mathrm{dequant}\left( c^{\mathrm{FP32}},\mathrm{X}^{\mathrm{Int8}} \right ) = \frac{\mathrm{X}^{\mathrm{Int8}}}{c^{\mathrm{FP32}}}=\mathrm{X}^{\mathrm{FP32}}$$
위와 같은 quantization은 제목의 이름에서 알 수 있는 것처럼 block 단위로 적용됩니다.
입력 텐서 $\mathrm{X}\in \mathbb{R}^{b \times h}$를 flatten하고 일정한 간격 $B$으로 나누어 n개의 segment를 갖게 됩니다.
즉 $n=(b\times h)/B$이고, n개의 block에 대해 각각 quantization을 수행하며, 각 quantization에서 발생한 quantization constant를 $c_i$로 표기합니다.
2.2. Low-rank Adapters
LoRA에 대해서는 워낙 많은 리뷰와 정리글이 있으니.. 간단히 수식만 짚고 넘어가겠습니다.
LoRA를 적용하지 않은 기존 모델의 경우, 입력 $X$를 가중치 $W$와 곱하여 출력 $Y$를 얻습니다.
이때 각 요소는 $X\in\mathbb{R}^{b\times h}$, $W\in\mathbb{R}^{h\times o}$, $Y\in\mathbb{R}^{h\times o}$이고 $Y=XW$가 됩니다.
여기에 LoRA를 적용하게 되면 $XW$는 freeze하게 되고, 학습 가능한 sub-layer를 추가하여 다음 식을 통해 $Y$를 얻습니다.
$$Y=XW+sXL_1L_2$$
이때 $s$는 스칼라값이고 $L_1\in\mathbb{R}^{h\times r}$, $L_2\in\mathbb{R}^{r\times o}$이므로 이를 $X$와 곱한 것은 결국 기존과 동일하게 $\mathbb{R}^{b\times o}$에 속하게 됩니다.
이보다 자세한 내용을 알아보고 싶으신 분들은 구글링, 또는 논문을 직접 참고해보시길 바랍니다!
2.3. Memory Requirement of Parameter-Efficient Finetuning
LoRA에 필요한 메모리의 양은 사용되는 adapter의 숫자와 크기에 의해 결정됩니다.
논문에서는 7B LLaMA 모델을 배치 사이즈 1로 FLAN v2에 LoRA로 학습하게 되는 경우, 업데이트 되는 LoRA의 weight는 기존 모델 weight의 0.2% 수준이며 input gradients의 memory footprint도 567MB에서 26MB 수준으로 감소됨을 언급하고 있습니다.
여기에 모델의 weight를 본 논문의 방식처럼 4-bit로 불러오게 되는 경우, 단 5,048MB의 메모리만을 차지하게 된다고 합니다.
모델이 차지하는 weight의 메모리의 크기가 다른 방식에 비해 획기적으로 작아, 더 많은 adapter를 활용하기 쉽다는 점을 강점으로 내세우고 있습니다.
이제 위 내용들을 바탕으로 하는 QLoRA에 대해 본격적으로 다뤄보겠습니다.
3. QLoRA Finetuning
QLoRA는 맨 위에서 언급했던 것처럼 (a) 4-bit NormalFloat (NF4), (b) Double Quantization, (c) Paged Optimizers 를 구성 요소로 갖습니다.
이전 챕터와 마찬가지로 각 요소를 하나씩 살펴보도록 하겠습니다.
3.1. 4-bit NormalFloat (NF4)
이 자료형은 Quantile Quantization에 기반을 두고 있으나, 각 quantization bin에 할당되는 자료의 개수를 동일하게 만들 수 있도록 normalize한 것으로 볼 수 있습니다.
quantile이란 한국어로는 '분위수'라고 번역되는데, 정의에 따르면 '확률 분포의 구간을 나누는 기준이 되는 수'를 뜻하며 우리에게는 사분위수가 가장 잘 알려져 있습니다.
(사분위수는 Q1=0.25 분위수, Q2=0.5 분위수, Q3=0.75 분위수로 해당 값보다 작은 값의 비중을 나타낸 것으로 이해할 수 있습니다. 즉, Q3의 값이 예를 들어 80이라면, 해당 분포에서는 80 이하에 해당하는 값이 전체의 75%를 차지한다는 뜻입니다)
어쨌든 Quantile Quantization은 입력 텐서의 이와 같은 분위수를 추정하여 양자화를 수행하는데, 그 비용이 크게 발생한다는 것이 문제점으로 꼽혀서 SRAM quantile과 같은 방법론들이 제시되기도 했다고 합니다.
본 논문에서는 사전학습된 neural network의 weights가 일반적으로 0을 중심으로 하면서 표준 편차 $\sigma$를 따르는 분포를갖는다는 점을 근거로 들어, 모든 weights를 특정 자료형의 범위에 정확하게 조정하는 것을 목표로 삼으며, 이때 임의의 구간 [-1, 1]로 조정한다고 합니다.
k-bit의 자료형을 만들기 위한 quantizatoin 과정은 다음 순서를 따릅니다.
(1) $N(0,1)$을 따르는 분포에서 $2^k+1$개의 분위수를 구합니다.
$2^k$개 대신 $2^k+1$개를 사용하는 이유는 $2^k$개의 구간을 얻기 위함입니다.
(2) 이 자료형을 [-1, 1] 범위로 normalize합니다.
(3) 입력 가중치 텐서에 absolute maximum scaling을 적용하여 [-1, 1] 범위로 normalize 합니다.
분위수를 사용하여 자료형을 만드는 것과 달리, 가중치 $W$는 2장에서 언급했던 Block-wise k-bit Quantization을 적용하는 것입니다.
즉, 양자화 상수를 입력 가중치 텐서에 곱하는 방식으로 quantization을 진행하는 것입니다.
이에 따라 (1), (2)를 통해 만든 k-bit 자료형과 (3)을 통해 변형한 입력 가중치 텐서의 범위가 동일해집니다.
마지막으로 $2^{k}+1$개의 분위수를 통해 만든 $2^k$개 간격 내에 존재하는 입력 가중치 텐서의 값들을 각 구간의 대표값으로 변환합니다.
이때 각 구간의 대표값은 다음 식을 통해 얻을 수 있습니다.
$$q_i=\frac{1}{2}\left(Q_X\left( \frac{i}{2^k+1} \right )+Q_X\left(\frac{i+1}{2^k+1} \right ) \right )$$
이렇게 quantization을 통해 획득한 k-bit 자료형은 $2^k$의 값을 표현할 수 있게 되는데, 비대칭적으로 값을 구성해야 합니다.
만약 $k=2$인 상황에서 4개의 숫자를 표현하는데 $(-2,-1,1,2)$라면 0이 생략되기 때문입니다.
따라서 $2^{k-1}$개를 0을 중심으로 좌측에서 구하고, $2^{k-1}+1$개를 0을 포함하여 값을 구성한 뒤 합치는 방식을 택한다고 합니다.
3.2. Double Quantization (DQ)
우리는 k-bit 자료형을 만들고, 가중치를 quantize 했습니다.
그런데 이때 가중치 $W$에 대해 정한 block size는 64입니다.
그래서 32-bit의 자료형을 quantize 하게 되는 경우 양자화 상수는 각 파라미터는마다 평균적으로 $32/64=0.5$ 비트를 더 차지하게 만듭니다.
DQ는 이 '양자화 상수'를 quantize하는 방법을 말합니다.
원래 자료형을 축소시킬 때 사용되었던 양자화 상수는 이후 dequantize에 사용되는데요,
이 상수를 quantize하게 되면 당연히도 그 표현력은 추가적으로 감소하게 됩니다. (정확하게 복원이 어려워진다는 뜻입니다)
그럼에도 저장 공간상의 이득이 크기 때문에 이러한 방법론을 사용하는 것이죠.
아래에서 식을 작성하겠지만 본 논문에서는 'FP32 자료형을 k-비트로 줄일 때 사용했던 양자화 상수'가 FP32 자료형인데, 이 양자화 상수를 (이를테면) FP8 자료형으로 quantize하는 예시를 언급하고 있습니다.
3.3. Paged Optimizers
CPU와 GPU 사이에서 page-to-page 전환이 자동적으로 이뤄지는 NVIDIA 통합 메모리를 사용합니다.
이는 GPU가 이따금씩 out-of-memory 위기에 놓일 때를 위한 것입니다.
저도 자세한 내용은 모르지만 GPU가 out-of-memory를 겪게 되면 연산을 CPU에 넘겨서 하다가, optimizer가 update하는 시점에는 다시 GPU로 paged back하는 방식이라고 합니다.
3.4. QLoRA
위에서 설명한 요소들을 조합하여 QLoRA를 single linear layer에 적용한 것을 식으로 표현한 것이 아래와 같습니다.
$$\mathrm{Y}^{\mathrm{BF16}}=\mathrm{X}^{\mathrm{BF16}}\mathrm{doubleDequant}\left(c_1^{\mathrm{FP32}}, c_2^{\mathrm{k-bit}}, \mathrm{W}^{\mathrm{NF4}}\right )+\mathrm{X}^{\mathrm{BF16}}L_1^{\mathrm{BF16}}L_2^{\mathrm{BF16}}$$
이때 $\mathrm{doubleDequant}(\cdot)$의 정의는 다음과 같습니다.
$$\mathrm{doubleDequant}(c_1^{\mathrm{FP32}},c_2^{\mathrm{k-bit}},\mathrm{W}^{\mathrm{k-bit}})=\mathrm{dequant}\left( \mathrm{dequant} \left(c_1^{\mathrm{FP32}}, c_2^{\mathrm{k-bit}} \right ), \mathrm{W}^{\mathrm{4bit}} \right )=\mathrm{W}^{\mathrm{BF16}}$$
저는 처음에 봤을 때 notation이 상당히 헷갈렸었는데, 첫 번째 quantization에 사용되는 양자화 상수는 $c_2$, 두 번째 quantization이 사용되는 양자화 항수는 $c_1$입니다.
즉, 가중치를 k-bit를 만들 때 사용되었던 양자화 상수 $c_2^{\mathrm{k-bit}}$를, 두 번째 양자화 상수 $c_1^{\mathrm{FP32}$를 이용해 FP32로 복원합니다.
이때 $c_2^{\mathrm{k-bit}}$는 두 번째 양자화를 통해 (예를 들어) FP8 자료형으로 축소된 상황이라고 가정한 것입니다.
(본 논문에서 특별한 언급이 없다면 $\mathrm{W}$는 NF4로, $c_2$는 FP8로 변형했다고 합니다)
이제 FP32로 복원된 양자화 상수를 이용해 4bit로 축소되어 있는 $\mathrm{W}^{\mathrm{4bit}}$를 BF16 자료형으로 복원합니다.
결과적으로 $\mathrm{W}^{\mathrm{BF16}}$이 됩니다.
QLoRA 방법론에 대한 모든 설명은 끝났습니다..!
사실 실험 및 결과에 대한 내용도 굉장히 많고 자세한 논문이긴한데, 이 포스팅에서는 다루지 않을 생각입니다.
컨셉을 이해하는 게 중요하다고 생각해서 이를 중점적으로 공부해야겠다 생각했었기 때문입니다 😅
극히 일부만 간단히 언급하자면 일반적으로 쓰이던 Float4 자료형으로 양자화한 결과와 NFloat4 자료형으로 양자화한 결과, 그리고 여기에 Double Quantization까지 적용한 결과를 비교하여 제시하고 있습니다.
NFloat4를 사용하는 것 자체는 Float4를 사용하는 것보다 우수한 실험 결과를 보여주었고요, 여기에 DQ를 적용하더라도 성능 하락폭은 거의 없다는 결과를 제시했습니다.
4. Gemma를 QLoRA로 fine-tuning 하기
마지막으로 이 QLoRA를 구글 코랩 환경에서 구현한 코드를 간단히 살펴보도록 하겠습니다.
사용한 모델은 최근 구글에서 공개한 Gemma-2B 입니다.
이 모델을 허깅페이스 허브에 업로드된 데이터셋 중에서 Alpaca 데이터셋의 형식을 따라 만든 'c-s-ale/dolly-15k-instruction-alpaca-format' 데이터셋에 fine-tuning 해보겠습니다.
참고로 이 데이터셋의 구성은 아래 이미지와 같습니다.
필요 패키지/라이브러리 설치 및 로드
!pip install -q trl peft datasets accelerate
!pip install -q -i https://pypi.org/simple/ bitsandbytes
각 라이브러리들이 뭔지 모르고 설치하면 또 나중에 어려움을 겪을 수 있으니 짧게만 설명해보겠습니다.
- trl - 원래는 Transformer Reinforcement Learning의 약자로 강화학습 관련된 라이브러리인데, Supervised Fine-Tuning Trainer를 지원하기 때문에 사용합니다.
- peft - Parameter Efficient Fine-Tuning의 약자로 LoRA를 지원합니다.
- datasets - 허깅페이스 허브에 등록되어 있는 데이터셋들을 다운로드받아 쓸 수 있도록 해줍니다.
- accelerate - 파이토치 코드로 작성된 모델의 분산/병렬 처리가 가능하도록 돕습니다.
- bitsandbytes - 양자화와 관련된 라이브러리로 모델을 4-bit로 불러올 때 사용합니다.
설치가 끝났으면 필요한 라이브러리/패키지를 한 번에 불러오겠습니다.
import torch
from datasets import load_dataset, DatasetDict
from transformers import (TrainingArguments, AutoTokenizer,
AutoModelForCausalLM, BitsAndBytesConfig)
from trl.trainer import SFTTrainer, ConstantLengthDataset
from peft import LoraConfig
학습용 데이터셋 구축하기
dataset = load_dataset("c-s-ale/dolly-15k-instruction-alpaca-format", split='train')
dataset
alpaca-format 이라는 이름을 붙여놓은 이 데이터셋은 instruction, category, input, output으로 구성되어 있습니다.
우리가 사용할 것은 category를 제외한 insturction, input, output 입니다.
첫 번째 내용들을 확인하면 다음과 같습니다.
print(f"# instruction\n{dataset['instruction'][0]}\n")
print(f"# input\n{dataset['input'][0]}\n")
print(f"# output\n{dataset['output'][0]}")
## instruction
#When did Virgin Australia start operating?
## input
#Virgin Australia, the trading name of Virgin Australia Airlines Pty Ltd, is an Australian-based airline. It is the largest airline by fleet size to use the Virgin brand. It commenced services on 31 August 2000 as Virgin Blue, with two aircraft on a single route. It suddenly found itself as a major airline in Australia's domestic market after the collapse of Ansett Australia in September 2001. The airline has since grown to directly serve 32 cities in Australia, from hubs in Brisbane, Melbourne and Sydney.
## output
#Virgin Australia commenced services on 31 August 2000 as Virgin Blue, with two aircraft on a single route.
모델에게 instruction finetuning을 위한 데이터를 전달할 때는 그냥 세 요소를 합쳐서 하나의 프롬프트를 구성하면 됩니다.
그러면 모델은 관련 질문이나 컨텍스트가 주어졌을 때 학습했던 방식과 유사하게 답변을 생성하게 됩니다.
고정된 프롬프트 양식에 맞게끔 각 내용들을 집어 넣어 프롬프트를 구성하는 함수 코드는 다음과 같습니다.
def generate_prompt(row):
return ("Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n"
"### Instruction\n"
"{instruction}\n"
"### Input\n"
"{input}\n"
"### Response\n"
"{output}\n").format_map(row)
이게 되게 이상하게 생긴 것처럼 보일 수 있기는 한데, 특히나 """ 로 텍스트를 작성하는 경우 왼쪽에 딱 붙여서 쓰지 않으면 사실 공백과 탭을 포함해서 글을 작성한 셈입니다.
따라서 우리 사람들이 보기 좋게 작성하는 게 아니라 모델이 보기 좋게 작성해야 합니다!!
위 코드를 데이터셋에 적용해야 하는데, 이후에 토크나이저를 불러오면서 함께 처리하겠습니다.
지금은 모델을 먼저 불러오도록 할게요.
모델 로드 및 LoRA 설정
# 모델을 불러올 때 사용
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4", # 모델 가중치 로드
bnb_4bit_compute_dtype=torch.float16, # 계산에 사용되는 자료형
)
device_map = 'auto' # gpu를 자동 할당
QLoRA에서는 모델의 가중치를 NF4 자료형으로 불러온다고 했습니다.
물론 실제 계산은 이를 dequantize하여 float16형으로 계산합니다.
참고로 device_map을 'auto'로 설정해주면 여러 대의 gpu가 있을 때 자동으로 메모리를 할당해줍니다.
가지고 있는 gpu의 종류가 다양하거나 특수한 상황이 아니라면 사실 알아서 처리해주는 게 엄청나게 편하긴 하더라고요..!
from huggingface_hub import login
login()
이제 모델을 불러와야 하는데요, 그 전에 허깅페이스 허브에 로그인을 먼저 해주어야 합니다.
왜냐하면 Gemma 모델도 Llama를 사용할 때와 마찬가지로 계정별로 모델 사용과 관련한 일부 제약들에 대한 동의를 제출해야 하기 때문입니다.
별 건 아니고 그냥 동의 동의 후 제출하면 바로 승인해줍니다.
어쨌든 위처럼 코드를 작성하고 해당 셀을 실행하면 다음 이미지와 같이 입력창이 보일 겁니다.
허깅페이스에 아이디가 있고 Gemma 사용까지 승인을 받았다면, 자신의 허깅페이스 토큰을 저기 빈 칸에 입력하고 로그인을 누르면 됩니다.
정상적으로 로그인이 되었다면 이런 메세지가 뜹니다!
base_model_name = 'google/gemma-2b'
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
quantization_config=bnb_config,
device_map=device_map,
trust_remote_code=True,
)
base_model.config.use_cache = False # 직전의 key/value attention cache를 사용하지 않음
base_model.config.pretraining_tp = 1 # linear layer에서 보다 정확한, 대신 느린 연산
모델을 불러오는 방법 자체는 다를 것이 없습니다.
다만 일부 config를 조정하는데, 설명은 코드의 주석을 참고해주시기 바랍니다.
저도 pretraining_tp의 역할이 명확히 와닿지는 않더라고요..!
만약 위에서 허깅페이스 허브에 로그인하지 않고 바로 모델을 불러오려고 시도한다면 다음과 같은 에러가 발생합니다.
따라서 사전에 Gemma 모델 사용 허가를 받고 반드시 허깅페이스 허브에 로그인 해주셔야 합니다!
# Trainer에 전달될 예정
peft_config = LoraConfig(
lora_alpha=8,
lora_dropout=0.1,
r=8, # trainable parameter 숫자에 가장 큰 영향을 줌
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"] # 어떤 가중치에 adapter를 적용할지 결정
)
print("model ready")
LoRA를 적용하기 위해서는 위와 같은 설정을 해주어야 합니다.
보통은 r과 alpha값을 다르게 쓸텐데, colab 환경에서 모델을 이정도로 세팅하지 않으면 메모리가 부족해서 코드가 돌아가지 않습니다..!
상황과 목적에 따라 target_modules도 조정해서 사용 가능합니다.
토크나이저 로드 및 데이터셋 구축 마무리
이제 토크나이저를 불러오고 아까 마무리 못한 데이터셋 구축을 이어나가겠습니다.
tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'right'
max_seq_length = 512
decoder 모델에서 볼 수 있는 특징이라면 pad_token을 지정해주어야 한다는 것입니다.
사전 학습을 할 때도 그냥 여러 글들을 쭉 이어 붙여 놓은 corpus에 노출되다 보니 padding 토큰이 없죠.
그래서 이를 문장 끝임을 알려주는 EOS 토큰으로 지정해줘야 합니다.
그리고 padding side, 이것이 굉장히 헷갈리는데요.
학습할 때는 sequence의 우측에 패딩을 삽입해야 합니다.
어차피 지금 위 코드에서 사용된 데이터는 일정 길이(512 토큰)로 잘라서 넘기는 방식이라 큰 의미도 없습니다만,
일반적으로 학습을 할 때 sequence의 좌측에 패딩을 넣게 되면 EOS 이후에 문장을 생성하도록 학습이 될 것입니다.
단, 추론을 할 때는 조금 다른 것이 입력에 response (output)을 제외하게 된다는 점입니다.
즉, 모델에게 [pad, pad, .. instruction, input] 을 입력해서 [output] 을 받아야 하는 것입니다.
이렇게 하지 않고 학습 때와 동일하게 padding을 우측에 주게 된다면 [instruction, input, pad, ... , pad] 이후에 [output]을 받아야 하는 상황으로 만들어 버리는 것입니다.
여튼 여기서는 right로 지정을 해주면 되겠습니다!
데이터셋이 제대로 만들어졌는지 확인해 볼수도 있는데요 아래 코드를 참고해주시기 바랍니다.
for sample in train_dataset:
break
sample
tokenizer.decode(sample['input_ids'])
Trainer 구축 및 학습
마지막으로 지금까지 모델, 데이터셋, 실험 세팅 등을 SFTTrainer에 한꺼번에 욱여넣고, 해당 인스턴스에 train 메서드를 사용하면 지정된 세팅에 따라 모델이 학습을 시작합니다.
batch_size = 2
gradient_accumulation_steps = 2
num_train_epochs = 1
training_args = TrainingArguments(
output_dir="./output/",
per_device_train_batch_size=batch_size,
fp16=True, # -> bf16=True
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.1,
gradient_accumulation_steps=gradient_accumulation_steps,
gradient_checkpointing=True,
do_eval=False,
evaluation_strategy="no",
num_train_epochs=num_train_epochs,
logging_strategy="steps",
logging_steps=10,
save_strategy="epoch",
gradient_checkpointing_kwargs={'use_reentrant': False},
)
인자는 본인이 원하는대로.. 잘 조절을 하면 되긴합니다..!
여기서 조금 특이한 점이라면 fp16=True 입니다.
사실 요즘 가장 많이 쓰이는 것은 bf16 자료형인데요, 이게 코랩에서는 지원이 안되더라고요..!
그래서 어쩔 수 없이 fp16 자료형으로 연산을 해야 합니다.
trainer = SFTTrainer(
model=base_model,
packing=True,
train_dataset=train_dataset,
peft_config=peft_config,
max_seq_length=max_seq_length,
tokenizer=tokenizer,
args=training_args,
)
print("trainer ready")
지금까지의 내용들을 종합하여 SFTTrainer에 담아주면 됩니다.
데이터셋, peft(LoRA), 토크나이저, 학습 args 등을 한 번에 전달해줍니다.
이제 이 트레이너의 train 메서드를 통해 학습을 진행하면 됩니다.
print("training start!")
trainer.train()
print("training finished!")
1 epoch으로 설정했을 뿐이긴 하지만 그래도 생각보다 얼마 걸리지도 않겠네요!
Loss가 초반부에 잘 떨어지는 것도 확인이 되고요, 아마 중간쯤에는 loss가 더 떨어지지 않아 학습이 중단될 것 같긴 하네요.
학습 중에 GPU를 얼마나 활용하고 있는지를 캡쳐해 보았습니다.
사실 처음에는 OOM이 발생해서 일부값을 수정하여 글을 작성한 것이고요, 본인의 필요나 가지고 있는 자원 및 환경에 따라 여러 값들을 조정해보시면 훨씬 좋은 결과들을 얻을 수 있을 것입니다.
참고로 Kaggle 노트북을 사용하면 주마다 30시간 동안 GPU를 무료 이용할 수 있는데요, T4는 두 장이나 주기 때문에 분산 처리도 가능하고 코랩과 다르게 중간에 끊길 일이 없어서 굉장히 좋습니다.
자원이 부족하신 분들은 특히나 반드시 사용해보시길 강추드립니다!
번외로.. 여기서는 다루지 않았습니다만, 모델 학습이 제대로 되었는지 확인하기 위해서는 보통 generate 함수를 사용합니다.
혹은 학습이 끝난 모델의 가중치를 저장하여 (여기서는 adapter만 학습했으므로 adapter 정보만 저장하게 됩니다) 이를 pipeline이라는 HuggingFace의 추론용 API를 활용하기도 합니다.
추론 및 검증을 위한 관련 내용들은 따로 찾아보시길 권장드립니다!
그럼 이상으로 QLoRA에 대한 소개, 사용 방법 등에 대한 글을 마치겠습니다.
잘못된 내용이나 부족한 점들은 댓글로 남겨주시면 감사하겠습니다.
긴 글 읽어주셔서 감사합니다 ☺️