所有课程

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 与训练框架结合使用

梯度检查点（再物化）

混合精度训练

模型并行策略

大规模优化算法

实践：实现梯度检查点

嵌套 pmap 与进阶分区

这部分内容有帮助吗？

参考文献

Parallelism in JAX: pmap, JAX Developers, 2024 (JAX Documentation) - 解释了 pmap 转换、轴命名和集体操作，是理解嵌套并行化的基础。
Distributed Arrays and shard_map in JAX, JAX authors, 2024 (JAX Documentation) - 详细介绍了高级数据分区、显式设备网格以及用于细粒度分布式数组管理的 jax.Array 系统，包括 shard_map。

© 2025 ApX Machine Learning用心打造