JAX 是什么、为什么需要 JAX、与 TensorFlow/PyTorch 对比、核心设计哲学。
安装 JAX (CPU/GPU/TPU)、JAX 与 NumPy 第一次亲密接触、设备数组基础。
jit 即时编译原理、使用与性能分析、静态参数与动态形状。
vmap 向量化映射、pmap 数据并行、grad 自动微分、三者组合使用。
XLA 编译器入门、HLO 与 StableHLO、设备抽象与后端选择。
CUDA 与 cuDNN 在 JAX 中的角色、GPU 内存管理、异步调度与流处理。
TPU 架构概览、v4/v5e/v5p 核心特性、JAX 在 TPU 上的最佳实践。
显存分析工具、梯度检查点、混合精度训练 (bfloat16/float16)、重计算与卸载。
tf.data 与 JAX 配合、多进程数据加载、流水线优化、缓存策略。
Flax 框架入门、nn.Module 与 nn.Sequential、参数初始化策略。
自定义训练循环、损失函数设计、优化器选择 (SGD/Adam/LAMB)、学习率调度。
pmap 数据并行实战、all-reduce 通信原语、梯度累积与同步。
模型并行与张量并行、流水线并行、FSDP (全分片数据并行)。
amp 与 autocast、损失缩放策略、精度与性能的权衡、实战案例。
JAX profiler 使用、NVIDIA Nsight 集成、TPU profiling、瓶颈定位。
手动算子融合、XLA 自动融合、自定义 cuda kernel、Triton 与 JAX。
padding 策略、scan 与 while_loop、动态批处理、ragged tensor 方案。
JAX 纯函数随机数设计、PRNG 序列管理、分布式随机数生成。
saved_model 导出、Serving 部署、量化与压缩、ONNX 转换。
GPT/LLaMA 在 JAX 中的实现、激活检查点、张量并行训练。
ViT 训练、扩散模型推理优化、混合精度与内存优化。
RL 环境集成、JAX 加速的 RL 算法、大规模并行环境模拟。
ODE 求解、物理模拟、JAX 与 NumPy 互操作、自定义 vjp。
vjp 与 jvp、自定义梯度函数、Hessian 计算、优化器二阶信息。
GPU/TPU/CPU 混合训练、云平台配置 (GCP/AWS)、容器化与 Kubernetes。
Optax 优化器库、Orbax 检查点管理、Einshape 张量操作、其他生态工具。
基准测试方法论、吞吐量与延迟测量、可重复性实验设计、报告生成。
JAX 2.0 新特性、Pallas 自定义 kernel、Pathways 系统、未来趋势。
从零构建分布式训练系统、端到端性能调优、项目总结与最佳实践。