趋近智
JAX 高级编程：性能、优化与规模化
章节 1: 高级 JAX 变换与控制流
核心转换回顾：jit, grad, vmap
精通 lax.scan 处理序列操作
使用 lax.cond 进行条件执行
使用 lax.while_loop 进行循环
结合控制流与变换
高级掩码技术
了解闭包和JAX暂存
实践：实现复杂的循环逻辑
章节 2: JAX 代码的性能优化
在 CPU、GPU 和 TPU 上分析 JAX 代码
理解 JAX 计算图 (jaxpr)
XLA 编译的作用
内存布局及其对性能的影响
避免重复编译
算子合并与操作优化
异步调度
实践：优化数值计算
章节 3: JAX 分布式计算
并行处理原理介绍
JAX 中的设备管理
使用 pmap 进行单一程序多数据 (SPMD) 运算
使用 pmap 实现数据并行
集合通信原语 (psum, pmean 等)
pmap 中轴名称的处理
嵌套 pmap 与进阶分区
多主机编程简介
实践：分布式数据并行训练
章节 4: 进阶自动微分方法
前向和反向模式自动微分回顾
雅可比向量积 (JVPs) 与 jax.jvp
向量-雅可比积 (VJPs) 与 jax.vjp
高阶导数
计算完整的雅可比矩阵和海森矩阵
使用 jax.custom_vjp 的自定义微分规则
jax.custom_jvp 自定义求导规则
通过控制流原语求导
处理不可微分函数
实践：实现自定义梯度
章节 5: 互操作性与自定义操作
JAX 与 NumPy 的集成
使用 DLPack 实现零拷贝数据共享
使用 jax.experimental.host_callback 调用外部 CPU/GPU 代码
使用 jax.pure_callback 进行无副作用调用
JAX 原语简介
定义自定义原语
实现抽象求值规则
为后端（CPU/GPU/TPU）实现转换规则
制定自定义原语的求导规则
实践：整合 C++ 函数
章节 6: 大规模模型训练方法
大型模型训练中的挑战概述
JAX 生态系统库（Flax, Haiku）简介
模型参数和状态的处理
将 pmap 与训练框架结合使用
梯度累积
梯度检查点（再物化）
混合精度训练
模型并行策略
大规模优化算法
实践：实现梯度检查点
梯度累积
这部分内容有帮助吗？
- Deep Learning, Ian Goodfellow, Yoshua Bengio, and Aaron Courville, 2016 (MIT Press) - 介绍了小批量梯度下降的概念，并讨论了近似大批量数据的方法。
- Language Models are Few-Shot Learners, Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei, 2020 Advances in Neural Information Processing Systems (NeurIPS) 33 DOI: 10.48550/arXiv.2005.14165 - 描述了超大型语言模型GPT-3的训练过程，其中梯度累积等技术对于内存管理至关重要。
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He, 2020 International Conference for High Performance Computing, Networking, Storage and Analysis (SC) DOI: 10.48550/arXiv.1910.02054 - 提出了一系列内存优化技术，包括利用或补充梯度累积以进行超大规模模型训练的方面。
- Large Model Training Best Practices, Flax Developers, 2024 (Flax) - Flax 官方指南，详细介绍了训练大型模型的策略，明确涵盖了梯度累积作为核心技术。