Minimind 系列笔记 01：模型结构精读

源码：jingyaogong/minimind（model/model_minimind.py，全文 287 行）。本文逐模块拆解关键片段，完整的源码放在各节的"完整源码"折叠块里，需要对照时展开即可。

完整源码：model_minimind.py（287 行全文）

python

import math, torch, torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
from transformers.modeling_outputs import MoeCausalLMOutputWithPast

class MiniMindConfig(PretrainedConfig):
    model_type = "minimind"
    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.use_moe = use_moe
        self.dropout = kwargs.get("dropout", 0.0)
        self.vocab_size = kwargs.get("vocab_size", 6400)
        self.bos_token_id = kwargs.get("bos_token_id", 1)
        self.eos_token_id = kwargs.get("eos_token_id", 2)
        self.flash_attn = kwargs.get("flash_attn", True)
        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
        self.hidden_act = kwargs.get("hidden_act", 'silu')
        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
        self.rope_theta = kwargs.get("rope_theta", 1e6)
        self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
        self.rope_scaling = {
            "beta_fast": 32, "beta_slow": 1, "factor": 16,
            "original_max_position_embeddings": 2048, "attention_factor": 1.0, "type": "yarn"
        } if self.inference_rope_scaling else None
        # MoE specific configs (ignored if use_moe = False)
        self.num_experts = kwargs.get("num_experts", 4)
        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return (self.weight * self.norm(x.float())).type_as(x)

def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
    freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
    if rope_scaling is not None:  # YaRN: f'(i) = f(i)((1-γ) + γ/s), γ∈[0,1] is linear ramp
        orig_max, factor, beta_fast, beta_slow, attn_factor = (
            rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
            rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
        )
        if end / orig_max > 1.0:
            inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
            low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
            ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
            freqs = freqs * (1 - ramp + ramp / factor)
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
    return freqs_cos, freqs_sin

def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
    q_embed = ((q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))).to(q.dtype)
    k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
    return q_embed, k_embed

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    bs, slen, num_key_value_heads, head_dim = x.shape
    if n_rep == 1: return x
    return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))

class Attention(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
        self.n_local_heads = config.num_attention_heads
        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 = config.head_dim
        self.is_causal = True
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn

    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=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)
        xq, xk = self.q_norm(xq), self.k_norm(xk)
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        past_kv = (xk, xv) if use_cache else None
        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
        if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
        else:
            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
            if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
        output = self.resid_dropout(self.o_proj(output))
        return output, past_kv

class FeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig, intermediate_size: int = None):
        super().__init__()
        intermediate_size = intermediate_size or config.intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

class MOEFeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.view(-1, hidden_dim)
        scores = F.softmax(self.gate(x_flat), dim=-1)
        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
        if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        y = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
            if mask.any():
                token_idx = mask.any(dim=-1).nonzero().flatten()
                weight = topk_weight[mask].view(-1, 1)
                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
            elif self.training:
                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
        if self.training and self.config.router_aux_loss_coef > 0:
            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
            self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
        else:
            self.aux_loss = scores.new_zeros(1).squeeze()
        return y.view(batch_size, seq_len, hidden_dim)

class MiniMindBlock(nn.Module):
    def __init__(self, layer_id: int, config: MiniMindConfig):
        super().__init__()
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        residual = hidden_states
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states), position_embeddings,
            past_key_value, use_cache, attention_mask
        )
        hidden_states += residual
        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
        return hidden_states, present_key_value

class MiniMindModel(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.dropout = nn.Dropout(config.dropout)
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        batch_size, seq_length = input_ids.shape
        if hasattr(past_key_values, 'layers'): past_key_values = None
        past_key_values = past_key_values or [None] * len(self.layers)
        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))
        # Recompute RoPE buffers lost during meta-device init (transformers>=5.x)
        if self.freqs_cos[0, 0] == 0:
            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
            self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
        position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
        presents = []
        for layer, past_key_value in zip(self.layers, past_key_values):
            hidden_states, present = layer(
                hidden_states,
                position_embeddings,
                past_key_value=past_key_value,
                use_cache=use_cache,
                attention_mask=attention_mask
            )
            presents.append(present)
        hidden_states = self.norm(hidden_states)
        aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
        return hidden_states, presents, aux_loss

