Files
llm-atlas/research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md
T
2026-07-30 07:53:12 +08:00

460 lines
16 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# Kimi K3 第五轮:Attention Residuals 梯度定义与深度扩展实验协议
> 协议 ID:`llm-atlas-k3-attnres-gradient-scale-v1`
> 冻结日期:2026-07-30(Asia/Shanghai)
> 状态:正式输出前预注册
> 前置定义审计:`research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md`
## 0. 目标与一句话研究问题
Round 04 在缩小模型中得到两个同时成立的结果:
- Full / Block AttnRes 的 2,000-step 验证 BPC 都优于 Baseline;
- 以“每个 Transformer block 的核心**参数**梯度 RMS”定义时,跨深度 CV 反而更高。
论文 Figure 5(c) 没有公开足以唯一恢复的梯度测量合同。Round 05 不猜作者的隐藏代码,而是
冻结一个可复现、与 Figure 5 的 output/gradient 并列叙述对齐的 activation-gradient 定义,
再问:
> 当深度从 16 增至 32、训练预算从 2,000 增至 8,000 step 时,Block AttnRes 是否比
> PreNorm Baseline 更一致地降低 post-MLP block-output activation gradient 的跨深度失衡?
这不是 K3 checkpoint forward,也不是论文 Figure 5 数值复画。
## 1. 一手来源与不可补写的空白
| 工件 | 固定 revision / checksum | 用途 |
|---|---|---|
| Attention Residuals GitHub | `85e22310fe5ee860b4a023de312d791de8a5a5e6` | 公式、Figure 5 / 8、模型尺度 |
| `Attention_Residuals.pdf` | SHA-256 `e5831b0d…a26b2f` | 论文一手图文 |
| WikiText-2 raw | `Salesforce/wikitext@b08601e04326c79dfdd32d625aee71d232d685c3` | 固定公开训练语料 |
| Round 04 protocol | `llm-atlas-k3-attnres-reduced-v1` | 公共主干与数据合同来源 |
官方仓库没有模型实现、训练脚本、checkpoint 或 Figure 5 原始数据。以下字段不能归因给论文:
- Figure 5 的确切 gradient tensor;
- norm / reduction;
- diagnostic batch;
- AMP / clipping 时点;
- 单点还是时间平均。
Grok Headless 只进行一次对抗式方法检查;其建议和错误都在前置审计中公开,不是事实来源。
## 2. 设计总览
```text
2 个深度:16 / 32 Transformer blocks
× 2 个 residual graph:PreNorm Baseline / Block AttnRes
× 3 个冻结 seed
× 8,000 training steps
× 32 windows/step
× 256 target bytes/window
= 12 个正式训练格
= 786,432,000 target bytes
```
只比较 Baseline 与 Block,因为论文 Figure 5 的训练动力学面板也是这两个结构的直接对照。
Round 04 的 Full AttnRes 结果保持公开,但本轮不增加一个与主问题无关的 6-run 分支。
## 3. 数据合同
继承 Round 04 的语料和预处理:
1. 按固定 parquet 行序读取 `text`;
2. 每行追加一个 `\n`;
3. UTF-8 编码,无 normalization、strip、去空行或大小写改写;
4. byte vocabulary `0..255`;
5. 每个窗口连续取 257 bytes,前 256 预测后 256。
固定拼接后 split:
| split | bytes | SHA-256 |
|---|---:|---|
| train | 10,951,563 | `0ca7d3e74dbe44564ea5942b85232f1bbcb525c9cd481cd5d28a87ee90e7e9b4` |
| validation | 1,148,008 | `a42356f6a8ff1d25daf25ec9db49e10a537c265581b61c74604bb63231dee719` |
| test | 1,292,014 | `bfe9eb16ab9987fb88bde4ea9a30a00f2a45db01dfc14bad78d05325789c4f12` |
训练第 `step`、第 `row` 的 window 起点:
```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 的 4 个结构格逐 step / row 使用完全相同的 token tensor。validation 64 windows、
diagnostic 16 windows,分别由标签 `validation-window` / `diagnostic-window` 与固定 index
生成,对全部结构与 seed 相同。
正式运行前 manifest 必须记录 parquet hash、split bytes/hash、全部 `3×8,000×32=768,000`
唯一训练窗口起点的 schedule hash、validation tensor hash 与 diagnostic tensor hash。
## 4. 模型合同
### 4.1 两个深度共享的结构
| 项 | 固定值 |
|---|---:|
| vocabulary | 256 bytes |
| context | 256 |
| `d_model` | 192 |
| heads / head dimension | 6 / 32 |
| `d_ff` | 768 |
| dropout | 0 |
| positional embedding | learned absolute,256 × 192 |
| norm | RMSNorm,`eps=1e-6` |
| attention | causal MHA;score 以 FP32 softmax |
| MLP | bias-free SwiGLU,`192→768`, `192→768`, `768→192` |
| embedding / readout | tied;final RMSNorm 后乘 token embedding |
全部 bias-free linear 与 embedding 初始化为 `N(0,0.02)`。attention output projection 与 MLP
down projection 的标准差为 `0.02 / sqrt(2×depth)`;普通 RMSNorm 为 1。
数值精度进一步固定为:attention / MLP 线性分支受 BF16 autocast;embedding、Baseline hidden
residual stream 与 Block aggregation partial 都以 FP32 累加。也就是说,Block 每个 BF16
branch output 在进入 `partial` 前显式转为 FP32。这样两种结构被捕获的 `h_l` 都是 FP32,
不会把 residual accumulator 精度差异混进梯度形状对比。
### 4.2 深度与 Block AttnRes 聚合
| Transformer depth | residual sublayers | aggregation groups | sublayers/group | Transformer blocks/group |
|---:|---:|---:|---:|---:|
| 16 | 32 | 8 | 4 | 2 |
| 32 | 64 | 8 | 8 | 4 |
Baseline 子层为:
```text
h ← h + f(RMSNorm(h))
```
Block AttnRes:
- embedding 永远是 source 0;
- 对已经完成的 aggregation-group sums 做跨组 softmax mixture;
- 组内 attention / MLP branch output 累加到 `partial`;
- 达到组边界时,保存完整 `partial` 为新 source,再开始下一组;
- output mixer 聚合 embedding + 8 个完整 group sums。
每个子层的 pseudo-query 为 `d_model` 向量,严格 zero-init;每个 source 的 key RMSNorm weight
严格 one-init。Baseline / Block 不要求总参数量或 residual-mixer FLOPs 相等,但同 seed /
depth 的 token/position embedding、attention、MLP、input norm、final norm 和 tied readout
必须逐 tensor SHA-256 exact。
## 5. 训练合同
| 项 | 固定值 |
|---|---:|
| seeds | `2026073001, 2026073002, 2026073003` |
| formal steps | 8,000 |
| batch | 32 |
| context | 256 |
| target bytes / run | 65,536,000 |
| optimizer | AdamW |
| betas / epsilon | `(0.9,0.95)` / `1e-8` |
| peak / min LR | `3e-4` / `3e-5` |
| warmup | 400 steps,linear |
| decay | cosine,step 400→8,000 |
| weight decay | `0.1` if `ndim>=2`,否则 `0` |
| parameter grad clip | global norm `1.0` |
| compute | BF16 autocast;FP32 optimizer state |
| compile | off / eager |
| device | one RTX 5090 |
| deterministic | deterministic algorithms;`CUBLAS_WORKSPACE_CONFIG=:4096:8` |
验证和诊断都发生在:
```text
step 0, 100, 500, 2,000, 4,000, 8,000
```
验证固定 64 windows,8 windows/eval batch,报告 token-mean CE nats 与
`bits_per_byte = CE / ln(2)`。诊断固定 16 windows,一次性输入,不改变 optimizer state。
计时合同:
- 前 20 个 training step 不计入;
- step 21–8,000 每步前后 CUDA synchronize;
- step 20 后 reset peak memory;
- 报告 mean / median / p95 step ms、peak allocated / reserved;
- timing、wall clock、hostname、GPU temperature 不进入 exact replay 字段。
## 6. 主梯度对象:精确到代码位置
### 6.1 `h_l` 定义
对 `l=1..depth`,统一在第 `l` 个 Transformer block 的 MLP branch 完成后捕获:
| 架构 | 精确定义 |
|---|---|
| Baseline | attention residual 与 MLP residual 都完成后的 hidden state |
| Block | MLP branch 已加入、本 aggregation partial 可能保存/reset **之前**的 partial |
所有 `h_l.shape = [16,256,192]`、dtype 为 FP32。实现必须为 diagnostic forward 返回独立
引用列表,不允许捕获 reset 后的零张量,不允许把 8 个 aggregation sources 当成 16 / 32 个
Transformer outputs。
### 6.2 diagnostic loss 与 gradient magnitude
```text
model.eval()
logits, h[1..depth] = forward(fixed_diagnostic_x, capture=true)
L = mean(cross_entropy(logits.float(), fixed_diagnostic_y))
backward(L) # 不使用 GradScaler,不执行 optimizer.step
m_l = sqrt(mean(float32(h_l.grad)² over batch×time×channel))
```
规则:
- forward 仍使用与训练一致的 BF16 autocast;
- logits 在 FP32 中计算 CE;
- loss 对全部 `16×256` targets 做算术平均,无 mask、无 label smoothing;
- backward 前 model / optimizer gradients 清零;
- activation gradient 在任何 parameter clipping 之前读取;
- diagnostic 不消耗训练数据,不进入 optimizer,不改变学习率或模型状态。
### 6.3 同位置 output magnitude
同一批 `h_l` 计算:
```text
o_l = sqrt(mean(float32(h_l)² over batch×time×channel))
```
它用于显示 Baseline 单调累积与 Block 组内周期,而不作为主 confirmatory endpoint。
## 7. 主指标与预注册判据
每个 diagnostic step 保存完整 `m_1..m_depth`。以下统计由未四舍五入的 float64 数组计算。
### 7.1 绝对尺度
```text
mean_grad = mean_l(m_l)
```
它防止“曲线更平只是全部梯度趋近于零”被 CV 隐藏。绝对尺度不设置优劣阈值,只公开。
### 7.2 归一化谱与 CV
```text
n_l = m_l / mean_grad
CV = population_std_l(m_l) / mean_grad
```
使用 population standard deviation(`ddof=0`)。网站必须同时显示 `m_l` 与 `n_l`,不能只显示
CV 排名。
### 7.3 前后四分位失衡
```text
q = depth / 4
Q_first = mean(m_1 .. m_q)
Q_last = mean(m_(depth-q+1) .. m_depth)
imbalance = abs(ln(Q_first / Q_last))
```
depth 16 时各取 4 层;depth 32 时各取 8 层。`imbalance=0` 才表示首尾一致;这个定义不会把
“早层偏大”和“晚层偏大”错误地都解释为越小越好。原始有符号比 `Q_first/Q_last` 仍公开。
### 7.4 seed 内配对对比
只在最终 step 8,000 做 confirmatory verdict:
```text
relative_CV_reduction
= (CV_baseline - CV_block) / CV_baseline
relative_imbalance_reduction
= (imbalance_baseline - imbalance_block) / imbalance_baseline
```
若 Baseline imbalance 精确为 0,则该 seed 的 relative imbalance reduction 定义为不可计算,
该 depth 自动不能得到“联合支持”;仍公开绝对差。
对每个 depth 分别判定:
- **joint directional support at this depth**:三个 seed 的 CV reduction 都 `>0`,其均值
`>=20%`;同时三个 seed 的 imbalance reduction 都 `>0`,其均值 `>=20%`。
- **joint directional concern at this depth**:三个 seed 的两项 reduction 都 `<0`,且两项
平均相对恶化都 `>=20%`。
- 其他:**mixed / inconclusive at this depth**。
总判定:
- 两个 depth 都 support:**scale-consistent directional support in this operationalization**;
- 两个 depth 都 concern:**scale-consistent directional concern in this operationalization**;
- 其他:**depth-dependent or inconclusive**。
不计算 p-value、population CI,不把 3 seeds 称为统计证明。
## 8. 必须公开的次要指标
### 8.1 全时间轨迹
六个预注册时点的以下数据必须全部公开,不能选择“最好看”的 checkpoint:
- validation BPC;
- absolute activation-gradient spectrum;
- normalized activation-gradient spectrum;
- activation CV;
- first/last quartile ratio 与 imbalance;
- output RMS spectrum;
- mean activation-gradient scale。
### 8.2 参数梯度
延续 Round 04 定义:每个 Transformer block 的 attention、MLP 与两个 input norm 的参数梯度
拼接后计算:
```text
parameter_grad_rms[l] = sqrt(sum(g²) / total_parameter_elements)
```
不含 embedding、final norm、LM head 与 AttnRes mixer 参数;clip 前读取。报告完整谱、CV 和
前后四分位失衡,但它们不进入 Round 05 主判定。
### 8.3 mixer / 成本
Block 同报:
- 各子层 softmax mixture 的 source-depth 分布;
- output mixer 分布;
- entropy 与 embedding / latest-complete-group mass;
- step time、peak allocated/reserved、参数量。
这些用于解释机制与成本,不改变 confirmatory verdict。
## 9. 运行前闸门
### 9.1 数据闸门
- split bytes/hash 与 Round 04 exact;
- 新 protocol 的 768,000-window schedule hash 落盘;
- validation / diagnostic tensor hash 落盘;
- 同 seed 四结构的至少 step 0 / 1 / 7,999 输入 tensor hash exact。
### 9.2 公共权重闸门
每个 depth / seed 的 Baseline 与 Block 公共参数逐 tensor exact;输出:
- 公共参数 tensor 数;
- 公共参数 element 数;
- name / shape / dtype / bytes 联合 hash;
- value bytes 联合 hash。
### 9.3 activation-gradient 闸门
四个 depth×architecture 格都必须通过:
1. 捕获数量严格等于 depth,shape 均为 `[16,256,192]`;
2. 全部 gradient 非 `None`、finite、storage 不别名;
3. 同输入把 loss 乘 2 后,每层 `m_l` 比值在 `2±1e-5`;
4. loss×2 前后 CV、normalized spectrum、quartile ratio 在 `1e-6` 绝对容差内;
5. 全新进程重复 smoke 的冻结字段 exact。
### 9.4 smoke
四个格都运行 20 training steps;冻结字段包括:
- protocol、architecture、depth、seed、device、dtype;
- data/schedule/tensor hashes;
- public-weight hashes;
- step 0 / 20 loss 与 validation;
- 全部预注册 diagnostic 数组;
- finite / alias / loss-scale checks。
smoke 不能写入 formal 目录。
## 10. 正式执行与独立 replay
12 个 formal grid 必须各自在全新进程中执行。目录键为:
```text
depth-{16|32}/{baseline|block}/seed-{2026073001|2026073002|2026073003}
```
完成后预注册 replay:
```text
depth-32 / block / seed-2026073001
```
replay 再从初始化训练完整 8,000 steps,不加载 formal checkpoint。比较时排除:
- wall time / step-time samples;
- peak memory;
- process ID / hostname;
- GPU 温度与驱动层瞬时字段;
- 文件路径和生成时间。
必须 exact 的字段:
- 数据与公共权重 hashes;
- 全部 validation CE/BPC;
- 全部 activation/output/parameter gradient 数组;
- mixer 数组与 summary;
- final model-state tensor hash;
- optimizer-state tensor hash;
- training-loss checkpoint 数组。
若 replay 不 exact,停止聚合并公开失败,不挑选另一 seed 替代。
## 11. 公开工件
正式结果完成后仓库必须包含:
```text
experiments/k3/attnres_gradient/
README.md
build_dataset.py
manifest.json
train.py
analyze.py
results/raw/*.json
results/compact.json
reproduction.json
research/
K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md
K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md
K3_ATTNRES_GRADIENT_SCALE_AUDIT.md
```
网站至少提供五个互相联动的视图:
1. 论文 Figure 5 的“已知 / 未定义”拆解;
2. activation gradient 绝对谱与 normalized spectrum;
3. depth 16 / 32、三 seed、六时间点对比;
4. output RMS 周期与 Block aggregation boundary;
5. activation gradient / parameter gradient 并排,以及判定、成本、哈希和声明边界。
图中必须能切换到所有负结果;不能只放均值、只放 final 或隐藏某个 seed。
## 12. 允许与禁止的结论
若达到支持条件,允许写:
> 在这个公开定义、两种缩小深度和 8,000-step byte-LM 合同中,Block AttnRes 方向一致地
> 降低了 post-MLP output activation gradient 的跨深度 CV 与首尾四分位失衡。
无论结果怎样,都禁止写:
- 复现了论文 Figure 5 的数值;
- 证明了论文未公开实现采用相同梯度定义;
- 证明了 Kimi K3 的真实梯度更健康;
- 证明 AttnRes 解决梯度消失、梯度爆炸或训练稳定性的全部问题;
- 从 3 seeds 推导总体显著性;
- 从 depth 16 / 32 外推至 48B、1T+400B Token 或 K3 2.8T 参数;
- 隐藏 activation 与 parameter gradient 方向不一致的结果。
## 13. 变更纪律
本文件提交后:
- 允许修复使实现符合本协议的 bug;
- 允许补充日志、注释、可视化和不改变数值的导出;
- 不允许看过 formal 结果后修改主指标、阈值、diagnostic step、seed、训练预算或 replay 格;
- 任何不得不改变实验合同的事项必须先停止、写入审计、升级 protocol ID,再重新执行全部 grid。