처음부터 구현하는 LoRA: LLM 파인튜닝 효율화

(lightning.ai)

1P by GN⁺ 2024-01-23 | ★ favorite | 댓글 1개

LoRA(Low-Rank Adaptation) 는 LLM 전체를 다시 학습하지 않고 작은 저랭크 행렬만 업데이트해 파인튜닝 비용을 줄이는 기법이며, 이 Studio는 LoRA 레이어를 직접 구현해 동작을 확인함
핵심은 일반 파인튜닝의 가중치 변화 ΔW를 두 작은 행렬 A, B의 곱으로 근사하는 것이며, 랭크 r이 작을수록 학습 파라미터와 표현 용량이 함께 줄어듦
5,000×10,000 가중치 행렬은 5,000만 파라미터를 갖지만, r=8 LoRA는 B 5,000×8과 A 8×10,000만 추가해 12만 파라미터로 400배 작아짐
DistilBERT 기반 IMDb 감성 분류에서 기본 LoRA는 Test acc 89.44% 를 기록해 마지막 두 레이어만 학습한 86.22%보다 높고, 전체 파인튜닝의 92.31%보다는 낮았음
하이퍼파라미터 탐색 후 LoRA는 약 50만 학습 파라미터로 Val acc 92.96%, Test acc 92.39%를 기록해, 6,695만 5,010개 파라미터를 학습한 전체 파인튜닝보다 약간 높은 정확도를 냄

LoRA가 줄이는 파인튜닝 비용

LoRA는 Low-Rank Adaptation의 약자로, LLM을 더 효율적으로 파인튜닝하기 위한 기법임
일반 파인튜닝은 딥러닝 모델의 모든 파라미터를 조정하지만, LoRA는 작은 저랭크 행렬 집합만 업데이트함
사전학습 LLM은 여러 작업에 쓸 수 있으나, 특정 데이터셋이나 작업에 맞추려면 파인튜닝이 유용함
모델이 커질수록 모든 레이어를 업데이트하는 방식은 계산 비용 부담이 커짐

ΔW를 작은 행렬 곱으로 근사

일반 파인튜닝에서는 가중치 행렬 W의 업데이트를 ΔW로 계산함
LoRA는 ΔW를 두 작은 행렬 A와 B의 곱으로 근사함
- PCA나 SVD에 익숙하다면 ΔW를 A와 B로 분해하는 방식에 가깝게 볼 수 있음
랭크 r은 LoRA의 하이퍼파라미터임
- 더 작은 r은 학습 파라미터 수를 줄이고 학습을 빠르게 하며 계산 요구량을 낮출 수 있음
- 동시에 작업별 정보를 포착하는 저랭크 행렬의 용량도 낮아짐
5,000×10,000 가중치 행렬 예시:
- 일반 업데이트 ΔW: 총 5,000만 파라미터
- r=8 LoRA: B 5,000×8, A 8×10,000
- 추가 파라미터: 80,000 + 40,000 = 120,000개
- 일반 파인튜닝 대비 400배 작음
실제 사용에서는 성능과 비용의 균형을 찾기 위해 여러 r 값을 실험해야 함

PyTorch로 LoRA 레이어 구현

기본 LoRALayer는 입력 차원, 출력 차원, 랭크, 스케일링 계수 alpha를 받음

class LoRALayer(torch.nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        super().__init__()
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.A = torch.nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.B = torch.nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x):
        x = self.alpha * (x @ self.A @ self.B)
        return x

in_dim은 LoRA를 적용할 레이어의 입력 차원이고, out_dim은 출력 차원임
rank는 A, B 행렬의 복잡도와 LoRA가 추가하는 파라미터 수를 제어함
alpha는 기존 모델 가중치에 LoRA가 주는 변화의 크기를 결정함
- alpha가 높으면 모델 동작을 더 크게 조정함
- alpha가 낮으면 더 미세한 변화가 됨
A는 작은 난수로 초기화하고 표준편차는 랭크의 제곱근으로 정함
- 초기 A 값이 너무 커지지 않게 하기 위한 선택임
B는 0으로 초기화함
- 학습 시작 전에는 B=0이므로 AB=0
- 역전파로 A와 B가 업데이트되기 전에는 LoRALayer가 원래 가중치에 영향을 주지 않음

Linear 레이어를 LinearWithLoRA로 교체

