본문 바로가기
TYLER SONGBlog
← Blog index
LLM Research9 min read

Project 1 학습 파이프라인: config 하나로 재현되는 baseline 학습 스크립트 만들기

Tiny GPT를 학습시키는 스크립트를 "설정 파일(config) 하나로 다시 돌릴 수 있게" 만드는 날이다. 학습률 스케줄(warmup+cosine), 혼합 정밀도(AMP), gradient clipping, 체크포인트, wandb 로깅을 각각 왜 쓰는지 이해하고, 200 step 스모크런으로 loss 감소와 재개를 확인한다. 이후 모든 실험이 비교되는

By 송민성

오늘의 목표

어제 만든 Tiny GPT 모델과 데이터 로더를 이용해, 같은 config를 넣으면 같은 방식으로 다시 학습되는 스크립트를 만든다. 그리고 소형 config(6층, d256)로 baseline을 한 번 끝까지 돌려서, 앞으로의 실험이 "무엇과 비교해서 좋아졌는지" 말할 수 있는 기준 숫자를 확보한다.

개념과 직관

학습 스크립트는 "모델 + 데이터 + 반복문"만 있으면 돌아간다. 하지만 실험을 여러 번 하려면 다음 다섯 가지가 필요하다. 하나씩 일상적인 비유로 보자.

1. config 기반 학습 (설정을 코드 밖으로 빼기) 레시피 카드와 요리사를 분리하는 것과 같다. 층 수, 학습률, 배치 크기 같은 숫자를 코드 안에 흩뿌려 두면, 나중에 "지난주 그 결과가 어떤 설정이었지?"를 알 수 없다. 숫자를 한 파일(config)에 모으고, 학습할 때 그 파일을 같이 저장하면 실험이 재현(reproducible, 같은 설정으로 같은 결과를 다시 얻을 수 있음)된다. 시드(seed, 난수 시작값)도 config에 넣는다.

2. 학습률 스케줄: warmup + cosine 학습률(learning rate, 한 번에 파라미터를 얼마나 크게 움직일지)은 일정하게 두지 않는 경우가 많다.

  • warmup(워밍업): 처음 몇십~몇백 step 동안 학습률을 0에서 목표값까지 서서히 올린다. 처음에는 모델이 무작위 상태라 큰 걸음을 내딛으면 휘청거리기 때문이다. 운동 전 준비운동과 같다.
  • cosine decay(코사인 감쇠): 그 후 코사인 곡선 모양으로 학습률을 천천히 0 근처까지 낮춘다. 막판에는 조심스럽게 걸어 정교하게 자리를 잡는다는 직관이다.

수식(텍스트):

text
warmup 구간 (step < warmup_steps):
    lr = max_lr * step / warmup_steps
그 이후:
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    lr = min_lr + 0.5 * (max_lr - min_lr) * (1 + cos(pi * progress))

3. AMP (Automatic Mixed Precision, 혼합 정밀도) 숫자를 32비트 대신 16비트(fp16/bf16)로 계산해 속도와 메모리를 아끼는 기법이다. 비유하면 소수점 아래 많은 자리까지 재는 정밀 저울 대신, 대부분은 간이 저울로 재고 중요한 합산만 정밀 저울을 쓰는 것이다. PyTorch는 torch.autocast가 연산별로 알아서 정밀도를 고른다. fp16은 표현 범위가 좁아 GradScaler(작은 기울기가 0으로 사라지지 않게 미리 키워 주는 도구)가 필요하고, bf16은 범위가 넓어 보통 필요 없다. GPU가 bf16을 지원하는지 확인하고 선택하자.

4. Gradient clipping (기울기 자르기) 가끔 기울기(gradient, 파라미터를 어느 방향으로 얼마나 바꿀지 알려 주는 값)가 비정상적으로 커지는 step이 생긴다. 이때 그대로 갱신하면 모델이 크게 망가진다. 기울기 전체 크기(norm)가 기준(예: 1.0)을 넘으면 비율로 줄여 준다. 급커브에서 속도 제한을 거는 안전장치다.

5. 체크포인트 (checkpoint)와 로깅

  • 체크포인트: 학습이 중간에 끊겨도 이어갈 수 있도록 상태를 저장한 파일. 모델 가중치만으로는 부족하다. 옵티마이저 상태, 현재 step, (사용한다면) GradScaler 상태, 난수 상태까지 같이 저장해야 "정말 이어서" 학습된다.
  • wandb(Weights & Biases): loss, 학습률, 기울기 norm 같은 값을 웹 대시보드에 자동 기록해 주는 실험 추적 도구. 계정이 부담되면 mode="offline"이나 CSV 파일 기록으로 대체해도 된다. 핵심은 "숫자를 남기는 습관"이다.

