DeepSeek V3 MoE 구조 리뷰
Apr 1, 2025
참고
DeepSeek V3 계열의 모델링 코드는 다음과 같다.
허깅페이스 플랫폼을 이용하는 모델링 코드는 3가지 버전이 모두 동일하다. 따라서 모델의 구조는 완전히 동일하다고 할 수 있다.
V3 공개시 Triton Kernel을 이용한 별도의 최적화된 인퍼런스 코드를 제공한 바 있다.
구조
설명
허깅페이스 구현을 살펴보면 DeepseekV3ForCausalLM > DeepseekV3Model > DeepseekV3DecoderLayer 구조로 되어 있으며, 기본적으로 Attention 후에 MLP(4번째부터는 MoE) 레이어 통과를 61회 반복하는 구조로 되어 있다.
DeepSeekMoE uses finer-grained experts and isolates some experts as shared ones.2
expert 1개의 크기는 0.044B이며, 256개 중 8개를 Topk 형태로 sigmoid 가중치만큼 합산하는 방식이다. 또한 동일한 크기의 Shared Expert 1개가 항상 함께 합산된다. 수식에는 \(N_s\)만큼으로 되어 있으나 코드 구현에는 1로 고정되어 있으며, 다이어그램에도 여러 개로 구성가능한 것으로 그려져 있지만 실제로는 1이다. 따라서 위 다이어그램에는 해당 부분에 빨간색으로 x 표시를 했다. expert의 수는 기존 V2보다 훨씬 더 늘어난 수치이며, Mixtral이나 Grok 같은 모델보다 월등히 더 촘촘 (finer-grained experts)하다.
Also for better representation of the input data, v3 increases the all-experts-activated layer from 1 to 3.1
이 설명에 따르면(논문에서는 이 내용을 찾을 수 없음) 표현력을 향상시키기 위해(better representation) 앞에 3개 레이어는 모두 활성화한다고 되어 있다. 그러면서 다음과 같이 257개가 모두 activated되는 것으로 계산했는데, 이 부분은 틀렸다. 실제로는 정확히 9배 더 큰 MLP이며, 따라서 원래 activated되어야 할 9개(8개+1개) 크기로 구성됐다고 보는게 맞다.
따라서 이 계산식은 틀렸다. 내가 계산한 바로는 \((61-3) \times 9 + 3 \times 9 = 549\) activated experts이며, expert 1개의 크기는 29.36M라고 했는데, 이 또한 내가 계산한 바로는 44M이 맞다. 계산해 보면, \(44M \times 549 = 24.15B\)가 나온다. 위에서는 37.96B라고 계산했는데, 원래 전체 크기가 37B인데 어텐션을 제외한 크기만 이미 37B를 넘어섰다는 점에서 계산이 완전히 잘못됐음을 알 수 있다. 내가 계산한 결과는 24.15B이며 여기에 어텐션과 기타 나머지를 포함하면 전체 37B가 될 것이다. 전체 계산 결과는 다음과 같다.
먼저 어텐션 레이어 1개의 파라미터는 0.187B이다.
- 어텐션 전체: \(187,107,328 \approx 0.187B\)
- 어텐션 Q 관련
- q_a_proj: \(7,168 \times 1,536 = 11,010,048\)
- q_a_layernorm: \(1,536\)
- q_b_proj: \(1,536 × (128 × 192) = 37,748,736\)
- 어텐션 KV 관련
- kv_a_proj_with_mqa: \(7,168 \times (512 + 64) = 4,128,768\)
- kv_a_layernorm: \(512\)
- kv_b_proj: \(512 \times (128 \times (192 - 64 + 128)) = 16,777,216\)
- 어텐션 O 관련
- o_proj: \((128 \times 128) \times 7,168 = 117,440,512\)
- 어텐션 Q 관련
MoE 구조에서 위 다이어그램에 계산해서 표시한 것처럼 1 expert가 0.044B이며, 3개 MLP 및 Activated Experts 전체는 24.15B이다.
어텐션과 MLP/MoE 앞에 진행되는 RMSNorm, 마지막 RMSNorm은 미미하므로 굳이 전체 계산에 포함하지 않아도 무방하다.
DeepseekV3DecoderLayer
의 RMSNorm: \(2 × 7,168 × 61 = 874,496\)DeepseekV3Model
의 RMSNorm: \(7,168\)
기타 임베딩과 언어 모델링 헤드(LM Head), MoE 게이트의 파라미터는 1.97B이다.
- 임베딩: \(129,280 \times 7,168 = 926,851,072 \approx 0.93B\)
- 언어 모델링 헤드: \(7,168 \times 129,280 = 926,851,072 \approx 0.93B\)
- MoE 게이트: \(58 \times 1,835,008 = 106,430,464 \approx 0.11B\)
이제 전체를 계산하면 다음과 같다.
- 모든 어텐션: 0.187B x 61개 레이어 = 11.4B
- 3개 MLP 및 나머지 모든 Activated Experts: 24.15B
- 기타: 1.97B
- 총합: 11.4B + 24.15B + 1.97B = 37.51B
논문2에서 제시한 37B와 정확하게 맞아 떨어진다.
기타
Preventing routing collapse
For MoE models, an unbalanced expert load will lead to routing collapse
MoE 모델의 쏠림을 방지하기 위해 학습시 bias를 포함하는 Auxiliary-Loss-Free Load Balancing을 사용한다.
DeepSeek-V3 keeps balanced expert load during training, and achieves better performance than models that encourage load balance through pure auxiliary losses.
이미 DeepSeek V2부터 routing collapse를 방지하기 위해 많은 노력을 한 흔적이 엿보인다.
Training hours
Decoding
최소 배포단위가 320개 GPU라고 하는데 256개 experts를 각 GPU에 분산하고, 나머지 58개 shared experts, 3개 MLP 등을 적절히 나머지 GPU에 분산한 거 같다.