DeepSeek V3, 순수 JAX 190줄로 구현
원제: Show HN: DeepSeek V3 in ~200 lines of pure Jax
왜 중요한가
최신 대형 언어모델 아키텍처를 최소한의 코드로 재현한 사례로, 연구자와 개발자의 DeepSeek V3 내부 구조 학습 진입 장벽을 크게 낮춘다.
개발자 Vaibhaav Tiwari가 DeepSeek V3의 핵심 아키텍처를 190줄의 순수 JAX 코드로 구현해 GitHub에 공개했다. MLA, DeepSeekMoE, aux-loss-free balancing, MTP를 모두 포함하며, MacBook Air에서 공개 도메인 탐정 소설 데이터셋으로 처음부터 사전학습을 진행했다.
Tiwari가 공개한 'deepseek-v3-jax' 저장소는 Flax, Equinox, PyTorch 등 고수준 프레임워크 없이 JAX만으로 DeepSeek V3 논문의 아키텍처를 재현한다. 모델 파라미터는 일반 딕셔너리 배열로 관리되며, 모든 구성 요소가 단일 파일 `deepseek_v3.py`에서 위에서 아래로 읽을 수 있는 함수 형태로 정리돼 있다.
구현 내용은 Multi-head Latent Attention(MLA), 분리형 RoPE, DeepSeekMoE, 보조 손실 없는 부하 분산, Multi-Token Prediction(MTP) 등을 포함한다. MLA는 토큰당 캐시 크기를 480개 숫자에서 144개로 줄여 메모리 효율을 높이는 핵심 기법이다. DeepSeekMoE는 항상 실행되는 공유 전문가 1개와 16개의 라우팅 전문가로 구성된다.
학습은 공개 도메인 탐정 소설 텍스트를 사용해 MacBook Air에서 처음부터 진행됐다. 저장소에는 모델 코드 외에도 데이터 준비, 학습, 샘플링, 테스트 스크립트가 함께 포함돼 있다. GitHub 별 142개, 포크 14개를 기록 중이며, 상세 설명은 작성자 블로그(vaibhaav.in)에 다이어그램과 함께 게재됐다.