개인 개발자의 998달러 3.8B LLM 학습

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

개인 개발자가 config-driven 학습 프레임워크 little-lm으로 3.8B 파라미터 LLM을 구축해 65.3B 토큰을 43시간 동안 학습했으며, 8× B200 사용 비용은 $998이었습니다. Muon과 AdamW의 조합, 사다리꼴 학습률 일정, ClimbMix 데이터셋, FP8, vocabulary padding을 적용해 1024 context 모델은 CORE 0.338을 기록했고, 2048 context 재학습에서는 0.3840까지 올랐습니다. 성능 차이의 상당 부분은 SQuAD와 boolq처럼 긴 prompt가 필요한 과제에서 context가 잘려 worked examples가 사라진 데서 비롯됐으며, 2048 context의 추가 점수는 주로 이 문제를 완화한 결과였습니다. Value embeddings는 파라미터를 19% 늘리면서 12,500 step 기준 CORE를 3.2% 높였지만, peak learning rate와 QK-norm 같은 핵심 설정은 별도 ablation을 거치지 않아 추가 검증이 남았습니다.

섹션별 상세

nanoGPT toy와 연구소급 학습 사이에는 개인이 수천 달러 규모로 의미 있는 LLM을 만들 수 있는 구간이 있으며, 이 프로젝트는 그 가능성을 직접 측정하는 데서 출발했습니다. little-lm은 YAML 파일 하나에 model, dataset, optimizer, schedule, callback을 지정하고 전역 registry에서 구성 요소를 이름으로 찾아 실행하는 구조라서 실험 변경을 코드 branch가 아닌 세 줄짜리 설정 차이로 표현합니다. 그 결과 3.848B 모델이 65.3B 토큰을 8× B200에서 43시간 동안 학습했고, $998의 비용으로 CORE 0.384를 기록해 비슷한 비용의 nanochat d32보다 높은 점수를 확보했습니다.

초기 858M Llama 실험은 FineWeb-Edu 16.4B 토큰 학습에 5.8일을 썼지만 PIQA 60.45%에 그쳐 GPT-2 124M의 약 63%보다 낮았습니다. 원인은 cosine decay가 전체 step의 약 70% 이후 학습률을 사실상 바닥에 붙였고, 858M 규모에 peak learning rate 2.5e-4가 보수적이었으며, 모든 파라미터에 AdamW를 적용했기 때문으로 정리됐습니다. 이후 warmup 5%와 평탄 구간을 거친 선형 cooldown, 행렬 파라미터용 Muon과 나머지 파라미터용 AdamW, ClimbMix, FP8와 vocabulary padding, 1024 context를 조합하면서 3.8B 학습의 마지막 step까지 eval loss가 계속 내려가는 상태를 만들었습니다.

최종 모델은 RMSNorm, RoPE, GQA, relu² MLP, QK-norm, logit softcap, layer별 learnable residual scalar, ResFormer-style value embeddings를 결합한 Llama-style decoder-only 구조입니다. query head 24개와 KV head 8개를 사용하고, 14개의 value-embedding table이 전체 3.848B 파라미터 가운데 721.2M, 즉 19%를 차지합니다. Value embeddings를 제거한 동일 조건 실험에서는 12,500 step 기준 loss가 2.1171에서 2.1075로, CORE가 0.3047에서 0.3147로 개선됐고 처리량은 477,908과 479,445 tokens/sec로 거의 같아 lookup 파라미터가 FLOPs보다 메모리와 optimizer state에 주로 비용을 남겼습니다.

학습 처리량은 FP8, vocabulary padding, fused linear cross-entropy, 비게이트 relu² MLP, bf16 master weights를 차례로 적용하며 크게 늘었습니다. 5090에서 baseline 26,144 tokens/sec가 최종 37,621 tokens/sec로 상승했고, vocabulary를 50,257에서 50,304로 채우면 tensor-core 경로가 열리며, fused loss는 full logits tensor를 만들지 않아 VRAM을 8GB 절약해 micro-batch 증가를 가능하게 했습니다. B200에서는 steady state 약 480,000 tokens/sec, SM activity 92%, B200당 약 1,047 TFLOP/s를 기록했으며, 로컬 shard 다운로드로 dataloader와 네트워크 대기를 피했습니다.

1024 context의 CORE 0.3384와 2048 context의 0.3840 차이는 모델의 일반적인 품질보다 context-dependent 평가 과제의 입력 절단에서 크게 발생했습니다. 1024 context에서는 SQuAD prompt의 100%, boolq의 99.8%, bigbench_language_id의 99.7%가 잘렸고, SQuAD의 10-shot 예시가 사라져 passage와 question만 남으면서 유창한 답변이 exact-token match에서 계속 0점 처리됐습니다. 2048 context에서는 SQuAD 점수가 0.0000에서 0.3114, boolq가 0.5131에서 0.7095로 올라갔지만 다른 19개 과제의 합산 변화는 +0.008에 그쳐, 긴 context의 추가 비용은 주로 측정 가능한 과제를 복원하는 데 쓰였습니다.

