9.9 KiB
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与 SGLang3c1717d9b6355e48855be7cc079db3ef9958f36d都把参数构造成 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. 研究问题与声明边界
本轮只问:
- 在核心 Transformer 权重、训练 Token、输入窗口与 optimizer 合同相同的缩小系统中, Full / Block AttnRes 的验证 bits-per-byte 相对 PreNorm 方向是否跨三个 seed 一致?
- 三种残差图的输入、分支输出和累计状态 RMS 怎样随深度变化?
- 固定诊断 batch 上,每个 Transformer block 的参数梯度 RMS 怎样分布?
- zero-init pseudo-query 经过训练后,是否出现邻近来源、embedding 和远层来源的可见分工?
- 额外残差运算带来多少实际 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. 数据合同
下载并校验:
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
预处理固定为:
- 按 parquet 行序读取
text; - 每行原文后追加一个
\n; - UTF-8 编码,不做 Unicode normalization、去空行、大小写或空白改写;
- 词表固定为 256 个原始 byte;
- 每个样本取连续
context+1=257bytes,前 256 预测后 256。
训练第 step、第 row 的起点:
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 |
| input/output embedding | tied |
三种 residual graph:
- PreNorm baseline:
h ← h + f(RMSNorm(h))。 - Full AttnRes:每个 attention / MLP 子层从 embedding 与所有早期子层输出中,
以
softmax(q_l^T RMSNorm(source))选出输入;当前分支输出成为下一个独立 source。 - Block AttnRes:32 个子层按顺序冻结为 8 块、每块 4 个子层;跨块保存 8 个 block sums,块内使用 partial sum,embedding 永远是 source 0。
AttnRes 每个子层增加:
- 一个
d_modelpseudo-query,初始化严格为 0; - 一个
d_modelRMSNorm key weight,初始化严格为 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)与:
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:
Δ_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 超参数或挑选架构。
正式运行顺序冻结为:
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. 冻结产物
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 和确定性复现账进入开源树。