parameter-golf 16MB 极限压缩实战：11L MLP3x + Int6 QAT + zstd-22 记录详解-平芜编程栈

【免费下载链接】parameter-golf

Train the smallest LM you can that fits in 16MB. Best model wins!

项目地址：https://gitcode.com/gh_mirrors/pa/parameter-golf

导读

本文解析 parameter-golf（OpenAI Model Craft Challenge）仓库中一个达到val_bpb = 1.1502（3 个种子均值）的经典提交记录：11 层 Transformer + 3x MLP 扩展 + Int6 量化感知训练 + zstd-22 压缩 + 滑动窗口评估。通过对照该记录目录下的 README.md、submission.json 与 train_gpt.py，读者可以掌握一条完整的"训练—量化—压缩—评估"技术链路：如何在 16MB 工件预算与 8×H100 十分钟训练预算下，用 Int6 量化省出的空间换取更深的网络，再用 zstd-22 与滑动窗口评估把每一点精度都榨干。

背景：在 16MB 与 10 分钟约束下做 L(N) 优化

parameter-golf 的挑战规则非常明确：训练一个压缩后工件不超过 16MB、在8×H100 上训练不超过 10 分钟的语言模型，按 FineWeb 验证集上的"与分词器无关的压缩率（bits per byte, BPB）"排名（见仓库根目录 README.md）。正如项目根 README 所述，这可以理解为一种 L(N) 优化：在参数量（N）固定的前提下，用数据、算力、步数、架构的自由度换取最低损失。因此，参赛记录的每一次精度提升，几乎都来自"压缩省空间 → 换取更多参数/更深网络 → 用更聪明的评估方式找回精度"的正循环。

本记录正是这一思路的典型样本：基线是 9 层 512 宽、2x MLP 扩展、int8 量化 + zlib 压缩；本记录通过11 层 + 3x MLP + Int6 QAT将模型做大，同时用zstd-22把压缩率再推高一档，最后用滑动窗口评估白捡约 0.034 BPB，最终跨 3 个种子稳定跑出 1.1502 BPB。

提交总览与关键结果

指标	数值
模型参数量	26,502,020（约 26.5M）
压缩后工件	均值 15,427,455 字节（约 15.4MB，3 个种子均在 16MB 内）
val_bpb（滑动窗口）	1.15015359（3 种子均值）
val_bpb（roundtrip，重载后复评）	1.18454250
val_loss（滑动窗口，nats）	1.94198447
相对基线提升	0.1307 nats
t 统计量	313.20（df=2，p << 0.001）
训练步数 / 单步耗时	~10,070 步 / 59.6 ms
硬件	8×H100 SXM（RunPod），600 秒训练

以上数据均可在 submission.json 中逐项核对，该文件是train_gpt.py运行时自动生成的官方记录。

架构：11 层 U-Net 变压器

在 Int6 量化与 zstd 压缩"腾出"的预算内，模型从基线的 9 层加宽到 11 层，具体规格为：

11 个 Transformer 块，模型维度 512；
GQA 注意力：8 个查询头、4 个 KV 头，RoPE 位置编码；
ReLU² MLP，3x 扩展：隐藏维度 512×3 = 1536（对应 train_gpt.py 中MLP类的torch.relu(self.fc(x))后接self.proj(x.square())）；
Tied embeddings：输入输出共享词嵌入，1024 BPE 词表（SentencePiece）；
U-Net 跳跃连接：5 个编码层 + 6 个解码层，解码层按逆序逐层叠加skip_weights缩放后的编码层残差；
26.5M 参数，zstd-22 压缩后约 15.4MB 工件。

从 train_gpt.py 的GPT类（约 L733 起）可以看到 U-Net 的实现方式：num_encoder_layers = num_layers // 2，前 5 层把中间激活压入skips栈，后 6 层按skips.pop()逆序叠加skip_weights[i]，与深度上的残差resid_mix一起构成双层跳连结构。这种结构让 11 层网络在深层信息传递上更高效，是小模型在有限参数下常用的"深度换宽度"手段。

Int6 QAT：训练时就把量化噪声"喂"给模型

本记录最核心的改动是Int6 量化感知训练（QAT）。思路是：既然导出时要把权重压到 int6，不如在训练中就用直通估计器（STE）模拟 int6 舍入误差，让模型提前适应低精度。

训练时的模拟量化

在 train_gpt.py 的CastedLinear.forward（约 L567 起）中，开启_qat后每个前向都会：

按每行绝对值的最大值amax计算缩放scale = amax / 127.0；
对权重做q_raw = (w_f / scale).round()；
若为 Int6 模式（_qat_int6），再做q = (q_raw / 4).round() * 4并clamp(-128, 124)——int6 等价实现：用 int8 存储但只保留 4 的倍数，即 64 个量化等级；
用 STE 技巧w = w + (w_q - w_f).detach()让前向使用量化权重、反向传播仍按原始权重计算梯度。