class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniMindConfig
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    def __init__(self, config: MiniMindConfig = None):
        self.config = config or MiniMindConfig()
        super().__init__(self.config)
        self.model = MiniMindModel(self.config)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        if self.config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight
        self.post_init()

    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, logits_to_keep=0, labels=None, **kwargs):
        hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])
        loss = None
        if labels is not None:
            x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
            loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
        return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states)

    @torch.inference_mode()
    def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
        input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
        past_key_values = kwargs.pop("past_key_values", None)
        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
        if streamer: streamer.put(input_ids.cpu())
        for _ in range(max_new_tokens):
            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
            outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
            attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
            logits = outputs.logits[:, -1, :] / temperature
            if repetition_penalty != 1.0:
                for i in range(input_ids.shape[0]):
                    seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
            if top_k > 0:
                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
                mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
            if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
            past_key_values = outputs.past_key_values if use_cache else None
            if streamer: streamer.put(next_token.cpu())
            if eos_token_id is not None:
                finished |= next_token.squeeze(-1).eq(eos_token_id)
                if finished.all(): break
        if streamer: streamer.end()
        if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
        return input_ids

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277

MiniMind 的价值在于"麻雀虽小五脏俱全"：它把 Qwen3 同款的现代 LLM 结构——RMSNorm、RoPE、GQA、SwiGLU、MoE、权重共享——全部用纯 PyTorch 手写了一遍，没有任何一行调用第三方模型库的抽象接口。读懂这 287 行，再去看 Llama / Qwen 的官方实现，会发现只是同样的东西换了排版。

本文逐模块过一遍代码，每个模块先看实现、再补原理。训练流程见 03-训练全流程，多模态结构见 04-多模态VLM与Omni。

一、全景：Config 与参数量

python

class MiniMindConfig(PretrainedConfig):
    model_type = "minimind"
    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.use_moe = use_moe
        self.vocab_size = 6400
        self.num_attention_heads = 8        # q_heads
        self.num_key_value_heads = 4        # kv_heads  ← GQA
        self.head_dim = hidden_size // num_attention_heads   # 96
        self.hidden_act = 'silu'
        self.intermediate_size = math.ceil(hidden_size * math.pi / 64) * 64   # 2432
        self.max_position_embeddings = 32768
        self.rope_theta = 1e6
        self.tie_word_embeddings = True     # 权重共享
        # MoE 专属
        self.num_experts = 4
        self.num_experts_per_tok = 1        # top-1 routing
        ...

主线 minimind-3 的 Dense 配置：

项	值	说明
params	64M	GPT-3 的 1/2700
vocab	6400	自训 BPE，见 02-Tokenizer与数据
layers / d_model	8 / 768	矮胖 vs 瘦高，见下
q_heads / kv_heads	8 / 4	GQA，KV Cache 省一半
head_dim	96	768 / 8
intermediate_size	2432	`ceil(768·π/64)·64`
max_pos / rope_theta	32768 / 1e6	支持 YaRN 外推

两个值得展开的细节：

为什么 intermediate_size 和 π 有关？ SwiGLU 论文（Shazeer 2020）的实验结论：把 FFN 的参数量固定时，激活维度取 $\frac{2}{3} \times 4 d \approx 2.67 d$ 效果最好。而 $\frac{2}{3} \cdot 4 = 8 / 3 \approx 2.67 \approx π$ ……更准确地说这是社区沿用的一个经验值：8/3 · d 再向上取整到 64 的倍数（对齐硬件 tile），768 × π / 64 = 37.7 → 38 × 64 = 2432。Llama 系列也是这个套路。

为什么 8 层 / 768 维？ MobileLLM 的系统研究：参数量固定时，深度比宽度重要，"瘦高"比"矮胖"更容易学到抽象概念。但"窄"有下限——d_model < 512 时词嵌入维度过窄的劣势会放大，且固定 q_heads 下 head_dim 偏小会伤害注意力质量。768/8 是训练效率、稳定性与效果的折中（MiniMind 作者做过 d_model × n_layers 的分配实验，结论与 MobileLLM 一致）。

二、RMSNorm：比 LayerNorm 少做一件事

python

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return (self.weight * self.norm(x.float())).type_as(x)

原理回顾。LayerNorm 对每个 token 的特征向量做三件事：

LayerNorm (x) = γ ⊙ \frac{x - μ}{\sqrt{σ^{2} + ϵ}} + β

RMSNorm（Root Mean Square Normalization）观察到：LayerNorm 的收益主要来自缩放不变性，减均值那一步可以省掉。于是只保留均方根归一化：

RMSNorm (x) = γ ⊙ \frac{x}{\sqrt{\frac{1}{d} \sum_{i = 1}^{d} x_{i}^{2} + ϵ}}

少算一次 mean、一次减法，大模型场景下这点开销的节省是可观的（每层 2 个 norm × 几十层 × 每秒上万 token）；
实验表明效果与 LayerNorm 相当，所以 Llama 之后的主流结构全线采用。

两处实现细节值得注意：

x.float() —— 归一化在 fp32 下算，再 .type_as(x) 转回训练精度。均值方差这种"对整个向量做 reduction"的统计量在 bf16/fp16 下误差大（平方和容易溢出/下溢），这是混合精度训练的标准防御姿势，后面 attention 的 softmax 同样如此。
eps 在分母里而不是分子，且默认 1e-5（config 传入 rms_norm_eps=1e-6），数值上避免除零。

