Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

기술 요약: 동적 희소성을 이용한 메모리 효율적 LLM 학습

문제 정의
동적 희소 훈련(Dynamic Sparse Training, DST)은 훈련 과정 전반에 걸쳐 희소한 토폴로지를 유지함으로써 심층 신경망의 효율성을 개선하는 패러다임을 제공한다. 그러나 LLM에 DST를 적용하는 것은 최적화 불안정성으로 인해 까다로운 것으로 밝혀졌다. 저자들은 표준 Adam 기반 옵티마이저를 DST에 나이브하게 적용할 경우, 토폴로지 업데이트(가지치기 및 재성장) 직후에 심각한 "손실 스파이크(loss spikes)"가 발생한다는 점을 지적한다.

그 근본 원인은 "콜드 스타트(cold-start)" 효과로 식별되었다: 새로 재성장된 파라미터들은 축적된 옵티마이저 상태(1차 및 2차 모멘트 추정치)가 부족하다. Adam에서 이러한 새로운 파라미터들에 대한 편향 수정(bias correction) 항은 성숙한 가중치에 비해 불균형적으로 큰 업데이트 크기를 초래한다. 이는 훈련 궤적을 방해하여 일시적인 손실 스파이크를 유발하고, 특히 모델 규모가 커지거나 토폴로지 업데이트가 공격적이 될 때 수렴을 저해한다. 또한, 기존의 DST 구현 방식은 활성화된 파라미터의 일부에 대해서만 옵티마이저 상태를 할당함에도 불구하고 전체 파라미터 공간에 대해 전체 크기의 옵티마이저 상태를 할당하기 때문에 실제 메모리 절감 효과를 실현하지 못하는 경우가 많다.

방법론: 희소 메모리 효율적 훈련 (SMET)
이러한 과제를 해결하기 위해, 저자들은 DST를 안정화하고 LLM의 메모리 소비를 줄이기 위해 설계된 프레임워크인 SMET을 제안한다. 이 방법은 세 가지 핵심 구성 요소로 이루어진다:

1. 옵티마이저 상태 웜업(Optimizer-State Warm-up): 콜드 스타트 효과를 완화하기 위해, SMET은 새로 재성장된 파라미터에 대해 두 단계 웜업 전략을 채택한다:

   • 암시적 웜업 (타임스텝 리셋): 재성장된 파라미터의 내부 타임스텝 카운터($t_i$)를 리셋한다. 이를 통해 Adam의 편향 수정 메커니즘을 재활성화하여 업데이트 규칙의 분모를 팽창시키고, 초기 업데이트 크기를 자연스럽게 제한한다.
   • 명시적 웜업 (학습률 스케일링): 재성장된 파라미터의 학습률이 $W$ 스텝 동안 0에서 목표치까지 선형적으로 증가하는 로컬 선형 웜업을 적용한다.
   • 이러한 메커니즘은 결합되어 업데이트 폭발을 방지하고, 토폴로지 변화 직후의 최적화 궤적을 부드럽게 만든다.

2. 밀도 인지 학습률 스케일링 (Density-Aware Learning Rate Scaling): 웜업으로 인해 도입된 보수적인 업데이트를 보상하고 희소 레이어의 감소된 유효 팬인(fan-in)을 고려하기 위해, SMET은 학습률을 스케일링한다. Maximal Update Parameterization ($\mu$P)에서 영감을 받아, 학습률은 현재 희소 밀도 $d$에 대해 $1/\sqrt{d}$로 스케일링된다. 이는 업데이트 크기가 밀집(dense) 훈련 환경과 유사하게 유지되도록 보장한다.

3. 메모리 효율적 최적화: SMET은 그래디언트와 옵티마이저 상태(모멘트 추정치 $m$ 및 $v$)가 활성화된(0이 아닌) 파라미터에 대해서만 저장되는 희소 저장 전략을 구현한다.

   • 그래디언트는 역전파 중에 활성 인덱스에서만 유지된다.
   • 활성 인덱스와 타임스텝을 추적하는 오버헤드를 관리하기 위해, 메타데이터 비용을 분산시키고자 파라미터를 크기 $B$의 블록으로 그룹화하는 블록 단위 저장 전략이 사용된다.
   • 이 접근 방식은 상태를 압축하지만 밀집 모델을 생성하는 방식과 달리, 훈련 전 과정 동안 희소성을 유지한다.

