@dataclass
class DeepSeekV3ModelArgs(BaseModelArgs):
    """
    配置模型参数

    Attributes:
        max_batch_size (int): 最大批量大小
        max_seq_len (int): 最大序列长度
        vocab_size (int): 词典大小
        dim (int): 嵌入维度
        inter_dim (int): MLP层的中间维度
        moe_inter_dim (int): MoE层专家的中间维度
        n_layers (int): Transformer层的数量
        n_dense_layers (int): 模型中密集层的数量, 前几层负载均衡收敛慢, 因此设置为密集层
        n_heads (int): 注意力头数
        n_routed_experts (int): MoE层中路由的专家数量
        n_shared_experts (int): MoE层中共享的专家数量
        n_activated_experts (int): MoE层中激活的专家数量
        n_expert_groups (int): 专家的分组数量
        n_limited_groups (int): 路由限制, 每次最多从限制的专家组里选择专家
        route_scale (float): 路由权重的缩放因子
        use_noaux_tc (bool): 是否使用无辅助损失的负载均衡策略
        bias_update_speed (float): 偏置更新速度
        use_seq_aux (bool): 是否使用序列级别的辅助损失
        seq_aux_alpha (float): 序列级别的辅助损失的权重
        q_lora_rank (int): query 的下投影维度【对应论文中的 d_c'】
        kv_lora_rank (int): key/value 的下投影维度【对应论文中的 d_c】
        qk_nope_head_dim (int): 没有位置编码的 q_t^C 和 k_t^C 的每个头的维度【对应论文中的 d_h】
        qk_rope_head_dim (int): 解耦的带 RoPE 的 q_t^R 和 k_t^R 的每个头的维度【对应论文中的 d_h^R】
        v_head_dim (int): value 的每个头的维度, 可以与 qk_nope_head_dim 不同【但在论文中也同样设定为 d_h】
        rope_theta (float): 旋转位置编码的基底【即 θ_d=b^(-2d/D) 中的 b】
        use_mtp (bool): 是否使用 MTP 策略
        mtp_loss_lambda (float): MTP 损失的权重
    """
    ...

class RMSNorm(nn.Module):
	"""RMS归一化"""
	...
    
def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
	"""应用RoPE位置编码"""
	...

class MLA(nn.Module):
    """
    多头潜在注意力: Multi-Headed Latent Attention Layer (MLA)

    Attributes:
        dim (int): 嵌入维度【对应论文中的 d】
        n_heads (int): 注意力头数【对应论文中的 n_h】
        q_lora_rank (int): query 的下投影维度【对应论文中的 d_c'】
        kv_lora_rank (int): key/value 的下投影维度【对应论文中的 d_c】
        qk_nope_head_dim (int): 没有位置编码的 q_t^C 和 k_t^C 的每个头的维度【对应论文中的 d_h】
        qk_rope_head_dim (int): 解耦的带 RoPE 的 q_t^R 和 k_t^R 的每个头的维度【对应论文中的 d_h^R】
        qk_head_dim (int): 最终执行注意力计算的 query 和 key 的每个头的维度【即 d_h + d_h^R】
        v_head_dim (int): value 的每个头的维度, 可以与 qk_nope_head_dim 不同【但在论文中也同样设定为 d_h】
        softmax_scale (float): 注意力计算的缩放因子【即 1/sqrt(d_h + d_h^R)】
    """
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.dim = args.dim
        self.n_heads = args.n_heads

        # 定义 query/key/value 维度
        self.q_lora_rank = args.q_lora_rank
        self.kv_lora_rank = args.kv_lora_rank
        self.qk_nope_head_dim = args.qk_nope_head_dim
        self.qk_rope_head_dim = args.qk_rope_head_dim
        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
        self.v_head_dim = args.v_head_dim

        # 低秩压缩 query
        self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False)
        self.q_norm = RMSNorm(self.q_lora_rank)
        self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False)
        
        # key / value 的维度变换
        self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
        self.kv_norm = RMSNorm(self.kv_lora_rank)
        self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False)

        # 输出的维度变换
        self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False)
        self.softmax_scale = self.qk_head_dim ** -0.5  # 注意力缩放因子
        
        # 缓存 c_t^KV 和 k_t^R
        self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False)
        self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False)

