3140억 파라미터(Parameter)의 대규모 AI 모델

Grok-1은 3140억 개의 파라미터를 가진 대규모 언어 모델(LLM)로, MoE(Mixture of Experts) 아키텍처를 사용한다. 특히, 각 토큰(Token)은 8개의 전문가(Expert) 중 2개를 활용하여 연산량을 줄이는 구조이다. RoPE(Rotary Positional Embedding)를 사용하며, 활성화 샤딩(Activation Sharding) 및 8비트 양자화(Quantization)를 지원하여 메모리 사용량을 최적화한다. 최대 시퀀스 길이(Maximum Sequence Length)는 8,192 토큰이다.

본 레포지토리는 Grok-1 모델을 JAX 프레임워크(Framework)를 사용하여 구현한 예제 코드를 제공한다. JAX는 자동 미분(Automatic Differentiation) 기능을 제공하여 모델 훈련 및 추론을 용이하게 한다. README에 따르면, MoE 레이어(Layer) 구현은 효율성을 최적화하기보다는 모델의 정확성 검증에 초점을 맞추었다고 설명한다. 이는 커스텀 커널(Custom Kernel) 없이 모델의 동작을 검증하기 위한 설계로 보인다.

모델 가중치는 토렌트(Torrent) 또는 허깅 페이스(Hugging Face) 허브(Hub)를 통해 다운로드할 수 있다. 허깅 페이스(Hugging Face)를 이용하는 경우, `huggingface-cli download` 명령어를 사용하여 가중치를 다운로드할 수 있다. 다운로드 받은 가중치는 `ckpt-0` 디렉토리에 위치해야 하며, 예제 코드는 이 디렉토리에서 모델을 로드한다. 대용량 모델(Large Model)이므로 충분한 GPU 메모리(GPU Memory)를 가진 환경에서 실행해야 한다.

MoE 아키텍처는 모델의 파라미터 수를 늘리면서도 계산 비용을 효율적으로 관리할 수 있는 장점이 있다. 각 토큰에 대해 모든 전문가를 사용하는 대신, 일부 전문가만 선택적으로 사용함으로써 계산량을 줄인다. 하지만, 라우팅(Routing) 메커니즘 구현의 복잡성, 전문가 간의 불균형(Imbalance), 그리고 통신 오버헤드(Communication Overhead) 등의 단점도 존재한다. 본 레포지토리의 구현은 이러한 트레이드오프(Trade-offs)를 고려하여 설계되었다고 볼 수 있다.