parameter-golf 记录解读：4096 词表 + 4×MLP + 高权重衰减的 16MB 模型配方（val_bpb 1.09785）

导读

本文围绕 parameter-golf 挑战中records/track_10min_16mb/2026-04-01_Vocab4096_MLPMult4_WD085/这份官方记录展开，它代表了一条重要的技术路线：把 SentencePiece 词表从 1024 扩到 4096、把 MLP 宽度从 3× 提到 4×，并配合更高的权重衰减（muon/embedding 均为 0.085）与一系列架构“减法”，在 8×H100 上 10 分钟内训练出val_bpb 1.09785（3-seed 平均，std 0.0004）的 16MB 以内模型。读完本文，你将掌握该配方的完整超参数表、数据与 tokenizer 的构建/下载流程、互质步长数据加载器与 GPTQ+Hessian 量化的实现细节，以及如何通过 train_gpt.py 一键复现。

记录背景：在 16MB 上限下"以压缩换容量"

Parameter Golf 的核心约束是：最终提交物 = 训练脚本代码字节数 + 压缩后模型字节数，合计不超过16,000,000 字节（十进制），且在 8×H100 SXM 上训练不超过 10 分钟，评估另行限时。因此模型大小不再是参数量，而是"量化+压缩后的磁盘体积"。

本记录的关键洞察来自作者 Kevin Clark 的实测：权重矩阵的可压缩性（量化压缩后体积 / 原始体积）与其 RMS（torch.sqrt(torch.mean(x**2))）高度相关，R² 接近 0.99。这意味着权重衰减可以直接压低权重的 RMS，从而压小压缩体积，把省出来的字节预算重新投入到增加模型参数量上。整份配方正是围绕"高 WD 压低压缩体积 → 扩大词表与 MLP → 用更多参数换更低 BPB"这条主线设计的。

该记录基于 2026-03-23_LeakyReLU_LegalTTT_ParallelMuon 的上一代记录（val_bpb 1.1194）演进而来，但做了一个大胆的方向调整：大幅删减上一代的技巧，而不是继续叠加。

三组改动：修复、简化与新增

原 README 将相对上一代记录的改动分为三组，本文逐一展开并与源码对照。

修复：滑动窗口评估的重复计分 bug

原评估代码在构造滑动窗口起点时使用了有缺陷的条件：

window_starts = [ws for ws in range(0, total_tokens, stride) if min(ws + seq_len, total_tokens) - ws >= 1]

该条件会让末尾的若干 token 被多个窗口重复计入损失与字节数。修复后的正确写法为：

window_starts = [ws for ws in range(0, total_tokens, stride) if ws + seq_len - stride < total_tokens]

作者实测该 bug 影响很小：在约 600 万 token 的验证集上只多加了约 2k 份重复贡献。修复逻辑在当前仓库的 train_gpt.py 的eval_val_sliding中已体现：窗口起点按h.eval_stride（默认 64）步进，context_size = seq_len - eval_stride，每个 token 只在其可获得的最大上下文窗口内被计分一次。

简化：对上一代技巧的全面"减法"

这一代果断移除了 6 项复杂机制：

• XSA 应用范围扩大：上一代只在最后 4 层使用 XSA（cross-subspace attention），本代改为所有 11 层全部启用（源码中xsa_last_n = 11，见 Hyperparameters）。
• 移除 Parameter Banking 与分布式 Muon：改为朴素 Muon + DDP，降低实现复杂度与通信开销。
• 移除测试时训练（TTT）：作者判断 0.1% 的额外 token 收益有限，且对长文档场景而言，把精力投向扩展序列长度更合理。
• 移除量化感知训练（QAT）：实测几乎没有收益。
• 移除 gated attention、value residuals。
• 移除 hash embeddings 与 smear gate：这两者主要用于缓解小词表的信息瓶颈，词表扩大到 4096 后必要性下降。

这些删除与"简化"方向一致：挑战的演化并非总是加法，去掉收益低于复杂度成本的机制同样是有效的优化。该记录在 README.md 中被列为"4096-Vocab + Larger Model + High WD + Simplifications"，正是对这一思路的总结。

新增：五处正向改动

