Yuqi Pan Yupeng Feng Jinghao Zhuang Siyu Ding et al

초록

주류의 Transformer 기반 대규모 언어 모델은 효율성 측면에서 중대한 한계에 직면해 있다. 학습 시 계산량은 시퀀스 길이의 제곱에 비례하여 증가하고, 추론 시 메모리 사용량은 선형적으로 증가함으로써 긴 컨텍스트 처리가 제한된다. NVIDIA 플랫폼 외부에서 대규모 모델을 구축하는 것도 안정적이고 효율적인 학습을 위한 도전 과제로 남아 있다. 이를 해결하기 위해 우리는 효율적인 긴 컨텍스트 학습과 추론을 위한 뇌에 영감을 받은 모델 세트인 SpikingBrain을 제안한다. SpikingBrain는 MetaX GPU 클러스터를 활용하며, 세 가지 핵심 방향에 초점을 맞춘다. (1) 모델 아키텍처: 적응형 스파이크 신경세포를 활용한 선형 및 하이브리드-선형 어텐션 아키텍처; (2) 알고리즘 최적화: 효율적인 변환 기반 학습 파이프라인과 전용 스파이크 인코딩 프레임워크; (3) 시스템 공학: MetaX 하드웨어에 맞춤화된 학습 프레임워크, 오퍼레이터 라이브러리 및 병렬 처리 전략.이러한 기술들을 활용해, 우리는 두 가지 모델을 개발했다. 하나는 선형 아키텍처를 채택한 SpikingBrain-7B, 다른 하나는 하이브리드-선형 MoE 아키텍처를 사용한 SpikingBrain-76B이다. 이 모델들은 NVIDIA 플랫폼 외부에서도 대규모 언어 모델 개발의 가능성을 입증한다. SpikingBrain는 오픈소스 Transformer 기반 기준 모델과 비슷한 성능을 달성하면서도, 지속적 사전학습에 약 1500억 토큰만을 사용한다. 제안된 모델은 긴 시퀀스 학습 효율성을 크게 향상시키며, (부분적으로) 일정한 메모리 사용과 이벤트 기반 스파이크 동작을 제공한다. 예를 들어, SpikingBrain-7B는 400만 토큰 시퀀스에 대해 첫 토큰 생성 시간(Time to First Token)에서 100배 이상의 속도 향상을 달성한다. 수백 대의 MetaX C550 GPU에서 수주간 안정적인 학습이 가능하며, 7B 모델은 모델 FLOPs 활용률(23.4%)을 달성한다. 제안된 스파이크 방식은 69.15%의 희소성(sparcity)을 달성하여 저전력 작동을 가능하게 한다. 종합적으로 본 연구는 뇌에 영감을 받은 메커니즘이 다음 세대의 효율적이고 확장 가능한 대규모 모델 설계를 이끌 수 있음을 보여준다.