def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
    """
    MLA 的前向传播

    Args:
        x (torch.Tensor): 输入 (batch_size, seq_len, dim)
        start_pos (int): 用于指定当前推理步骤的起始位置，即从序列的哪个位置开始计算
        freqs_cis (torch.Tensor): 预先计算的复数 RoPE 矩阵
        mask (Optional[torch.Tensor]): 掩码

    Returns:
        torch.Tensor: 输出 (batch_size, seq_len, dim)
    """

    bsz, seqlen, _ = x.size()
    end_pos = start_pos + seqlen

    # -------------------------- query 部分 --------------------------
    q = self.wq_b(self.q_norm(self.wq_a(x)))
    q = q.view(bsz, seqlen, self.n_heads, self.qk_head_dim)
    q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
    q_pe = apply_rotary_emb(q_pe, freqs_cis)

    # ----------------------- key / value 部分 -----------------------
    kv = self.wkv_a(x)
    kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
    k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)

    # --------------------------- 矩阵吸收 ---------------------------
    wkv_b = self.wkv_b.weight
    wkv_b = wkv_b.view(self.n_heads, -1, self.kv_lora_rank)
    q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
    
    # -------------------------- 注意力实现 --------------------------
    if self.training:  # 训练阶段不使用缓存
        kv = self.kv_norm(kv)
        k_pe = k_pe.squeeze(2)
        scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) + torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale
    else:  # 推理阶段使用缓存
        self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)
        self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)
        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)
    
    # ----------------------- 计算输出+矩阵吸收 -----------------------
    if self.training:
        x = torch.einsum("bsht,btc->bshc", scores, kv)
    else:
        x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
    x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
    x = self.wo(x.flatten(2))
    return x

q = self.wq_b(self.q_norm(self.wq_a(x)))
q = q.view(bsz, seqlen, self.n_heads, self.qk_head_dim)
q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
q_pe = apply_rotary_emb(q_pe, freqs_cis)

kv = self.wkv_a(x)
kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)

wkv_b = self.wkv_b.weight
wkv_b = wkv_b.view(self.n_heads, -1, self.kv_lora_rank)
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])

if self.training:  # 训练阶段不使用缓存
    kv = self.kv_norm(kv)
    k_pe = k_pe.squeeze(2)
    scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) + torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale
else:  # 推理阶段使用缓存
    self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)
    self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)
    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)

if self.training:
    x = torch.einsum("bsht,btc->bshc", scores, kv)
else:
    x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
x = self.wo(x.flatten(2))