실험에는 문서 경계별 flex attention mask, Liger RMSNorm과 RoPE kernel, streaming dataset, nanochat-style initialization 같은 선택도 포함됐지만 복잡도나 효과가 충분하지 않아 일부를 되돌렸습니다. 문서별 attention mask는 best-fit packing과 unconditional causal attention으로 대체됐고, Liger RoPE는 microbenchmark에서 2.2배 빨랐어도 end-to-end 처리량 변화가 없었으며, streaming은 로컬 shard보다 2~3% 느리고 간헐적인 네트워크 지연을 일으켰습니다. 앞으로는 value embeddings에 투입한 721M 파라미터의 대체 용도, 동일 wall-clock 기준 context 비교, commonsense_qa 하락 원인, optimizer sharding과 추가 데이터 탐색이 남아 있습니다.

용어 해설

Muon Optimizer: — 행렬 파라미터의 gradient를 Newton-Schulz 직교화로 갱신하는 최적화 기법입니다. 나머지 파라미터에는 AdamW를 함께 적용하며, 이 글에서는 step당 계산 비용이 늘어도 전체 수렴 속도를 높이는 방식으로 활용됐습니다.
사다리꼴 학습률 스케줄(Trapezoidal Learning-Rate Schedule): — 초기 warmup 뒤 일정한 최고 학습률을 유지하고, 학습 후반부에 선형으로 낮추는 일정입니다. cosine decay처럼 후반부 대부분을 지나치게 낮은 학습률로 보내지 않아 마지막 단계까지 손실을 줄이는 데 초점을 둡니다.
Grouped-Query Attention: — 여러 query head가 더 적은 수의 key·value head를 공유하는 Attention 구조입니다. 이 모델은 query head 24개와 KV head 8개를 사용해 KV 상태와 메모리 부담을 줄이는 구성을 택했습니다.
Value Embeddings: — 일부 층에 별도의 어휘 임베딩 테이블을 추가해 Transformer 내부의 value 경로에 개념 관련 편향을 제공하는 구성입니다. 3.8B 모델에서는 14개 테이블이 전체 파라미터의 19%를 차지했지만 lookup 방식이라 처리량 손실은 거의 없었습니다.
CORE 벤치마크(CORE): — 언어 이해와 상식 관련 22개 과제를 묶어 모델 성능을 측정하는 평가 지표입니다. 무작위 기준을 중심으로 점수를 산출하므로 점수가 낮은 구간에서는 작은 logit 변화도 상대적인 점수 차이를 크게 만들 수 있습니다.

코드 예제

yaml

model:
  hidden_size: 3072
  intermediate_size: 12288 # 4x, non-gated
  num_hidden_layers: 28
  num_attention_heads: 24
  num_key_value_heads: 8 # 3:1 GQA
  head_dim: 128
  hidden_act: relu2
  gated_mlp: false
  qk_norm: true
  logit_softcap: 15.0
  layer_scale: true
  value_embeddings: true # 14 tables, alternating layers
  tie_word_embeddings: false
  rope_theta: 10000.0
  rms_norm_eps: 1.0e-6
  vocab_pad_to: 64 # 50257 -> 50304
  max_position_embeddings: 2048
dtype: bf16
engine:
  compile: true
  fp8: true
precision: bf16
total_batch_size: 2293760 # 20 x 2048 x 7 grad_accum x 8 GPUs
loss: LigerFusedLinearCrossEntropyLoss(softcap=15.0)
optimizer: # composite, one group per parameter class
  matrix: Muon lr=0.02 momentum=0.95 wd=0.0
  embeddings: AdamW lr=0.1414 betas=(0.8, 0.995) eps=1e-10 wd=0.001
  lm_head: AdamW lr=0.002828 betas=(0.8, 0.96) eps=1e-10 wd=0.01
  value_embeds: AdamW lr=0.0707 betas=(0.8, 0.995) eps=1e-10 wd=0.01
  scalars: AdamW lr=0.005 betas=(0.8, 0.95) eps=1e-10 wd=0.05
scheduler:
  trapezoidal:
    warmup_ratio: 0.05
    warmdown_ratio: 0.50
    final_lr_frac: 0.05
data:
  dataset: nvidia/Nemotron-ClimbMix (karpathy/climbmix-400b-shuffle shards)
  tokenizer: gpt2 (tiktoken)
  block_size: 2048
  packing: best-fit, BOS-aligned
  batch_size: 20 per rank
  num_workers: 11
trainer:
  max_steps: 32000 # stopped at ~28,000 -> 65.3B tokens
  eval_every: 4000 # must divide max_steps or the final CORE is skipped

little-lm의 최종 3.8B 모델 학습에 사용한 전체 YAML 설정으로, 모델 구조와 optimizer, 데이터셋, batch 구성, 평가 주기를 한곳에 고정합니다.

text

The capital of France is Paris. It is the largest city in France and the second largest city in Europe The french revolution happened in 1789 and 1799, and was a time of great change in france At the center of the milky way there is a supermassive black hole. It is called Sagittarius A* (pronounced Electrons orbit around the nucleus of an atom in a series of energy levels. The energy levels are numbered Newton discovered the laws of motion and gravity. He also discovered the law of universal gravitation. Newton's

학습 종료 후 생성된 텍스트 예시로, 사실 단위가 이어지지만 문장 경계와 완결성이 불안정한 모델의 출력 특성을 드러냅니다.

기술

little-lm
nanochat
Muon
AdamW
FP8
torch._scaled_mm
LigerFusedLinearCrossEntropyLoss
ClimbMix
DistributedDataParallel
B200
RTX 5090
H100

활용 사례

개인 연구자의 소규모 LLM 사전학습
제한된 GPU 예산에서의 학습 설정 비교
학습 throughput과 capability-per-dollar 최적화
context length에 민감한 언어 이해 benchmark 평가