⚡ JAX 全栈性能分析调优实战
📘 30章 · 从入门到调优专家
🎯 目录 · 点击跳转
30/30
1
CH01
JAX初探
JAX是什么？为什么需要JAX？与NumPy、TensorFlow、PyTorch对比
➤
2
CH02
环境搭建
安装JAX (CPU/GPU/TPU) · 验证安装 · VS Code + Jupyter 配置
➤
3
CH03
JAX核心概念
函数式编程 · 不可变状态 · 纯函数 · JIT编译原理
➤
4
CH04
JAX基础操作
jnp.array · 设备内存管理 · PRNGKey · 向量化操作
➤
5
CH05
自动微分入门
grad() · 高阶导数 · jacfwd/jacrev · vmap与pmap
➤
6
CH06
JIT编译深入
jit()装饰器 · 静态参数 · 动态形状 · 缓存与重编译
➤
7
CH07
性能分析工具概览
jax.profiler · perfetto · nsys (NVIDIA Nsight)
➤
8
CH08
基准测试与计时
block_until_ready() · timeit · 自定义基准框架
➤
9
CH09
内存分析
设备内存分配 · 泄漏检测 · XLA内存优化
➤
10
CH10
计算图可视化
xla_computation · HLO导出 · TensorBoard集成
➤
11
CH11
性能瓶颈定位
jax.profiler追踪 · 火焰图 · 热点函数识别
➤
12
CH12
向量化优化
vmap陷阱 · 批量大小 · 内存带宽与计算强度
➤
13
CH13
并行计算策略
pmap多设备 · 数据并行 · 模型并行 · 流水线并行
➤
14
CH14
编译优化技巧
融合操作 · 代数简化 · 布局优化 · 常量折叠
➤
15
CH15
数据类型与精度
bfloat16/float16 · 混合精度 · jnp.float0
➤
16
CH16
控制流优化
lax.cond/while_loop/scan · 静态与动态循环
➤
17
CH17
随机数生成优化
PRNG链式调用 · 并行随机 · 确定性随机数
➤
18
CH18
自定义VJP规则
自定义梯度 · implicit differentiation · 检查点策略
➤
19
CH19
XLA编译器深入
HLO与MHLO · 编译流水线 · AOT编译 · 持久化缓存
➤
20
CH20
TPU性能调优
TPU架构 · 矩阵单元 · 数据加载 · 编译预热
➤
21
CH21
GPU性能调优
CUDA内核融合 · 共享内存 · warp调度 · NCCL通信
➤
22
CH22
CPU性能调优
多线程 · SIMD向量化 · 缓存友好 · MKL后端
➤
23
CH23
数据加载与预处理
tf.data集成 · 内存映射 · 预取 · 数据混洗
➤
24
CH24
大规模训练优化
梯度累积/压缩 · ZeRO优化器 · 重计算
➤
25
CH25
模型推理优化
量化 · 剪枝 · 蒸馏 · JIT推理 · 批处理推理
➤
26
CH26
分布式训练调试
NCCL超时 · 拓扑感知 · 负载不均 · 梯度同步
➤
27
CH27
性能回归测试
CI集成 · 性能基准 · 阈值告警 · 历史趋势
➤
28
CH28
端到端案例
ResNet训练调优 · Transformer推理 · GAN生成调优
➤
29
CH29
高级调试技巧
NaN/Inf追踪 · 断点调试 · 打印张量 · 反编译HLO
➤
30
CH30
未来趋势与生态
JAX 2.0展望 · 社区生态 · 与TF/PyTorch互操作 · 贡献指南
➤