逐内核剖析 JAX vs PyTorch:同为 6.5ms 为何 PyTorch 更快
RobertTLange · x · 2026-10-21
JAX 老兵 RobertTLange(evosax 作者)用一手实验重新审视「JAX vs PyTorch」之争。
- 背景：他多年偏爱 JAX,jit + v/pmap 让进化策略在大规模种群评估上体验极佳。在 Noise benchmark 零阶(无反向传播)赛道上，他用 JAX + Pallas 内核创下纪录(#29:8×H100 上 975 次更新，val loss < 4.2)。
- 关键发现：benchmark 维护者把他的方案原样移植到 PyTorch(相同估计器、配方、任务划分与集合通信)，在同一节点上反而更快——PyTorch 6.35s vs JAX 6.8–7.6s。该 benchmark 此后规定新条目必须用 PyTorch。
- 他随后对两者做逐内核 profiling:在 GPU 层面两者基本是同一个程序，典型单次更新 PyTorch 耗时 6.49ms,与 JAX 几乎持平，差异来自框架层面的开销分布。
「Infra」频道最新
- GitHub 出现故障：Actions、PR 与 Webhook 全面延迟 — DanielLockyer · 2026-10-21
- NVIDIA 联手六大封测厂扩产：基板产能 3 年翻 2.5 倍 — nvidia · 2026-10-21
- 延迟数小时后 Node.js 24.22.1 全平台 Docker 镜像终于发布 — DanielLockyer · 2026-10-21
- 美光谈存储协同设计：与客户从共研到量产的 AI 存储五段路线 — BenBajarin · 2026-10-21
- 孟菲斯投票通过暂停新建数据中心 135 天 — Polymarket · 2026-10-21
- 美光发布 AI 级 NAND 设计：高 IOPS 与高带宽两条产品线瞄准 KV cache — BenBajarin · 2026-10-21