LoRA는 보통 신경망의 Linear/피드포워드 레이어에 적용됨
기존 forward가 두 Linear 레이어를 순서대로 호출한다면, LoRA 적용 후에는 각 Linear 출력에 LoRA 출력을 더함

def forward(self, x):
    x = self.linear_1(x) + self.lora_1(x)
    x = F.relu(x)
    x = self.linear_2(x) + self.lora_2(x)
    return logits

기존 PyTorch 모델을 수정할 때는 각 Linear 레이어를 LinearWithLoRA로 바꾸는 방식이 단순함

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )

    def forward(self, x):
        return self.linear(x) + self.lora(x)

LinearWithLoRA는 원래 Linear 레이어와 새 LoRALayer를 함께 보유함
사전학습 모델의 Linear 레이어를 LinearWithLoRA로 교체하면 LoRA를 장착한 뒤 파인튜닝할 수 있음

DistilBERT로 IMDb 분류 실험

실습 예제는 생성 텍스트보다 정확도 평가가 쉬운 텍스트 분류를 사용함
모델은 Hugging Face transformers의 사전학습 DistilBERT를 사용함

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2)

새 LoRA 가중치만 학습하기 위해 모든 모델 파라미터의 requires_grad를 False로 설정함

for param in model.parameters():
    param.requires_grad = False

DistilBERT에는 6개 Transformer 레이어가 있으며, 각 레이어 안에 Linear 레이어가 있음
- attention에는 q_lin, k_lin, v_lin, out_lin이 있음
- FFN에는 lin1, lin2가 있음
- 출력 쪽에는 pre_classifier, classifier 두 Linear 레이어가 있음

선택적으로 LoRA를 적용하는 설정

기본 LoRA 설정은 attention의 query와 value 가중치 행렬에만 LoRA를 적용함

lora_r = 8
lora_alpha = 16
lora_dropout = 0.05
lora_query = True
lora_key = False
lora_value = True
lora_projection = False
lora_mlp = False
lora_head = False

루프를 돌며 DistilBERT의 각 Transformer 레이어에서 선택된 Linear 레이어를 LinearWithLoRA로 교체함
- lora_query=True이면 q_lin 교체
- lora_key=True이면 k_lin 교체
- lora_value=True이면 v_lin 교체
- lora_projection=True이면 out_lin 교체
- lora_mlp=True이면 ffn.lin1, ffn.lin2 교체
- lora_head=True이면 pre_classifier, classifier 교체
교체 후 모델 출력에서 q_lin과 v_lin 등이 LinearWithLoRA로 바뀐 것을 확인할 수 있음

기본 LoRA와 일반 파인튜닝 비교

기본 LoRA 설정으로 IMDb Movie Reviews 분류를 학습한 결과:
- Train acc: 92.15%
- Val acc: 89.98%
- Test acc: 89.44%
마지막 두 출력 레이어만 파인튜닝한 결과:
- Train acc: 86.68%
- Val acc: 87.26%
- Test acc: 86.22%
- 학습 파라미터 수: 592,130개
기본 LoRA는 마지막 두 레이어만 학습하는 방식보다 Test acc가 높았고, 학습 파라미터 수는 147,456개로 더 적음
전체 레이어를 전통적인 방식으로 파인튜닝한 결과:
- Train acc: 96.41%
- Val acc: 92.80%
- Test acc: 92.31%
- 학습 파라미터 수: 66,955,010개
전체 파인튜닝은 기본 LoRA보다 Test acc가 약 2% 높지만, LoRA 설정보다 약 450배 많은 파라미터를 업데이트함

LoRA 하이퍼파라미터 탐색

LoRA 성능은 lora_r, lora_alpha, 적용 대상 레이어 설정에 따라 달라질 수 있음
03_finetune-lora.py는 하이퍼파라미터를 명령줄 인자로 받음

python 03_finetune-lora.py --lora_alpha 32 --lora_r 16

다른 LoRA 적용 대상을 함께 켤 수도 있음

python 03_finetune-lora.py \
--lora_alpha 32 \
--lora_r 16 \
--lora_query True \
--lora_key True \
--lora_value True \
--lora_projection True \
--lora_mlp True \
--lora_head True

