기술 요약: GNMR – 저정밀도 대규모 언어 모델 학습을 위한 런타임 안정성 제어

1. 문제 정의

저정밀도 실행 경로(예: FP8, 양자화된 활성화 함수, 혼합 정밀도 레시피)를 사용하여 대규모 언어 모델(LLM)을 학습시키는 것은 계산, 메모리 및 통신 비용을 줄이는 데 필수적입니다. 그러나 이러한 경로는 수치적 불안정성 위험을 초래하며, 이는 모델 전체나 학습 타임라인에 걸쳐 균일하게 나타나지 않습니다. 대신, 저정밀도 섭동(perturbation)은 특정 연산자나 레이어에서 짧은 시간 동안 발생하는 일시적이고 국소적인 불안정성으로 나타나는 경우가 많습니다.

그래디언트 클리핑(gradient clipping), 전역 그래디언트 노이즈 분석, 또는 정적 정밀도 레시피와 같은 기존의 안정성 접근 방식은 이 문제의 런타임 제어(runtime control) 측면을 해결하는 데 한계가 있습니다. 이러한 방식들은 특정 불안정 유닛을 식별할 만큼 세밀하지 못하거나, 비용이 많이 드는 전역 통계에 의존하거나, 과도한 오버헤드 없이 비정상 상태(non-stationary)의 위험에 동적으로 적응할 수 없습니다. 핵심 과제는 다음을 수행할 수 있는 컨트롤러를 설계하는 것입니다:

1. 특정 연산자의 국소적, 일시적 불안정성 감지.
2. 필요할 때만 복구 동작(고충실도 경로로 전환) 트리거.
3. 저정밀도 학습의 이점을 상쇄하지 않도록 엄격한 오버헤드 예산(메모리 및 계산) 하에서 작동.
4. 특정 커널 구현이나 하드웨어 레시피로부터 제어 로직을 분리하여 **백엔드에 무관(backend-agnostic)**하게 유지.

2. 방법론: GNMR 컨트롤러

저자들은 경량 런타임 안정성 컨트롤러인 **GNMR (Gradient Norm-to-Mean Ratio)**을 제안합니다. GNMR은 연산자별 그래디언트 통계량을 모니터링하고, 위험도가 높은 유닛을 저비용 실행 경로에서 고충실도 복구 경로로 라우팅하는 방식으로 작동합니다.

2.1 위험 모니터링 (이중 타임스케일)

GNMR은 전체 프로세스 통계량이나 비용이 많이 드는 역전파(backward pass) 없이도 불안정성을 감지하기 위해 두 가지 이력 정규화된 위험 점수를 활용합니다:

• GNMR (장기 관점): 유닛의 현재 그래디언트 노름($\|g_t\|$)을 누적된 역사적 평균과 비교합니다. 이는 유통되는 일반적인 행동에서 벗어나는 유닛을 식별합니다.
  $$\text{GNMR}_t(\theta) := \frac{\|g_t\|}{\frac{1}{t-1}\sum_{\tau=1}^{t-1} \|g_\tau\|}$$
• $\Delta$-GNMR (단기 윈도우): 현재 GNMR 값을 최근 윈도우 평균 GNMR 값과 비교합니다. 이는 위험의 급격한 단기 스파이크를 감지합니다.
  $$\Delta\text{-GNMR}_t(\theta) := \frac{\text{GNMR}_t(\theta)}{\frac{1}{\delta}\sum_{\tau=1}^{\delta} \text{GNMR}_{t-\tau}(\theta)}$$

이 점수들은 무차원(dimensionless)이며 서로 다른 이질적인 모듈 간에도 비교 가능하므로, 컨트롤러가 고유한 스케일에 관계없이 유닛의 위험 순위를 매길 수 있게 합니다.

2.2 예산 기반 복구 정책

컨트롤러는 위험 점수를 제한된 동작으로 변환합니다:

• 후보 선택: 위험 점수가 임계값($\alpha_t, \beta_t$)을 초과하거나 현재 "잠금(locked)" 상태인 유닛들이 후보 풀을 형성합니다.
• **활성 예산 ($maxO$):** 컨트롤러는 후보 풀에서 가장 위험도가 높은 상위$maxO$개의 유닛을 선택하여 복구 경로(더 높은 정밀도)에서 실행합니다. 그 외의 복구 가능한 유닛들은 저비용 경로에 머뭅니다.
• 잠금 메커니즘: 선택된 유닛에는 잠금 카운터($T_{lock}$)가 할당됩니다. 유닛이 계속 위험한 상태라면 후보 풀에 유지되며, 안정화되면 잠금 카운터가 감소합니다. 이는 빈번한 전환(thrashing)을 방나지 않으면서도 단계당 활성 복구 유닛의 하드 캡(hard cap)을 준수합니다.

2.3 이론적 배경

