DeepSeek モデルアーキテクチャ徹底解説

# ブロック単位量子化の擬似コード
# 順伝播: FP8行列乗算
for each block (128x128) in activation_matrix:
scale = max(abs(block)) / 448 # E4M3の最大値は448
block_fp8 = round(block / scale) # FP8に量子化
output_block = block_fp8 @ weight_fp8 # FP8行列乗算
output_block = output_block * scale # 逆量子化

# 逆伝播: 勾配はBF16精度で計算
gradient = backward(output, loss) # BF16勾配

# GRPOのコアロジック
# 各質問に対して、G個の候補回答を生成
group_responses = [model.generate(question) for _ in range(G)]
group_rewards = [reward_model(question, r) for r in group_responses]

# グループ内相対優位性（平均を引き、標準偏差で割る）
mean_reward = mean(group_rewards)
std_reward = std(group_rewards)
advantages = [(r - mean_reward) / std_reward for r in group_rewards]

# ポリシー勾配更新
loss = -mean(advantages * log_prob_ratio - beta * kl_divergence)
model.update(loss)

# DeepSeekMoE フォワード伝播（DeepSeek-V3 オープンソースコードに基づく） class DeepSeekMoE(nn.Module): def __init__(self): self.num_experts = 256 # ルーティング専門家の数 self.num_shared_experts = 8 # 共有専門家の数 self.top_k = 8 # Top-K ルーティング self.n_routed_experts = 256 self.routed_scaling_factor = 1.0 # ゲーティングネットワーク self.gate = nn.Linear(hidden_size, self.n_routed_experts, bias=False) # ルーティング専門家（各専門家は FFN） self.experts = nn.ModuleList([ ExpertFFN() for _ in range(self.num_experts) ]) # 共有専門家 self.shared_experts = SharedExpertFFN() def forward(self, hidden_states): identity = hidden_states orig_shape = hidden_states.shape # 1. ゲーティングルーティング gate_logits = self.gate(hidden_states) # Top-K 専門家を選択 weights, indices = torch.topk(gate_logits, self.top_k, dim=-1) weights = F.softmax(weights, dim=-1, dtype=torch.float32) # 2. ルーティング専門家の計算 routed_output = self.moe_forward(hidden_states, weights, indices) # 3. 共有専門家の計算（常に活性化） shared_output = self.shared_experts(identity) # 4. 出力を結合 final_output = routed_output + shared_output return final_output def moe_forward(self, hidden_states, topk_weights, topk_ids): # トークンを割り当てられた専門家ごとにグループ化 outputs = torch.zeros_like(hidden_states) for expert_idx in range(self.num_experts): # この専門家に割り当てられたトークンを見つける expert_mask = (topk_ids == expert_idx).any(dim=-1) if not expert_mask.any(): continue expert_input = hidden_states[expert_mask] expert_output = self.experts[expert_idx](expert_input) # 重み付き累積 weight = topk_weights[expert_mask][topk_ids[expert_mask] == expert_idx] outputs[expert_mask] += expert_output * weight.unsqueeze(-1) return outputs

# MLA マルチヘッド潜在アテンション（DeepSeek-V3 オープンソースコードに基づく） class MultiHeadLatentAttention(nn.Module): def __init__(self): self.q_lora_rank = 1536 # Query 低ランク分解次元 self.kv_lora_rank = 512 # KV 結合圧縮次元 self.qk_rope_head_dim = 64 # RoPE 部分の次元 self.qk_nope_head_dim = 128 # 非 RoPE 部分の次元 self.v_head_dim = 128 self.num_heads = 128 # Query 低ランク分解 self.q_a_proj = nn.Linear(hidden_size, self.q_lora_rank, bias=False) self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank) self.q_b_proj = nn.Linear(self.q_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.qk_rope_head_dim)) # KV 結合圧縮 self.kv_a_proj_with_mqa = nn.Linear( hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False) self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank) self.kv_b_proj = nn.Linear( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)) def forward(self, hidden_states, freqs_cis, kv_cache=None): # 1. Query 低ランク分解 q = self.q_a_proj(hidden_states) q = self.q_a_layernorm(q) q = self.q_b_proj(q) q = q.view(-1, self.num_heads, self.qk_nope_head_dim + self.qk_rope_head_dim) q_nope, q_rope = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_rope = apply_rotary_emb(q_rope, freqs_cis) # RoPE エンコーディング # 2. KV 結合圧縮 kv_a = self.kv_a_proj_with_mqa(hidden_states) c_kv, k_rope = kv_a.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_rope = apply_rotary_emb(k_rope.unsqueeze(1), freqs_cis) # RoPE エンコーディング # 3. 低次元潜在ベクトル c_KV をキャッシュ（推論時） if kv_cache is not None: kv_cache.append(c_kv) # 4. 上投影で Key と Value を復元 c_kv = self.kv_a_layernorm(c_kv) kv = self.kv_b_proj(c_kv) kv = kv.view(-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) # 5. Query と Key の各部分を連結 q = torch.cat([q_nope, q_rope], dim=-1) k = torch.cat([k_nope, k_rope.expand(-1, self.num_heads, -1)], dim=-1) # 6. 標準アテンション計算 attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt( self.qk_nope_head_dim + self.qk_rope_head_dim) attn_weights = F.softmax(attn_weights, dim=-1) output = torch.matmul(attn_weights, v) return output

# MTP マルチトークン予測モジュール（DeepSeek-V3 オープンソースコードに基づく） class MultiTokenPredictionModule(nn.Module): def __init__(self): self.hidden_size = 7168 # 入力融合層 self.embedding_norm = nn.RMSNorm(hidden_size) self.hidden_norm = nn.RMSNorm(hidden_size) self.input_proj = nn.Linear(hidden_size * 2, hidden_size, bias=False) # MTP Transformer ブロック self.mtp_block = TransformerBlock() # 出力投影（共有 LM Head） self.output_norm = nn.RMSNorm(hidden_size) self.shared_head = nn.Linear(hidden_size, vocab_size, bias=False) def forward(self, hidden_states, input_ids): # 1. 現在のトークンの埋め込みを取得 token_emb = self.embedding_norm(embedding(input_ids)) # 2. メインモデルの隠れ状態とトークン埋め込みを融合 hidden_norm = self.hidden_norm(hidden_states) combined = torch.cat([hidden_norm, token_emb], dim=-1) combined = self.input_proj(combined) # 3. MTP Transformer ブロックを通過 mtp_output = self.mtp_block(combined) # 4. 出力予測 mtp_output = self.output_norm(mtp_output) logits = self.shared_head(mtp_output) return logits # トレーニング時に使用 hidden = transformer(input_ids) # メインモデルのフォワード logits_main = lm_head(hidden) # 標準の next-token 予測 logits_mtp = mtp_module(hidden, input_ids) # MTP 追加予測 loss_main = cross_entropy(logits_main[:, :-1], targets[:, 1:]) loss_mtp = cross_entropy(logits_mtp[:, :-2], targets[:, 2:]) total_loss = loss_main + 0.3 * loss_mtp # 加重組み合わせ