03_gridsearch.py는 사용 가능한 모든 GPU에서 다음 그리드를 실행함
- alpha_values = [1, 4, 8, 16, 32, 64]
- rank_values = [1, 2, 4, 8, 16, 32]
- lora_query = ["True"]
- lora_key = ["False", "True"]
- lora_value = ["True"]
- lora_projection = ["False", "True"]
- lora_mlp = ["False", "True"]
- lora_head = ["False", "True"]
스크립트는 Visual Studio Code, 명령줄 터미널, Job으로 실행할 수 있으며, Job은 완료 후 자동 종료됨
결과는 results.txt에 저장됨

그리드 탐색에서 나온 최고 설정

results.txt 기준 가장 좋은 하이퍼파라미터 구성은 다음과 같음

lora_r: 8
lora_alpha: 1
lora_query: True
lora_key: False
lora_value: True
lora_projection: False
lora_mlp: True
lora_head: False

이 설정의 결과:
- Val acc: 92.96%
- Test acc: 92.39%
이 LoRA 설정은 학습 파라미터가 약 500k이며, 전체 파인튜닝의 66M 파라미터보다 훨씬 적음
정확도는 전체 파인튜닝의 Val acc 92.80%, Test acc 92.31%보다 약간 높음

실행 환경과 추가 자료

Studio 상단의 Run을 클릭하면 코드를 포함한 환경을 복제할 수 있음
Studio 복제 후에는 추가 설치, 다운로드, 설정 단계 없이 코드 파일을 실행할 수 있음
관련 노트북과 스크립트:
- 00_lora-layer.ipynb: LoRA 레이어 구현
- 01_finetune-last-layers.ipynb: 마지막 레이어 파인튜닝
- 02_finetune-with-lora.ipynb: LoRA 파인튜닝
- 03_finetune-lora.py: LoRA 하이퍼파라미터 인자 실행
- 03_gridsearch.py: LoRA 하이퍼파라미터 그리드 탐색
- 04_finetune-all-layers.ipynb: 전체 레이어 파인튜닝
추가 자료:
- Parameter-Efficient LLM Finetuning With Low-Rank Adaptation (LoRA): LoRA 개념과 실습 예제를 더 길게 소개함
- Practical Tips for Finetuning LLMs Using LoRA (Low-Rank Adaptation): LoRA 파인튜닝 관련 질문과 실무 팁을 제공함

GN⁺ 2024-01-23 [-]

Hacker News 의견들

기법 흐름은 Maxime Labonne의 LLMs 101로 따라가고 있음: https://github.com/mlabonne/llm-course#4-supervised-fine-tun...
LoRA != LoRa라서 계속 헷갈림. 이미 있던 약어를 재사용한 게 싫음
- 나도 마찬가지임. 본업이 기계 학습인데도, 아니 어쩌면 그래서인지, 문맥이 거의 없는 곳에서 이 약어를 볼 때마다 한 번씩 멈칫함
  HN 첫 페이지처럼 두 의미가 모두 자연스러운 곳에서는 특히 그렇다
- “Low-Rank Adaptation” 말고 다른 뜻이 뭐임? 차이를 검색하기도 어렵다
- 사람들이 너무 전문화돼서 자기 버블 밖에서 무슨 일이 일어나는지 신경 쓰지 않으면 이런 일이 생김
- 소프트웨어 쪽 사람들이 하드웨어 관련 이름을 가져다 붙이는 흐름이 싫음
- 지금까지 서로 관련 없는 두 기술이 같은 약어를 쓰게 된 건 아쉬움
컴퓨터 과학 분야에서 “이 숫자들, 즉 하이퍼파라미터가 결과에 정확히 어떻게 영향을 주는지 모르니 여러 값을 넣어보고 제일 잘 되는 걸 쓰자” 같은 말을 한다는 게 아직도 이상하게 느껴짐
- “여러 값을 넣어보고 제일 잘 되는 걸 쓰자”는 값 찾기에 몬테카를로 시뮬레이션을 쓰는 것과 비슷하지 않나?
  가끔 최적/정답 대신 국소 최댓값에 걸릴 수 있지만 그래도 작동함. 닫힌 형태의 공식으로 풀 수 없으니 수십억 번쯤 무작위 표본을 뽑아 원하는 값을 찾는 식이고, LLM도 같다는 뜻은 아니지만 이런 접근은 꽤 자주 씀
