핵심 요약
대규모 언어 모델(LLM) 사전 학습 시 발생하는 대역폭 제한과 수치적 불안정성은 효율적인 스케일업을 방해하는 주요 요소다. Character.AI는 노암 샤지어(Noam Shazeer)의 주도하에 Squinch라는 6비트 그래디언트 압축 알고리즘을 포함한 5가지 핵심 최적화 기법을 개발했다. 이 기법들은 통신 비용 절감, 양자화 오류 방지, 증류 데이터 저장 효율화 등을 통해 제한된 자원에서도 고성능 모델을 학습할 수 있게 한다. 현재 회사는 오픈소스 모델 기반의 사후 학습(Post-training)으로 전환했으나 공개된 기술들은 여전히 효율적인 AI 시스템 구축에 중요한 시사점을 제공한다.
배경
트랜스포머 아키텍처, 분산 학습(Distributed Training), 양자화(Quantization), 모델 증류(Distillation)
대상 독자
대규모 언어 모델 사전 학습 및 최적화 인프라를 설계하는 엔지니어 및 연구원
의미 / 영향
이 기술들은 하드웨어 자원이 제한된 환경에서도 대규모 모델을 효율적으로 학습할 수 있는 실질적인 방법론을 제시한다. 특히 노암 샤지어의 초기 아이디어들이 실제 프로덕션 수준에서 어떻게 구현되었는지 보여줌으로써 효율적인 AI 시스템 설계를 위한 중요한 기술적 자산이 된다.
섹션별 상세
Squinch는 통신 대역폭이 부족한 환경에서 분산 학습 효율을 높이기 위해 설계된 6비트 그래디언트 압축 알고리즘이다. 각 블록은 8개의 그래디언트 값을 48비트로 압축하며 부호와 크기를 정밀하게 캡처하여 bfloat16 수준의 정확도를 유지한다. 트랜스포머 그래디언트의 특성에 맞춰 동적 범위를 튜닝함으로써 통신 비용을 획기적으로 낮추면서도 학습 충실도 손실을 최소화했다. 이는 네트워크 대역폭이 병목인 크로스 도메인 학습이나 희소 MoE 모델 환경에서 특히 유용하다.
Attention Z-Reg는 어텐션 로짓(Attention Logits)의 수치 범위를 제어하여 bfloat16의 고정밀 구간을 최대한 활용하는 정규화 기술이다. 로짓의 합인 Z값이 0에 가깝게 유지되도록 시프트하여 값이 커질수록 정밀도가 떨어지는 bfloat16의 수치적 한계를 극복한다. 별도의 손실 함수가 아니라 최적화 과정에서 그래디언트에 직접 추가되는 방식으로 구현되어 연산 효율성을 확보했다. 이 기법은 어텐션 로짓뿐만 아니라 선형 모델 로짓에도 적용되어 전반적인 학습 안정성을 높인다.
Dynamic Clamping은 양자화 인식 학습(QAT) 과정에서 활성화 값이 0으로 수렴하여 정보가 손실되는 현상을 방지한다. 가중치의 RMS(Root Mean Square) 값에 따라 클램핑 범위를 동적으로 조정하여 고정된 임계값을 사용할 때 발생하는 양자화 오류를 줄인다. 특히 FFN(Feed-Forward Network)의 업프로젝션 상태에서 입력값의 스케일에 맞춰 범위를 설정함으로써 학습 안정성과 모델 정확도를 동시에 개선했다. 이는 가중치가 매우 작아질 때 발생하는 수치적 클러터링 문제를 효과적으로 해결한다.
Visibility Mask는 배치 내 아이템 간의 관계를 visibility_start와 visibility_limit이라는 두 개의 텐서로 표현하여 복잡한 어텐션 범위를 관리하는 API다. 트리 구조의 문서 관계나 채팅 데이터를 네이티브하게 표현할 수 있으며 여러 독립적인 청크를 하나로 묶어 학습 시스템의 효율을 높인다. 양방향 어텐션 지원은 물론 추론 시 빔 서치(Beam Search)나 페이지드 어텐션(Paged Attention) 환경에서도 유연하게 작동한다. 이를 통해 데이터 패킹 효율을 극대화하고 다양한 어텐션 패턴을 단일 구조로 처리한다.
Gumbel Softmax를 활용한 증류 최적화는 교사 모델의 방대한 출력 확률 분포를 저장하는 비용을 줄이기 위해 Gumbel Top-k 샘플링 기반의 서브샘플링 기법을 사용한다. 기대값을 보존하면서 샘플링된 토큰의 가중치를 재계산하여 학생 모델이 편향되지 않도록 설계했다. 오프라인 증류 시 필요한 저장 공간과 대역폭을 대폭 절감하면서도 교사 모델의 확률 분포를 정확히 모방할 수 있다. 이 알고리즘은 대규모 어휘집을 가진 모델의 증류 학습 시 발생하는 스토리지 병목 문제를 해결하는 데 핵심적인 역할을 한다.
</> 코드 예제 포함
실무 Takeaway
- 분산 학습 시 네트워크 대역폭이 병목이라면 Squinch와 같은 블록 단위 6비트 압축을 통해 bfloat16 대비 성능 저하 없이 통신 비용을 75% 이상 절감할 수 있다.
- bfloat16 사용 시 수치적 정밀도 유지를 위해 로짓의 범위를 0 근처로 강제하는 Z-Reg 기법을 적용하여 대규모 모델의 학습 안정성을 확보해야 한다.
- 양자화 인식 학습(QAT)에서 고정된 클램핑 임계값은 정보 손실을 유발하므로 가중치 스케일에 연동된 동적 클램핑(Dynamic Clamping) 도입이 필수적이다.
언급된 리소스
AI 분석 전체 내용 보기
AI 요약 · 북마크 · 개인 피드 설정 — 무료