본 논문은 예산 기반 복구가 학습 목적 함수의 1단계 매끄러움 경계(one-step smoothness bound)를 줄여준다는 경계 수준의 분석을 제공합니다. 이론적으로 고정된 예산 하에서의 최적 전략은 가장 큰 "복구 격차(recovery gaps)"(저비용 경로와 고충실도 경로 사이의 섭동 페널티 차이)를 가진 유닛을 복구하는 것입니다. GNMR과 $\Delta$-GNMR은 온라인에서 이 오라클 선택(oracle selection)을 근사하기 위한 경량화된 프록시 역할을 합니다.

3. 주요 기여

본 논문은 네 가지 주요 기여를 설명합니다:

1. 정식화: 저정밀도 안정성을 제약된 온라인 제어 문제로 프레임화하여, 위험 식별 및 복구 결정을 백엔드 구현 세부 사항으로부터 분리했습니다.
2. 이중 타임스케일 모니터: GNMR 및 $\Delta$-GNMR의 설계를 통해 지속적인 편차와 급격한 변화를 모두 포착하는 실시간, 무차원, 국소적 위험 신호를 제공합니다.
3. 예산 기반 정책: 복구 오버헤드에 대한 하드 캡($maxO$)을 강제하고 잠금 구간을 사용하여 활성 집합을 안정화함으로써 저비용 실행을 보장하는 폐루프 제어 정책을 제시합니다.
4. 컨트롤러 수준의 증거: 세 가지 서로 다른 설정(활성화 양자화 스트레스, DeepSeek 스타일의 혼합 정밀도 사전 학습, LLaMA-2 13B 미세 조정)에 대한 종합적인 평가를 통해, 새로운 커널을 도입하거나 기본 학습 레시피를 변경하지 않고도 GNMR이 안정성을 개선함을 입증했습니다.

4. 실험 결과

저자들은 세 가지 시나리오에서 GNMR을 평가했습니다:

• 활성화 양자화 스트레스 (LLaMA-2 사전 학습):

  • 설정: 프로젝션 연산자의 저장된 활성화를 4-bit 또는 8-bit로 양자화함.
  • 결과: 정적 저비용 경로(4-bit)는 수렴에 실패하거나 높은 퍼플렉시티(perplexity)를 생성했습니다. GNMR이 제어하는 4-bit/8-bit 복구는 다양한 모델 크기(60M ~ 13B)에 대해 고정된 8-bit 및 BF16 베이스라인과 유사한 퍼플렉시티를 달성했습니다.
  • 희소성: 복구 신호는 매우 희소했으며, 임계값 초과 발생률은 대상 블록의 0.2% 미만이었습니다. 이는 복구가 일시적인 위험에 대해서만 필요함을 확인시켜 줍니다.

• DeepSeek 스타일의 혼합 정밀도 사전 학습:

  • 설정: DeepSeek-V3 스타일의 정밀도 계층(6-bit/8-bit 동적)을 사용하는 LLaMA-2 모델에서 평가.
  • 결과: GNMR이 제어하는 동적 정밀도는 학습 전반에 걸쳐 고정된 고정밀 참조 모델(8-bit 및 16-bit)의 성능에 근접했습니다.
  • 오버헤드: 단계별 시간 및 처리량(throughput)은 저비용 경로와 거의 동일하게 유지되어, 엔지니어링 오버헤드가 미미함을 보여주었습니다.

• LLaMA-2 13B 미세 조정 스트레스:

  • 설정: LoRA 어댑터를 사용한 미세 조정 시 고정된 INT8 vs. GNMR 제어형 INT8/BF16 비교.
  • 결과: GNMR 방식은 저비용 실행을 유지하면서도 MMLU에서 BF16과 대등한 성능을 보였고, HellaSwag에서는 이를 능가하는 고충실도 다운스트림 성능을 달 achievement 했습니다.

5. 의의 및 주장

본 논문은 GNMR을 저정밀도 학습의 "런타임 제어" 공백을 메우는 백엔드에 무관한 컨트롤러로 포지셔닝합니다. GNMR의 의의는 다음과 같습니다:

• 저비용 실행 보존: 엄격한 예산 하에서 복구 대상을 희소한 부분 집합으로 제한함으로써, GNMR은 전체 고정밀도 학습의 메모리 및 계산 페널티를 피합니다.
• 제어와 구현의 분리: 컨트롤러는 새로운 수치 형식, 커널 또는 옵티마이저 변경 없이 기존의 저비용 및 복구 경로 위에서 작동합니다.
• 비정상 상태 대응: 정적 안전장치와 달리, GNMR은 학습 위험의 비정상적 특성에 적응하여 정적 레시피가 놓칠 수 있는 일시적 불안정성을 식별하고 복구합니다.

저자들은 GNMR이 국소적인 수치적 위험을 동적으로 감지하고 복구하는 메커니즘을 제공함으로써, 저비용 실행의 효율성을 희생하지 않으면서도 공격적인 저정밀도 학습 전략의 생존 가능성을 지원한다고 결론짓습니다.