코드로 직접 해보기

라이브러리 기능을 쓰기 전에, 학습률 스케줄은 직접 함수로 만들어 원리를 확인한다.

1단계: config 파일 (`config.py` 또는 yaml)

python
from dataclasses import dataclass, asdict @dataclass class Config: n_layer: int = 6 d_model: int = 256 n_head: int = 4 block_size: int = 256 batch_size: int = 32 max_lr: float = 3e-4 min_lr: float = 3e-5 warmup_steps: int = 20 max_steps: int = 200 # 스모크런. baseline은 더 크게 grad_clip: float = 1.0 seed: int = 1337 ckpt_path: str = "ckpt.pt" use_amp: bool = True

dataclass는 값들을 묶어 두는 간단한 상자다. asdict(cfg)로 딕셔너리로 바꿔 체크포인트와 wandb에 같이 저장한다.

2단계: 학습률 스케줄을 직접 구현

python
import math def get_lr(step, cfg): if step < cfg.warmup_steps: return cfg.max_lr * (step + 1) / cfg.warmup_steps progress = (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps) coeff = 0.5 * (1 + math.cos(math.pi * progress)) return cfg.min_lr + coeff * (cfg.max_lr - cfg.min_lr) cfg = Config() for s in [0, 10, 19, 20, 100, 199]: print(s, round(get_lr(s, cfg), 6))

먼저 이 값들이 "올라갔다가 내려가는지" 눈으로 확인한 뒤 학습에 붙인다.

3단계: 학습 루프 (AMP + clipping + 로깅) 어제 만든 model, get_batch(x, y 배치를 돌려주는 함수)가 있다고 가정한다. 모델은 logits, loss = model(x, y)를 돌려준다고 두자. 아니라면 본인 인터페이스에 맞게 고친다.

python
import torch, random, numpy as np def set_seed(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) device = "cuda" if torch.cuda.is_available() else "cpu" set_seed(cfg.seed) model = model.to(device) opt = torch.optim.AdamW(model.parameters(), lr=cfg.max_lr, weight_decay=0.1) amp_dtype = torch.bfloat16 use_amp = cfg.use_amp and device == "cuda" start_step = 0 for step in range(start_step, cfg.max_steps): lr = get_lr(step, cfg) for g in opt.param_groups: g["lr"] = lr x, y = get_batch("train") x, y = x.to(device), y.to(device) opt.zero_grad(set_to_none=True) with torch.autocast(device_type=device, dtype=amp_dtype, enabled=use_amp): logits, loss = model(x, y) loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) opt.step() if step % 10 == 0: print(f"step {step} loss {loss.item():.4f} lr {lr:.2e} gnorm {grad_norm:.2f}") # wandb.log({"loss": loss.item(), "lr": lr, "grad_norm": grad_norm}, step=step)

여기서 loss.backward()는 기울기를 계산해 주는 PyTorch 함수라고만 알면 된다. 원리는 이후 Day에서 다룬다. bf16을 쓰면 GradScaler는 필요 없다. fp16을 쓴다면 scaler를 추가해야 한다.

4단계: 체크포인트 저장과 재개

python
import os def save_ckpt(path, model, opt, step, cfg): torch.save({ "model": model.state_dict(), "opt": opt.state_dict(), "step": step, "cfg": asdict(cfg), }, path) def load_ckpt(path, model, opt): ck = torch.load(path, map_location=device) model.load_state_dict(ck["model"]) opt.load_state_dict(ck["opt"]) return ck["step"] + 1 # 다음 step부터 시작

학습 루프에서 if step % 50 == 0: save_ckpt(...)로 저장하고, 시작 전에 if os.path.exists(cfg.ckpt_path): start_step = load_ckpt(...)로 불러온다.

5단계: 스모크런 → baseline 스모크런(smoke run)은 "전원을 켜 보고 연기가 나는지 확인하는" 짧은 시험 가동이다. max_steps=200으로 에러 없이 끝나는지, loss가 내려가는지만 본다. 통과하면 같은 스크립트로 max_steps와 데이터 양만 늘려 소형 config(6층, d256) baseline을 돌린다. 이때 바꾸는 것은 config 값뿐이어야 한다.

오늘의 실험