- 공학적으로 만든 것과 발견한 것의 차이처럼 느껴짐
  지금까지 업계의 대부분은 공학적으로 설계된 것이었고, LLM은 발견된 것에 가깝다
- 그런 상향식 땜질은 Dijkstra 본인이 관찰했듯 미국에서 컴퓨터 과학이 시작된 방식과도 비슷함: https://www.cs.utexas.edu/users/EWD/transcriptions/EWD06xx/E...
  이상적으로는 이론적 기반이 필요하지만, 이론을 만들거나 검증할 만큼의 데이터를 끌어내려면 때로는 무작위 탐색이 필요함
- Minsky와 다른 사람들이 퍼셉트론이 비선형 함수를 모델링하지 못한다며 폄하한 탓도 큼. LLM은 현대 CPU와 GPU 없이는 어차피 나오기 어려웠겠지만, 그렇다고 더 나은 이론적 기반을 미리 갖출 수 없었다는 뜻은 아님
  우리는 있어야 할 위치보다 몇 년 뒤처져 있음. 1990년대 게임 업계에서 일할 때는 신경망이 잘해야 막다른 길이고 나쁘게는 사기라는 게 “상식”이었음. 몇몇 권위자가 모두를 말린 탓에 시간을 이렇게 많이 잃은 건 정말 아쉽고, 이번에는 그런 일이 반복되지 않게 해야 함
- Stable Diffusion 설정을 연구하는 느낌이 이와 비슷함. 금방 알게 되는 건 추측이 많이 섞여 있다는 것임
언제 미세 조정을 하고 언제 RAG를 써야 하는지 아직 명확하지 않음
예전에는 미세 조정이 주로 모델의 행동을 바꾸는 용도라고 생각했는데, 최근에는 일부 회사가 지식 추가에도 미세 조정을 쓰는 것처럼 보임. 미세 조정의 주요 사용처가 궁금함
- 주요 사용처는 여전히 행동 변화라고 봄. 지시 미세 조정, 분류용 미세 조정 등이 그렇다
  가중치에 지식을 추가하는 건 사전 학습으로 하는 게 가장 좋음. 또는 생성 중 질의할 외부 데이터베이스나 문서가 있다면 질문처럼 RAG를 쓰면 됨. 참고로 NeurIPS 2023 LLM Efficiency Challenge에서 24시간 안에 GPU 1개로 “최고” LLM을 미세 조정한 모든 우승자는 LoRA 또는 QLoRA(양자화 LoRA)를 사용했음
- 추가 데이터가 간결하지 않거나 문맥이 필요하면 미세 조정이 RAG보다 나음
  문맥이 너무 많거나 초점이 흐리면 프롬프트 추종성이 희석될 수 있고, RAG는 모델이 더 높은 차원의 토큰 연관을 배우게 해주지 못함. 그래서 필요한 내용을 보강 자료에서 운 좋게 뽑아와야 하는데, 그러면 그다지 고급 검색 엔진보다 낫지 않음. 정부나 대기업 내부 문서처럼 공개 데이터셋에 잘 나타나지 않는 자체 미세 방언을 가진 전문 말뭉치를 다룰 때 특히 문제가 됨
- 내가 이해한 바로는 미세 조정은 비정상적으로 효과적임 [0]. 문맥 내 학습은 기본 모델이 얼마나 강력한지와 RAG를 어떻게 구성하는지, 즉 질의 처리·임베딩 검색·결과 순위화 등에 크게 의존하기 때문임 [1]
  읽은 논문에 따르면 미세 조정은 새 도메인 지식을 추가하거나 특정 지식을 강화할 수 있고, RAG는 강화에만 제한됨. 다만 다른 트레이드오프를 가진 두 기법이 비슷한 수준의 능력을 보이기도 함 [2]
  
  [0] Fast.ai: Can Models learn from one sample, https://www.fast.ai/posts/2023-09-04-learning-jumps/ / https://archive.is/eJMPR
  
  [1] LlamaIndex: Advanced RAG, https://blog.llamaindex.ai/a-cheat-sheet-and-some-recipes-fo... / https://archive.is/qtBXX
  
  [2] Microsoft: RAG vs Fine-tuning: Pipelines, Tradeoffs, and a Case Study, https://arxiv.org/html/2401.08406v2#S6 / https://archive.is/UQ8Sa#S6
