동적 희소성을 활용한 메모리 효율적 LLM 학습: 안정성에서 실질적 스케일링까지

동적 희소 학습 (Dynamic Sparse Training, DST)은 심층 신경망 (Deep Neural Networks)의 학습 및 추론 효율성을 향상시키기 위한 유망한 패러다임을 제공합니다. 그러나 우리는 대규모 언어 모델 (Large Language Model, LLM) 학습에서 DST가 토폴로지 (Topology) 업데이트 이후 손실 급증 (Loss Spikes)으로 나타나는 최적화 불안정성 (Optimization Instability) 문제를 겪을 수 있음을 발견했습니다. 본 연구에서는 표준 Adam 기반 옵티마이저 (Optimizer)를 단순하게 사용할 경우, 새로 재생성된 파라미터 (Parameters)에 대해 콜드 스타트 (Cold-start) 문제가 발생하여 과도하게 큰 업데이트를 초래하고 학습 역학 (Training Dynamics)을 방해한다는 것을 보여줍니다. 이 문제를 해결하기 위해, 우리는 옵티마이저 웜업 (Optimizer Warm-up)을 통해 DST를 안정화하고 밀도 인식 학습률 스케일링 (Density-aware Learning-rate Scaling)을 통해 학습 진행을 개선하는 SMET (Sparse Memory-Efficient Training)를 제안합니다. SMET은 활성 파라미터 (Active Parameters)에 대해서만 그래디언트 (Gradients)와 옵티마이저 상태 (Optimizer States)를 저장함으로써 메모리 소비를 더욱 줄입니다. 우리는 SMET 하에서의 업데이트 동작에 대한 이론적 분석을 제공하여 최적화 안정성이 향상되었음을 보여줍니다. 광범위한 실험을 통해 SMET이 LLM의 안정적이고, 확장 가능하며, 메모리 효율적인 희소 사전 학습 (Sparse Pre-training)을 가능하게 함을 입증하였으며, 이는 희소 학습 (Sparse Training)이 밀집 학습 (Dense Training)의 실질적인 대안이 될 수 있는 길을 열어줍니다. 우리의 코드는 다음에서 공개적으로 사용할 수 있습니다: https://github.com/QiaoXiao7282/SMET.