关键点：脚本并非真正实现 6-bit 位宽，而是**"int8 容器 + 4 的倍数"，64 级离散值正好等价于 6 比特信息量，同时规避了非对齐位宽的工程复杂度。在 quantize_float_tensor_int6（导出阶段）中可以看到完全一致的舍入逻辑，保证训练模拟与导出量化严格对齐**——这正是 QAT 有效的先决条件。

后训练量化导出

训练结束后，quantize_state_dict_int8（约 L385 起）按如下规则导出：

2D 矩阵：per-row int8（或 int6），每行一个 fp16 缩放因子，按 99.99984 分位数裁剪后再量化，降低离群值影响；
非 2D 张量与小于 65,536 元素的小张量：直接以fp16透传（INT8_KEEP_FLOAT_MAX_NUMEL = 65_536），避免元数据反客为主；
tied embedding 单独走 fp16 透传（fp16_embed开关）：输入/输出头共享的 1024×512 矩阵精度敏感，用 fp16 保留质量；
层范围由INT6_LAYER_START/END控制，本记录设为 0–10，即11 层全部块权重走 int6。

导出的格式标记为int8_clean_per_row_v1，保存量化值、缩放、原始 dtype 与 per-row 元数据；最终写出final_model.int8.ptz。

zstd-22：比 zlib 再省约 1.5MB

量化后的 state_dict 先用torch.save序列化，再压缩。脚本默认回退方案是zlib.compress(..., level=9)；当安装了zstandard且USE_ZSTD=1时，改用zstd_mod.ZstdCompressor(level=args.zstd_level)（train_gpt.py）。

zstd-22 的收益在 README 中被明确量化：比 zlib 节省约 1.5MB。对一个 16MB 上限的挑战而言，这 1.5MB 恰好是"11 层 + 3x MLP"能塞进预算的关键，压缩环节因此从"最后收尾"变成了"架构决策的输入"。运行前提：pip install zstandard。

值得留意的是，README 中报出的 val_bpb（1.1502）是滑动窗口评估口径，而 roundtrip 口径（重载量化工件后在相同脚本内复评）为 1.18454——两个口径相差约 0.034 BPB，这正是滑动窗口评估贡献的"免费精度"，详见下文。

优化器与学习率：低 LR + 解耦权重衰减 + 动量预热

本记录对优化器做了三处针对性调整，全部指向"让权重更适合量化"：

解耦权重衰减（WD=0.04）

Muon 与 AdamW 均启用 0.04 的解耦权重衰减。其中 Muon 侧的实现并非标准权重衰减，而是在opt.step()之后手动执行p.mul_(1.0 - muon_lr * muon_weight_decay)（train_gpt.py）。较小的权重范数直接缩小了量化的动态范围，让 per-row 缩放更稳、舍入误差更小——README 明确称其"keeps weights small and quantization-friendly"。

三组学习率

参数组	学习率	优化器
矩阵参数（Transformer 块内 2D 权重）	`MATRIX_LR=0.025`	Muon
标量/向量/控制张量	`SCALAR_LR=0.025`	AdamW
Tied embedding（tok_emb 权重）	`TIED_EMBED_LR=0.035`	AdamW

从 train_gpt.py 的优化器拆分逻辑可见：矩阵参数走 Muon（内含zeropower_via_newtonschulz5的 Newton-Schulz 正交化，约 L116），标量与控制张量走 AdamW，embedding 独立一组 AdamW。三组学习率都比基线（0.04）显著调低，README 将其归为第 10 项改动。

Muon 动量预热（0.92 → 0.99）

动量从 0.92 起步，在 1500 步内线性爬升到 0.99。实现位于训练主循环（约 L1256-L1259）：frac = min(step / muon_momentum_warmup_steps, 1.0)，muon_momentum = (1-frac)*start + frac*target，每个 step 写入优化器参数组。训练初期动量低 → 更新幅度更保守，配合 3000 步的 warmdown（WARMDOWN_ITERS=3000）共同构成"慢启动、稳收尾"的节奏，为后期量化友好的小权重收敛创造条件。

滑动窗口评估：每 64 token 计一次分

验证时，脚本支持两种口径（train_gpt.py）：

普通 roundtrip：按 train_seq_len=1024 切块，每块独立前向，val_bpb = 1.18454；
滑动窗口：stride=64，窗口起点range(0, total_len - seq_len, stride)，每窗口只对最后 64 个 token（score_start = seq_len - stride）计分。

滑动窗口让每个 token 都以"前面最多 1024 token 的上下文"被评估，相当于把模型的实际有效上下文喂满，得到更低的 loss——本记录中约 0.034 BPB 的提升即来源于此。评估实现eval_val_sliding在 rank 间交错分配窗口位置（rank_positions = positions[rank::world_size]），全程torch.inference_mode+ bf16 autocast，3 个种子的滑动评估合计约 88 秒，仍在 10 分钟评估预算内（README 注明"Sliding window eval takes ~88s on 8xH100"）。

注意：滑动窗口并非"作弊"，而是改变评估口径；提交时须同时诚实报告 roundtrip 口径（见 submission.json 的val_bpb_roundtrip字段），挑战对口径有合规要求。

