Search

AI/ DeepSeek-V3

DeepSeek-V3

DeepSeek-V3은 제약된 하드웨어를 사용하여 저렴한 비용으로 최상위 모델과 견주는 성능(오픈소스 기준 최상위 성능)의 모델로 세상에 큰 인상을 남겼다. 이를 위해 여러 방법들이 사용되었지만, 특히 chip level 최적화까지 수행한 엔지니어링 노력이 저렴한 비용을 가능하게 했고, V3 모델에 이르기까지 쌓인 DeepSeek 연구원들의 연구 경험이 모델의 높은 성능을 달성할 수 있었으리라 생각 됨.
아래에는 DeepSeek-V3의 주요한 부분들을 요약한 내용이며, 더 상세한 내용은 논문(아래 링크)이나 기타 자료 참조.

Architecture

DeepSeek의 Transformer Block은 최근 경향을 따라 RMSNorm - Attention - RMSNorm - FFN으로 구성된다. 아래 그림 참조.
여기서 RMSNorm(Root Mean Square Normalization)은 이름 그대로 제곱 평균의 제곱근 RMS=1Ki=1Kxi2\text{RMS} = \sqrt{{1\over K} \sum_{i=1}^K x_i^2}을 이용하여 normalization을 수행하는 것으로, 선택적으로 scaling factor γ\gamma와 분모로 사용되는 RMS가 00이 되는 것을 막기 위한 작은 값 ϵ\epsilon을 추가하여 아래처럼 계산된다.
xi=γ(xi1Ki=1Kxi2+ϵ)x_i' = \gamma\cdot \left({x_i \over \sqrt{{1\over K}\sum_{i=1}^K x_i^2}+\epsilon}\right)
Attention과 FFN Layer에 대해서는 이하 설명 참조

Multi-Head Latent Attention