三、RoPE：旋转位置编码

3.1 动机：让注意力天然感知"相对距离"

注意力本身是置换不变的——不打位置信息，"猫咬狗"和"狗咬猫"对它完全一样。经典方案有两类：

绝对位置编码（原始 Transformer / GPT-2）：把位置 $m$ 编码成向量直接加到 token embedding 上。简单，但外推性差，且注意力只能间接地从加性偏置里"猜"相对距离。
相对位置编码（T5 等）：在 attention score 里直接加 $m - n$ 的偏置。表达相对距离直接，但实现侵入 attention 内部、且 KV Cache 不友好。

RoPE（Rotary Position Embedding）的思路非常漂亮：把 $q, k$ 向量视为复数，按位置 $m$ 旋转 $m θ$ 角度。这样：

⟨ R_{m} q, R_{n} k ⟩ = ⟨ q, R_{n - m} k ⟩

两个位置各自旋转后再做内积，结果只依赖相对位移 $m - n$ ——绝对位置操作，相对位置效果，attention 结构零侵入，KV Cache 完全兼容。

3.2 实现

python

def precompute_freqs_cis(dim, end=32*1024, rope_base=1e6, rope_scaling=None):
    # 每两个维度共用一个频率：freq_i = 1 / theta^(2i/d)
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()          # (end, dim//2) 位置×频率
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin                    # (end, dim)

def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    def rotate_half(x):
        return torch.cat((-x[..., x.shape[-1]//2:], x[..., :x.shape[-1]//2]), dim=-1)
    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed, k_embed

三个细节：

频率递减：第 $i$ 对维度的频率是 $θ^{- 2 i / d}$ 。低维转得快（感受局部位置），高维转得慢（感受全局位置）——类似不同周期的钟表指针组合出唯一的"时刻"。rope_base=1e6（Llama 2 用 1e4）把高频分量压得更慢，直接利好长上下文。
为什么 cos/sin 要 cat 一遍？ 这是"旋转"的实数化写法。对二维子空间 $[x_{1}, x_{2}]$ 旋转 $ϕ$ 角，标准形式是 $(\begin{matrix} \cos ϕ & - \sin ϕ \\ \sin ϕ & \cos ϕ \end{matrix}) (\begin{matrix} x_{1} \\ x_{2} \end{matrix})$ 。把整个向量按"前后两半"重排后（rotate_half 把 $[a, b]$ 变成 $[- b, a]$ ，即 90° 旋转），乘上"复制两半"的 cos/sin，恰好等价于每对维度独立做二维旋转。HF Llama 系列就是这个写法。
只在 embedding 层算一次位置：MiniMindModel.forward 里把 freqs_cos[start_pos : start_pos+seq_len] 切片后传给每一层——所有层共享同一套位置信息，不用每层重复计算。

3.3 YaRN：免训练的长文本外推

训练长度 2048，推理想上 32768？RoPE 的麻烦在于没见过的位置角度会让注意力崩掉。YaRN 的思路：不同频率分量区别对待——高频（转得快、早就"见过"各种角度）保持不动；低频（转得慢、外推时进入陌生角度区）按插值压缩：

f^{'} (i) = f (i) \cdot ((1 - γ) + \frac{γ}{s}), γ \in [0, 1]

其中 $s$ 是扩展倍数（factor=16）， $γ$ 是随频率变化的线性 ramp：高频 $γ \to 0$ （不缩），低频 $γ \to 1$ （缩到 $1 / s$ ，即把"未来"的 32768 个位置压回模型见过的 2048 范围内插值）。

python

if rope_scaling is not None:
    # beta_fast / beta_slow：标定"多快算高频/多慢算低频"的边界
    low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim//2 - 1)
    ramp = torch.clamp((torch.arange(dim//2).float() - low) / max(high - low, 0.001), 0, 1)
    freqs = freqs * (1 - ramp + ramp / factor)

inv_dim(b) 是从"波长"反解维度下标的解析式：波长 $λ (i) = 2 π / f (i)$ ， $β$ 指定边界波长，inv_dim 算出对应的维度索引。中间频率在 $[0, 1 / s]$ 之间线性过渡。

MiniMind 把 YaRN 做成推理时开关（inference_rope_scaling），训练完全不动——这就是"免训练外推"。切换时只需用带 scaling 的参数重算一遍 cos/sin 表。

四、Attention：GQA 与 KV Cache

4.1 Q/K/V 投影与 QK-Norm

python

self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False)  # 768→768
self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)  # 768→384
self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False)  # 768→384
self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False)
self.q_norm = RMSNorm(head_dim, eps=config.rms_norm_eps)
self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps)

两个知识点：

QK-Norm：对 Q/K 的每个 head 维度再做一次 RMSNorm（注意是对 head_dim=96 归一化，不是对整个 hidden）。动机是抑制 attention logits 的数值暴涨（大训练时长下个别 head 的 score 会异常大，softmax 进入饱和区、梯度消失）。Qwen3 等新模型都加了。
bias=False：所有线性层全部去 bias。现代 LLM 的惯例——实验表明去掉对效果几乎无影响，还省参数、省显存、省一点计算。

4.2 GQA：为什么 4 个 KV head 够 8 个 Q head 用

多头注意力（MHA）里 Q/K/V 都是 8 个 head。但推理时 KV Cache 要为每个 KV head 存完整的序列：8 头 × 96 维 × fp16 × 每层每 token = 8×96×2×8层 = 12KB/token。上下文越长这块越肥，decoding 完全变成显存带宽瓶颈。

**分组查询注意力（GQA）**的观察：多个 Q head 期待的 K/V 信息高度重叠，何必人手一份？让 4 个 KV head 服务 8 个 Q head，每 2 个 Q head 共享一组 KV：

python

self.n_rep = num_attention_heads // num_key_value_heads    # 2

def repeat_kv(x, n_rep):    # (B, S, 4, 96) → (B, S, 8, 96)
    bs, slen, n_kv, head_dim = x.shape
    if n_rep == 1: return x
    return x[:, :, :, None, :].expand(bs, slen, n_kv, n_rep, head_dim) \
                              .reshape(bs, slen, n_kv * n_rep, head_dim)

KV head 数减半 → KV Cache 直接省一半，质量几乎无损（MQA 是极端版：1 个 KV head 服务所有 Q，省得最多但掉点）。Qwen3 / Llama 3 全线采用。

4.3 KV Cache：增量解码的关键

训练时整段序列一次 forward，无所谓 cache。推理自回归时如果每生成一个 token 都重算全部历史，复杂度是 $O (n^{2})$ 次 token 计算的重复浪费。解法：把每层的 K、V 缓存下来，新 token 只算自己的 $q_{t}, k_{t}, v_{t}$ 拼上去：

python

xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
xq, xk = self.q_norm(xq), self.k_norm(xk)
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
if past_key_value is not None:
    xk = torch.cat([past_key_value[0], xk], dim=1)   # 沿 seq 维拼接
    xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None

注意 RoPE 在拼 cache 之前施加——cos/sin 表按 start_pos（缓存长度）切片，新 token 拿到的是自己在全序列里的绝对位置角度，缓存里的旧 K 已经转过，无需重转。

模型主 forward 里 start_pos = past_key_values[0][0].shape[1] 就是从缓存反推当前位置——一行代码同时服务了位置切片和增量拼接两件事。

4.4 Flash Attention：两条计算路径

python

self.flash = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn

if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) \
        and (attention_mask is None or torch.all(attention_mask == 1)):
    output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=..., is_causal=True)
