跳转至

Read this page in English

KITE/SST:公开小模型等算量诊断

这是 KITE/SST 的缩小机制检查,不是论文 MoE 规模或原文 benchmark 的复现。运行脚本:scripts/benchmark_kite_equal_flops.py。

采用公开 WikiText-2 原始文本,独立 train、validation、test 文件,按 UTF-8 字节做 32-token next-token 任务。固定种子 42、43、44。宽度 32 的 2 层 Prefiller 经两阶段扩成 2+2 层 SST;对照为 4 层 dense 模型。两者初始 embedding、输出头和前两层权重相同,但参数量分别为 64,704 与 68,800。

SST 训练 30+30 步;对照按 PyTorch profiler 可识别的前后向 FLOPs 匹配为 17+28 步。两者阶段边界均重启 AdamW,学习率 0.001。对照 FLOPs / SST FLOPs = 0.9993;计数不含 optimizer、LayerNorm 和内存移动,因而不是严格硬件算量相等。test 不参与方法或超参选择。

Seed SST val NLL Dense val NLL SST test NLL Dense test NLL SST 末 token ms Dense 末 token ms
42 3.7512 4.0244 3.7661 4.0508 0.279 0.346
43 3.5698 3.7940 3.6004 3.8177 0.279 0.343
44 3.5245 3.7784 3.5059 3.7568 0.281 0.356
PYTHONPATH=src python scripts/benchmark_kite_equal_flops.py \
  --data-root data --seeds 42,43,44 \
  --source-steps 30 --continuation-steps 30

仅能得出:在这个固定小任务和近似 FLOPs 预算下,SST 的 NLL 与本机末 token 延迟优于该 dense 对照。这里的“本机”不是 GPU 延迟,不能外推为论文 67B 扩展性或 GPU 吞吐。

A100 CUDA 推理微基准

另用 scripts/benchmark_kite_gpu_latency.py 测量宽度 32、随机初始化的 2+2 SST 与 4 层 dense 对照,固定种子 42、43、44,每个配置预热后重复 20 次。输入是程序生成的 byte token,不是 WikiText-2 训练后的权重,也不产生语言质量指标。decode-one 是完整前缀之后生成一个 token 的耗时;SST 的 prefill-last-only 仅计算与末 token 相关的 prefiller 输出,和全序列 prefill 分列报告。末 token logits 与完整 prefill 的最大绝对误差为 2.46e-7。

上下文 token SST 完整 prefill SST 末 token prefill SST decode-one Dense prefill Dense decode-one
32 1.421 1.417 1.435 1.374 1.447
128 1.423 1.422 1.438 1.378 1.454
256 1.441 1.431 1.455 1.406 1.457

单位为 ms,表中是三种子各自重复测量中位数的再中位数;脱敏 GPU 运行凭证记录了固定版本和命令。在此极小模型、短上下文的 launch-bound 微基准中,没有测出明确的 CUDA 延迟优势。它只关闭“小模型真实 CUDA prefill/decode 是否能跑”的验证项;原论文规模、长上下文、吞吐和真实训练权重仍待验证,不能据此宣称 KITE 论文效果已复现。