LoRA 线性层模块 将权重更新分解为两个低秩矩阵的乘积: ΔW = B @ A 前向传播: y = W @ x + (B @ A) @ x * scaling 优势: - 大幅减少可训练参数（rank << min(in_features, out_f

混合专家模型模块 结构： 1. Router: 计算每个 token 对每个专家的得分 2. Top-K 选择: 选择得分最高的 K 个专家 3. 专家计算: 被选中的专家对 token 进行处理 4. 加权融合: 根据路由得分加权融合专家输出

多头注意力模块 支持自注意力（Self-Attention）和交叉注意力（Cross-Attention）： - 自注意力：Q = K = V = x_query - 交叉注意力：Q = x_query, K = V = x_context Args:

旋转位置编码模块 通过将查询和键向量按照位置进行旋转，使注意力分数包含相对位置信息。 公式: f(x, m) = x * cos(m*θ) + rotate_half(x) * sin(m*θ) Args: head_dim: 旋转编码的维度（通

缩放点积注意力模块 计算查询和键的点积，除以缩放因子后应用 softmax 得到注意力权重， 最后用注意力权重对值进行加权求和。 Args: dropout_p: Dropout 概率，默认 0.0

标准前馈神经网络模块 由两个线性层组成，中间使用 ReLU 激活函数。 intermediate_dim 通常是 model_dim 的 4 倍。 Args: model_dim: 输入/输出维度 intermediate_dim

分组查询注意力模块 在 GQA 中： - Q 有 num_heads 个头 - K 和 V 只有 num_kv_heads 个头（num_kv_heads < num_heads） - 多个 Q 头共享同一个 KV 头（通过 repeat_kv 实现）

层归一化模块 公式: LayerNorm(x) = (x - mean) / sqrt(var + eps) * gamma + beta Args: model_dim: 归一化的特征维度 eps: 数值稳定性常数，防止除零，默认 1e

多头潜在注意力模块 通过低秩投影将 Q 和 KV 压缩到潜在空间，显著减少 KV 缓存的显存占用。 同时结合 RoPE 位置编码保持位置感知能力。 Args: model_dim: 模型隐藏维度 num_heads: 注意力头数

预训练损失模块 计算因果语言模型的交叉熵损失。 通过预测下一个词来训练语言模型。 Args: ignore_index: 忽略的标签索引，不计入损失计算，默认 -100

RMS 归一化模块 公式: RMSNorm(x) = x / RMS(x) * gamma 其中 RMS(x) = sqrt(mean(x^2) + eps) 相比 LayerNorm，RMSNorm 不计算均值，计算量更小。 Args:

监督微调损失模块 在 SFT 阶段，我们通常只希望计算 response 部分的损失， 而不计算 prompt 部分的损失。该模块支持通过 prompt_lengths 来屏蔽 prompt。 Args: 无

SwiGLU 前馈神经网络模块 结构: output = Down(Gate(x) * Up(x)) 其中 Gate 使用 SiLU (Swish) 激活函数，实现门控机制。 相比传统 FFN 的 ReLU 激活，SwiGLU 的门控机制可以更好地捕捉复杂的非线