完整复现命令与全部环境变量

运行前需安装zstandard：

pip install zstandard

本记录的标准复现命令（来自 README.md）：

RUN_ID=submission \ SEED=1337 \ NUM_LAYERS=11 \ MLP_MULT=3 \ MATRIX_LR=0.025 \ SCALAR_LR=0.025 \ TIED_EMBED_LR=0.035 \ FP16_EMBED_EXPORT=1 \ INT6_LAYER_START=0 \ INT6_LAYER_END=10 \ QAT_ENABLED=1 \ QAT_INT6=1 \ MUON_WEIGHT_DECAY=0.04 \ ADAM_WEIGHT_DECAY=0.04 \ MUON_MOMENTUM=0.99 \ MUON_MOMENTUM_WARMUP_START=0.92 \ MUON_MOMENTUM_WARMUP_STEPS=1500 \ WARMDOWN_ITERS=3000 \ USE_ZSTD=1 \ EVAL_STRIDE=64 \ torchrun --standalone --nproc_per_node=8 train_gpt.py

以上环境变量在 train_gpt.py 的Hyperparameters类（约 L38-L107）中逐一定义，除本记录用到的取值外，还有一批默认值与可选项，整理如下：

环境变量	本记录取值	默认值	作用
`NUM_LAYERS`	11	9	Transformer 层数（U-Net 编码/解码层自动分配）
`MLP_MULT`	3	2	MLP 隐藏层扩展倍数（hidden = MLP_MULT × model_dim）
`MODEL_DIM`/`NUM_HEADS`/`NUM_KV_HEADS`	512 / 8 / 4	同左	模型宽度与 GQA 头数
`VOCAB_SIZE`	1024	1024	BPE 词表大小，须与 tokenizer 匹配
`FP16_EMBED_EXPORT`	1	0	tied embedding 导出为 fp16（保留质量）
`INT6_LAYER_START`/`INT6_LAYER_END`	0 / 10	-1 / -1	int6 量化的层范围（-1 表示关闭）
`QAT_ENABLED`/`QAT_INT6`	1 / 1	0 / 0	开启 QAT 模拟量化 / 使用 int6 模式
`MATRIX_LR`/`SCALAR_LR`/`TIED_EMBED_LR`	0.025 / 0.025 / 0.035	0.04 / 0.04 / 0.05	矩阵、标量、embedding 三组学习率
`MUON_WEIGHT_DECAY`/`ADAM_WEIGHT_DECAY`	0.04 / 0.04	0.0 / 0.0	解耦权重衰减
`MUON_MOMENTUM`	0.99	0.95	Muon 最终动量
`MUON_MOMENTUM_WARMUP_START`/`_STEPS`	0.92 / 1500	0.85 / 500	动量预热起点与步数
`WARMDOWN_ITERS`	3000	1200	学习率线性退火步数
`USE_ZSTD`/`ZSTD_LEVEL`	1 / 22	0 / 22	使用 zstd 压缩（否则回退 zlib level 9）
`EVAL_STRIDE`	64	0	滑动窗口步长（0 表示关闭滑动评估）
`EVAL_SEQ_LEN`/`EVAL_BATCH_SIZE`	0 / 32	0 / 32	滑动评估窗口长度与批大小
`SEED`	1337 / 42 / 123	1337	随机种子（3 种子取均值）
`ITERATIONS`	20000（受墙钟截断）	20000	目标训练步数
`MAX_WALLCLOCK_SECONDS`	600.0	600.0	训练墙钟上限，到时即停

训练脚本自带WARMUP_STEPS=20步编译预热：先完整跑一小段前反向，把torch.compile的编译成本消化掉，再恢复初始权重与优化器状态，确保计时从真实训练起点开始——这是 10 分钟计时公平性的关键实现（train_gpt.py）。

数据路径与分词器同样由环境变量控制（DATA_PATH、TOKENIZER_PATH，默认指向./data/datasets/fineweb10B_sp1024与./data/tokenizers/fineweb_1024_bpe.model），脚本仅支持 SentencePiece.model分词器，并校验词表大小一致后才开始训练。

多种子统计：0.1307 nats 的提升不是噪声

本记录用 3 个种子（1337 / 42 / 123）验证稳定性：

Seed	slide_loss (nats)	slide_bpb	rt_bpb	Artifact（字节）
1337	1.94265607	1.15055135	1.18484075	15,360,260
42	1.94207795	1.15020896	1.18456681	15,556,813
123	1.94121940	1.14970047	1.18421993	15,365,293
Mean	1.94198447	1.15015359	1.18454250	15,427,455
Std	0.00072288

（数据来源：README.md 与 submission.json）

三个种子 BPB 的标准差仅约 0.0004，均值相对基线提升 0.1307 nats，t 统计量高达 313.20（df=2，p << 0.001），统计显著性毋庸置疑；三个工件事务大小在 15.36MB～15.56MB 之间，全部落在 16MB 上限之内。这份记录因此成为 2026-03-19 时点参数高尔夫赛道上"深度 × 量化 × 压缩"组合的可靠基准。