1. 词表 1024 → 4096：使用仓库内的 data/download_hf_docs_and_tokenize.py 构建 sentencepiece tokenizer 并预分词。词表模型仅增长约 50KB，最终产物仍低于 16MB。更大词表意味着同等序列长度下模型看到更多上下文、每步训练数据也更多。
2. 更大但更强正则化的模型：muon 权重衰减 0.04 → 0.085，新增 embedding 权重衰减 0.085；Adam 权重衰减从 0.04 降到 0.02（标量参数无需低幅值）；mlp_mult3 → 4；学习率 0.025 → 0.02（大模型通常需要更小 LR）。
3. 互质步长数据加载器：来自 #726，通过在数据文件间跳跃采样，避免同一文档的序列出现在相同/邻近的 minibatch 中。
4. GPTQ Hessian 感知量化：基于 #1060，从训练预算中预留时间做 Hessian 计算。
5. 更高效的字节打乱 + brotli 压缩、sigmoid 门控 UNet 跳跃连接：均来自 #1089；同时将qk_gain_init从 1.5 提升到 4（#1125）。

完整超参数表（默认值，均可被环境变量覆盖）

train_gpt.py将全部超参数定义为类属性，运行时通过os.environ.get读取，未设置时使用以下默认值（对应本记录实际采用的配置）：

其中与上一代对比最值得关注的三项：mlp_mult=4.0、muon_wd=0.085、embed_wd=0.085，它们共同体现了"用高权重衰减压低压缩体积、换取更大模型容量"的核心思路。注意train_batch_tokens与val_batch_tokens的语义是全局 token 数，实际每个 rank 的批量由world_size × grad_accum_steps分摊。

模型架构：11 层 UNet + XSA + Value Embedding

从源码看，模型是一个带跳跃连接的类 UNet 解码器结构（GPT 类）：

• 前 5 层为 encoder（num_layers // 2），逐层将隐藏态压入skips列表；后 6 层为 decoder，从skips弹出并与当前隐藏态按skip_weights加权融合。
• sigmoid 门控跳跃连接：skip_gates参数经torch.sigmoid后，用torch.lerp(scaled_skip, x, g)在"跳过连接"与"当前隐藏态"之间做插值，门控值为 0 时完全走 skip，为 1 时完全走主干。
• XSA（cross-subspace attention）：在CausalSelfAttention中，对注意力输出 y 减去其在 V 上投影的残差（_xsa_efficient），把每个 head 的输出正交化到与 V 无关的方向上，用于抑制注意力坍缩；use_xsa通过xsa_last_n=11在所有层开启。
• Value Embedding：一个共享的ValueEmbedding(vocab_size, 128, kv_dim)，token id 经小嵌入投影后加到 V 上（v = v + v_embed），默认注入第 9、10 层，为深层注入词级先验。
• MLP 激活：F.leaky_relu(self.fc(x), negative_slope=0.5).square()，即 LeakyReLU(0.5)²，延续上一代记录的激活选择。
• Partial RoPE：rope_dims=16，只对 head 维度的前 16 维施加旋转，其余维度原样通过（apply_rotary_emb的x_pass分支）。
• 层归一化缩放：ln_scale=True时，注意力与 MLP 输入乘以1 / sqrt(layer_idx + 1)的递减因子。
• Logit softcap：30 * tanh(logits / 30)，限制 logit 幅值以利于 int6 量化。

权重初始化上，tie_embeddings=True时嵌入以std=0.005正态初始化，线性层矩阵（≥64×64）用正交初始化，proj投影层与注意力输出层零初始化，保证残差结构在训练初期稳定。

优化器：Muon 矩阵优化 + 三组 AdamW 分工

优化器设计（Optimizers 类）按参数形态分流：

• Muon：负责所有 2D 矩阵参数（注意力与 MLP 权重，排除控制张量），lr=0.02、wd=0.085、momentum=0.99。其 step 内部用 Newton–Schulz 五次迭代（zeropower_via_newtonschulz5，系数 a=3.4445, b=-4.7750, c=2.0315）做正交化，并支持分布式下的跨 rank 分块计算 +all_reduce汇总——但本记录去掉了上一代的 Parameter Banking，直接使用标准 DDP + Muon。
• AdamW（embed）：嵌入权重，lr=0.03、wd=0.085。
• AdamW（scalar）：所有标量/控制张量（attn_scale、mlp_scale、resid_mix、q_gain、skip_weights、skip_gates、ve_layer_scales等，由CONTROL_TENSOR_NAME_PATTERNS定义），lr=0.02、wd=0.02。

训练中 Muon 的 momentum 从 0.92 线性 warmup 到 0.99（前 1500 步），学习率按lr_mul在warmdown_frac=0.667时开始线性衰减到min_lr。此外全程维护 EMA（ema_decay=0.997）权重，训练结束后加载 EMA 状态作为最终模型——这与上一代"EMA + SWA"的配方相比又做了一次减法（本代不再使用 SWA）。