class MLP(nn.Module):
    """
    前馈层, 前馈方法与Llama3相同

    Attributes:
        w1 (nn.Module): 实现 input-to-hidden 的转换
        w2 (nn.Module): 实现 hidden-to-output 的转换
        w3 (nn.Module): 实现 input-to-hidden 的转换
    """
    def __init__(self, dim: int, inter_dim: int):
        """
        MLP 初始化

        Args:
            dim (int): 嵌入维度
            inter_dim (int): 隐藏层维度
        """
        super().__init__()
        self.w1 = nn.Linear(dim, inter_dim, bias=False)
        self.w2 = nn.Linear(inter_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, inter_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        MLP 前向传播

        Args:
            x (torch.Tensor): 输入张量 (batch_size, seq_len, dim)

        Returns:
            torch.Tensor: 输出张量 (batch_size, seq_len, dim)
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Gate(nn.Module):
    """
    即 Router, MoE 中的门控网口, 用于动态路由

    Attributes:
        dim (int): 嵌入维度
        topk (int): 每个输入激活的专家数量
        n_groups (int): 专家的分组数量
        topk_groups (int): 选中 topk 个分组
        route_scale (float): 路由权重的缩放因子
        n_routed_experts (int): 路由的专家数量
        weight (torch.nn.Parameter): 门控的可学习权重参数
        bias (Optional[torch.nn.Parameter]): 门控的偏置项
        original_scores (Optional[torch.Tensor]): 原始的亲和度得分, 形状为 (batch_size * seq_len, n_routed_experts)
    """
    def __init__(self, args: DeepSeekV3ModelArgs):
        """
        门控网络初始化

        Args:
            args (ModelArgs): 模型配置参数
        """
        super().__init__()
        self.dim = args.dim
        self.topk = args.n_activated_experts
        self.n_groups = args.n_expert_groups
        self.topk_groups = args.n_limited_groups
        self.route_scale = args.route_scale
        self.n_routed_experts = args.n_routed_experts
        self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
        self.bias = nn.Parameter(torch.empty(args.n_routed_experts), requires_grad=False) if args.use_noaux_tc else None
        self.original_scores = None
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """
        初始化参数
        """
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        门控网络的前向传播

        Args:
            x (torch.Tensor): 输入张量 (batch_size * seq_len, dim), 在输入前已经调整好形状

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: 路由权重和选择的专家索引, 形状均为 (batch_size * seq_len, topk)
        """
        # -------------------------- 计算亲和度分数 --------------------------
        scores_logits = F.linear(x, self.weight, None)
        scores = scores_logits.sigmoid()
        self.original_scores = scores
        scores_for_topk = scores
        if self.bias is not None:
            scores_for_topk = scores_for_topk + self.bias
        
        # -------------------------- 专家分组 --------------------------
        if self.n_groups > 1:
            scores_view = scores_for_topk.view(x.size(0), self.n_groups, -1)
            if self.bias is None:
                group_scores = scores_view.amax(dim=-1)
            else:
                group_scores = scores_view.topk(2, dim=-1)[0].sum(dim=-1)
            indices_groups = group_scores.topk(self.topk_groups, dim=-1)[1]
            mask = torch.ones(x.size(0), self.n_groups, dtype=torch.bool, device=x.device)
            mask.scatter_(dim=1, index=indices_groups, value=False)
            scores_for_topk = scores_view.masked_fill(mask=mask.unsqueeze(-1), value=float("-inf")).flatten(1)
        
        # -------------------------- 计算权重和索引 --------------------------
        _, indices = torch.topk(scores_for_topk, self.topk, dim=-1)
        weights = scores.gather(dim=1, index=indices)
        weights_sum = weights.sum(dim=-1, keepdim=True)
        weights = weights / (weights_sum + 1e-6)
        weights = weights * self.route_scale
            
        return weights.type_as(x), indices

scores_logits = F.linear(x, self.weight, None)
scores = scores_logits.sigmoid()

self.original_scores = scores
scores_for_topk = scores

if self.bias is not None:
    scores_for_topk = scores_for_topk + self.bias

scores_view = scores_for_topk.view(x.size(0), self.n_groups, -1)

if self.bias is None:
    group_scores = scores_view.amax(dim=-1)
else:
    group_scores = scores_view.topk(2, dim=-1)[0].sum(dim=-1)

indices_groups = group_scores.topk(self.topk_groups, dim=-1)[1]

mask = torch.ones(x.size(0), self.n_groups, dtype=torch.bool, device=x.device)
mask.scatter_(dim=1, index=indices_groups, value=False)

scores_for_topk = scores_view.masked_fill(mask=mask.unsqueeze(-1), value=float("-inf")).flatten(1)

_, indices = torch.topk(scores_for_topk, self.topk, dim=-1)

weights = scores.gather(dim=1, index=indices)
weights_sum = weights.sum(dim=-1, keepdim=True)
weights = weights / (weights_sum + 1e-6)
weights = weights * self.route_scale

class MoE(nn.Module):
    """
    Mixture-of-Experts (MoE) 混合专家模块

    Attributes:
        dim (int): 嵌入维度
        n_routed_experts (int): 路由专家数量
        n_activated_experts (int): 每个输入激活的专家数
        gate (nn.Module): 门控机制
        experts (nn.ModuleList): 专家列表
        shared_experts (nn.Module): 共享专家
        use_seq_aux (bool): 是否使用序列级别的辅助损失
        seq_aux_alpha (float): 序列级别的辅助损失的权重
        bias_update_speed (float): 偏置更新速度
    """
    def __init__(self, args: DeepSeekV3ModelArgs):
        """
        MoE 初始化

        Args:
            args (ModelArgs): 模型配置参数
        """
        super().__init__()
        self.dim = args.dim
        self.n_routed_experts = args.n_routed_experts
        self.n_activated_experts = args.n_activated_experts
        self.gate = Gate(args)
        self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) for _ in range(self.n_routed_experts)])
        self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)
        self.use_seq_aux = args.use_seq_aux
        self.seq_aux_alpha = args.seq_aux_alpha
        self.bias_update_speed = args.bias_update_speed

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        MoE 前向传播

        Args:
            x (torch.Tensor): 输入张量 (batch_size, seq_len, dim)

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: 输出张量 (batch_size, seq_len, dim), 序列级辅助损失, 全局负载情况
        """
        # -------------------- 变量准备 --------------------
        shape = x.size()
        bsz, seqlen = shape[:2]
        x = x.view(-1, self.dim)
        weights, indices = self.gate(x)
        y = torch.zeros_like(x)
        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts)
        global_counts = None

        # -------------------- 无辅助损失负载均衡策略 --------------------
        if self.gate.bias is not None and self.training:
            global_counts = counts.clone()
            if dist.is_available() and dist.is_initialized():
                dist.all_reduce(global_counts, op=dist.ReduceOp.SUM)
            avg_count = sum(global_counts).float() / self.n_routed_experts
            is_distributed_and_master = dist.is_initialized() and dist.get_rank() == 0
            is_not_distributed = not dist.is_initialized()
            if is_distributed_and_master or is_not_distributed:
                for i, count in enumerate(global_counts):
                    error = avg_count - count
                    self.gate.bias.data[i] += self.bias_update_speed * torch.sign(error)
            if dist.is_available() and dist.is_initialized():
                dist.broadcast(self.gate.bias.data, src=0)

        # -------------------- 序列级别的辅助损失 --------------------
        if self.use_seq_aux and self.training:
            scores_for_seq_aux = self.gate.original_scores.view(bsz, seqlen, -1)
            scores_for_seq_aux = scores_for_seq_aux / scores_for_seq_aux.sum(dim=-1, keepdim=True)
            P_i = scores_for_seq_aux.mean(dim=1) 
            f_i = F.one_hot(indices.view(bsz, -1), num_classes=self.n_routed_experts)
            f_i = f_i.sum(dim=1)
            f_i = (f_i * self.n_routed_experts) / (self.n_activated_experts * seqlen)
            seq_aux_loss = (f_i * P_i).sum() * self.seq_aux_alpha
        else:
            seq_aux_loss = None

        # -------------------- 计算专家输出 --------------------
        for i in range(self.n_routed_experts):
            if counts[i] == 0:
                continue
            expert = self.experts[i]
            idx, top = torch.where(indices == i)
            y[idx] += expert(x[idx]) * weights[idx, top, None]
        z = self.shared_experts(x)
 
        return (y + z).view(shape), seq_aux_loss, global_counts