else:
    scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
    if self.is_causal:
        scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf")).triu(1)
    if attention_mask is not None:
        scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
    output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv

手写路径正是教科书公式 $softmax (\frac{Q K^{⊤}}{\sqrt{d_{k}}}) V$ ，但有三处巧思：

causal mask 只盖住最后 seq_len 列：scores 的形状是 (B, H, q_len, kv_len)。有 cache 时 kv_len > q_len（历史部分天然 causal，不该遮）。[:, :, :, -seq_len:] 精确地把上三角 -inf 只加在"新 token 之间"，历史区域原样保留。这是手写 attention 对 KV Cache 的关键配合，很多复现翻车就在这里。
padding mask 用 -1e9 而不是 -inf：-inf 在极端情况（整行被 mask）下 softmax 会产生 NaN，-1e9 平滑兜底。
softmax 强制 fp32：.float() 进 .type_as() 出，理由同 RMSNorm。

SDPA（scaled_dot_product_attention）路径的条件为什么这么苛刻？——因为 is_causal=True 在 PyTorch 里的语义是"query 的第 i 个位置只能看 kv 的前 i 个"。带 KV Cache 的解码步 q_len=1 而 kv_len>1，这个语义会把历史也遮掉，所以 decode 步必须走手写路径；同理带 padding mask 时 SDPA 需要显式传 attn_mask，这里为了简化只在"无 padding 的训练 prefill"场景用 SDPA。SDPA 内部会自动选择 FlashAttention-2 / mem-efficient / math 后端——这就是"Flash Attention"的落地方式：不手写 online-softmax 分块算法，而是让 PyTorch 调度现成 kernel。

完整源码：Attention 类（含 flash 条件判断与两条路径）

python

