전이 학습 FLOPs의 1% 미만으로 Llama-3.2-3B와 유사한 성능을 내는 2.3B MoE (360M 활성) 하이브리드 Mamba-2

우리는 Llama-3.2-3B와 몇 점 차이로 근접하는, 2.3B MoE (360M 활성) 하이브리드 Mamba-2를 그 사전 학습 FLOPs의 <1%만 사용하여 사전 학습했습니다.

전용 클러스터는 없었습니다. 이 작업은 단일 코드베이스에서 H100, A100, V100 (네, V100도), 그리고 TPU v5p/v6e 사이를 이동하며 진행되었습니다.

Rigel을 소개합니다.

수치입니다. lm-eval-harness로 제로샷(Zero-shot) 측정했으며, 모두에게 동일한 설정이 적용됩니다.
Rigel은 훨씬 더 많은 활성 파라미터를 가진 모델에 비해 훨씬 적은 학습 FLOPs만으로 근접한 성능을 보여줍니다.

Rigel은 3.5T 토큰으로 학습되었습니다.

Rigel 아키텍처:

• 하이브리드 Mamba-2 MoE (활성 360M / 총 2.3B)
• 40개 레이어: 30개의 Mamba-2 + 10개의 GQA (3:1), SWA (4,096 윈도우) + GQA + 모든 어텐션 레이어의 XSA
• 모든 FFN은 세밀한 MoE 구조: 128개 전문가(experts), 상위 2개 활성
• NoPE, 연결된 임베딩(tied embeddings), 하이퍼파라미터 전송을 위한 μP

여섯 가지 혼합 단계(mixture phases)를 사용하여 사전 학습되었습니다:

우리는 클러스터가 없었기 때문에 규칙은 '가지고 있는 컴퓨팅 자원으로 작업한다'였습니다.

작업의 대부분은 TPU v6e와 A100에서 진행되었고, 일부는 BF16을 지원하지 않는 32GB V100이었습니다 (FP16 + 손실 스케일링). 우리는 V100용 MoE 커널과 TPU용 Pallas Mamba-2 커널을 작성하고 torch-xla를 실행하여 GPU와 TPU 모두에 적용했습니다.