DeepSeek-V3는 일반적인 Multi-Head Attention 대신 Multi-Head Latent Attention이라는 별도의 방법을 사용하는데, 이것은 attention 연산 시에 메모리 사용량이 큰 KV 캐시를 줄이기 위해 key와 value를(과 query도) 저차원으로 압축하는 것이 핵심이다. 이것은 DeepSeek-V2에서 이미 적용 및 (철저한) 검증이 되었다고 함.
ctKV=WDKVht[kt,1C;kt,2C;...;kt,nhC]=ktC=WUKctKVktR=RoPE(WKRht)kt,i=[kt,iC;ktR][vt,1C;vt,2C;...;vt,nhC]=vtC=WUVctKV\begin{align} {\mathbf{c}_t^{KV}} &= W^{DKV}\mathbf{h}_t \\ [\mathbf{k}_{t,1}^C;\mathbf{k}_{t,2}^C;...;\mathbf{k}_{t,n_h}^C] &= \mathbf{k}_t^C = W^{UK}\mathbf{c}_t^{KV} \\ \mathbf{k}_t^{R} &= \text{RoPE}(W^{KR}\mathbf{h}_t) \\ \mathbf{k}_{t,i} &= [\mathbf{k}_{t,i}^C;\mathbf{k}_t^R] \\ [\mathbf{v}_{t,1}^C;\mathbf{v}_{t,2}^C;...;\mathbf{v}_{t,n_h}^C] &= \mathbf{v}_t^C = W^{UV}\mathbf{c}_t^{KV} \end{align}
여기서 attention block에 대한 tt-번째 입력 벡터 ht\mathbf{h}_t에 down-projection 가중치 행렬 WDKVW^{DKV}를 곱해 압축된 latent 벡터 ctKV\mathbf{c}_t^{KV}를 생성한다. 이렇게 생성된 ctKV\mathbf{c}_t^{KV}는 다시 attention 연산을 위해 원래로 되돌려야 하는데 이때 up-projection 가중치 행렬 WUKW^{UK}WUVW^{UV}를 이용해서 ctKV\mathbf{c}_t^{KV}를 각각 key에 대한 벡터 ktC\mathbf{k}_t^C(식 (2))와 value에 대한 벡터 vtC\mathbf{v}_t^{C}(식 (5))로 복구한다.
추가로 sequence 정보를 처리하기 위해서는 각 token에 position 정보가 반영되어야 하는데, 위와 같이 데이터를 latent로 압축하면 position 정보가 날아가는 문제가 발생한다. 이를 해결하기 위해 추가로 입력 ht\mathbf{h}_t에 대해 별도의 가중치 행렬 WKRW^{KR}을 이용을 곱한 다음 Rotary Positional Embedding(RoPE)(AI/ Rotary Positional Embedding (RoPE))를 적용하여 position 정보를 반영한 ktR\mathbf{k}_t^R를 구하고 이것을 앞서 복구한 ktC\mathbf{k}_t^C와 결합하여 최종적으로 position 정보가 반영된 kt\mathbf{k}_t를 구한다. (value는 key에 해당하는 값이므로 position 정보가 불필요하여 RoPE를 처리하지 않는다)
query는 KV 캐시로 저장되지 않기 때문에 저차원으로 압축하는 것이 의미가 없게 생각될 수 있지만 DeepSeek 연구원들은 학습 하는 동안 activation 메모리를 줄일 수 있기 때문에 query에 대해서도 유사한 과정을 수행한다.
ctQ=WDQht[qt,1C;qt,2C;...;qt,nhC]=qtC=WUQctQ[qt,1R;qt,2R;...;qt,nhR]=qtR=RoPE(WQRctQ)qt,i=[qt,iC;qt,iR]\begin{align} \mathbf{c}_t^Q &= W^{DQ}\mathbf{h}_t \\ [\mathbf{q}_{t,1}^C;\mathbf{q}_{t,2}^C;...;\mathbf{q}_{t,n_h}^C] &= \mathbf{q}_t^C = W^{UQ}\mathbf{c}_t^Q \\ [\mathbf{q}_{t,1}^R;\mathbf{q}_{t,2}^R;...;\mathbf{q}_{t,n_h}^R] &= \mathbf{q}_t^R =\text{RoPE}(W^{QR}\mathbf{c}_t^Q) \\ \mathbf{q}_{t,i} &= [\mathbf{q}_{t,i}^C;\mathbf{q}_{t,i}^R] \end{align}
key를 처리한 것과 유사하게 입력 ht\mathbf{h}_t을 latent vector로 압축시키는 down-projection 가중치 행렬 WDQW^{DQ}를 이용하고 다시 up-projection 가중치 행렬 WUQW^{UQ}를 이용해 qtC\mathbf{q}_t^C를 복구한다. 또한 key와 마찬가지로 token에 position 정보가 필요하므로 가중치 행렬 WQRW^{QR}을 압축된 입력 ctQ\mathbf{c}_t^Q(입력을 사용한 key와 달리 query에서는 압축된 벡터를 사용하는 것에 주의)에 곱한 다음 RoPE에 적용하여 position 정보를 반영한 qtR\mathbf{q}_t^R을 구한 다음 qtC\mathbf{q}_t^C와 결합하여 최종적으로 qt\mathbf{q}_t를 구한다.
위와 같이 복구된 query (qt,i\mathbf{q}_{t,i}), keys (kj,i\mathbf{k}_{j,i}), values(vj,iC\mathbf{v}_{j,i}^C)를 이용하여 출력 ut\mathbf{u}_{t}은 아래와 같이 일반적인 attention 연산과 동일하게 계산된다.
ot,i=j=1tSoftmaxj(qt,ikj,idh+dhR)vj,iCut=WO[ot,1;ot,2;...;ot,nh]\begin{align} \mathbf{o}_{t,i} &= \sum_{j=1}^t \text{Softmax}_j \left({\mathbf{q}_{t,i}^\top\mathbf{k}_{j,i} \over \sqrt{d_h + d_h^R}} \right)\mathbf{v}_{j,i'}^C \\ \mathbf{u}_t &= W^O[\mathbf{o}_{t,1};\mathbf{o}_{t,2};...;\mathbf{o}_{t,n_h}] \end{align}

Sample Code

실제 DeepSeek-V3의 git에 공개된 Multi-Head Latent Attention 구현의 forward 부분 코드. 원래 코드에는 일반적인 attention을 위한 부분도 존재하는데, 아래는 보기 쉽게 해당 부분을 삭제하고 latent attention을 수행하는 부분만 남김
실제 수식과 차이가 있는 부분도 있는데, 대표적으로 위의 수식 상에서는 마치 up-projection 된 query와 key에 대해 연산을 수행하는 것처럼 나오는데, 실제 코드 상에서는 거꾸로 query를 down-projection한 후에 연산하고 최종 단계에서 up-projection을 수행한다. 아마도 작은 크기에서 연산하는 것이 성능에 더 낫기 때문일 것으로 생각 됨.
추가로 논문 상에서는 positional embedding을 수행할 때 query에 대해서는 압축된 것을 사용하고, key는 입력된 것을 사용하는데, 실제 코드 구현상에서는 query가 입력 크기로 rope가 수행되고, key는 압축된 크기로 rope를 수행함
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]): bsz, seqlen, _ = x.size() end_pos = start_pos + seqlen q = self.wq_b(self.q_norm(self.wq_a(x))) q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim) q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) # kv의 positional embedding과 no-positional embedding과 연산하기 위해 분리 q_pe = apply_rotary_emb(q_pe, freqs_cis) # 분리된 q_pe에만 rope embedding. 즉 query는 down-norm-up 된 것에 대해 rotary embedding을 적용함. kv = self.wkv_a(x) # kv를 down-projection하는 W^{KV} kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) # down projection 후에 key에 대한 positional embedding을 위해 kv와 kv_pe로 split k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis) # 분리된 k_pe를 이용하여 rope embedding. 즉 key는 down-projected 된 것에 대해 rotary embedding을 적용함 # 논문과 코드의 구현이 반대로 되어 있다. 논문 상에서는 q_pe는 입력을 압축 시킨 것을 사용하고 k_pe는 입력을 그대로 사용하는데, 코드 상에서는 q_pe가 원본으로 돌린 후에 사용하고, k_pe가 압축된 것을 사용함 # kv를 up-projection하는 행렬의 설정 부분 wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size) wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank) # 논문 상으로는 wkv_b를 이용해서 kv를 up-projection한 후에 q와 곱해야 하는데, 실제 코드 상에서는 q_nope을 down-projection해서 down-projection된 kv와 곱한다. # 이것은 수학적으로 동등하고, 큰 것끼리 곱하는 것보다 작은 것끼리 곱하는 것이 성능상 유리하기 때문일 것으로 추정됨 q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim]) self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv) # latent로 압축된 kv를 정규화해서 kv_cache에 저장 self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2) # latent로 압축된 것 중 positional embedding이 된 것을 별도의 cache에 저장 # 최종적으로 softmax에 들어가기 전 score를 구하는 부분. # self.softmax_scale은 1 / \sqrt{d_k + d_h^R}에 해당한다. 그 외의 나머지 부분이 q, k, v의 내적에 해당한다. # 그 외의 나머지 부분이 q, k, v의 내적에 해당하며, 이때 positinoal embedding이 반영된 부분과 그렇지 않은 부분끼리 내적한 다음 최종 합산하여 전체 식을 구성한다. scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) + torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale if mask is not None: scores += mask.unsqueeze(1) scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x) # 일반적인 attention의 마지막에 softmax()를 통해 구한 확률 분포에 value를 곱하는 부분 x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos]) # 지금까지 latent 상에서 연산이 이루어졌으므로 최종적으로 up-projection해서 원래 크기로 돌린다. x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:]) # 원래 크기로 돌아온 multihead 값을 마지막으로 out 행렬에 통과시켜 하나의 값으로 합산하고 반환한다. x = self.wo(x.flatten(2)) return x
Python
복사