shape = x.size()
bsz, seqlen = shape[:2]
x = x.view(-1, self.dim)
weights, indices = self.gate(x)
y = torch.zeros_like(x)
global_counts = None

counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts)

global_counts = counts.clone()
if dist.is_available() and dist.is_initialized():
    dist.all_reduce(global_counts, op=dist.ReduceOp.SUM)
avg_count = sum(global_counts).float() / self.n_routed_experts

is_distributed_and_master = dist.is_initialized() and dist.get_rank() == 0
is_not_distributed = not dist.is_initialized()
if is_distributed_and_master or is_not_distributed:
    for i, count in enumerate(global_counts):
        error = avg_count - count
        self.gate.bias.data[i] += self.bias_update_speed * torch.sign(error)
if dist.is_available() and dist.is_initialized():
    dist.broadcast(self.gate.bias.data, src=0)

scores_for_seq_aux = self.gate.original_scores.view(bsz, seqlen, -1)
scores_for_seq_aux = scores_for_seq_aux / scores_for_seq_aux.sum(dim=-1, keepdim=True)
P_i = scores_for_seq_aux.mean(dim=1)

f_i = F.one_hot(indices.view(bsz, -1), num_classes=self.n_routed_experts)
f_i = f_i.sum(dim=1)
f_i = (f_i * self.n_routed_experts) / (self.n_activated_experts * seqlen)
seq_aux_loss = (f_i * P_i).sum() * self.seq_aux_alpha

for i in range(self.n_routed_experts):
    if counts[i] == 0:
        continue
    expert = self.experts[i]
    idx, top = torch.where(indices == i)
    y[idx] += expert(x[idx]) * weights[idx, top, None]

z = self.shared_experts(x)