class Attention(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
        self.n_local_heads = config.num_attention_heads
        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 = config.head_dim
        self.is_causal = True
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn

    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=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)
        xq, xk = self.q_norm(xq), self.k_norm(xk)
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        past_kv = (xk, xv) if use_cache else None
        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
        if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
        else:
            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
            if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
            if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
        output = self.resid_dropout(self.o_proj(output))
        return output, past_kv

五、FeedForward：SwiGLU

python

class FeedForward(nn.Module):
    def __init__(self, config, intermediate_size=None):
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
        self.up_proj   = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]    # SiLU

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

对应公式：

FFN (x) = (SiLU (x W_{g a t e}) ⊙ x W_{u p}) W_{d o w n}

相比经典 FFN $ReLU (x W_{1}) W_{2}$ ，SwiGLU 多了一条不带激活的门控分支 $x W_{u p}$ ，与激活后的主分支逐元素相乘——"让网络自己学哪些维度该放行"。SiLU 即 $x \cdot σ (x)$ ，平滑、非单调、处处可导，比 ReLU 在深层网络里梯度表现更稳。

三个矩阵（gate/up/down）比经典两矩阵多 50% 参数，所以配合上文 1.33 倍缩放的 intermediate_size（ $π d \approx 8 d / 3$ ）补偿，总参数量保持在与 4d-FFN 相当的水平。

六、MoE：稀疏激活的专家混合

6.1 路由与加权

python

class MOEFeedForward(nn.Module):
    def __init__(self, config):
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)      # 路由器
        self.experts = nn.ModuleList([
            FeedForward(config, intermediate_size=config.moe_intermediate_size)
            for _ in range(num_experts)])

    def forward(self, x):
        x_flat = x.view(-1, hidden_dim)                       # (B*S, d)
        scores = F.softmax(self.gate(x_flat), dim=-1)         # 每个专家的分
        topk_weight, topk_idx = torch.topk(scores, k=num_experts_per_tok, dim=-1, sorted=False)
        if norm_topk_prob:
            topk_weight = topk_weight / (topk_weight.sum(-1, keepdim=True) + 1e-20)
        y = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
            if mask.any():
                token_idx = mask.any(dim=-1).nonzero().flatten()
                weight = topk_weight[mask].view(-1, 1)
                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))

MoE 的核心思想：与其让一个巨大的 FFN 处理所有 token，不如拆成 N 个"专家"，每个 token 只挑 top-k 个走。minimind-3-moe 是 4 专家 / top-1：总参 198M，激活仅 64M——参数容量翻 3 倍，计算量不变。

实现要点：

按专家分桶而非按 token 循环：mask.any(dim=-1) 找出被分给专家 $i$ 的 token，index_add_ 把加权输出累加回去。top-1 下一个 token 只进一个桶；top-k 的话 mask 在 expert 维有多个 True，topk_weight[mask] 平铺后正好与 token 索引一一对应（同一 token 会出现 k 次，各带自己的权重）。
norm_topk_prob：top-k 权重归一化（除以权重和），让不同 token 的输出量级一致。top-1 时无影响（单元素归一化=1），这是为 Qwen3-MoE 兼容预留的。
没有 shared expert：Qwen3-MoE 有一个所有 token 必经的共享专家，MiniMind 为保持极简去掉了（README 明确说明对齐"去除 shared expert"的 Qwen3-MoE 风格）。

6.2 负载均衡 aux loss

路由是 argmax 类离散操作，不可导。如果放任自流，路由器会很快"偏心"——少数专家吃掉几乎所有 token，其他专家荒废（赢了通吃，强者愈强）。经典解法是 Switch Transformer 的辅助负载均衡损失：

python

if self.training and router_aux_loss_coef > 0:
    load = F.one_hot(topk_idx, num_experts).float().mean(0)    # 每个专家实际接到的 token 占比
    self.aux_loss = (load * scores.mean(0)).sum() * num_experts * router_aux_loss_coef

L_{a u x} = α \cdot N \sum_{i = 1}^{N} f_{i} \cdot p_{i}

$f_{i}$ ：专家 $i$ 实际接收的 token 比例（不可导，one-hot 统计）；
$p_{i}$ ：路由器给专家 $i$ 的平均概率（可导）。

均匀分布时 $f_{i} = p_{i} = 1 / N$ ， $L_{a u x} = α \cdot N \cdot N \cdot \frac{1}{N^{2}} = α$ ，取最小值。梯度通过 $p_{i}$ 流回路由器，逼着它把 token 摊匀。router_aux_loss_coef=5e-4，训练脚本里直接 loss = res.loss + res.aux_loss 加进总损失。

6.3 两个工程细节

哑梯度保 DDP 不挂：

python

elif self.training:
    y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())