주요 기여

1. 불안정성 분석: 본 논문은 DST를 이용한 LLM 사전 훈련에서의 손실 스파이크 현상을 식별하고 이론적으로 분석하며, 이를 Adam 옵티마이저 내에서 새로 재성장된 가중치의 콜드 스타트 효과 때문이라고 규명하였다.
2. 안정화 프레임워크: 저자들은 토폴로지 업데이트를 안정화하고 훈련 성능을 향상시키기 위해 옵티마이저 상태 웜업과 밀도 인지 학습률 스케일링을 통합한 SMET을 제안한다.
3. 메모리 효율적 구현: 활성 파라미터에 대해서만 그래디언트와 옵티마이저 상태를 저장하는 희소 최적화 구현을 개발하였으며, 이는 훈련 전 과정 동안 희소성을 유지하면서 비구조적(unstructured) 및 블록 단위 희소 패턴을 모두 지원한다.
4. 체계적 평가: 다양한 모델 규모와 희소도 수준에 걸친 DST에 대한 종합적인 평가를 제공하며, SMET이 기존 DST 방법들을 능가하고 밀집 훈련 성능에 근접함을 입증한다.

실험 결과
C4 및 OpenWebText 데이터셋으로 사전 훈련된 LLaMA 기반 모델(60M에서 1B 파라미터 범위)에 대해 실험을 수행하였다.

• 성능: SMET은 다른 메모리 효율적 베이스라인(LoRA, ReLoRA, GaLore) 및 기존 DST 방법(SET, RigL)보다 일관되게 우수한 성능을 보인다. 더 큰 모델 규모(예: 1B 파라미터)에서 SMET은 높은 희소도 수준(예: 밀도 $d=0.1$)에서도 밀집 훈련에 근접한 퍼플렉시티(perplexity)를 달et한다.
• 안정성: 해당 방법은 베닐라(vanilla) DST에서 관찰되는 손실 스파이크를 효과적으로 제거하여 매끄러운 훈련 곡선을 생성한다.
• 메모리 효율성: 활성 파라미터에 대해서만 상태를 저장함으로써, SMET은 그래디언트와 옵티마이저 상태의 메모리 점유율을 크게 줄인다. 1B 모델에서 밀도 $d=0.4$일 때 상당한 메모리 절감이 관찰된다.
• 확장성: 모델 크기가 커질수록 희소 훈련과 밀집 훈련 사이의 성능 격차가 좁아지며, 이는 희소 훈련이 대규모 모델에 특히 유효할 수 있음을 시사한다.
• 일반화: SMET의 이점은 서로 다른 재성장 전략(무작위 및 그래디언트 기반)과 서로 다른 모델 아키텍처(LLaMA 및 Qwen) 전반에 걸쳐 일반화된다.

의의 및 주장
본 논문은 SMET이 LLM을 위한 실질적인 밀집 훈련의 대안으로서 희소 훈련의 길을 열어준다고 주장한다. 동적 희소 훈련에 내재된 최적화 불안정성을 해결하고 희소 상태 저장을 통해 실제 메모리 절감을 실현함으로써, SMET은 안정적이고 확장 가능하며 메모리 효율적인 사전 훈련을 가능하게 한다. 저자들은 SMET이 훈련 안정성을 개선하고 옵티마이저 메모리를 줄이지만, 추론 속도를 완전히 실현하기 위해서는 활성화 메모리 문제와 미세한(fine-grained) 희소성에 대한 효율적인 하드웨어 지원이 여전히 과제로 남아 있다고 언급한다. 그럼에도 불구하고, 본 연구는 적절한 안정화와 메모리 관리가 동반된다면 동적 희소성이 파운데이션 모델로 성공적으로 확장될 수 있음을 보여준다.