互质步长数据加载器：打破文档聚集偏差

DistributedTokenLoader 是本记录数据侧的核心改动（来自 PR #726）。其工作机制可概括为：

1. 把每个训练 shard 按seq_len切成若干 token 块（block），shard 内以随机相位起步。
2. 为每个 shard 的块游标挑选一个与块数互质的步长（_pick_coprime_stride，通过math.gcd(s, n) == 1保证），从而以循环方式遍历全部块且不产生短周期重复。
3. 每步先按剩余块数加权（alpha = 0.90 - 0.40 * progress，随训练进度从 0.90 降到 0.50）采样一组 shard，再以互质步长交错取出窗口，最后按 rank 切分给各进程。

互质步长的价值在于：同一文档的连续序列块会以"跳跃"的方式分散到不同的 minibatch，避免模型在相邻 minibatch 中反复看到同一篇文档的高相关片段，相当于一种廉价的文档级去相关。源码中每个 shard 的 RNG 由文件路径字节流哈希播种，保证跨进程可复现。

数据与 Tokenizer 准备：两条路径任选

路径一：直接下载 sp4096 缓存数据

rm -f data/manifest.json MATCHED_FINEWEB_REPO_ID=kevclark/parameter-golf \ python3 data/cached_challenge_fineweb.py --variant sp4096 --train-shards 143

务必先删除data/manifest.json：下载脚本会在本地缓存 manifest，而默认仓库（willdepueoai/parameter-golf，见 cached_challenge_fineweb.py）的旧 manifest 不包含 sp4096 条目，若不删除会导致变体找不到。该脚本按 manifest 中的files_train数量逐个拉取fineweb_train_*.bin与fineweb_val_*.bin，并把 tokenizer 的.model/.vocab一并落地到data/tokenizers/。

路径二：从零重建 tokenizer 与数据集

若想验证数据流水线，可先写入 4096 词表的规格文件，再调用导出脚本：

cat > data/tokenizer_specs_4096.json << 'EOF' [ { "name": "sp_bpe_4096", "kind": "sentencepiece_bpe", "vocab_size": 4096, "tokenizer_train_docs": 5000000 } ] EOF python3 data/download_hf_docs_and_tokenize.py \ --output-root data \ --tokenizer-config data/tokenizer_specs_4096.json \ --skip-byte

对照 download_hf_docs_and_tokenize.py 的源码，这条链路会依次完成：从 HF 数据集仓库拉取docs_selected.jsonl及其 sidecar → 用 SentencePiece Trainer 训练 BPE tokenizer（model_type="bpe"、byte_fallback=True、character_coverage=0.999、add_dummy_prefix=False、split_digits=True、normalization_rule_name="nmt_nfkc"，pad/bos/eos/unk 分别为 0/1/2/3）→ 对全部文档逐条编码 → 按SHARD_SIZE=1e8token 切分写盘。写入格式为 256×int32 头 +uint16token 流（DATAFILE_MAGIC=20240520、DATAFILE_VERSION=1），与加载端load_data_shard的校验逻辑严格对应。默认先取前 50,000 篇文档作为固定的验证集（NUM_VAL_DOCS=50000），后续全部作为训练集。

注意：预分词脚本要求训练数据存在docs_selected.jsonl；若只想快速复现训练，走路径一即可，无需本地重建。

运行环境与复现命令

本记录的脚本直接导入flash_attn_interface，必须使用 Flash Attention 3（Hopper 架构），官方环境为 PyTorch 2.11.0+cu130。安装命令：

pip install torch --index-url https://download.pytorch.org/whl/cu130 pip install --no-cache-dir \ "https://download.pytorch.org/whl/cu130/flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl" pip install -r requirements.txt

requirements.txt中除 torch 与 flash-attn-3 外仅需brotli、huggingface-hub、numpy、sentencepiece、tqdm（requirements.txt）。

在 8×H100 上复现记录的命令：

RUN_ID=1337 SEED=1337 \ torchrun --standalone --nproc_per_node=8 train_gpt.py

脚本从data/datasets/fineweb10B_sp4096/读取训练/验证分片，从data/tokenizers/fineweb_4096_bpe.model加载 tokenizer（路径由Hyperparameters中的datasets_dir/tokenizer_path模板拼接）。torchrun会注入RANK/WORLD_SIZE/LOCAL_RANK，脚本据此判断是否进入分布式模式，并在main()中要求8 % world_size == 0，保证grad_accum_steps = 8 // world_size为整数。