DDP 要求所有 rank 的 backward 触达完全相同的参数集合。某个 batch 里恰好有专家没接到任何 token 时，它的参数没有梯度 → DDP all-reduce 等待超时直接挂死。这一行 0 * sum(...) 让没活干的专家也产生一个全零梯度，凑齐"全家桶"。（主 forward 里 aux_loss += sum(p.sum() for p in self.audio_proj.parameters()) * 0 是同一招，见 04 篇。）

MoE 训练反而更慢？ README 的实测：4-expert/top-1 的 MoE 比 dense 慢约 50%。原因是原生 PyTorch 按专家分桶 → 每个桶一次小矩阵乘 → kernel 启停和调度开销远大于一次大 GEMM。工业级靠 Triton 融合 kernel / DeepSpeed-MoE / Megatron 的 expert parallel 解决；MiniMind 选择保留原生 PyTorch 普适性，这是清醒的折中。"MoE 推理更快"和"这个 MoE 实现训练更慢"并不矛盾。

完整源码：MOEFeedForward 类

python

class MOEFeedForward(nn.Module):
    def __init__(self, config: MiniMindConfig):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts = nn.ModuleList([FeedForward(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)])
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.view(-1, hidden_dim)
        scores = F.softmax(self.gate(x_flat), dim=-1)
        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
        if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        y = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
            if mask.any():
                token_idx = mask.any(dim=-1).nonzero().flatten()
                weight = topk_weight[mask].view(-1, 1)
                y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
            elif self.training:
                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())   # 哑梯度保 DDP
        if self.training and self.config.router_aux_loss_coef > 0:
            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
            self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
        else:
            self.aux_loss = scores.new_zeros(1).squeeze()
        return y.view(batch_size, seq_len, hidden_dim)

七、Block：Pre-Norm 残差

python

class MiniMindBlock(nn.Module):
    def forward(self, hidden_states, position_embeddings, ...):
        residual = hidden_states
        hidden_states, present = self.self_attn(
            self.input_layernorm(hidden_states), position_embeddings, ...)
        hidden_states += residual
        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
        return hidden_states, present

h \leftarrow h + Attn ({RMSNorm}_{1} (h)), h \leftarrow h + FFN ({RMSNorm}_{2} (h))

Pre-Norm（norm 在残差分支内）vs Post-Norm（原始 Transformer，norm 在残差相加后）：Pre-Norm 的恒等通路让梯度可以无衰减地直通底层，深网络训练稳定得多——这是 GPT-2 之后的事实标准。代价是最终输出"没有归一化过"，所以主干末端补了一个 self.norm(hidden_states) 再进 lm_head。

一个微小但讲究的细节：hidden_states += residual 是原地加法，mlp 分支是普通加法——前者省一次内存分配，这种抠门在 287 行的教学代码里依然随处可见。

八、组装：主干与 CausalLM

8.1 主干

python

class MiniMindModel(nn.Module):
    def __init__(self, config):
        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(num_hidden_layers)])
        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
        freqs_cos, freqs_sin = precompute_freqs_cis(...)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    def forward(self, input_ids, attention_mask, past_key_values, use_cache, **kwargs):
        ...
        aux_loss = sum(l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward))
        return hidden_states, presents, aux_loss

persistent=False 的含义：cos/sin 表可以从参数确定性重建（纯 arange + cos/sin），不作为模型状态存进 checkpoint——省 4MB 权重文件体积，也让不同 max_position 配置互不干扰。

8.2 词嵌入权重共享（Weight Tying）

python

class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    def __init__(self, config):
        self.model = MiniMindModel(self.config)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.model.embed_tokens.weight = self.lm_head.weight

输入嵌入矩阵 $E \in R^{6400 \times 768}$ 和输出投影 $W \in R^{6400 \times 768}$ 职责互逆（one-hot→dense vs hidden→vocab logits），论文 tied embeddings 的结论是共享后小模型效果更好：6400×768 ≈ 4.9M 参数，对 64M 的模型占了近 8%，省下这部分还顺带充当正则（输入输出在同一个语义空间中对齐）。大模型上这条逐渐不再必要，但小模型依然是标配。

8.3 Loss：shift 藏在模型里

python

def forward(self, input_ids, ..., labels=None, logits_to_keep=0, **kwargs):
    hidden_states, past_key_values, aux_loss = self.model(...)
    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
    logits = self.lm_head(hidden_states[:, slice_indices, :])
    if labels is not None:
        x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
        loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)

语言模型的 next-token 目标：位置 $i$ 的 logits 应预测位置 $i + 1$ 的 token。logits[..., :-1] 对齐 labels[..., 1:] 完成 shift。shift 放模型里做而不是 dataloader 里，是 HF 风格——数据侧只管给"整段 token + 整段 label"，训练/推理共用一个 forward。

