LabHub
开始
学习 学习路径 课程

MiniMind — 亲手从头到尾训练一个小型语言模型

LoRA 不动权重,只加一条旁路

在 LabHub 中继续学习

한국어 원문으로 표시합니다.

한 줄 요약

LoRA 는 학습된 가중치 W 를 얼려 두고, 그 옆에 작은 행렬 둘(A: 입력→r, B: r→출력)을 달아 B·A 만 학습한다. B 를 0 으로 시작하므로 처음엔 모델이 조금도 바뀌지 않고, 학습이 끝나면 W + B·A 로 합쳐 곁가지 없이 쓸 수 있다. MiniMind 의 model_lora.py 는 이것을 60줄로 구현한다. 이 모듈에서는 SFT 모델에 LoRA 를 얹어 말투를 바꾸고, 기반 가중치가 한 비트도 바뀌지 않았는지를 직접 확인한다.

왜 이게 필요했나

LoRA 논문은 GPT-3 175B 를 예로 든다 — 과제마다 전체 미세조정한 모델을 따로 두면 과제 하나에 1,750억 파라미터가 필요하다. 사전학습 가중치를 얼리고 저랭크 행렬만 학습하면 학습하는 파라미터가 1만 분의 1 로, GPU 메모리가 3분의 1 로 줄고, 어댑터와 달리 추론 지연도 더하지 않는다고 보고했다. MiniMind README 도 같은 쓰임을 든다 — 기반 모델의 일반 능력은 두고 의료 같은 도메인이나 '나는 누구' 같은 자기 인식을 LoRA 로 얹는다. 자료가 충분하면 전체 SFT 도 되지만, 그때는 도메인 자료에 과적합해 일반 능력을 잃지 않게 섞는 일이 따로 필요하다.

어떻게 동작하나

MiniMind 의 구현은 세 부분이다.

class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        self.A = nn.Linear(in_features, rank, bias=False)   # 정규분포(std 0.02)
        self.B = nn.Linear(rank, out_features, bias=False)  # 0
    def forward(self, x): return self.B(self.A(x))

def apply_lora(model, rank=16):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
            module.lora = LoRA(...);  module.forward = 원래 forward(x) + lora(x)

어디에 붙나. 입력과 출력 차원이 같은 Linear 에만 붙는다. 우리 모델(128차원, Q 헤드 4×32)에서는 q_proj(128→128)와 o_proj(128→128)가 해당하고, k_proj·v_proj(128→64)와 FFN(128↔448)은 빠진다. 층 4개 × 2 = 8곳, 곳마다 A(8×128)+B(128×8) = 2,048개라 학습하는 파라미터는 16,384개, 전체의 1.6% 다. 논문은 어느 행렬에 붙일지 고르는 문제를 따로 다루는데, MiniMind 는 이 간단한 규칙으로 정했다.

시작점. B 가 0 이라 B·A = 0 이다. 얹은 직후의 출력은 기반 모델과 완전히 같다. 학습은 '기반 모델에서 얼마나 벗어날까' 를 0 에서부터 배운다.

배율. 논문은 ΔW·x 를 α/r 로 곱해 rank 를 바꿔도 학습률을 다시 고르지 않아도 되게 한다. MiniMind 의 구현은 이 배율 없이 더한다 — 그래서 rank 를 바꾸면 학습률도 다시 봐야 한다.

얼리기와 저장. train_lora.py 는 이름에 lora 가 없는 파라미터를 전부 requires_grad=False 로 두고, 옵티마이저에는 LoRA 파라미터만 넘긴다. 저장은 save_lora 가 LoRA 가중치만 fp16 으로 떼어 저장한다 — 수십 KB 다. 쓸 때는 기반 모델에 apply_loraload_lora, 또는 merge_lora(W + B·A)로 합친 모델 하나를 만든다.

현장에서 만나는 모습

하나의 기반 모델에 고객사·과제마다 다른 LoRA 를 두면, 저장은 기반 하나 + 작은 파일 여럿이고 서빙은 요청마다 LoRA 만 갈아 끼우면 된다. 이때 기반 가중치가 몰래 바뀌면 모든 LoRA 가 동시에 틀어진다 — 얼리기를 빼먹은 학습 스크립트 하나가 그렇게 만든다. 그래서 학습이 끝나면 기반 가중치를 원본과 비트 단위로 견줘 보는 것이 싸고 확실한 확인이다. 반대로 한 가지 용도로만 쓸 거라면 합쳐서(merge) 배포해 추론 때 곁가지 계산을 없앤다.

MiniMind 원본과 이 코스가 다른 점

MiniMind 의 기본 rank 는 16 이고, train_lora.py 는 lr 1e-4 로 10 에폭 돈다. 이 코스는 rank 8, lr 5e-3, 150걸음이다. 모델이 작고 과제(말투 하나)가 단순해 rank 를 줄여도 충분하고, 짧은 시간에 효과를 보려고 학습률을 키웠다. 알파 배율이 없는 구현이라 rank 와 학습률은 함께 움직인다 — rank 를 두 배로 늘리면 B·A 의 크기가 커지는 만큼 같은 학습률이 사실상 더 큰 걸음이 된다. 또 MiniMind 는 LoRA 를 SFT 모델(full_sft) 위에 얹는 것을 기본으로 삼는다. 사전학습 모델에 바로 얹으면 대화 형식부터 배워야 해서, 16,384 개 파라미터로는 부족하다.

다음 실습에서 할 것

SFT 기준 모델에 rank 8 LoRA 를 얹어 어디에 붙는지와 얹은 직후 출력이 그대로인지 확인한다. LoRA 만 학습하도록 얼린 뒤 '…다냥.' 말투 자료로 150걸음 학습하고, 기반 가중치가 그대로인지, 떼어 둔 질문에서 말투가 얼마나 바뀌었는지, 합친 모델이 LoRA 를 얹은 모델과 같은 출력을 내는지 잰다.