本文从 PJRT、IFRT、SPMD、MPMD、单控制器和多控制器等概念出发,介绍 Google Pathways 的系统动机与核心设计。
《How to Scale Your Model》完整中文翻译,覆盖 Roofline 分析、TPU/GPU 架构、分片矩阵乘法、Transformer 并行策略、LLaMA 训练与推理、性能分析等全部内容
JAX Pallas 文档完整中文版,覆盖快速入门、Grid/BlockSpec、软件流水线、TPU 内核编写(矩阵乘法、稀疏计算、分布式、SparseCore)、Mosaic GPU 内核编写(Blackwell 矩阵乘法、集合通信)、设计笔记等全部内容
在 GKE 上对 TPU v6e 2×2 集群进行 ICI 带宽实测,验证 device_put 跨 mesh 走 ICI 直连不经 Host,实测单条 Link 单向 93 GB/s、psum 4-device 聚合 445 GB/s,并分析标称 800 GB/s 与实际带宽的差距
编译器通过一行@jax.jit 将8行代码就优化到了250 VLIW bundles. 本文通过HLO/LLO追踪了 TPU 编译器的优化之旅
SGLang 的 JAX 后端实现,支持在 TPU 上运行高效推理。