실험 A: 200 step에서 loss가 감소하는가

  • 설정: 위 config 그대로 200 step.
  • 관찰: step 0의 loss를 기록한다. 어휘 크기를 V라 할 때, 아무것도 모르는 모델의 loss는 대략 ln(V) 근처여야 정상이다(예: V=50000이면 약 10.8, 문자 단위 V=65이면 약 4.17). 처음 값이 이와 크게 다르면 버그 신호다. 그 뒤 loss가 단조롭게 내려가진 않더라도 전체적으로 줄어드는지 본다.
  • 기록: step 0 / 50 / 100 / 200의 loss, lr, grad_norm을 표로 남긴다.

실험 B: 체크포인트 재개가 진짜 이어지는가

  1. max_steps=200으로 돌리되 step 100에서 일부러 중단(Ctrl+C)한다. 마지막 저장이 step 100 근처인지 확인한다.
  2. 같은 명령으로 다시 실행해 "step 101부터" 시작하는지 확인한다.
  3. 비교 대조군: 중단 없이 200 step을 한 번에 돌린 결과와 step 150 부근 loss를 비교한다. 완전히 같지는 않을 수 있다(데이터 샘플링의 난수 상태를 저장하지 않았다면 배치가 달라진다). 크게 튀지 않고 이어지는지가 확인 포인트다.
  4. 반례 실험: 옵티마이저 상태를 일부러 불러오지 않고 재개해 보면, 재개 직후 loss가 잠깐 튀는 것을 관찰할 수 있다.

실험 C(선택): warmup을 0으로 바꾸기 warmup_steps=0, 학습률을 크게(예: 3e-3) 두면 초반 loss가 불안정해지는지 본다. 변경은 config 한 줄만 바꾸고, wandb 런 이름에 설정을 적어 구분한다.

기록 방법: 실험마다 config 파일, git 커밋 해시(있다면), 결과 표, 한 줄 해석("warmup을 빼니 step 15에서 grad_norm이 치솟음")을 노트에 남긴다.

왜 이걸 배우나 (LLM / Post-training 연결)

Pretraining이든 SFT든, 대형 학습은 거의 항상 "warmup+decay 스케줄, 혼합 정밀도, gradient clipping, 중단 후 재개 가능한 체크포인트, 지표 로깅"의 조합 위에서 돌아가므로 오늘 만든 골격이 그대로 확장된다. Post-training 실험에서는 수많은 설정을 비교해야 하기 때문에, config로 재현되는 baseline이 없으면 개선 여부 자체를 주장할 수 없다.

초보자가 막히는 지점

  • loss가 `nan`이 된다: 학습률이 너무 크거나 fp16 오버플로가 흔한 원인이다. 학습률을 낮추고, fp16 대신 bf16을 쓰고, grad_norm을 로그로 찍어 어느 step에서 폭발하는지 확인한다.
  • `zero_grad`를 빼먹음: 기울기가 누적되어 loss가 이상하게 움직인다. 루프 안에서 매 step 호출한다.
  • clipping 위치 실수: clip_grad_norm_은 backward() 뒤, opt.step() 앞에 둔다. 순서가 바뀌면 효과가 없다.
  • 재개했는데 step이 0부터 시작: start_step을 range에 반영하지 않았거나 ck["step"]을 저장하지 않은 경우다.
  • 스케줄을 step마다 적용하지 않음: 옵티마이저 생성 시 lr만 넣고 끝내면 스케줄이 작동하지 않는다. 매 step param_groups의 lr을 갱신하는지 확인한다.
  • CPU에서 AMP 에러: 개발 환경에 GPU가 없으면 use_amp를 꺼서 로직 확인부터 하고, 속도는 나중에 본다.
  • wandb 로그인 막힘: wandb offline이나 CSV 로깅으로 먼저 우회하고, 파이프라인 검증을 우선한다.
  • GPU 메모리 부족(OOM): batch_size나 block_size를 줄인다. 결과 비교를 위해 바꾼 값은 반드시 config에 기록한다.

알고리즘 (병행 1시간)

LeetCode Medium 2문제(DP). 한 문제당 30분 안팎으로 시간을 잡는다.

풀이 후 문제마다 다음을 기록한다.

  • 접근: 상태(dp[i]가 무엇을 뜻하는가)를 한 문장으로 정의하고, 점화식과 초기값을 적는다.
  • 시간복잡도 / 공간복잡도: 반복문 구조에서 도출하고, 공간을 1차원이나 변수 몇 개로 줄일 수 있는지도 생각한다.
  • 틀린 이유: 틀린 제출이나 막힌 지점을 남긴다(예: 경계 인덱스, 초기값 누락, 상태 정의 모호함). 30분 안에 못 풀면 풀이를 본 뒤, 다음날 아무것도 보지 않고 다시 풀어 본다.