class Block(nn.Module):
    """
    Transformer Block, 包括 Attention 和 Feed-Forward 部分

    Attributes:
        attn (nn.Module): 注意力层 (MLA)
        ffn (nn.Module): 前馈网络 (MoE)
        attn_norm (nn.Module): 注意力层的 Layer Normalization
        ffn_norm (nn.Module): 前馈网络的 Layer normalization
    """
    def __init__(self, layer_id: int, args: DeepSeekV3ModelArgs):
        """
        初始化 Transformer Block

        Args:
            layer_id (int): Transformer 的层索引
            args (ModelArgs): 模型配置参数
        """
        super().__init__()
        self.attn = MLA(args)
        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)
        self.attn_norm = RMSNorm(args.dim)
        self.ffn_norm = RMSNorm(args.dim)

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        Transformer Block 的前向传播

        Args:
            x (torch.Tensor): 输入 (batch_size, seq_len, dim)
            start_pos (int): 用于指定当前推理步骤的起始位置，即从序列的哪个位置开始计算
            freqs_cis (torch.Tensor): 预先计算的复数 RoPE 矩阵
            mask (Optional[torch.Tensor]): 掩码

        Returns:
            Tuple[torch.Tensor, Optional[torch.Tensor]]: 输出 (batch_size, seq_len, dim)
        """
        x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask)
        if isinstance(self.ffn, MoE):
            h, seq_aux_loss, global_counts = self.ffn(self.ffn_norm(x))
            return x + h, seq_aux_loss, global_counts
        x = x + self.ffn(self.ffn_norm(x))
        return x, None, None

class MTP(nn.Module):
    """
    多token预测(Multi-Token Prediction, MTP)

    Attributes:
        args (ModelArgs): 模型配置参数
        embed (nn.Module): 嵌入层
        head (nn.Module): 输出投影
        h_norm (nn.Module): 对上一个 MTP 模块或主模型输出的 hidden state 应用的 Layer Normalization
        x_norm (nn.Module): 对本 MTP 模块输入应用的 Layer Normalization
        output_norm (nn.Module): 对本 MTP 模块输出应用的 Layer Normalization
        linear_proj (nn.Module): 线性投影层，用于将 MTP 模块的两个输入拼接后进行投影
        transformer_block (nn.Module): Transformer Block
    """
    def __init__(self, args: DeepSeekV3ModelArgs, embed: nn.Module, head: nn.Module):
        """
        初始化 MTP

        Args:
            args (ModelArgs): 模型配置参数
            embed (nn.Module): 嵌入层, 来自于 Transformer 共享
            head (nn.Module): 输出投影, 来自于 Transformer 共享
        """
        super().__init__()
        self.embed = embed
        self.head = head
        self.h_norm = RMSNorm(args.dim)
        self.x_norm = RMSNorm(args.dim)
        self.output_norm = RMSNorm(args.dim)
        self.linear_proj = nn.Linear(2 * args.dim, args.dim, bias=False)
        self.transformer_block = Block(0, args)
    
    def forward(self, x: torch.Tensor, h: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        MTP 的前向传播

        Args:
            x (torch.Tensor): 输入 token_ids (batch_size, mtp_seq_len)
            h (torch.Tensor): 上一个 MTP 模块或主模型的输出 (batch_size, mtp_seq_len, dim)
            start_pos (int): 用于指定当前推理步骤的起始位置，即从序列的哪个位置开始计算
            freqs_cis (torch.Tensor): 预先计算的复数 RoPE 矩阵

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: 输出 Logits (batch_size, vocab_size) 和给下一个 MTP 使用的 h (batch_size, mtp_seq_len, dim)
        """
        seqlen = x.size(1)
        mask = None

        x = self.embed(x)
        x = self.x_norm(x)
        h = self.h_norm(h)
        x = self.linear_proj(torch.cat([x, h], dim=-1))

        if seqlen > 1:
            mask = torch.full((seqlen, seqlen), float("-inf"), device=x.device).triu_(1)
        x, _, _ = self.transformer_block(x, start_pos, freqs_cis, mask)

        h = x
        x = self.output_norm(x)
        logits = self.head(x)

        return logits, h

