SFT Training Loop 직접 작성: TRL 없이 Qwen 0.5B를 파인튜닝하며 배우는 유효 배치·LR·warmup
라이브러리(TRL 등) 없이 직접 짠 학습 루프로 Qwen 0.5B를 SFT(지도 미세조정)하는 날이다. 질문-답변 데이터에서 "답변 부분만" 손실을 계산하는 법, 메모리가 모자랄 때 쓰는 gradient accumulation, 학습률과 warmup을 어떻게 잡는지 배운다. 학습률 1e-5, 5e-5, 1e-4를 비교하면서 VRAM과 손실 곡선도 기록한다
By 송민성
오늘의 목표
TRL 같은 SFT 전용 라이브러리 없이, Hugging Face(HF)에서 받은 Qwen 0.5B 모델과 내가 직접 작성한 학습 루프만으로 SFT를 끝까지 돌린다. 그 과정에서 VRAM(GPU 메모리) 사용량을 기록하고, 학습률(learning rate, LR) 세 가지를 비교해 본다.
개념과 직관
SFT(Supervised Fine-Tuning, 지도 미세조정) 는 이미 사전학습(pretraining)된 모델에게 "이런 질문엔 이렇게 답해라"라는 정답 예시를 보여주며 추가로 학습시키는 일이다. 사전학습 모델이 책을 많이 읽은 학생이라면, SFT는 모범 답안집으로 "질문에 답하는 방식"을 연습시키는 과정이다. 학습 방식 자체는 지금까지 해온 것과 같다. 다음 토큰을 맞히고, 틀린 정도(손실, loss)를 계산하고, 파라미터를 조금 고친다.
Full fine-tuning 은 모델의 모든 파라미터를 업데이트하는 방식이다. 0.5B(5억 개) 모델이라 한 장의 GPU에서 가능하지만 메모리를 꽤 쓴다. 파라미터 외에 기울기(gradient)와 옵티마이저 상태(AdamW라면 파라미터마다 값 2개를 더 저장)도 메모리에 올라가기 때문이다. 대략 "모델 크기의 서너 배 이상"이 필요하다고 감각을 잡아두자. 정확한 수치는 오늘 직접 측정해서 기록한다.
1) 답변 부분만 손실 계산하기 (label masking)
학습 데이터는 보통 "질문(prompt) + 답변(response)"을 이어 붙인 한 줄의 토큰열이다. 그런데 모델이 질문까지 외우도록 학습할 필요는 없다. 그래서 질문 토큰의 정답 라벨을 -100으로 바꿔 손실 계산에서 제외한다. PyTorch의 cross_entropy는 -100인 위치를 기본으로 무시한다. HF 모델도 labels를 주면 내부에서 이 규칙으로 손실을 계산한다.
토큰: [질문 토큰들........][답변 토큰들....]
라벨: [-100 -100 ... -100][답변 토큰들....]2) 유효 배치와 gradient accumulation
배치(batch)는 한 번 업데이트할 때 같이 보는 샘플 묶음이다. 큰 배치가 학습에 안정적이지만 GPU 메모리가 모자라면 못 쓴다. 이때 gradient accumulation(기울기 누적, 작은 배치 여러 번의 기울기를 더해 두었다가 한 번에 업데이트하는 것) 을 쓴다.
비유하면, 한 번에 32권을 못 들어서 8권씩 4번 나르고, 4번 다 나른 뒤에야 "정리 완료" 도장을 찍는 것이다.
유효 배치 크기(effective batch size) = 마이크로 배치 × 누적 횟수 × (GPU 개수)
예: 마이크로 배치 4 × 누적 8 = 유효 배치 32주의할 점은 손실을 누적 횟수로 나눠야 한다는 것이다. 그래야 기울기의 크기가 "유효 배치 32의 평균"과 같아진다.
loss_for_backward = loss / 누적 횟수3) 학습률과 warmup
학습률은 한 번 업데이트할 때 파라미터를 얼마나 크게 움직일지 정하는 값이다. SFT의 학습률은 사전학습보다 작게 잡는다. 이미 잘 학습된 지식을 큰 보폭으로 망가뜨리지 않기 위해서다. 비유하면 완성된 그림을 수정할 때는 굵은 붓이 아니라 가는 붓을 쓰는 것과 같다.
Warmup(워밍업, 학습 초반에 LR을 0에서 목표값까지 서서히 올리는 것) 은 초반의 불안정한 업데이트를 줄여준다. 옵티마이저 상태가 아직 쌓이지 않은 초반에 큰 LR로 시작하면 모델이 크게 흔들릴 수 있기 때문이다. 그 뒤에는 보통 LR을 서서히 줄이는 스케줄(예: cosine decay)을 함께 쓴다.
step < warmup_steps: lr = base_lr × step / warmup_steps
그 이후: lr = base_lr × 0.5 × (1 + cos(pi × 진행률)) # cosine decay4) 에폭
에폭(epoch)은 전체 데이터를 한 바퀴 도는 횟수다. SFT는 데이터가 많지 않아서 보통 1~3 에폭 정도로 짧게 돈다. 너무 많이 돌리면 학습 데이터만 외우는 과적합(overfitting)이 생긴다. 오늘은 작은 데이터로 2 에폭을 기준으로 한다.
코드로 직접 해보기
사전 준비: pip install torch transformers datasets (GPU 환경 권장). 모델 이름은 Hugging Face에서 받을 수 있는 Qwen 0.5B 계열 Instruct가 아닌 Base 체크포인트를 쓰는 것을 권한다(예: Qwen/Qwen2.5-0.5B). 정확한 이름은 HF 모델 페이지에서 직접 확인하자.
1단계: 모델·토크나이저 불러오기
import torch, math, time, random
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL = "Qwen/Qwen2.5-0.5B" # HF에서 이름을 확인하고 사용
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(MODEL)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16).to(device)
model.train()bfloat16은 메모리를 절반으로 줄이는 16비트 숫자 형식이다. GPU가 지원하지 않으면 torch.float32로 바꾸되 VRAM이 두 배가 된다.
2단계: 데이터 만들고 라벨 마스킹하기
raw = [
{"q": "대한민국의 수도는?", "a": "서울입니다."},
{"q": "1+1은?", "a": "2입니다."},
# 실제로는 수백~수천 개를 준비 (직접 만들거나 공개 SFT 데이터셋 일부 사용)
]
def encode(ex, max_len=256):
prompt = f"질문: {ex['q']}\n답변: "
answer = ex["a"] + tok.eos_token
p_ids = tok(prompt, add_special_tokens=False)["input_ids"]
a_ids = tok(answer, add_special_tokens=False)["input_ids"]
ids = (p_ids + a_ids)[:max_len]
labels = ([-100] * len(p_ids) + a_ids)[:max_len]
return ids, labels
def collate(batch):
L = max(len(x[0]) for x in batch)
ids = torch.full((len(batch), L), tok.pad_token_id)
lab = torch.full((len(batch), L), -100)
att = torch.zeros((len(batch), L), dtype=torch.long)
for i, (a, b) in enumerate(batch):
ids[i, :len(a)] = torch.tensor(a)
lab[i, :len(b)] = torch.tensor(b)
att[i, :len(a)] = 1
return ids.to(device), lab.to(device), att.to(device)
data = [encode(ex) for ex in raw]패딩(길이를 맞추려고 채운 토큰) 위치의 라벨도 -100으로 둬서 손실에서 빠지게 한다. attention_mask는 "진짜 토큰이 있는 곳"을 알려주는 표시다.
3단계: LR 스케줄 함수
def get_lr(step, total, base_lr, warmup):
if step < warmup:
return base_lr * (step + 1) / warmup
progress = (step - warmup) / max(1, total - warmup)
return base_lr * 0.5 * (1 + math.cos(math.pi * progress))4단계: 학습 루프 (gradient accumulation 포함)
def train(base_lr, epochs=2, micro_bs=4, accum=8, warmup_ratio=0.05):
# 매 실험마다 모델을 새로 불러와야 공정하다 (아래 '막히는 지점' 참고)
m = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16).to(device)
m.train()
opt = torch.optim.AdamW(m.parameters(), lr=base_lr, weight_decay=0.0)
steps_per_epoch = len(data) // (micro_bs * accum)
total = max(1, steps_per_epoch * epochs)
warmup = max(1, int(total * warmup_ratio))
torch.cuda.reset_peak_memory_stats()
step, log = 0, []
for ep in range(epochs):
random.shuffle(data)
for i in range(0, len(data) - micro_bs * accum + 1, micro_bs * accum):
opt.zero_grad(set_to_none=True)
step_loss = 0.0
for k in range(accum):
s = i + k * micro_bs
ids, lab, att = collate(data[s:s + micro_bs])
out = m(input_ids=ids, attention_mask=att, labels=lab)
(out.loss / accum).backward()
step_loss += out.loss.item() / accum
for g in opt.param_groups:
g["lr"] = get_lr(step, total, base_lr, warmup)
torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
opt.step()
step += 1
log.append(step_loss)
print(f"step {step}/{total} loss {step_loss:.4f} lr {g['lr']:.2e}")
peak = torch.cuda.max_memory_allocated() / 1e9 if device == "cuda" else 0
print(f"peak VRAM: {peak:.2f} GB")
return m, log, peakclip_grad_norm_은 기울기 전체 크기가 1.0을 넘으면 비율로 줄여 한 번의 업데이트가 폭주하는 것을 막는다(gradient clipping).
데이터가 micro_bs × accum = 32개보다 적으면 steps_per_epoch이 0이 되어 아무것도 돌지 않는다. 오늘은 최소 수백 개 이상의 샘플을 준비하자.
5단계: 학습 전후 생성 확인
@torch.no_grad()
def ask(m, q):
m.eval()
x = tok(f"질문: {q}\n답변: ", return_tensors="pt").to(device)
y = m.generate(**x, max_new_tokens=50, do_sample=False)
m.train()
return tok.decode(y[0][x["input_ids"].shape[1]:], skip_special_tokens=True)학습 전 모델과 학습 후 모델에 같은 질문을 던져 답변 형식이 바뀌었는지 눈으로 확인한다. 숫자(loss)만 보지 말고 실제 출력을 꼭 읽어보자.
오늘의 실험
질문: 학습률이 달라지면 손실 곡선, 안정성, 결과물이 어떻게 달라지는가?
- 바꾸는 것:
base_lr∈ {1e-5, 5e-5, 1e-4} - 고정하는 것: 데이터, 에폭(2), 마이크로 배치(4), 누적(8), warmup 비율, 시드(seed, 난수 시작값)
- 관찰하는 것:
1. 손실이 내려가는 속도와 마지막 손실 값 2. 손실이 갑자기 튀거나 nan이 되는지(발산) 3. 학습 후 같은 질문 5개에 대한 답변 품질 4. 최대 VRAM(GB)
기록은 아래와 같은 표로 남긴다. 값은 직접 측정한 것만 채운다.
| LR | 첫 5스텝 평균 loss | 마지막 5스텝 평균 loss | 발산 여부 | peak VRAM | 샘플 답변 메모 |
|---|---|---|---|---|---|
| 1e-5 | |||||
| 5e-5 | |||||
| 1e-4 |
손실 로그 리스트를 matplotlib으로 그려 세 곡선을 한 그래프에 겹쳐 보면 비교가 쉽다. 마지막에는 "내 데이터와 이 설정에서는 어떤 LR이 가장 나았고, 그 이유를 어떻게 추측하는지"를 두세 줄로 적는다. 손실이 가장 낮은 LR이 항상 최선은 아니다. 답변이 이상해졌거나 학습 데이터만 외웠다면 손실이 낮아도 좋은 결과가 아니다.
왜 이걸 배우나 (LLM / Post-training 연결)
SFT는 사전학습된 모델이 "지시를 따르는 모델"이 되는 첫 단계이고, 이후 선호 학습 같은 post-training 단계도 대부분 SFT 모델에서 출발한다. 라벨 마스킹, 유효 배치, LR·warmup 설정은 LLM Post-training 실무에서 매번 마주치는 요소라서, 라이브러리 내부를 직접 짜 본 경험이 있으면 학습이 이상할 때 원인을 좁혀 나갈 수 있다. 이어지는 LoRA/QLoRA에서는 오늘 측정한 full fine-tuning의 VRAM이 비교 기준이 된다.
초보자가 막히는 지점
- CUDA out of memory:
micro_bs를 줄이고accum을 늘려 유효 배치를 유지한다.max_len을 줄이는 것도 효과가 크다. 그래도 안 되면 다른 프로세스가 GPU를 잡고 있는지 확인한다. - 손실이 `nan`: LR이 너무 크거나
bfloat16환경이 맞지 않는 경우가 많다. LR을 낮추고 gradient clipping이 켜져 있는지 확인한다. - 손실을 `accum`으로 안 나눔: 누적 횟수만큼 기울기가 커져 사실상 LR이
accum배 커진 것처럼 동작한다. - `zero_grad` 위치 오류: 마이크로 배치 안쪽에서 호출하면 누적이 사라진다. 업데이트 한 번당 한 번, 안쪽 루프 시작 전에 호출한다.
- 실험 간 모델 재사용: 이전 실험에서 학습된 모델로 다음 LR을 돌리면 비교가 무효가 된다. 매번 원본 체크포인트에서 새로 불러온다.
- 라벨이 전부 `-100`:
max_len으로 잘라냈을 때 답변이 통째로 잘리면 손실 계산 대상이 사라진다. 잘린 샘플 비율을 점검한다. - EOS 토큰을 빼먹음: 답변 끝에 EOS가 없으면 모델이 멈출 줄 모르고 계속 생성한다.
- 손실만 보고 판단: 숫자가 좋아도 출력이 이상할 수 있다. 반드시 생성 결과를 읽어본다.
알고리즘 (병행 1시간)
LeetCode Medium 2문제 (DP, 동적계획법)
DP는 큰 문제를 작은 문제로 쪼개고, 한 번 푼 작은 문제의 답을 저장해 재사용하는 방식이다. 문제는 DP 태그가 붙은 Medium 중에서 고른다. 예를 들어 1차원 점화식으로 풀 수 있는 문제 하나와, 2차원 표가 필요한 문제 하나를 고르면 균형이 좋다.
각 문제마다 아래를 기록한다.
- 접근: "dp[i]가 무엇을 뜻하는가"를 한 문장으로 정의하고, 점화식과 초기값을 적는다.
- 시간·공간 복잡도: 상태 개수 × 상태당 계산량으로 구하고, 공간을 줄일 수 있는지도 생각해 본다.
- 틀린 이유: 초기값 오류, 인덱스 범위(off-by-one), 점화식 누락 등 어디서 틀렸는지 구체적으로 적는다. 풀지 못했다면 어느 지점에서 막혔는지, 풀이를 본 뒤 무엇을 이해했는지를 남긴다.