M07 · 按 stage 加载与打包权重

shards_for_stage：

weight_map = {
  "thinker.model.embed_tokens.weight":      "model-00001.safetensors",
  "thinker.model.layers.0.mlp.gate.weight": "model-00002.safetensors",
  "talker.model.codec_embedding.weight":    "model-00002.safetensors",
  "talker.code_predictor.lm_head.0.weight": "model-00003.safetensors",
  "code2wav.decoder.0.conv.weight":         "model-00003.safetensors",
  "thinker.lm_head.weight":                 "model-00004.safetensors",
}
shards_for_stage(weight_map, "thinker")  -> ["model-00001…", "model-00002…", "model-00004…"]
shards_for_stage(weight_map, "talker")   -> ["model-00002…", "model-00003…"]
shards_for_stage(weight_map, "code2wav") -> ["model-00003…"]

同一个 shard 可同时属于两个 stage；结果去重并排序。thinker.lm_head.* 也属于 thinker stage。

select_stage(state, "thinker", "thinker.model.")：keys thinker.model.norm.weight、thinker.model.layers.0.input_layernorm.weight、thinker.lm_head.weight、talker.model.norm.weight、thinkers.model.x → {"layers.0.input_layernorm.weight", "norm.weight"}。值原样保留，只改 key。

load_thinker_text：多一个 thinker.model.rotary_emb.inv_freq → LoadError；num_layers 比实际多 → 那层什么都没有 → LoadError；thinker.lm_head.weight 和 talker.* 存在不算错误。

关键问题

1. 判断 key 属于哪个 stage 该用哪个前缀？为什么 "thinker"（不带点）不够安全？
2. 去前缀时 "thinker.model.norm.weight" 变成什么？"thinker.lm_head.weight" 会怎样？
3. 收集第 i 层为什么必须用 "layers.{i}."（带结尾点）？i=1 时不带点会误收哪些层？
4. qkv 的行数是多少？为什么沿 dim 0 拼而不是 dim 1？
5. w13 是先对每个 expert cat(gate, up) 再 stack，还是先 stack 所有 gate 再和 up 拼？
6. 如何判断 expert id “恰好是 0..E-1 且完整”？E 从哪里得知？
7. “多余 key”怎么检测？需要记录哪些 key 已被使用？

细节

• stage 级前缀（"thinker."）vs strip 前缀（"thinker.model."）：用 startswith + k[len(strip):]。
• nn.Linear.weight 是 [out, in]，多个共享输入的 Linear 合并就是在 out 维（行）上拼。
• PackedLayer 的 9 个字段与 HF key 一一对应；mlp.gate.weight 是 router。
• load_thinker_text 返回的第一个元素是 {"embed_tokens": ..., "norm": ...}（key 里没有 .weight）。

• shards_for_stage：一行集合推导，sorted(set)。
• pack_layer：切出 sub；小助手 take(name)：不存在 raise LoadError，存在就记入 used 并返回。expert key 匹配成功时也加入 used。
• 收集 experts：experts.setdefault(int(id), {})[proj] = v；检查非空、连续、三种 proj 齐全；E = len(experts)。
• 打包：qkv = cat([q, k, v], 0)；w13 = stack([cat([gate_e, up_e], 0) for e in range(E)])；w2 = stack([down_e ...])。按 e = 0..E-1 的顺序。
• 顺序陷阱：多余 key 检查必须在所有 take 之后。
• load_thinker_text：select_stage(state, "thinker", "thinker.model.") → 对 range(num_layers) 调 pack_layer → 检查非 layers. 开头的 key 集合恰好是 {embed_tokens.weight, norm.weight}。

自检：打包后是否等价

x = torch.randn(5, hidden)
q, k, v = (x @ L.qkv.T).split([H*D, KVH*D, KVH*D], dim=-1)   # q 应等于 x @ q_proj.T
g, u = (x @ L.w13[e].T).chunk(2, dim=-1)
y = (F.silu(g) * u) @ L.w2[e].T                              # 应等于 expert e 的 SwiGLU 输出

这也正是 M08 把 PackedLayer 还原成 M02/M03 权重的方式：qkv 按行 split，w13 按第 1 维对半切。