From fd6448ce6153cac2699c697118255680b4325d7b Mon Sep 17 00:00:00 2001 From: Zhenyu Wang Date: Fri, 15 Aug 2025 14:45:05 +0800 Subject: [PATCH] feat: add explanation --- model/model_minimind.py | 10 +++++++++- trainer/train_pretrain.py | 9 ++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 4b43610..4eab595 100644 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -133,6 +133,7 @@ class Attention(nn.Module): self.n_local_kv_heads = self.num_key_value_heads self.n_rep = self.n_local_heads // self.n_local_kv_heads self.head_dim = args.hidden_size // args.num_attention_heads + # 把所有头权重都映射到同一个维度，方便并行运算 self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) @@ -151,6 +152,7 @@ class Attention(nn.Module): attention_mask: Optional[torch.Tensor] = None): bsz, seq_len, _ = x.shape xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x) + # 分割并行计算后的结果，还原至各个注意力头 xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) @@ -164,6 +166,7 @@ class Attention(nn.Module): xv = torch.cat([past_key_value[1], xv], dim=1) past_kv = (xk, xv) if use_cache else None + # GQA时可能需要多个Q 共用 KV 矩阵 xq, xk, xv = ( xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), @@ -179,13 +182,18 @@ class Attention(nn.Module): output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True) else: + # 计算注意力分数 scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) + # 添加三角掩码，防止模型看到未来的信息（在训练和推理阶段都 mask）。这是实现 kv cache 的关键 scores = scores + torch.triu( torch.full((seq_len, seq_len), float("-inf"), device=scores.device), diagonal=1 ).unsqueeze(0).unsqueeze(0) # scores+mask if attention_mask is not None: + # attention mask 标记了输入token，并把有意义的 token 值设置为 1 + # 掩盖 pad 部分的注意力分数，防止模型关注空白token + # 掩盖的结果就是把 pad 部分的分数设置为极小值，从而在 softmax 时忽略这些 token extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) extended_attention_mask = (1.0 - extended_attention_mask) * -1e9 scores = scores + extended_attention_mask @@ -384,7 +392,7 @@ class MiniMindModel(nn.Module): start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0 hidden_states = self.dropout(self.embed_tokens(input_ids)) - + # 在模型顶部计算一次位置编码，之后传给每个 Attention 层 position_embeddings = ( self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length] diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index f0cfc58..da25375 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -47,17 +47,24 @@ def train_epoch(epoch, wandb): res.logits.view(-1, res.logits.size(-1)), Y.view(-1) ).view(Y.size()) + # 实际句子长度通常小于max_length，pad部分的loss_mask是0，不需要计算损失 loss = (loss * loss_mask).sum() / loss_mask.sum() + # 添加辅助损失，用于训练专家选择 loss += res.aux_loss + # 梯度累积（/累计步数） loss = loss / args.accumulation_steps + # 放大 Loss，因为 loss 高精度而梯度低精度，防止极小的梯度丢失。 + # 相当于链式法则 ∂L_scaled/∂W₁ = ∂(k×L)/∂W₁ = k × ∂L/∂W₁，k是缩放因子 scaler.scale(loss).backward() if (step + 1) % args.accumulation_steps == 0: + # 缩小梯度 scaler.unscale_(optimizer) + # 梯度裁剪，防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - scaler.step(optimizer) + # 更新缩放因子 scaler.update() optimizer.zero_grad(set_to_none=True)