两个实用设计：

ignore_index=-100（cross_entropy 默认值）：label 为 -100 的位置不计损失。Pretrain 用它 mask 掉 padding，SFT 用它 mask 掉 prompt 部分——数据构造见 02-Tokenizer与数据。
logits_to_keep：只对序列最后 N 个位置算 logits。RLHF 场景生成完一整段后只需要 response 部分的 logprob，提前把前缀的 logits 省掉，能省一大块显存和 lm_head 的计算（6400 维的输出投影不小）。

九、generate：解码采样策略全手动实现

MiniMindForCausalLM.generate 用 30 行手写了完整的采样循环，是理解"LLM 推理到底在干嘛"的最好教材。每个策略单独拆开：

9.1 温度

python

logits = outputs.logits[:, -1, :] / temperature

softmax 前把 logits 除以 $T$ ： $T \to 0$ 趋近 argmax（贪心）， $T > 1$ 拉平分布更随机。数学上等价于对分布做 $p_{i}^{1 / T}$ 归一化。默认 0.85——比 1.0 稍保守。

9.2 重复惩罚

python

seen = torch.unique(input_ids[i])
score = logits[i, seen]
logits[i, seen] = torch.where(score > 0, score / rp, score * rp)

对已出现过的 token 惩罚。注意正负分开处理：logit > 0 时除以 rp（>1 即压低），logit < 0 时乘 rp（让负的更负）。对负 logit 也做除法反而会加重它（除以 >1 的数让 -3.0 变 -2.7，奖励了它），所以必须分象限。默认 rp=1.0 即关闭。

9.3 Top-k

python

logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')

只保留分数前 k 的候选，其余 -inf。一行核心逻辑：拿 topk 结果的最后一个值（第 k 名的分数）当阈值，低于它的全部枪毙。

9.4 Top-p（nucleus）

python

sorted_logits, sorted_indices = torch.sort(logits, descending=True)
mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0     # ① 右移一位
logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')

top-k 的候选数量固定，但有时前 5 个就占 99% 概率（该砍），有时 50 个才占 60%（不该砍那么多）。top-p 改为动态砍：按概率从高到低排序，累计到 $p$ 为止的集合保留。

三行代码里藏着一个经典 off-by-one：cumsum > p 会把"恰好跨过 p 的那个 token"也标成 True（它的累计概率包含它自己，砍掉它可能 cumsum 就到不了 p 了）。所以第 ③ 行把 mask 右移一位——跨过阈值的那位保留，之后的全砍。mask[..., 0] = 0 保证概率最高的 token 永远存活（分布不至于整体塌缩）。

最后一行 scatter 是因为 mask 是在排序后的空间里算的，要按 sorted_indices 映射回原始索引位置再置 -inf。

9.5 循环与终止

python

next_token = torch.multinomial(torch.softmax(logits, dim=-1), 1) if do_sample \
             else torch.argmax(logits, dim=-1, keepdim=True)
...
if eos_token_id is not None:
    finished |= next_token.squeeze(-1).eq(eos_token_id)      # 或上本轮是否 EOS
    if finished.all(): break

batch 解码的细节：已完成的序列强制继续输出 eos（torch.where(finished, eos, next_token)）占位，等所有序列都 finished 才 break——避免 batch 内快慢不一导致的索引错乱。streamer.put(...) 每步把 token 喂给流式输出器，WebUI 的打字机效果就来自这里。

另外值得注意：这个 generate 支持 num_return_sequences > 1（GRPO 一个 prompt 采样 N 个回答就靠它，见 03 篇），以及 return_kv 把 KV Cache 带出来给 Agentic RL 多轮拼接复用。

完整源码：generate 解码循环

python

# https://github.com/jingyaogong/minimind/discussions/611
@torch.inference_mode()
def generate(self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=2, streamer=None, use_cache=True, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs):
    input_ids = kwargs.pop("input_ids", inputs).repeat(num_return_sequences, 1)
    attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
    past_key_values = kwargs.pop("past_key_values", None)
    finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
    if streamer: streamer.put(input_ids.cpu())
    for _ in range(max_new_tokens):
        past_len = past_key_values[0][0].shape[1] if past_key_values else 0
        outputs = self.forward(input_ids[:, past_len:], attention_mask, past_key_values, use_cache=use_cache, **kwargs)
        attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
        logits = outputs.logits[:, -1, :] / temperature
        if repetition_penalty != 1.0:
            for i in range(input_ids.shape[0]):
                seen = torch.unique(input_ids[i]); score = logits[i, seen]; logits[i, seen] = torch.where(score > 0, score / repetition_penalty, score * repetition_penalty)
        if top_k > 0:
            logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
        if top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
            mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
            logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
        next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
        if eos_token_id is not None: next_token = torch.where(finished.unsqueeze(-1), next_token.new_full((next_token.shape[0], 1), eos_token_id), next_token)
        input_ids = torch.cat([input_ids, next_token], dim=-1)
        past_key_values = outputs.past_key_values if use_cache else None
        if streamer: streamer.put(next_token.cpu())
        if eos_token_id is not None:
            finished |= next_token.squeeze(-1).eq(eos_token_id)
            if finished.all(): break
    if streamer: streamer.end()
    if kwargs.get("return_kv"): return {'generated_ids': input_ids, 'past_kv': past_key_values}
    return input_ids

