wuyang
277 lines
10 KiB
Markdown
277 lines
10 KiB
Markdown
# Kimi K3 第四轮：Attention Residuals 独立小模型复现协议
|
||
|
||
> 协议 ID：`llm-atlas-k3-attnres-reduced-v1`
|
||
> 冻结日期：2026-07-30（Asia/Shanghai）
|
||
> 状态：正式输出前预注册
|
||
> 目标：用可在单张 RTX 5090 上完整训练、复跑和检查的小模型，验证 Attention
|
||
> Residuals 的运算合同与早期训练动力学；不冒充 K3 checkpoint forward 或原论文规模复现。
|
||
|
||
## 1. 为什么本轮不直接“跑 K3 第一层”
|
||
|
||
截至本协议冻结时，官方开放工件仍是：
|
||
|
||
- `moonshotai/Kimi-K3@9f62e4e9fffbd0a83ddd60e1c209d828994b3569`；
|
||
- K3 GitHub `MoonshotAI/Kimi-K3@7c5be9599120d7993748de66a76128614f15f210`；
|
||
- Hugging Face remote code 把 KDA `A_log` 初始化为 `[num_heads]=[96]`；
|
||
- checkpoint header 中每个 KDA `A_log` 仍为 `[128]`；
|
||
- 当前 vLLM `fa2a2589bd2a1fce0851df7fd42ffb54b6195f04` 与 SGLang
|
||
`3c1717d9b6355e48855be7cc079db3ef9958f36d` 都把参数构造成 head 轴 96，并沿该轴
|
||
sharding；没有公开的 128→96 转换合同。
|
||
|
||
因此，裁剪、重复、平均或把 `[128]` 改解释成 channel 参数都会引入未公开假设。本轮不做。
|
||
官方也没有发布可加载的 K3 reduced checkpoint；Attention Residuals 论文所述 48B/3B 模型权重
|
||
仍未公开。
|
||
|
||
## 2. 一手来源
|
||
|
||
| 工件 | 固定 revision / checksum | 本轮用途 |
|
||
|---|---|---|
|
||
| Kimi K3 HF | `9f62e4e9…b3569` | 说明完整 checkpoint 的执行边界 |
|
||
| Kimi K3 GitHub | `7c5be959…f210` | 报告与部署入口 |
|
||
| Attention Residuals GitHub | `85e22310fe5ee860b4a023de312d791de8a5a5e6` | 官方论文与公式 |
|
||
| `Attention_Residuals.pdf` | SHA-256 `e5831b0d…a26b2f` | Full / Block AttnRes 运算合同 |
|
||
| WikiText-2 raw | `Salesforce/wikitext@b08601e04326c79dfdd32d625aee71d232d685c3` | 固定公开语料 |
|
||
|
||
Grok Headless 只承担正式输出前的对抗式方法检查，不是事实来源，也不参与结果解释。
|
||
|
||
## 3. 研究问题与声明边界
|
||
|
||
本轮只问：
|
||
|
||
1. 在核心 Transformer 权重、训练 Token、输入窗口与 optimizer 合同相同的缩小系统中，
|
||
Full / Block AttnRes 的验证 bits-per-byte 相对 PreNorm 方向是否跨三个 seed 一致？
|
||
2. 三种残差图的输入、分支输出和累计状态 RMS 怎样随深度变化？
|
||
3. 固定诊断 batch 上，每个 Transformer block 的参数梯度 RMS 怎样分布？
|
||
4. zero-init pseudo-query 经过训练后，是否出现邻近来源、embedding 和远层来源的可见分工？
|
||
5. 额外残差运算带来多少实际 step time 与 peak allocated memory？
|
||
|
||
本轮不能回答：
|
||
|
||
- K3 2.8T checkpoint 的真实 hidden state、router load 或 cache trace；
|
||
- 原论文 194M–528M activated / 38.7B–119B Token scaling law；
|
||
- AttnRes 的 benchmark 能力、长上下文质量或生产吞吐；
|
||
- KDA、MLA、MoE 与 AttnRes 的联合因果贡献；
|
||
- 任意总体显著性、p-value 或跨数据集普适性。
|
||
|
||
## 4. 数据合同
|
||
|
||
下载并校验：
|
||
|
||
```text
|
||
wikitext-2-raw-v1/train-00000-of-00001.parquet
|
||
wikitext-2-raw-v1/validation-00000-of-00001.parquet
|
||
wikitext-2-raw-v1/test-00000-of-00001.parquet
|
||
```
|
||
|
||
预处理固定为：
|
||
|
||
1. 按 parquet 行序读取 `text`；
|
||
2. 每行原文后追加一个 `\n`；
|
||
3. UTF-8 编码，不做 Unicode normalization、去空行、大小写或空白改写；
|
||
4. 词表固定为 256 个原始 byte；
|
||
5. 每个样本取连续 `context+1=257` bytes，前 256 预测后 256。
|
||
|
||
训练第 `step`、第 `row` 的起点：
|
||
|
||
```text
|
||
z = first 8 bytes of SHA256(
|
||
protocol_id + "\0train-window\0" + seed + "\0" + step + "\0" + row
|
||
)
|
||
start = uint64_be(z) mod (len(train_bytes) - 257)
|
||
```
|
||
|
||
因此同一 seed 的三种架构逐 step、逐 row 使用完全相同的 token tensor，不依赖 Python、NumPy
|
||
或 CUDA RNG 的消费顺序。
|
||
|
||
验证窗口固定 64 条，诊断窗口固定 16 条，分别用标签 `validation-window` 与
|
||
`diagnostic-window`、固定 index 取 SHA-256 起点；它们对全部 seed / 架构相同。
|
||
|
||
manifest 必须记录：
|
||
|
||
- dataset revision、三个 parquet SHA-256；
|
||
- 拼接后 split byte length 与 SHA-256；
|
||
- 训练 seed×step×row 的总 window-schedule hash；
|
||
- validation / diagnostic tensor hash。
|
||
|
||
## 5. 模型合同
|
||
|
||
所有模型：
|
||
|
||
| 项 | 固定值 |
|
||
|---|---:|
|
||
| vocabulary | 256 bytes |
|
||
| context | 256 |
|
||
| Transformer blocks | 16 |
|
||
| residual sublayers | 32（每块 attention + MLP） |
|
||
| `d_model` | 192 |
|
||
| heads | 6 |
|
||
| head dimension | 32 |
|
||
| `d_ff` | 768 |
|
||
| dropout | 0 |
|
||
| positional embedding | learned absolute, 256 × 192 |
|
||
| norm | RMSNorm, `eps=1e-6` |
|
||
| attention | causal multi-head softmax，score 用 FP32 softmax |
|
||
| MLP | bias-free SwiGLU，two `192→768` projections + one `768→192` projection |
|
||
| input/output embedding | tied |
|
||
| final readout | RMSNorm 后乘 tied token embedding |
|
||
|
||
三种 residual graph：
|
||
|
||
1. **PreNorm baseline**：`h ← h + f(RMSNorm(h))`。
|
||
2. **Full AttnRes**：每个 attention / MLP 子层从 embedding 与所有早期子层输出中，
|
||
以 `softmax(q_l^T RMSNorm(source))` 选出输入；当前分支输出成为下一个独立 source；
|
||
32 个子层结束后用同型 output mixer 聚合 embedding + 32 个输出。
|
||
3. **Block AttnRes**：32 个子层按顺序冻结为 8 块、每块 4 个子层；跨块保存 8 个 block
|
||
sums，块内使用 partial sum，embedding 永远是 source 0；最后用同型 output mixer 聚合
|
||
embedding + 8 个完整 block sums。
|
||
|
||
AttnRes 每个子层增加：
|
||
|
||
- 一个 `d_model` pseudo-query，初始化严格为 0；
|
||
- 一个 `d_model` RMSNorm key weight，初始化严格为 1。
|
||
|
||
Full / Block 各自再增加一组相同结构的 output mixer。所有 bias-free linear 与 embedding 使用
|
||
`N(0, 0.02)`；attention output projection 与 MLP down projection 按
|
||
`0.02 / sqrt(2 × 16)` 缩放初始化；普通 RMSNorm 为 1。
|
||
|
||
核心 token / position embedding、attention、MLP、输入 norm 与 tied output weights 在同一 seed
|
||
的三种架构中必须逐 tensor SHA-256 exact。AttnRes 额外参数与 mixer 运算如实报告，不通过
|
||
无作用 dummy 参数伪造“完全等容量”。
|
||
|
||
本实验匹配：
|
||
|
||
- 核心 Transformer 参数；
|
||
- optimizer steps；
|
||
- 每步训练 Token；
|
||
- 逐 Token 数据顺序。
|
||
|
||
本实验**不匹配**：
|
||
|
||
- residual mixer FLOPs；
|
||
- step wall time；
|
||
- peak activation memory。
|
||
|
||
后二者正是次要观测指标，不能被用于支持“同算力下更好”。
|
||
|
||
## 6. 优化与运行合同
|
||
|
||
| 项 | 固定值 |
|
||
|---|---:|
|
||
| seeds | `2026073001, 2026073002, 2026073003` |
|
||
| formal steps | 2,000 |
|
||
| batch | 32 |
|
||
| tokens / run | 16,384,000 target bytes |
|
||
| optimizer | AdamW |
|
||
| betas | `(0.9, 0.95)` |
|
||
| epsilon | `1e-8` |
|
||
| peak LR | `3e-4` |
|
||
| min LR | `3e-5` |
|
||
| warmup | 100 steps，linear |
|
||
| decay | cosine，step 100→2,000 |
|
||
| weight decay | `0.1` for parameters with `ndim >= 2`; otherwise `0` |
|
||
| grad clip | global norm `1.0` |
|
||
| compute | BF16 autocast，FP32 optimizer state |
|
||
| device | one RTX 5090 |
|
||
| compile | off / eager |
|
||
| RNG | deterministic algorithms，`CUBLAS_WORKSPACE_CONFIG=:4096:8` |
|
||
|
||
验证发生在 step `0, 100, 250, 500, 1000, 1500, 2000`。每次使用固定 64 条窗口；
|
||
以 8 条一个 eval batch，报告 token-mean cross entropy（nats）与：
|
||
|
||
```text
|
||
bits_per_byte = cross_entropy_nats / ln(2)
|
||
```
|
||
|
||
计时：
|
||
|
||
- 前 20 个训练 step 不进入 wall-time 统计；
|
||
- step 21–2,000 每步前后 CUDA synchronize；
|
||
- step 20 后 reset peak memory stats；
|
||
- 同报 mean / median / p95 step ms、peak allocated 与 peak reserved。
|
||
|
||
## 7. 预注册指标
|
||
|
||
### 7.1 主指标
|
||
|
||
每个 seed 在 step 2,000 的验证 `bits_per_byte`：
|
||
|
||
```text
|
||
Δ_full = Full AttnRes − PreNorm
|
||
Δ_block = Block AttnRes − PreNorm
|
||
```
|
||
|
||
只报告三个 paired delta、mean、min–max。
|
||
|
||
预注册解释规则：
|
||
|
||
- 三个 seed 全同为负，且 mean delta ≤ `−0.010 BPC`：**directional support in this reduced protocol**；
|
||
- 三个 seed 全同为正，且 mean delta ≥ `+0.010 BPC`：**directional concern in this reduced protocol**；
|
||
- 其他情况：**inconclusive at this budget**。
|
||
|
||
不报告 p-value、population CI 或“复现了论文提升”。
|
||
|
||
### 7.2 次要指标与精确定义
|
||
|
||
- `branch_output_rms[l]`：第 l 个 attention / MLP 分支输出在 batch×time×channel 上的 RMS；
|
||
- `layer_input_rms[l]`：进入第 l 个分支 norm 之前的 mixture / residual state RMS；
|
||
- `stream_state_rms[l]`：baseline 加法后的 state；Full 为新 source 输出集合的整体 RMS；
|
||
Block 为当前 partial sum 的 RMS；
|
||
- `parameter_grad_rms[block]`：固定 diagnostic batch 上、clip 前，该 block 所有核心参数梯度
|
||
拼接后的 `sqrt(sum(g²)/numel)`；不含 embedding、LM head 与 AttnRes mixer 参数；
|
||
- `depth_weights[l, source]`：固定 diagnostic batch 上 softmax 权重对 batch×time 求均值；
|
||
- `source_entropy[l]`：同一权重先逐 token 计算 entropy，再对 batch×time 求均值；
|
||
- `step_ms` 与 CUDA memory：按 §6 固定协议。
|
||
|
||
这些指标是描述性机制账，不进入主判定。
|
||
|
||
## 8. Smoke、正式运行与复现
|
||
|
||
正式输出前先跑：
|
||
|
||
- 三种架构 × seed 1；
|
||
- 20 steps；
|
||
- batch 4；
|
||
- 同一 window schedule；
|
||
- 断言无 NaN、loss finite、common-weight hashes exact；
|
||
- 在全新进程重跑相同 smoke，比较数据 hash、初始 common hash、逐 eval loss 和最终 common
|
||
checkpoint hash。
|
||
|
||
smoke 只决定实现是否可运行，不用于改变 formal 超参数或挑选架构。
|
||
|
||
正式运行顺序冻结为：
|
||
|
||
```text
|
||
seed 2026073001: baseline → full → block
|
||
seed 2026073002: baseline → full → block
|
||
seed 2026073003: baseline → full → block
|
||
```
|
||
|
||
全部 9 个 run 完成后，在全新进程复跑 `seed 2026073001 / block / 2,000 steps`，比较：
|
||
|
||
- manifest / dataset / schedule hashes；
|
||
- common initial weight hash；
|
||
- step 0–2,000 eval history；
|
||
- final common parameter hash；
|
||
- final mixer parameter hash；
|
||
- diagnostic tensors与 depth-weight hashes。
|
||
|
||
计时不要求 exact；数值字段要求 exact，若底层确定性限制导致非 exact，必须报告最大差异，
|
||
不能降低字段范围后宣称复现。
|
||
|
||
## 9. 冻结产物
|
||
|
||
```text
|
||
research/K3_ATTNRES_REDUCED_PROTOCOL.md
|
||
research/K3_ATTNRES_REDUCED_AUDIT.md
|
||
experiments/k3/attnres/build_dataset.py
|
||
experiments/k3/attnres/train.py
|
||
experiments/k3/attnres/analyze.py
|
||
experiments/k3/attnres/README.md
|
||
experiments/k3/attnres/manifest.json
|
||
src/data/k3-attnres-reduced.json
|
||
src/data/k3-attnres-reduced-compact.json
|
||
```
|
||
|
||
原始 parquet 与训练 checkpoint 留在本机 cache，不进入公开仓库；manifest、原始指标 JSON、
|
||
分析结果、代码、环境、checksum 和确定性复现账进入开源树。
|