처음부터 만드는 LLM 32e부: 학습률 스케줄링과 최적화 개입

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

핵심 요약

대형 언어 모델 학습 시 고정된 학습률을 사용하면 손실 함수 최소값 근처에서 수렴하지 못하고 진동하는 문제가 발생한다. 이를 해결하기 위해 초기 그래디언트 폭주를 막는 선형 웜업과 부드러운 수렴을 유도하는 코사인 감쇠 기법을 도입했다. DeepSeek의 하이퍼파라미터 스케일링 법칙을 적용해 현재 연산 예산에 맞는 최적 피크 학습률을 0.0014로 도출했다. PyTorch의 SequentialLR을 활용해 두 스케줄러를 체이닝하고 체크포인트 저장 시 스케줄러 상태를 포함함으로써 학습의 연속성을 확보했다.

배경

Python 및 PyTorch 기본 지식, 경사 하강법(Gradient Descent) 및 옵티마이저 개념, LLM 학습 파이프라인에 대한 이해

대상 독자

LLM을 처음부터 학습시키거나 학습 파이프라인을 최적화하려는 머신러닝 엔지니어

의미 / 영향

학습률 스케줄링은 LLM 학습의 수렴 속도와 최종 성능에 결정적인 영향을 미친다. 특히 연산 예산에 따른 스케일링 법칙을 적용함으로써 시행착오 비용을 줄이고 효율적인 모델 학습이 가능해진다.

섹션별 상세

고정 학습률의 한계와 스케줄링의 필요성을 분석했다. 학습 초기에는 큰 보폭으로 손실 지형을 탐색해야 하지만, 최적점 근처에서는 보폭을 줄여야 정밀한 수렴이 가능하다. 특히 V자형 손실 곡선에서는 고정 학습률이 최소값 주변을 무한히 진동하게 만들 수 있음을 확인했다.

코사인 감쇠(Cosine Decay) 스케줄링을 도입했다. Chinchilla 논문의 권장 사항에 따라 학습률을 피크값의 1/10 수준까지 코사인 곡선을 그리며 낮추는 방식을 채택했다. 이는 초기 탐색과 후기 미세 조정 단계를 자연스럽게 연결하여 최적의 성능을 이끌어낸다.

학습 초기 안정성을 위한 선형 웜업(Linear Warmup)을 적용했다. 모델 초기화 직후의 무작위 에러로 인한 급격한 파라미터 업데이트를 방지하기 위해, 전체 학습 스텝의 약 5% 구간 동안 학습률을 0에서 피크값까지 서서히 높이는 과정을 추가했다.

DeepSeek 논문의 공식을 활용해 최적 피크 학습률을 산출했다. 연산 예산(FLOPs)과 학습률 사이의 상관관계를 정의한 스케일링 법칙 공식을 통해, 현재 프로젝트의 예산인 5e18 FLOPs에 적합한 학습률이 약 1.4e-3임을 도출하고 이를 실험에 적용했다.

python

C = 5e18 # Compute budget in FLOPs
eta_opt = 0.3118 * C**-0.1250
# Result: 0.00143385376262387

DeepSeek 논문의 공식을 활용하여 연산 예산에 따른 최적 학습률을 산출하는 코드

PyTorch의 SequentialLR을 이용한 구현 및 체크포인트 관리 방법을 정립했다. LinearLR과 CosineAnnealingLR을 순차적으로 연결하여 구현했으며, 학습 중단 후 재개 시 스케줄러의 state_dict를 함께 로드해야 학습률이 초기화되는 문제를 방지할 수 있음을 확인했다.

python

peak_lr = 0.0014
warmup_period = 1600
decay_period = 32000

optimizer = torch.optim.AdamW(
    model.parameters(), lr=peak_lr, weight_decay=0.1
)

warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.00001, end_factor=1.0, total_iters=warmup_period
)

cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=decay_period, eta_min=peak_lr / 10
)

scheduler = torch.optim.lr_scheduler.SequentialLR(
    optimizer, 
    schedulers=[warmup_scheduler, cosine_scheduler], 
    milestones=[warmup_period]
)

PyTorch의 SequentialLR을 사용하여 선형 웜업과 코사인 감쇠를 결합한 스케줄러 구현 예시

실무 Takeaway

학습 초기 불안정성을 제어하기 위해 전체 스텝의 5% 내외를 선형 웜업 구간으로 설정하여 학습률을 점진적으로 높여야 한다.
DeepSeek의 하이퍼파라미터 스케일링 공식을 사용하면 대규모 실험 없이도 연산 예산에 기반한 합리적인 피크 학습률을 추정할 수 있다.
PyTorch에서 스케줄러를 체이닝할 때는 SequentialLR을 사용하고, 체크포인트 저장 시 반드시 스케줄러의 상태 정보도 포함해야 한다.

언급된 리소스

문서Build a Large Language Model (from Scratch)

논문DeepSeek-V3 Technical Report