JAX 提速 2.27 倍:速度跑赛道把 nanGPT 训到 4.179 仅需 6.76 秒
RobertTLange · x · 2026-10-21
Robert Lange(前 Google Brain/Elite 研究员)为 qlabs-eng/noise 的速度跑(speedrun)赛道提交了 JAX 版 PR #29:在 8×H100 80GB 上,把验证 loss 训到 4.1792 只用 6.76 秒(从第一次更新到最后一次更新计时)。
- 相比同一 recipe 的 PR #25(15.35 秒)提速 2.27 倍,相比赛道纪录 #7(83.3 秒)提速 12.3 倍。
- 稳定性:同 recipe 再跑 5 个种子,平均 6.74±0.18 秒,val loss 4.162±0.016,五个种子全部低于 4.2。
- 方法要点:估计器和 recipe 完全不变——所有信号仍是同一前向测量,只通过 moment kernels(每 GPU 一个 head)把计算本身做快;改动集中在 train.py 和一个新 kernel 文件。该分支堆叠在 #25 之上,需先合 #25。
对 JAX/内核优化和训练加速感兴趣的读者值得看 diff。
「Infra」频道最新
- Super Micro 财测大超预期：FY27 Q1 营收指引 185–195 亿美元 — firstadopter · 2026-10-21
- PyTorch 大会 Demo Theater 下午场：IBM、Modal、Arm 等 5 场演示日程 — PyTorch · 2026-10-21
- SPCX 算力达 1.7 GW，Q4 冲 2.5 GW 可期 — JOBhakdi · 2026-10-21
- 开发者开源 LLM 成本计算器：接 OpenRouter 实时报价对比 250+ 模型 — zackku · 2026-10-21
- 马斯克放话：2027 年目标 13GW 算力，单集群 44 万块 Vera Rubin — scaling01 · 2026-10-21
- VectorDash 用 Rust 重写:Windows 本地搜索+剪贴板+AI 助手合一 — andy2na · 2026-10-21