手写 Mojo Metal 内核在 M4 Max 上训练 GPT-2，快过 PyTorch MPS 1.71 倍
ulmentflam · reddit · 2026-07-22
- 作者把 Karpathy 的 llm.c 移植到 Mojo，并补上了 Metal 后端，让 GPT-2 124M 能在 M4 Max 上训练，训练时不依赖 PyTorch 或 CPython。
- 实测里，Mojo/Metal 的 bf16 路径达到 8138 tok/s，相比 PyTorch MPS bf16 快 1.71 倍；但 MLX 仍然更快，最高到 10077 tok/s。
- 文章拆解了主要瓶颈：matmul 占了大约 70% 的 step 时间；作者的 Metal bf16 matmul 只比 fp32 快约 1.1 倍，而 MLX 的 bf16 能吃到 tensor cores，接近 2 倍 提升。
- 正确性方面，项目用梯度对比、10 步 loss 轨迹和 235 个等价性测试做门禁；仓库里还有一个 FineWeb checkpoint，在 HellaSwag 上拿到 29.53%。
- 作者强调核心 kernel 和 trainer 代码都靠手写完成，并表示 Mojo 在跨硬件可移植性上没有想象中“编译器自动兜底”，反而需要更多按设备/厂商分支。
「Infra」频道最新
- Primis 探索算力对冲基础设施，想把价格做稳 — dolos_diary · 2026-07-22
- PyTorch 基金会发布六个托管项目季度更新 — PyTorch · 2026-07-22
- Puri.li 开放免费网页搜索 API，索引已达 1.85 亿页 — skillplayed · 2026-07-22
- 开源编码 Agent Octomind 推出云端机器和 21 模型接入 — donk8r · 2026-07-22
- Atome LM 在 18 项 MCU 任务上持平或胜出，体积小 5 到 70 倍 — themoroccanship · 2026-07-22
- 算力短缺或让亚马逊微软谷歌反而提价增厚利润 — RihardJarc · 2026-07-22