趋近智
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 与训练框架结合使用
梯度累积
梯度检查点（再物化）
混合精度训练
模型并行策略
大规模优化算法
实践：实现梯度检查点
梯度检查点（再物化）
这部分内容有帮助吗？
- Training Deep Nets with Sublinear Memory Cost, Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin, 2016 arXiv preprint arXiv:1604.06174 DOI: 10.48550/arXiv.1604.06174 - 本文介绍了梯度检查点（也称为激活检查点或重新物化）技术，以减少深度神经网络训练期间的内存消耗。
- jax.checkpoint, JAX developers, 2024 -
jax.checkpoint（别名为jax.remat）转换的官方文档，提供了在JAX中进行内存优化的实际使用细节和示例。 - The Reversible Residual Network: Backpropagation Without Storing Activations, Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse, 2017 Advances in Neural Information Processing Systems 30 (NIPS 2017) (Curran Associates Inc.) - 本文介绍了可逆残差网络，这是一种通过在反向传播过程中允许激活的完美重构来进一步优化内存的架构设计，它扩展了重新计算的原理。
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He, 2020 SC20: International Conference for High Performance Computing, Networking, Storage and Analysis (IEEE) DOI: 10.1109/SC41405.2020.00024 - 本文介绍了ZeRO，一套用于大规模分布式模型训练的内存优化技术，其框架中整合了激活检查点策略。