运行流程（train_and_eval）为：warmup 若干步后进入主训练循环，按墙钟时间（扣除预留的 GPTQ 时间）决定停止步数；结束时加载 EMA 权重，先做一次 pre-quantization 评估，再执行serialize（量化+压缩落盘）→deserialize（读回重建）→run_evals（int6 往返评估 + 滑动窗口评估）。日志写入logs/{RUN_ID}.txt，量化产物为final_model.int6.ptz。

量化与压缩：GPTQ 感知 Hessian + int6/int8 混合 + brotli

Hessian 收集

collect_hessians在训练结束后、量化前，用gptq_calibration_batches=64批校准数据跑一遍前向，通过 forward hook 对每个参数量大于 65536 的CastedLinear（mlp/attn 类）累加H = X^T X，得到逐层的二阶 Hessian 近似。这段计算从 600 秒预算中预留gptq_reserve_seconds=10秒。

GPTQ 主路径

gptq_quantize_weight实现了 Frantar et al. (ICLR 2023) 的 GPTQ，并带三处工程细节：

1. 死列清零与阻尼：对角为 0 的列置 1，再按0.01 × diag 均值加阻尼，保证 Cholesky 分解数值稳定；
2. actorder：按 Hessian 对角从大到小重排列，逐块（block_size=128）量化；
3. Cholesky 误差补偿：用cholesky_inverse+ 上三角 Cholesky 得到 H⁻¹ 的三角因子，量化误差沿块内列序反向传播到后续列（W_block[:, j:] -= err * Hinv_block[j, j:]）。

量化方式为int6（clip_range=31），并对 5 档分位（0.9990 / 0.9995 / 0.9999 / 0.99999 / 1.0）做逐行 clip 搜索，取重建 MSE 最小者；若 Cholesky 失败则回退到纯分位搜索的quantize_int6_per_row。小参数（≤65536 元素）与控制张量原样 FP16/FP32 透传，其余张量走 int8 逐行量化（quantize_float_tensor，clip 分位 99.99984%）。

压缩

量化结果经torch.save后，先做字节打乱（_byte_shuffle，按 stride=2 转置字节流，带BSHFmagic 头），再用brotli quality=11压缩（也可通过COMPRESSOR=lzma切换）。最终提交大小 = 压缩后模型字节 + 脚本代码字节，本记录三个种子均控制在 ~15.9MB（见submission.json的bytes_total与bytes_code=68206）。

结果复现与统计显著性

submission.json 记录了三种子完整结果：

值得注意的细节：滑动窗口评估（stride=64，每个 token 用尽可能长的前缀上下文）显著优于 int6 往返的整序列评估，因此提交的val_bpb 1.09785采用滑动窗口口径（3-seed 平均，std=0.0004）。这与挑战规则"允许任意序列长度评估"一致——上下文越长，长程依赖建模得越好。该分数在同日排行榜中位列4096-Vocab + Larger Model + High WD + Simplifications，作为 PR #1218 的记录（见根目录 README.md 排行榜）。

经验总结：这份配方教给我们什么

1. 压缩视角的参数量预算：在 16MB 硬约束下，权重衰减不只是正则化工具，更是"压缩体积调节阀"——RMS 与可压缩性的强相关性（R²≈0.99）意味着可以用 WD 精确控制模型体积，把预算花在词表、宽度等更有收益的地方。
2. 减法有时优于加法：TTT、QAT、hash embeddings、smear gate、gated attention、Parameter Banking 的移除说明，机制复杂度与收益不成正比时，删掉它们反而能腾出时间与字节预算。
3. 词表大小是免费午餐：1024 → 4096 的 tokenizer 只增加约 50KB，却让模型在相同序列长度下看到更多上下文、每步获得更多训练信号。
4. 量化管线与训练管线协同设计：GPTQ 的 Hessian 预留时间、int6 逐行 clip 搜索、字节打乱 + brotli，每一步都以"提交物字节数"为最终目标，而非仅看训练 loss。

若要继续深入，可对照阅读该记录的 train_gpt.py（1627 行完整训练/量化/评估链路）、上一代基线 2026-03-23_LeakyReLU_LegalTTT_ParallelMuon/README.md，以及数据侧工具 data/download_hf_docs_and_tokenize.py 与 data/cached_challenge_fineweb.py。