Character.ai가 공개하는 대규모 트랜스포머 학습 효율화 기술 5가지

핵심 요약

Character.ai는 오픈소스 모델 기반으로 전환하기 전, 대규모 트랜스포머 학습의 속도와 효율성을 높이기 위해 다양한 독자적 기술을 개발했다. 노암 샤지어(Noam Shazeer)가 주도한 이 연구는 제한된 네트워크 대역폭 환경에서 6비트 그래디언트 압축(Squinch)을 통해 성능 저하 없이 학습을 가능하게 했다. 또한 정밀도 정규화, 동적 클램핑, 가시성 마스크, Gumbel Softmax 기반 증류 최적화 등을 통해 학습 안정성과 효율성을 극대화했다. 현재는 대규모 사전 학습 대신 오픈소스 모델의 사후 학습(Post-training) 및 강화학습(RL)에 이러한 최적화 역량을 집중하고 있다.

배경

트랜스포머(Transformer) 아키텍처, 분산 학습(Distributed Training) 원리, 양자화(Quantization) 및 지식 증류(Knowledge Distillation) 개념

대상 독자

대규모 언어 모델 학습 인프라 및 최적화를 담당하는 ML 엔지니어

의미 / 영향

Character.ai가 공유한 기술들은 하드웨어 자원이 제한된 환경에서도 고성능 모델을 학습시킬 수 있는 실질적인 해법을 제시한다. 특히 그래디언트 압축과 효율적인 어텐션 마스킹 기법은 분산 학습 효율을 극대화하려는 업계의 노력에 중요한 참고 자료가 된다.

섹션별 상세

Squinch는 노암 샤지어가 발명한 6비트 그래디언트 압축 알고리즘으로, 그래디언트를 블록 단위로 양자화하여 통신 대역폭을 획기적으로 줄인다. 당시 최신 시스템의 4분의 1 수준 대역폭에서도 bfloat16 수준의 정확도를 유지하며 분산 학습을 수행할 수 있게 했다. 8개의 그래디언트 값을 48비트 표현으로 인코딩하며, 트랜스포머 그래디언트의 특성에 맞춰 동적 범위를 튜닝하여 통신 비용을 최소화한다. 네트워크 대역폭이 병목인 환경이나 희소 MoE 모델 학습에서 특히 유용하다.

Attention Z-Reg는 어텐션 로짓(Logits)의 수치 범위를 제어하여 bfloat16 표현의 고정밀 구간을 활용하도록 돕는 정규화 기술이다. 로짓의 합인 'Z' 값이 0에 가깝게 유지되도록 시프트하여, 값이 커질수록 정밀도가 떨어지는 bfloat16의 한계를 극복한다. 이는 손실 함수가 아니라 최적화 과정에서 그래디언트에 직접 추가되는 방식으로 구현되어 학습 안정성을 높인다. 수치적 해상도가 높은 구간에서 연산이 이루어지므로 모델의 최종 성능 유지에 기여한다.

Dynamic Clamping은 양자화 인식 학습(QAT) 과정에서 작은 활성화 값이 0으로 붕괴되는 현상을 방지하기 위한 기술이다. 고정된 클램핑 범위를 사용하는 대신, 가중치의 RMS(Root Mean Square) 값에 따라 클램핑 범위를 실시간으로 계산한다. 이를 통해 FFN 네트워크의 양자화 오차를 크게 줄이고 모델의 정확도와 학습 안정성을 개선한다. 가중치가 작아질 때 활성화 값이 0 근처에 몰려 발생하는 정보 손실 문제를 효과적으로 해결한다.

Visibility Mask는 배치 내 아이템 간의 상호 관계를 두 개의 텐서로 압축하여 표현하는 효율적인 어텐션 API이다. 이 방식은 트리 구조의 문서 관계나 채팅 데이터를 네이티브하게 표현할 수 있으며, 여러 독립적인 청크를 하나로 묶어 학습 시스템의 효율을 높인다. 인과적 마스킹뿐만 아니라 양방향 어텐션과 빔 서치 추론 등 다양한 구조를 유연하게 지원한다. 특히 복잡한 컨텍스트 관계를 가진 대화형 AI 시스템 구축에 최적화된 구조를 제공한다.

Gumbel Softmax를 활용한 증류 최적화는 교사 모델의 방대한 출력 확률 분포를 저장하는 비용을 줄이기 위해 도입됐다. 전체 어휘 사전을 저장하는 대신 Gumbel Top-k 샘플링을 통해 일부만 추출하되, 기대값을 보존하여 학생 모델의 편향을 방지한다. 이 알고리즘은 오프라인 증류 과정에서 저장 공간과 대역폭 비용을 실질적으로 절감하면서도 교사 모델의 분포를 정확히 반영한다. 대규모 어휘 사전을 가진 모델의 지식 증류 시 발생하는 데이터 관리 문제를 해결하는 실용적인 해법이다.

</> 코드 예제 포함

실무 Takeaway

네트워크 대역폭이 제한된 환경에서는 Squinch와 같은 블록 단위 6비트 그래디언트 압축을 통해 통신 병목을 해결하고 학습 속도를 높일 수 있다.
bfloat16의 수치적 특성을 고려하여 로짓 범위를 0 근처로 유지하는 Attention Z-Reg를 적용하면 대규모 모델 학습의 수치적 안정성을 확보할 수 있다.
양자화 학습 시 가중치 스케일에 맞춘 Dynamic Clamping을 적용하여 활성화 값의 소실을 막고 양자화 모델의 성능 저하를 최소화해야 한다.

언급된 리소스

GitHubpipelining-sft

GitHubOvi

논문Categorical Reparameterization with Gumbel-Softmax