class DeepSeekV3Model(BaseModel):
    """
    Transformer 整体模型

    Attributes:
        max_seq_len (int): 最大序列长度
        embed (nn.Module): 嵌入层
        layers (torch.nn.ModuleList): Transformer Blocks 列表
        norm (nn.Module): 在最后一个 Transformer Block 后应用 Layer Normalization
        head (nn.Module): 输出投影
        freqs_cis (torch.Tensor): 预先计算的复数 RoPE 矩阵
        use_mtp (bool): 是否使用 MTP
        mtp_loss_lambda (float): MTP 损失的权重
        MTP (torch.nn.ModuleList): MTP 模块
    """
    model_name = "mini_deepseekv3"

    def __init__(self, args: DeepSeekV3ModelArgs):
        """
        初始化 Transformer

        Args:
            args (ModelArgs): 模型配置参数
        """
        super().__init__()
        self.max_seq_len = args.max_seq_len
        self.embed = nn.Embedding(args.vocab_size, args.dim)
        self.layers = torch.nn.ModuleList()
        for layer_id in range(args.n_layers):
            self.layers.append(Block(layer_id, args))
        self.norm = RMSNorm(args.dim)
        self.head = nn.Linear(args.dim, args.vocab_size, bias=False)
        self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False)
        self.use_mtp = args.use_mtp
        self.mtp_loss_lambda = args.mtp_loss_lambda
        self.mtp =MTP(args, self.embed, self.head)
    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor = None, start_pos: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播

        Args:
            input_ids (torch.Tensor): 输入张量, 内容为 token_ids (batch_size, seq_len)
            targets (torch.Tensor): 目标张量, 内容为 token_ids (batch_size, seq_len)
            start_pos (int): 起始位置, 默认为 0

        Returns:
            torch.Tensor: 输出 Logits (batch_size, vocab_size), Loss
        """
        # ----------------------- 变量初始化 -----------------------
        seqlen = input_ids.size(1)
        h = self.embed(input_ids)
        freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
        mask = None
        main_loss = 0.0
        total_seq_aux_loss = 0.0
        mtp_loss = 0.0
        all_global_counts = []

        # ----------------------- 主模型部分 -----------------------
        if seqlen > 1:
            mask = torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device).triu_(1)
        for layer in self.layers:
            h, seq_aux_loss, global_counts = layer(h, start_pos, freqs_cis, mask)
            if seq_aux_loss is not None:
                total_seq_aux_loss += seq_aux_loss
            if global_counts is not None:
                all_global_counts.append(global_counts)
        
        h_for_mtp = h
        h = self.norm(h)
        logits = self.head(h)
        if self.training:
            main_loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="mean")

        # ----------------------- MTP 部分 -----------------------
        if self.use_mtp and self.training:
            mtp_logits, _ = self.mtp(input_ids[:, 1:], h_for_mtp[:, :-1], 0, self.freqs_cis[0:seqlen-1])
            targets_for_mtp = targets[:, 1:]
            mtp_loss = F.cross_entropy(mtp_logits.reshape(-1, mtp_logits.size(-1)), targets_for_mtp.reshape(-1), reduction="mean")
        
        # ----------------------- 计算总损失 -----------------------
        loss = main_loss + total_seq_aux_loss + self.mtp_loss_lambda * mtp_loss

        return logits, loss, (main_loss, total_seq_aux_loss, self.mtp_loss_lambda * mtp_loss, all_global_counts)

if seqlen > 1:
    mask = torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device).triu_(1)

for layer in self.layers:
    h, seq_aux_loss, global_counts = layer(h, start_pos, freqs_cis, mask)
    if seq_aux_loss is not None:
        total_seq_aux_loss += seq_aux_loss
    if global_counts is not None:
        all_global_counts.append(global_counts)

h_for_mtp = h
h = self.norm(h)
logits = self.head(h)
if self.training:
    main_loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="mean")

if self.use_mtp and self.training:
    mtp_logits, _ = self.mtp(input_ids[:, 1:], h_for_mtp[:, :-1], 0, self.freqs_cis[0:seqlen-1])
    targets_for_mtp = targets[:, 1:]
    mtp_loss = F.cross_entropy(mtp_logits.reshape(-1, mtp_logits.size(-1)), targets_for_mtp.reshape(-1), reduction="mean")

loss = main_loss + total_seq_aux_loss + self.mtp_loss_lambda * mtp_loss

return logits, loss, (main_loss, total_seq_aux_loss, self.mtp_loss_lambda * mtp_loss, all_global_counts)

max_batch_size: int = 16
max_seq_len: int = 512
vocab_size: int = -1  # 加载模型时传入
dim: int = 768
inter_dim: int = 3072
moe_inter_dim: int = 512
n_layers: int = 12
n_dense_layers: int = 3
n_heads: int = 12

# moe
n_routed_experts: int = 8
n_shared_experts: int = 1
n_activated_experts: int = 2
n_expert_groups: int = 4
n_limited_groups: int = 2
route_scale: float = 1.
use_noaux_tc: bool = True
bias_update_speed: float = 0.001
use_seq_aux: bool = True
seq_aux_alpha: float = 0.0001

# mla
q_lora_rank: int = 384  # 源码中若 q_lora_rank=0, 则不使用下投影，这里我们使用下投影，并略去若 q_lora_rank=0 的逻辑
kv_lora_rank: int = 256  # 论文中 d_c = 4 * d_h
qk_nope_head_dim: int = 64  # d_h
qk_rope_head_dim: int = 32  # 论文中 d_h^R = d_h / 2
v_head_dim: int = 64

# RoPE
rope_theta: float = 10000.0

# MTP
use_mtp: bool = True
mtp_loss_lambda: float = 0.0001