DeepSeek MoE with Auxiliary-Loss-Free Load Balancing

Basic Architecture

일반적으로 FFN 레이어는 연산 비용이 매우 크기 때문에 MoE 같은 방법을 통해 일부 파라미터만 활성화하도록 하여 연산 비용을 아끼는데, DeepSeek는 FFN 레이어에 아래와 같은 MoE 방식을 적용한다. 이를 DeepSeek MoE라고 명명함.
여기서는 일반적인 MoE 방식과 달리 routed expert에 추가로 항상 활성화되는 shared expert를 사용하고, routed expert 사이의 load balancing을 위한 auxiliary loss를 사용하지 않는다는 특징이 있다. (MLA와 마찬가지로 이 기법도 DeepSeek-V2에서 도입되고 검증되었음)
우선 아래의 식 (12)에서 FFNi(s)\text{FFN}_i^{(s)}는 항상 활성화 되는 shared expert로 보다 general한 feature를 학습하게 되고 FFNi(r)\text{FFN}_i^{(r)}은 task에 따라 활성화 될 수도 있고 안 될 수도 있는 routed expert로 task에 특정한 feature를 학습하게 된다. 덕분에 general feature는 항상 활성화된 expert에서 학습하고, task-specific feature는 해당 task에 대해서만 활성화되는 expert에서 학습해서 효율성을 높임.
한편 식 (12)의 gi,tg_{i,t}는 gating 메커니즘으로 입력 ut\mathbf{u}_tii-번째 expert에 대한 임베딩 벡터 ei\mathbf{e}_i와의 내적(후 sigmoid를 통해 확률로 변환)을 구하고, 전체 expert 중에 상위 kk개 안에 드는 expert만 활성화 시키는 방식으로 MoE를 구현한다. 이를 통해 특정 expert가 특정 task에 대해 활성화 되도록 학습이 됨.
ht=ut+i=1NsFFNi(s)(ut)+i=1Nrgi,tFFNi(r)(ut)gi,t=gi,tj=1Nrgj,tgi,t={si,tsi,tTopk({sj,t1jNr},Kr)0otherwisesi,t=Sigmoid(utei)\begin{align} \mathbf{h}_t' &= \mathbf{u}_t + \sum_{i=1}^{N_s} \text{FFN}_i^{(s)}(\mathbf{u}_t) + \sum_{i=1}^{N_r} g_{i,t}\text{FFN}_i^{(r)}(\mathbf{u}_t) \\ g_{i,t} &= {g_{i,t}' \over \sum_{j=1}^{N_r} g_{j,t}'} \\ g_{i,t}' &= \begin{cases} s_{i,t} & s_{i,t} \in \text{Topk}(\{s_{j,t}|1 \le j \le N_r \}, K_r) \\ 0 & \text{otherwise} \end{cases} \\ s_{i,t} &= \text{Sigmoid}(\mathbf{u}_t^\top \mathbf{e}_i) \end{align}
추가로 입력 ut\mathbf{u}_t를 출력에 추가하여 모델이 입력을 출력으로 변환 하는 것을 직접 학습하지 않고 입력과 출력 사이(’입력’ - ‘출력’)의 잔차(residual)을 학습하도록 유도한다.

Auxiliary-Loss-Free Load Balancing

한편 일반적인 MoE 모델의 경우 expert 사이의 load balancing이 잡히지 않으면 즉, 특정 expert는 매번 활성화되고, 특정 expert는 거의 활성화되지 않으면, expert를 병렬화하는 의미를 없앤다. 이를 제어하기 위해 즉, expert 사이의 load balancing을 잡기 위해 보조적적인 loss를 사용하는게 일반적인데, auxiliary loss가 자칫 모델 성능을 손상 시킬 수 있기 때문에 DeepSeek 연구진은 auxiliary-loss-free load balancing 전략을 개척한다. 구체적으로 아래와 같이 각 expert에 대한 affinity score si,ts_{i,t}에 바이어스 항 bib_i를 도입한다.
gi,t={si,tsi,t+biTopk({sj,t+bj1jNr},Kr)0otherwise\begin{align} g_{i,t}' = \begin{cases} s_{i,t} & s_{i,t} + b_i \in \text{Topk}(\{s_{j,t} + b_j|1 \le j \le N_r\}, K_r)\\ 0 & \text{otherwise} \end{cases} \end{align}
이 bias 항은 학습 가능한 파라미터이며 expert에 과부하가 걸리면 bias를 줄여서 score가 낮아지게 해서 expert가 덜 선택 되게 하고, 과소부하가 걸리면 bias를 커져서 score를 높아지게 해서 expert가 더 선택되게 하여 전체적으로 expert가 고르게 선택될 수 있게 학습된다. 이러한 auxiliary-loss-free 접근은 결과적으로 모델 성능을 높이고 load balancing을 잘 달성할 수 있었다고 함.
참고로 이 bias가 추가된 식은 오직 routing에서만 사용되며, FFN 출력에 곱해지는 gating value에는 원래의 affinity score 식(14)가 사용된다.

Complementary Sequence-Wise Auxiliary Loss

DeepSeek는 기본적으로 auxiliary-loss-free 전략을 사용하지만 임의의 단일 시퀀스 내에서 극단적인 불균형이 발생하는 것을 방지하기 위해 보완적으로 아래의 sequence-wise balance loss를 추가로 사용한다.
Lbal=αi=1NrfiPifi=NrKrTt=1T1 (si,tTopk({sj,t1jNr},Kr)si,t=si,tj=1Nrsj,tPi=1Tt=1Ts1,t\begin{align}\mathcal{L}_\text{bal} &= \alpha \sum_{i=1}^{N_r} f_i P_i \\ f_i &= {N_r\over K_r T}\sum_{t=1}^T \mathbf{1} \ (s_{i,t} \in \text{Topk}(\{s_{j,t}|1 \le j \le N_r\}, K_r) \\ s_{i,t}' &= {s_{i,t} \over \sum_{j=1}^{N_r} s_{j,t}} \\ P_i &= {1\over T} \sum_{t=1}^T s_{1,t}' \end{align}
여기서 하이퍼파라미터 α\alpha는 balance factor로 극단적으로 작은 값이 설정되며, 1()1(\cdot)은 내부의 값이 만족될 때 11, 그렇지 않으면 00을 반환하는 indicator function으로 1(si,tTopk({sj,t1jNr},Kr)1(s_{i,t} \in \text{Topk}(\{s_{j,t}|1 \le j \le N_r\}, K_r)는 expert ii가 토큰 tt에서 Topk에 몇 번 들어갔었는지를 합산한다.

Node-Limited Routing

일반적으로 MoE에서 각 expert는 여러 GPU에 걸쳐 분산될 수 있다. 따라서 각 GPU가 각 expert에 접근할 수 있도록 all-to-all 커뮤니케이션을 구성하게 된다. 이때 모든 GPU가 모든 expert에 접근하면 특히 GPU 클러스터가 커질수록 커뮤니케이션 비용이 매우 커지기 때문에, DeepSeek 연구진은 이를 해결하기 위해 각 토큰이 최대 MM개의 node에만 보낼 수 있게 설정함. 이것은 각 노드에 분포된 expert의 가장 높은 KrM{K_r \over M} affinity score의 합을 따라 선택되도록 함. 이를 통해 DeepSeek의 MoE 학습 프레임워크가 full computation-communication overlap에 가까운 수준을 달성할 수 있었다고 함.

No Token-Dropping

기존 MoE 모델에서 load balancing이 잘 못 잡히면, 특정 expert에 너무 많은 token이 몰려서 해당 expert에 대한 GPU에 과부하가 걸리기 때문에 GPU를 보호하기 위해 일부 token을 drop하는 기법을 사용하는데, DeepSeek-V3은 효율적인 load balancing 전략으로 이런 token-dropping 기법을 사용하지 않을 수 있었다고 함.

Sample Code

Gating 메커니즘에 대한 코드는 아래 참조
class Gate(nn.Module): # ... 생략 def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: # weight는 routed expert의 수만큼의 행을 갖는 가중치 행렬이다. # 입력에 대해 weight를 linear 연산하여 입력에 대한 expert별 score를 구한다. scores = linear(x, self.weight) if self.score_func == "softmax": scores = scores.softmax(dim=-1, dtype=torch.float32) else: scores = scores.sigmoid() original_scores = scores # bias는 load balancing을 위한 용도이므로 routing 계산 단계에서만 더해준다. # bias는 expert가 과도하게 선택되면 해당 expert의 선택이 줄어들게 만들고 expert가 과소선택되면 해당 expert의 선택이 증가되도록 제어 한다. if self.bias is not None: scores = scores + self.bias # routing group이 존재하면 routing을 2단계로 처리한다. # 우선 해당 expert를 group 수로 나누고, 가장 점수가 높은 group을 제외한 나머지 group은 mask를 씌워 제외한다. # 그 후 선택된 group 내에서 topk expert를 선별하여 routing 한다. if self.n_groups > 1: scores = scores.view(x.size(0), self.n_groups, -1) if self.bias is None: group_scores = scores.amax(dim=-1) # bias가 None이면 amax() else: group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1) # bias가 있으면 topk() indices = group_scores.topk(self.topk_groups, dim=-1)[1] mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False) scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1) # bias가 추가된 scores에 대해 topk를 취하여 가장 큰 값을 갖는 인덱스를 추출한다. indices = torch.topk(scores, self.topk, dim=-1)[1] # original_scores에서 indices에 해당하는 값을 추출한다. weights = original_scores.gather(1, indices) # sigmoid 함수를 사용하는 경우 각 행의 합이 1이 되도록 정규화한다. if self.score_func == "sigmoid": weights /= weights.sum(dim=-1, keepdim=True) # route_scale를 곱하여 값을 조정한다. weights *= self.route_scale # 최종적으로 weights와 indices를 반환한다. return weights.type_as(x), indicesㅍ
Python
복사
위의 gating 메커니즘을 활용한 MoE는 아래와 같이 구현된다.
class MoE(nn.Module): # ... 생략 def forward(self, x: torch.Tensor) -> torch.Tensor: shape = x.size() x = x.view(-1, self.dim) # gate에서 입력에 적합한 가중치와 index를 가져온다. weights, indices = self.gate(x) # index를 이용하여 사용할 routed expert에 대한 카운트를 설정한다. y = torch.zeros_like(x) counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() # 각 routed expert에 대한 카운트가 0이 아니면 해당 expert를 사용한다. for i in range(self.experts_start_idx, self.experts_end_idx): if counts[i] == 0: continue expert = self.experts[i] # 여기서 idx는 batch index이고 top은 expert index이다. idx, top = torch.where(indices == i) y[idx] += expert(x[idx]) * weights[idx, top, None] # shared expert는 항상 사용한다. z = self.shared_experts(x) # 만약 분산 처리가 이루어지는 경우, 모든 노드에서 계산된 y의 합을 구해야 한다. if world_size > 1: dist.all_reduce(y) # 최종적으로 모든 계산된 값을 합산하여 출력한다. return (y + z).view(shape)
Python
복사

Multi-Token Prediction

일반적으로 Transformer는 현재 입력된 토큰에 대해 11개의 출력을 생성하는데, DD개의 토큰을 추가로 예측하도록 하는 것이 Multi-Token Prediction(MTP)이다. 이것은 정답이 주어지는 상황에서만 사용 가능하므로 inference에는 사용하지 않고 오직 학습 단계에서만 활용된다.
MTP를 처음 제안한 Gloeckle et al(2024)의 작업은 독립적인 출력 헤드를 사용하여 병렬적으로 DD개의 추가 토큰을 예측 했던 반면 DeepSeek에서는 순차적으로 DD개의 토큰을 예측하도록 해서 causal chain을 유지했음.
MTP를 구현하기 위해 기본적인 Transfomer의 구조를 갖는 Main Model과 별개로 해당 모델의 Embedding Layer, Output Head, Transformer Block를 재사용하여 DD개의 MTP 모듈을 생성한다. 아래 그림 참조.
여기서 각 MTP 모듈은 이전(혹은 Main Model)의 Output Head를 통과하기전의 출력과 해당 시점의 정답 토큰을 입력으로 받아 추가 토큰을 예측한다. 여기서 이전 모듈의 출력은 RMSNorm을 통과시키고, 정답 토큰은 embedding layer를 통과시킨 후 RMSNorm을 통과시키고 둘을 concatenate하고 해당 MTP 모듈의 독립적인 Linear Projection을 수행한 후 Main Model과 공유되는 Transformer Block을 통과시켜 출력을 생성한다. 여기서 생성된 출력은 다음 MTP 모듈을 위해 사용될 수 있고, 이 출력을 output head를 통과시키면 해당 MTP 모듈에 대한 loss를 계산할 수도 있다.
이 설명에 대한 수식은 아래 참조. 여기서 hik1\mathbf{h}_i^{k-1}은 이전 MTP 모듈(혹은 Main Model)의 output head를 통과시키기 전의 출력이고, tikt_{i_k}는 해당 시점의 정답 토큰이고 hik\mathbf{h}_i'^{k}는 해당 시점의 MTP 모듈의 Linear Projection의 결과(transformer block의 입력)이고, hik\mathbf{h}_i^{k}는 Transformer block을 통과한 결과로 output head의 입력, Pi+k+1kP_{i+k+1}^k는 output head의 결과로 해당 MTP 모듈의 cross-entropy loss 계산에 쓰인다.
hik=Mk[RMSNorm(hik1);RMSNorm(Emb(ti+k))]h1:Tkk=TRMk(h1:Tkk)Pi+k+1k=OutHead(hik)\begin{align}\mathbf{h}_i'^k &= M_k[\text{RMSNorm}(\mathbf{h}_i^{k-1}); \text{RMSNorm}(\text{Emb}(t_{i+k}))] \\ \mathbf{h}_{1:T-k}^k &= \text{TRM}_k(\mathbf{h}_{1:T-k}'^k) \\ P_{i+k+1}^k &= \text{OutHead}(\mathbf{h}_i^k) \end{align}
각 MTP 모듈에 대한 cross-entropy loss는 LMTPk\mathcal{L}_\text{MTP}^k은 아래처럼 계산된다.
LMTPk=CrossEntropy(p2+k:T+1k,t2+k:T+1)=1Ti=2+kT+1logPik[ti]\begin{align} \mathcal{L}_\text{MTP}^k = \text{CrossEntropy}(p_{2+k:T+1}^k, t_{2+k:T+1}) = -{1\over T} \sum_{i=2+k}^{T+1} \log P_i^k[t_i] \end{align}
전체 MTP 모듈의 손실은 아래와 같이 평균내고 가중치 factor λ\lambda를 곱해 사용한다.
LMTP=λDk=1DLMTPk\begin{align} \mathcal{L}_\text{MTP} = {\lambda\over D} \sum_{k=1}^D \mathcal{L}_\text{MTP}^k \end{align}
이러한 절차는 전반적으로 모델의 학습을 개선할 수 있다고 함.

Training Framework

DeepSeek의 성능은 앞선 방법들과 이전에 사용했던 것보다 훨씬 많은 대규모 데이터셋을 활용한 덕분으로 생각되는데, DeepSeek의 놀라운 비용은 아래의 엔지니어링 기법 덕분이라고 생각 됨.

Dual Pipe and Computation-Communication Overlap

대규모 모델을 단일 GPU에 띄울 수 없기 때문에 여러 병렬화 방법을 사용하는데, 그 중 Data Parallelism(DP)은 여러 GPU에 동일한 모델을 띄우고 Datset을 병렬화하여 각 GPU에 넣고 계산한 뒤 결과를 취합하는 방식이고, Pipeline Parallelism(PP)는 대규모 모델을 layer 별로 나누어 여러 GPU에 띄우고, 각 GPU 별로 연산을 전파하는 식으로 처리하는 방식이고, Tensor Parallelism은 대규모 모델에 대해 각 GPU가 모델의 모든 layer를 가지지도록 행렬 계산 방식에 의해 파라미터를 분리한 다음, 입력을 모든 GPU에서 처리한 후에, 최종적으로 결과를 취합하는 방식을 말한다.
GPU 분산 환경에서 시퀀스 모델을 처리할 때, 필연적으로 앞선 GPU의 연산이 끝나야 다음 GPU에서 연산을 할 수 있기 때문에 뒤에 존재하는 GPU에 대해 Idle 상태인 상황이 발생할 수 밖에 없는데, 이렇게 GPU가 놀고 있는 시간을 Pipeline Bubble이라 한다.
현재의 GPU 연산에서 가장 큰 병목은 연산이 아니라 메모리 IO 접근이나 다른 칩 간의 커뮤니케이션에서 발생하기 때문에 DeekSeek 연구진은 computation과 communication에 대한 최적화를 위해 computation 단계와 comnunication 단계가 1:1이 되도록 다음과 같이 4개의 chunk를 구성함
attention - computation
all-to-all dispatch(MoE routing) - communication
MLP - computation
all-to-all combine(MoE 출력 결합) - communication
이렇게 구분한 chunk를 기준으로 Idle로 존재하는 GPU를 최소화 하기 위해 아래 그림과 같이 GPU 연산 단계를 overlapping하는 Dual Pipe 전략을 사용하여 pipeline bubble을 최소화함

Efficient Implementation of Cross-Node All-to-All Communication

또한 DualPipe의 계산적 성능을 최대한 보장하기 위해 연산과 통신에 모두 쓰이는 SM(streaming multiprocessor)에 대해 통신에 전담되는 SM의 수를 20개로 줄여 GPU의 연산과 통신이 최대한 overlap 되도록 했음. 이를 위헤 CUDA의 PTX(parallel thread excution) 명령어도 커스터마이징하였다고 함.

Extremely Memory Saving with Minimal Overhead

통신과 함께 GPU 병목의 주요 원인인 memory IO를 줄이기 위해 DeepSeek 연구진은 아래와 같은 방법들을 사용함
Recomputation of RMSNorm and MLA Up-Projection
역전파하는 동안 모든 RMSNorm 연산과 MLA up-projection을 재계산하여 출력 활성화를 지속적으로 저장할 필요성을 없앤다. 이 전략은 작은 오버헤드만으로 activation 저장을 위한 메모리 요구량을 상당히 줄인다.
Exponential Moving Average in CPU
모델 파라미터에 대한 Exponential Moving Average(EMA)를 보존하는데, EMA 파라미터는 CPU 메모리에 저장되고 각 학습 단계 이후에 비동기적으로 업데이트 됨
Shared Embedding and Output Head for Multi-Token Prediction
DualPipe 전략을 사용하여 모델의 가장 얕은 레이어(embedding layer를 포함하여)와 가장 깊은 레이어(출력 헤드를 포함하여)를 동일한 PP rank에 배포한다. 이 배열은 MTP 모듈과 메인 모델 사이의 공유된 임베딩과 출력 헤드의 파라미터의 gradient를 물리적 공유할 수 있게 한다. 이 물리적 공유 메커니즘은 추가로 메모리 효율성을 강화한다.

FP8 Training

Mixed Precision Framework

일반적으로 FP8 format은 pre-trained 모델을 더 작은 크기로 만드는 양자화를 위해 사용되는 형식인데, pre-trained 단계부터 양자화 format을 사용하려는 시도들이 있었고, DeepSeek-V3도 유사하게 FP8 mixed precision 방식을 이용하여 대규모 학습을 진행 했음.
이를 통해 기존 BF16 보다 두 배 빠른 연산 속도를 얻으면서 메모리도 절감할 수 있었는데, 학습 단계에서 정밀도에 민감한 부분들이 존재하기 때문에 mixed 형식으로 일부 정밀도에 민감한 영역에 대해서는 FP32나 BF16 등을 사용하고 그 외의 부분에 대해 FP8 format을 사용함. 그렇게 해서 학습 안정성을 가지면서도 효율적인 학습을 수행할 수 있었음.
참고로 FP8로 실행된 연산은 아래와 같고
Fprop(forward pass)
Dgrad(activation backward pass)
Wgrad(weight backward pass)
정밀도에 민감한 아래의 컴포넌트들에 대해서는 원본 정밀도(BF16 또는 FP32)를 유지함
임베딩 모듈
출력 헤드
MoE gating 모듈
normalization 연산기
attention 연산기
추가로 수치적 안정성을 보장하기 위해 아래의 요소들에 대해서는 더 높은 정밀도를 사용했다고 함
master 가중치
가중치 gradients
optimizer state

Improved Precision from Quantization and Multiplication

mixed precision FP8 프레임워크를 사용하기 때문에 양자화 단계가 필요한데 이때 다음과 같은 fine-grained quantization을 수행함.
우선 입력 벡터에 대해 채널별로 1×Nc1 \times N_c (여기서 Nc=128N_c=128) 형태로 tile basis로 group화 하고, 가중치 행렬에 대해서는 Nc×NcN_c \times N_c 형태로 block basis로 group화를 수행함. 그 후에 각 group 내의 최대 절댓값(Max Absolute Value)을 이용해 값이 [1.0,1.0][-1.0, 1.0] 내로 들어오도록 스케일링함. 이러한 미세한 양자화는 outlier에 더 견고하게 만들어서 효율적인 학습이 가능했다고 함.
한편 low-precision GEMM 연산이 종종 underflow 이슈에 시달려서 연구원들이 확인해 보니 H800 GPU에서 GEMM 연산을 가속화 하는데 쓰이는 Tensor core의 레지스터 정밀도가 14비트로 제한되어있었다고 함. CUDA core의 레지스터는 FP32 정밀도로 데이터를 저장하는 반면, Tensor core는 레지스터 정밀도가 14비트 밖에 안되서 내부 차원의 크기가 커질 수록 누적되는 값의 오차가 증가하는 문제가 있었음.
때문에 일정한 간격(NcN_c)마다 Tensor core의 중간 결과를 CUDA core의 레지스터에 복사해서 정밀도를 높이는 작업을 수행 함(이것은 애초에 Tensor core의 구조상 높은 정밀도를 유지하기 어렵기 때문에 최신 GPU에서도 유사한 기법을 사용해 줘야 한다고 함). 덕분에 실행 속도는 감소 됐지만 정밀도를 개선할 수 있었다고 함.
추가로 양자화를 online으로 했다는 내용과 FP8 format을 이전 연구에서는 E4M3와 E5M2를 사용했었는데 DeepSeek-V3에서는 E4M3로 동일했다는 내용도 있음
또한 Optimizer State와 Activation에 대해 더 낮은 정밀도로 압축해서 메모리 소비와 커뮤니케이션 오버헤드를 줄였다고 함.
이하 추론 및 배포에서 부하 관리 전략과 DeepSeek 연구진이 Hardward Vendor들에게 한 제안이나 학습 데이터와 하이퍼파라미터 설정, Evaluation에 대한 내용은 생략. 해당 내용은 아래 논문 내용 참조.

Reinforcement Learning

reasoning 능력에 대해서는 DeepSeek-R1 시리즈가 더 낫지만, DeepSeek-V3도 reasoning 능력을 가진 모델이다. 다만 reasoning 능력을 발전시키기 위한 방법이 R1 시리즈와 유사하므로 관련 내용은 아래의 DeepSeek-R1 내용 참조.
참고로 DeepSeek-R1 시리즈는 DeepSeek-V3를 기준으로 만들어진 모델이지만, DeepSeek-V3의 reasoning 능력을 발달시키는데 DeekSeek-R1을 이용해 만들어진 데이터를 학습에 사용했다는 점이 흥미롭다. 즉 두 모델은 서로에게 영향을 주었다.

참고