little-lm 3.8B는 무작위 초기화부터 653억 토큰을 사전학습해 CORE 0.384를 달성한 개인 프로젝트로, B200 8대에서 43시간 동안 학습하는 데 998달러가 들었음
- 초기 실험의 낮은 성능을 개선하기 위해
Muon, 사다리꼴 학습률 스케줄, ClimbMix 데이터, FP8 등을 적용했으며, 최종 학습에서는 마지막 단계까지 평가 손실이 감소함
- 처리량 최적화에서는 개별 연산 속도보다
전체 학습 속도와 메모리 사용량이 중요했음. 융합 손실 함수는 같은 배치에서 6% 느렸지만, 절약한 메모리로 배치를 키우자 전체 처리량이 높아짐
- 컨텍스트를 1024에서 2048토큰으로 늘리자 CORE가
0.3384에서 0.3840으로 상승했지만, 상승분의 83%는 프롬프트 잘림에 영향을 받던 SQuAD와 BoolQ에서 나왔음
- 구성요소 제거 실험에서는
가치 임베딩이 거의 같은 처리량으로 성능을 높인 반면, GQA는 학습 속도가 빨라도 같은 손실에 도달하는 비용에서는 이점이 없었음
개인 규모의 사전학습과 모델 구성
little-lm은 저녁 시간에 개발하고 RTX 5090에서 디버깅한 뒤 임대한 B200에서 학습을 마친 프로젝트로, Andrej Karpathy의 nanochat에서 큰 영향을 받았음
-
최종 모델은
38억 4800만 파라미터이며, 컨텍스트 길이에 따른 학습 결과는 다음과 같음
-
1024토큰에서는 573억 토큰을 35.9시간 동안 학습했으며, 비용은 820달러, CORE는 0.3384였음
-
2048토큰에서는 653억 토큰을 43시간 동안 학습했으며, 비용은 998달러, CORE는 0.3840이었음
-
비교 대상으로 GPT-2 1.5B의 CORE는
0.2565, nanochat d26 약 561M은 약 0.258, d32 약 1B는 0.310이었음
-
nanochat d32는 H100 8대로 약 33시간, 약 1000달러를 사용했음
-
little-lm은 더 큰 모델이며 B200을 사용했으므로 동일 하드웨어 비교는 아님. 실험 당시에는 B200의 작업량 대비 비용이 더 유리했음
-
학습 프레임워크는
YAML 설정 기반이며, 구성요소를 전역 레지스트리에 등록하고 이름으로 선택함
-
옵티마이저나 데이터셋 교체를 설정 한 줄로 처리하고, 실험을 별도 코드 분기 대신 작은 YAML 변경으로 표현할 수 있었음
-
관심사 분리, 명확한 인터페이스, 교체 가능한 구성요소에 들인 초기 비용은 수렴 문제를 디버깅할 때 회수할 수 있었음
-
모델은
Llama 계열 구조에 RMSNorm, RoPE, GQA, relu² MLP, QK-norm, 로짓 소프트캡, 층별 학습 가능한 잔차 스칼라, ResFormer 방식의 가치 임베딩을 사용함
-
디코더는 28개 층이며 쿼리 헤드 24개, KV 헤드 8개로 구성됨
-
토큰 임베딩과 가중치를 공유하지 않는 LM 헤드는 각각 1억 5450만 파라미터, 디코더 층은 28억 1870만 파라미터임
-
가치 임베딩은 한 층씩 건너뛰며 배치한
vocab × kv_dim
테이블 14개로, 7억 2120만 파라미터이자 전체의 약 19%를 차지함
초기 실패에서 바꾼 학습 방식
- 초기에는
858M Llama 모델을 단일 A100에서 약 5.8일 동안 FineWeb-Edu 164억 토큰으로 학습했으나, PIQA 정확도가 약 60.45%에 그쳤음
- GPT-2 124M의 약 63%보다 낮았고, 생성문도 반복적이거나 의미가 불분명했음
- AdamW 학습률
2.5e-4
, 5% 워밍업, 0까지 감소하는 코사인 스케줄, 누적 배치 256, 컨텍스트 2048을 사용했음
학습률 스케줄과 최대 학습률이 주요 개선 대상이었음
- 코사인 스케줄에서는 전체 단계의 약 70% 이후 곡선이 평평해졌으며, 너무 낮아진 학습률 때문에 마지막 30%의 연산이 거의 기여하지 못했을 가능성이 있었음
- 858M 규모에서
2.5e-4
는 보수적인 최대 학습률이었음
-
이후 5% 워밍업, 일정 학습률 유지, 마지막 50% 구간에서 최대값의 5%까지 선형 감소하는 사다리꼴 스케줄을 사용함
-
행렬 파라미터에는
Muon, 나머지에는 AdamW를 적용함
-
Newton-Schulz 직교화 때문에 그래디언트 누적 횟수가 적은 실험에서는 옵티마이저 단계 비용이 약 25% 늘었지만, 7회 누적에서는 추가 비용이 약 4%로 줄어듦
-
전체 실행 시간 기준 수렴은 더 빨랐으며, 개인적으로 진행한 비전과 자연어 처리 프로젝트에서도 Muon으로 수렴이 개선됐음
FineWeb-Edu를 ClimbMix로 교체하면서 수렴 속도가 크게 향상됐으며, Karpathy의 실험에서도 같은 경향이 있었음
FP8과 짧은 컨텍스트로 처리량을 높임
torch._scaled_mm
과 세 GEMM 모두에 대한 동적 텐서 단위 스케일링으로 FP8 학습을 구현함
- 어휘 크기를 50,257에서 64의 배수인 50,304로 패딩했으며, 패딩 자체의 추가 개선은 작았음
- 컨텍스트를 2048에서 1024로 줄이면 같은 메모리에서 배치를 대략 두 배로 키울 수 있었고, 설정에 따라 처리량이 약 10% 개선됨
전체 학습 처리량을 높인 최적화
-
1024토큰 학습은 안정 상태에서
초당 약 48만 토큰을 처리했으며, 정기 CORE 평가를 포함한 전체 실행 시간은 35.9시간이었음
-
2500단계에서 평가 손실 2.3278, CORE 0.2389였고, 25,000단계에서는 각각 1.9868과 0.3384에 도달함
-
안정 상태의 SM 활성도는 92%, SM 점유율은 40%였음
-
B200 한 대당 지속 성능은 약 1047 TFLOP/s로, 밀집 연산 FP8 최고 성능 대비 MFU 약 25%였음. bf16 최고 성능 기준으로는 약 50%이며, 모든 선형 층이 FP8로 실행되는 것은 아님
-
단일 노드의 3.8B 모델에는
DistributedDataParallel을 사용했으며, 그래디언트 통신은 병목이 아니었음
-
RTX 5090의 858M 모델에서 bf16과 컴파일을 사용한 기본 처리량은
초당 26,144토큰이었고, 최적화 후 37,621토큰으로 약 44% 증가함
-
FP8과 어휘 패딩을 함께 적용한 개선폭은 약 33%였음
-
Liger의
FusedLinearCrossEntropyLoss
는 LM 헤드 행렬곱과 손실 계산을 합치고 내부적으로 청크 처리해 전체 (B*T, vocab)
로짓 텐서를 만들지 않음
- 배치 6에서는 처리량이 34,724에서 32,952토큰/초로 약 6% 줄었지만, 메모리는 27,852에서 19,630MiB로 감소함
- 배치를 10으로 키우자 메모리 28,872MiB에서 37,621토큰/초에 도달함
- 같은 배치만 비교하면 손해처럼 보여도, 확보한 메모리를 활용하면 전체 처리량은 개선됐음
비게이트 MLP와 bf16 마스터 가중치도 큰 영향을 줌
-
작은 모델에서 SwiGLU를 relu²로 바꾸면 행렬곱이 3회에서 2회로 줄어 처리량이 183,035에서 214,173토큰/초로 증가하고 메모리 약 6GB를 절약함
-
SwiGLU의 중간 차원 배율 2.75를 relu²에 그대로 적용하면 학습이 나빠졌으며, 비게이트 구조에서는 4배를 사용함
-
1.5B 설정에서 옵티마이저 마스터 가중치를 fp32 대신 bf16으로 유지하자 메모리가 27% 감소하고 처리량이 초당 64만에서 140만 토큰으로 약 2.2배 증가함
-
다만 4000단계 CORE가 0.23에서 0.22로 낮아지는 품질 저하가 있었음
-
같은 코드와 FP8을 사용한 150M 모델에서
B200은 RTX 5090보다 2.59배 빨랐음
-
처리량은 각각 477,440과 184,662토큰/초였으며, 더 큰 메모리로 배치를 추가 확대하는 효과는 포함하지 않은 비교임
효과가 없었거나 되돌린 시도
문서 경계 마스킹은 실험 규모에서 모델 품질을 개선하지 못했음
- flex attention과 토큰별 문서 ID로 다른 문서에 대한 어텐션을 차단했지만, 복잡도에 상응하는 이득이 없었음
- 약 10줄의 best-fit 패킹과
F.scaled_dot_product_attention(..., is_causal=True)
로 대체함
- Karpathy도 BOS 정렬 패킹에서 문서 간 정보 누출의 악영향이 크지 않다고 확인했지만, 결과는 학습 데이터 분포의 영향을 받음
Liger RoPE와 RMSNorm은 제거함
- RoPE는 마이크로벤치마크에서 2.2배 빨랐으나 전체 처리량에는 측정 가능한 차이가 없었음
- RMSNorm은 PyTorch 2.9의
F.rms_norm
보다 느렸으며, 측정 시간은 0.41ms 대 0.25ms였음
nanochat 방식 초기화는 초반 손실만 조금 낮췄으며, 약 1500단계 이후에는 기존 방식과 곡선이 겹쳤음
, 선형 가중치 균등분포, 출력 투영 0 초기화, LM 헤드 N(0, 0.001)
을 사용함
초기화 대비 측정 가능한 품질 차이는 없었지만, 더 원칙적인 초기화 방식이라는 이유로 유지함
데이터 스트리밍은 초기 탐색에는 편리했지만, 유료 장시간 학습에서는 네트워크 정체와 중단 위험이 있었음
- 로컬 샤드는 네트워크가 정상일 때도 처리량을 2~3% 높였음
- 몇 시간 이상 학습할 때는 시작 전에 몇 분을 들여 데이터를 내려받는 방식을 선택함
가치 임베딩과 GQA 제거 실험
가치 임베딩은 동일 설정으로 12,500단계, 약 290억 토큰까지 비교함
-
사용 시 3.848B, 평가 손실 2.1075, CORE 0.3147이었고, 제거 시 3.128B, 손실 2.1171, CORE 0.3047이었음
-
처리량은 각각 479,445와 477,908토큰/초로 사실상 같았음
-
조회 연산이므로 메모리와 옵티마이저 상태는 늘지만 FLOPs는 거의 늘지 않음
-
파라미터 약 19% 증가로
손실은 0.46%, CORE는 3.2% 개선됐음
-
해당 구간의 손실 감소 속도로 환산하면 약 1200단계, 전체 25,000단계의 약 5%를 추가 학습한 효과에 해당함
-
CORE의 상대 개선폭은 손실의 약 7배였고, 학습이 진행되며 차이는 줄었음
-
CORE는 정확도 기반이라 작은 로짓 변화로 정답 여부가 뒤집힐 수 있고, 무작위 기준선으로 중심화하므로 낮은 점수 구간에서 상대 차이가 증폭됨
GQA 실험은 더 저렴하게 충분히 학습할 수 있는 1.5B 모델에서 MHA, 2:1 GQA, 3:1 GQA를 비교함
-
MHA는 KV 헤드 12개로 손실 2.1433, CORE 0.2857을 기록했으며, 처리량은 초당 123.5만 토큰, 학습 시간은 9.4시간이었음
-
2:1 GQA는 KV 헤드 6개로 손실 2.1652, CORE 0.2618을 기록했으며, 처리량은 초당 141.4만 토큰, 학습 시간은 8.2시간이었음
-
3:1 GQA는 KV 헤드 4개로 손실 2.1762, CORE 0.2812를 기록했으며, 처리량은 초당 138만 토큰, 학습 시간은 8.4시간이었음
-
GQA는
처리량이 약 1215% 높지만 평가 손실이 1.01.5% 나빴음
-
MHA는 메모리 부족을 피하려고 마이크로배치를 줄였으며, 이것도 처리량 차이에 영향을 줌
-
MHA가 각 GQA의 최종 손실에 도달하는 데는 약 7.9시간과 7.4시간이 걸려, 이 규모에서는 GQA가 학습 비용을 줄이지 못했음
-
GQA의 실질적인 장점은 추론 시 KV 캐시가 2~3배 작아진다는 점임
CORE만으로 작은 차이를 순위화하기는 어려웠음
- 작은 성능 차이를 판단할 때는
CORE의 변동성도 고려해야 함
- 손실이 더 나쁜 3:1 GQA가 2:1보다 CORE는 높았고, 체크포인트 사이 점수 변동도 약 0.01이었음
- 3:1이 2:1보다 약간 느린 이유는 확정하지 못했으며, 노드 간 편차일 가능성이 있음
컨텍스트 길이가 CORE 측정에 미친 영향
- CORE 22개 과제 중
SQuAD, BoolQ, bigbench_language_id는 1024토큰에서 거의 모든 프롬프트가 잘렸음
- 잘림 비율은 각각 100%, 99.8%, 99.7%였음
- SQuAD 점수는 학습 초반 0.1478에서 최종 0으로 계속 감소함
- BoolQ는 10,000단계에서 0.6294로 정점을 찍은 뒤 0.5131로 내려갔고, 언어 식별은 우연 수준에 머물렀음
SQuAD의 평가 형식과 잘림 방식이 점수 하락을 설명함
- DCLM 평가 묶음의 SQuAD는 예제 10개를 먼저 주는 10-shot 과제이며, 평가 프롬프트 중앙값은 1998토큰이었음
- 평가 하네스는 길이 초과 시 마지막
max_seq_len
토큰만 유지함
-
끝에 있는 약 169토큰의 실제 시험 지문과 질문은 남았지만, 기대 출력 형식을 알려주는 예제들이 잘렸음
-
정답과 정확히 일치해야 점수를 얻으므로, 학습이 진행되며 유창한 문장을 생성할수록 초반의 우연한 짧은 정답이 사라졌음
-
컨텍스트를
2048토큰으로 늘릴 때 메모리 사용량을 유지하도록 마이크로배치를 절반으로 줄여, 옵티마이저 단계당 토큰 수는 동일하게 유지함
-
임대 시간을 아끼기 위해 약 28,000단계에서 중단했으므로 학습률 감소를 끝까지 완료하지 못했으며, 해당 결과는 하한값으로 봄
-
20,000단계의 평가 손실은 두 실행 모두 소수점 셋째 자리까지 2.016으로 같았지만, CORE는 0.034 차이가 났음
-
최종 CORE 상승폭
0.0456의 83%는 SQuAD와 BoolQ에서 나왔음
-
SQuAD는 0에서 0.3114로 상승했고, 잘림은 100%에서 47%로 감소함
-
BoolQ는 0.5131에서 0.7095로 상승했고, 잘림은 99.8%에서 3.2%로 감소함
-
BoolQ의 무작위 기준선은 다수 클래스인 “yes” 비율 약 0.62이므로, 원점수 상승 0.196은 중심화 후 약 0.517로 확대됨
-
두 과제를 제외한 나머지 20개 과제의 합산 기여는 약 0.008로, 약 14% 더 많은 학습 토큰만으로 얻을 법한 수준이었음
언어 식별은 잘림이 99.7%에서 14%로 줄어도 점수가 약 0.005만 올랐으며, 이 모델에 특히 어려운 과제였음
-
반대로 commonsense_qa는 0.072, cs_algorithms는 0.031 하락함
-
2048토큰은 이 실험에서 일반적인 품질 향상보다
평가 가능한 과제를 늘리는 선택에 가까웠음
-
처리량은 초당 48만에서 43.7만 토큰으로 약 9% 감소함
-
후속 실험에서 컨텍스트 확대에 따른 CORE 상승은 1.5B 모델 약 0.014, 3.8B 모델 약 0.046이었음
-
긴 컨텍스트로 잘렸던 과제를 평가할 수 있게 되더라도, 실제 점수 상승폭은 모델이 그 과제를 얼마나 잘 수행하는지에 따라 달라짐
검증하지 못한 설정과 후속 과제
최대 학습률, 스케줄 대안, QK-norm은 별도로 제거 실험을 하지 않았음
- AdamW 학습률은 nanochat의
sqrt(768/d_model)
규칙을 따랐고, Muon 학습률 0.02도 그대로 가져옴
- 코사인에서 사다리꼴 스케줄로 바꿨지만 더 효율적인 스케줄이 있을 수 있으며, QK-norm은 기본 활성 상태를 유지함
- 적은 예산에서는 기존 실험을 재사용할 수 있지만, 다른 모델과 데이터 규모에도 결과가 적용된다는 가정이 남음
가치 임베딩의 파라미터 재배분과 동일 시간의 컨텍스트 비교가 필요함
- 가치 임베딩 제거 실험은 해당 7억 2100만 파라미터를 다른 구조에 쓰는 경우와 비교하지 않았음
- 2048토큰 실행은 컨텍스트뿐 아니라 학습 시간도 늘었으므로, 평가 문제는 확인했지만 동일 시간 대비 품질 문제는 해결하지 못함
- 긴 프롬프트가 필요 없는 commonsense_qa가 0.072 하락한 이유도 미해결임
옵티마이저 샤딩은 추가 처리량 개선 후보임
- 현재 DDP와 단일 GPU용 Muon 조합은 모든 랭크가 전체 옵티마이저 상태를 보유하고 같은 Newton-Schulz 갱신을 중복 계산함
- nanochat은 DDP 래퍼 없이 옵티마이저 내부에서 ZeRO-2 샤딩을 수행하고 reduce-scatter, 계산, all-gather를 겹쳐 실행함
- 메모리 절감은 배치 확대로 이어질 수 있지만, Muon은 전체 그래디언트 행렬이 필요하므로 행렬을 랭크별로 분할하는 것만으로는 중복 직교화가 사라지지 않음
- 랭크마다 완전한 행렬을 나눠 맡기는 방식은 가능성이 있지만, 이 프로젝트에서는 탐색하지 않았음
CORE와 ClimbMix 데이터 분석에 충분한 시간을 쓰지 못했으며, 같은 연산 예산에서 성능을 더 높일 후속 탐색 대상으로 남음
- 실행 설정은
GPT-2 토크나이저와 BOS 정렬 best-fit 패킹을 사용하며, 데이터는 nvidia/Nemotron-ClimbMix
의 karpathy/climbmix-400b-shuffle
샤드임
-
생성 예시로 프랑스 수도와 혁명, 은하 중심의 Sagittarius A*, 전자의 에너지 준위, 뉴턴의 운동과 중력 법칙에 관한 문장 시작 부분을 공개함
-
2019년 연구실의 최전선 결과였던
GPT-2 1.5B의 CORE 점수를 7년 뒤 개인이 저녁 시간과 시간제 임대 하드웨어, 998달러로 크게 넘어섰음
-
최전선의 발전과 함께, 과거 연구실이 필요했던 일부 작업을 개인 엔지니어도 수행할 수 있게 됐음
댓글과 토론