- 이들은 자기회귀 모델임. 앞부분으로부터 뒤에 올 요소를 예측할 수 있는 새로운 유형의 시퀀스가 있고, 그 방식이 모델이 전에 본 것과 다르다면 미세 조정이 타당해 보임
  특정 데이터 상황에서 무엇을 할지 정하는 기준으로는 꽤 모호하지만, 대략적인 휴리스틱으로는 충분할 수 있음. 지식 추가가 여기에 들어가는지는 실험 없이는 취향의 문제일지도 모름
좋은 글임. 이 분야 사람은 아니지만 원 논문을 읽었을 때는 LoRA가 마지막 밀집층에만 적용되고, 모든 층에 독립적으로 적용되지는 않는다고 이해했음. 내가 잘못 읽었을 수도 있음
링크의 구현이 왜 이런지 좀 파보니 QLoRA에서 이 방식을 썼고 흥미로운 효과가 있는 것 같음. QLoRA의 결정에 대한 메모를 추가하면 좋을 듯함. 다만 왜 작동하는지는 잘 모르겠고, 초보자 관점에서는 마지막 층에 LoRA를 적용하는 건 말이 되지만 각 선형층마다 반복해서 적용하는 근거는 감이 안 잡힘. 직관을 설명해 줄 수 있나?
- 기계 학습의 많은 것처럼 어떤 층을 쓸지는 이론보다 경험적 증거에 더 좌우됨. 일반적인 LoRA 학습 파이프라인에서는 기본 모델을 고정하고 LoRA 층만 조정함
  LoRA 층으로 바꾸는 층이 많을수록 최적화의 자유도가 커짐. 일부 미세 조정 방식은 마지막 층만 미세 조정하라고 권하는데, 입력의 “가장 고차” 표현을 담고 있다고 추정하기 때문임. 다른 방식은 모든 층을 미세 조정함. 대부분 데이터와 문제에 따라 달라지고, LoRA는 이 관례를 그대로 반영함
Axolotl의 “처음부터”가 아니라 설정에서 시작하는 접근을 선호함. Axolotl은 Mistral, Llama 2 미세 조정을 지원하고, 샘플 패킹, FlashAttention, xFormers 같은 최신 기법도 많이 지원함
나는 LoRA를 처음부터 배우기보다 미세 조정 데이터를 모으고 큐레이션하는 데 집중해서 데이터 중심 미세 조정을 함
- 우리(Lightning AI)의 lit-gpt 라이브러리도 같은 일을 함
  https://github.com/Lightning-AI/lit-gpt
이름 짓기가 어렵긴 함. 처음에는 이게 “long range”의 LoRa나 IoT 센서 통신인 LoRaWAN 이야기인 줄 알았음
미세 조정에 가장 많이 쓰이는 라이브러리는 뭐가 있나? “처음부터”가 아닌 방식 기준임
- https://github.com/OpenAccess-AI-Collective/axolotl
와, 나도 처음엔 당연히 LoRa 이야기인 줄 알았음
LoRA의 성능 비용은 어느 정도인가?
- 학습 중에는 역전파로 일부 매개변수만 업데이트하므로 전체 미세 조정보다 효율적임
  추론 중에는 두 가지가 가능함. 순전파 중 LoRA 값을 동적으로 더하면 이론상 약간 느려질 수 있지만, 고객별로 작은 가중치 세트를 따로 유지하고 싶다면 장점이 되기도 함. 큰 기본 모델 하나만 돌리고 고객별 LoRA 가중치를 즉석에서 적용할 수 있기 때문임. LoRA 가중치를 기본 모델에 다시 병합하면 기본 모델과 정확히 같은 성능을 낼 수 있음

답변달기

처음부터 구현하는 LoRA: LLM 파인튜닝 효율화

LoRA가 줄이는 파인튜닝 비용

ΔW를 작은 행렬 곱으로 근사

PyTorch로 LoRA 레이어 구현

Linear 레이어를 LinearWithLoRA로 교체

DistilBERT로 IMDb 분류 실험

선택적으로 LoRA를 적용하는 설정

기본 LoRA와 일반 파인튜닝 비교

LoRA 하이퍼파라미터 탐색

그리드 탐색에서 나온 최고 설정

실행 환경과 추가 자료

함께 보면 좋은 글 β

댓글과 토론

Hacker News 의견들