"""模型主体：Transformer Block、MTP 多 token 预测头、完整模型。""" import math from typing import Optional, Tuple, List import torch import torch.nn as nn import torch.nn.functional as F from .layers import RMSNorm, build_rope_cache from .mla import MLA from .moe import MoE, DenseFFN def _shift_left(x: torch.Tensor, k: int, fill: int): """y[:, i] = x[:, i+k]，越界位置填 fill。""" if k == 0: return x.clamp_min(0) if fill == 0 else x y = torch.full_like(x, fill) if k < x.size(1): y[:, :-k] = x[:, k:] if fill == 0: y = y.clamp_min(0) return y class Block(nn.Module): """Pre-Norm 残差块：x + MLA(norm(x)) → x + FFN(norm(x))。""" def __init__(self, cfg, layer_id: int): super().__init__() self.attn_norm = RMSNorm(cfg.dim, cfg.norm_eps) self.attn = MLA(cfg) self.ffn_norm = RMSNorm(cfg.dim, cfg.norm_eps) # V3 的做法：靠前的层用稠密 FFN，后面的层才换成 MoE，训练早期更稳 self.is_moe = layer_id >= cfg.n_dense_layers self.ffn = MoE(cfg) if self.is_moe else DenseFFN(cfg) def forward(self, x, cos, sin, start_pos=0, mask=None): x = x + self.attn(self.attn_norm(x), cos, sin, start_pos, mask) h, aux = self.ffn(self.ffn_norm(x)) return x + h, aux class MTPModule(nn.Module): """Multi-Token Prediction 模块（DeepSeek-V3）。 普通语言模型每个位置只预测下一个 token，训练信号比较稀疏。 MTP 再挂一个轻量模块，让位置 i 同时预测 i+2： h'_i = Block( W · [ RMSNorm(h_i) ; RMSNorm(Emb(t_{i+1})) ] ) p(t_{i+2}) = Head(RMSNorm(h'_i)) 嵌入层和输出头与主干共享，所以额外开销只有一层 Block。 好处有二：训练时数据效率更高（模型被迫规划更远）； 推理时这个头天然就是个 draft model，可以做自投机解码，实测约 1.8x 加速。 """ def __init__(self, cfg, layer_id: int): super().__init__() self.h_norm = RMSNorm(cfg.dim, cfg.norm_eps) self.e_norm = RMSNorm(cfg.dim, cfg.norm_eps) self.proj = nn.Linear(2 * cfg.dim, cfg.dim, bias=False) self.block = Block(cfg, layer_id) def forward(self, h, emb, cos, sin, start_pos=0, mask=None): z = self.proj(torch.cat([self.h_norm(h), self.e_norm(emb)], dim=-1)) return self.block(z, cos, sin, start_pos, mask) class DeepSeekFlashMini(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.dim) self.layers = nn.ModuleList([Block(cfg, i) for i in range(cfg.n_layers)]) self.norm = RMSNorm(cfg.dim, cfg.norm_eps) self.head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False) if cfg.tie_embeddings: self.head.weight = self.embed.weight self.mtp = nn.ModuleList([MTPModule(cfg, cfg.n_layers + i) for i in range(cfg.n_mtp)]) self.mtp_norm = RMSNorm(cfg.dim, cfg.norm_eps) if cfg.n_mtp > 0 else None cos, sin = build_rope_cache(cfg.qk_rope_head_dim, cfg.max_seq_len, cfg.rope_theta, cfg.rope_scaling) self.register_buffer("rope_cos", cos, persistent=False) self.register_buffer("rope_sin", sin, persistent=False) self.apply(self._init_weights) # 残差分支的输出投影按深度缩放初始化，防止深层激活方差爆掉 for name, p in self.named_parameters(): if name.endswith("wo.weight") or name.endswith("w_down.weight"): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layers)) self._cache_ready = False @staticmethod def _init_weights(m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02) # ================= 前向 ================= def _rope(self, start_pos: int, T: int): if start_pos + T > self.cfg.max_seq_len: raise ValueError( f"位置 {start_pos + T} 超出上下文上限 max_seq_len={self.cfg.max_seq_len}。" f"请加大配置里的 max_seq_len（并考虑设置 rope_scaling 做外推）。") return self.rope_cos[start_pos:start_pos + T], self.rope_sin[start_pos:start_pos + T] def _mask(self, T: int, start_pos: int, device): """带 cache 的增量解码时，为多 token 输入构造对齐的因果 mask。""" if T == 1: return None if not self._cache_ready: return None # 训练路径：交给 SDPA 的 is_causal S = start_pos + T return torch.ones(T, S, dtype=torch.bool, device=device).tril(S - T) def forward_trunk(self, tokens: torch.Tensor, start_pos: int = 0): """主干前向，返回 (最后一层隐状态 h, logits, 路由辅助损失)。""" B, T = tokens.shape cos, sin = self._rope(start_pos, T) mask = self._mask(T, start_pos, tokens.device) h = self.embed(tokens) aux_total = h.new_zeros(()) for layer in self.layers: h, aux = layer(h, cos, sin, start_pos, mask) aux_total = aux_total + aux logits = self.head(self.norm(h)) return h, logits, aux_total def mtp_forward(self, h: torch.Tensor, next_tokens: torch.Tensor, start_pos: int = 0, depth: int = 0): """MTP 第 depth 层：给定主干隐状态和「下一个 token」，预测再下一个。""" B, T = next_tokens.shape cos, sin = self._rope(start_pos, T) mask = self._mask(T, start_pos, next_tokens.device) emb = self.embed(next_tokens) h2, aux = self.mtp[depth](h, emb, cos, sin, start_pos, mask) logits = self.head(self.mtp_norm(h2)) return h2, logits, aux def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, start_pos: int = 0): """训练入口。targets[i] 应为 tokens[i+1]（由 dataset 准备好）。""" h, logits, aux = self.forward_trunk(tokens, start_pos) if targets is None: return {"logits": logits} loss_main = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100) out = {"logits": logits, "loss_main": loss_main, "aux_loss": aux} # ---- MTP：第 d 个模块用 t_{i+1+d} 当输入，监督目标是 t_{i+2+d} ---- loss_mtp = logits.new_zeros(()) if self.cfg.n_mtp > 0: cur_h = h for d in range(self.cfg.n_mtp): inp = _shift_left(targets, d, fill=0) # t_{i+1+d} tgt = _shift_left(targets, d + 1, fill=-100) # t_{i+2+d} cur_h, mtp_logits, aux_d = self.mtp_forward(cur_h, inp, start_pos, d) aux = aux + aux_d loss_mtp = loss_mtp + F.cross_entropy( mtp_logits.reshape(-1, mtp_logits.size(-1)), tgt.reshape(-1), ignore_index=-100) loss_mtp = loss_mtp / self.cfg.n_mtp out["loss_mtp"] = loss_mtp out["aux_loss"] = aux out["loss"] = loss_main + self.cfg.mtp_loss_weight * loss_mtp + aux return out # ================= KV cache 管理 ================= def setup_cache(self, max_batch: int, max_seq_len: int, device=None, dtype=None): device = device or self.embed.weight.device dtype = dtype or self.embed.weight.dtype for m in self.modules(): if isinstance(m, MLA): m.setup_cache(max_batch, max_seq_len, device, dtype) self._cache_ready = True def clear_cache(self): for m in self.modules(): if isinstance(m, MLA): m.clear_cache() self._cache_ready = False def set_attn_impl(self, impl: str): for m in self.modules(): if isinstance(m, MLA): m.attn_impl = impl # ================= 统计 ================= def num_params(self) -> dict: cfg = self.cfg total = sum(p.numel() for p in self.parameters()) if cfg.tie_embeddings: total_unique = total else: total_unique = total # 激活参数量：每个 token 实际参与计算的那部分 inactive = 0 for layer in self.layers: if layer.is_moe: per_expert = sum(p.numel() for p in layer.ffn.experts[0].parameters()) skipped = cfg.n_routed_experts - cfg.n_activated_experts inactive += per_expert * skipped mtp_params = sum(p.numel() for p in self.mtp.parameters()) return { "total": total_unique, "activated": total_unique - inactive - mtp_params, # 推理时 MTP 可不用 "mtp": mtp_params, "embedding": self.embed.weight.numel(), } def kv_cache_bytes_per_token(self, dtype_size: int = 2) -> dict: """对比 MLA 与等规模 MHA 的 KV cache 开销。""" cfg = self.cfg mla = (cfg.kv_lora_rank + cfg.qk_rope_head_dim) * cfg.n_layers * dtype_size mha = 2 * cfg.n_heads * cfg.v_head_dim * cfg.n_layers * dtype_size return {"mla": mla, "mha_equivalent": mha, "ratio": mha / mla} def expert_load_stats(self) -> List[float]: return [layer.ffn.gate.load_stats()["imbalance"] for layer in self.layers if layer.is_moe] @torch.no_grad() def update_expert_bias(self): """训练循环每个 optimizer step 后调用一次。""" speed = self.cfg.bias_update_speed if speed <= 0: return for m in self.modules(): if isinstance(m, MoE): m.gate.update_bias(speed)