十、小结

把整条前向链路串起来：

text

input_ids
  → Embedding(6400, 768)                     [与 lm_head 共享权重]
  → 8 × [ RMSNorm → RoPE+GQA Attention → 残差 → RMSNorm → SwiGLU/MoE → 残差 ]
  → RMSNorm
  → lm_head → logits → shift + CE(+aux_loss)

对照大模型，你会认出几乎所有现代组件：Qwen3 同款的 GQA + QK-Norm + SwiGLU + YaRN + weight tying，Switch Transformer 的 aux loss，HF 风格的 shift-in-model loss 与 logits_to_keep。区别只在规模——而这正是 MiniMind 的立意：结构没有秘密，秘密都在训练里。

下一篇 02-Tokenizer与数据：模型吃进来的 token 是怎么来的，chat template 和 loss mask 又是怎么设计的。

数据结构

数据库

离散数学

数字电路

概率论

数学建模

软件工程

面向对象程序设计

Python 与数据分析

计算机组成

操作系统

计算机网络

编译原理

图形学

自然语言处理

Linux 实验

分布式系统

Internet原理与技术

BlogWeb

Stars Saying

黑马商城

苍穹外卖

黑马点评

黑马商城单体版

JavaWeb AI Tlias

RuoYi DKD

LinkQwen

InterX

LeetCode Agent Lab

00-预备知识

01-PR1108-逐代码段解析

02-PR1135-Cube化逐代码段解析

03-两个Bug深度剖析

04-开发心得与数值性能实测

01-PR1169-模型推理性能优化

02-PR1194-固定Shape及分档推理适配

NumPy

PyTorch

Python 八股文

LeetCUDA

LeetGPU

Triton

Minimind 系列笔记 01：模型结构精读 ​

一、全景：Config 与参数量 ​

二、RMSNorm：比 LayerNorm 少做一件事 ​

三、RoPE：旋转位置编码 ​

3.1 动机：让注意力天然感知"相对距离" ​

3.2 实现 ​

3.3 YaRN：免训练的长文本外推 ​

四、Attention：GQA 与 KV Cache ​

4.1 Q/K/V 投影与 QK-Norm ​

4.2 GQA：为什么 4 个 KV head 够 8 个 Q head 用 ​

4.3 KV Cache：增量解码的关键 ​

4.4 Flash Attention：两条计算路径 ​

五、FeedForward：SwiGLU ​

六、MoE：稀疏激活的专家混合 ​

6.1 路由与加权 ​

6.2 负载均衡 aux loss ​

6.3 两个工程细节 ​

七、Block：Pre-Norm 残差 ​

八、组装：主干与 CausalLM ​

8.1 主干 ​

8.2 词嵌入权重共享（Weight Tying） ​

8.3 Loss：shift 藏在模型里 ​

九、generate：解码采样策略全手动实现 ​

9.1 温度 ​

9.2 重复惩罚 ​

9.3 Top-k ​

9.4 Top-p（nucleus） ​

9.5 循环与终止 ​

十、小结 ​

Minimind 系列笔记 01：模型结构精读

一、全景：Config 与参数量

二、RMSNorm：比 LayerNorm 少做一件事

三、RoPE：旋转位置编码

3.1 动机：让注意力天然感知"相对距离"

3.2 实现

3.3 YaRN：免训练的长文本外推

四、Attention：GQA 与 KV Cache

4.1 Q/K/V 投影与 QK-Norm

4.2 GQA：为什么 4 个 KV head 够 8 个 Q head 用

4.3 KV Cache：增量解码的关键

4.4 Flash Attention：两条计算路径

五、FeedForward：SwiGLU

六、MoE：稀疏激活的专家混合

6.1 路由与加权

6.2 负载均衡 aux loss

6.3 两个工程细节

七、Block：Pre-Norm 残差

八、组装：主干与 CausalLM

8.1 主干

8.2 词嵌入权重共享（Weight Tying）

8.3 Loss：shift 藏在模型里

九、generate：解码采样策略全手动实现

9.1 温度

9.2 重复惩罚

9.3 Top-k

9.4 Top-p（nucleus）

9.5 循环与终止

十、小结