🚀 JAX 模型部署与推理加速
实战课程 · 30章完整目录
⚡ JIT / XLA
🧠 Flax / Haiku
📦 Triton / Serving
🎯 从入门到集群
📘 共30章 · 点击卡片跳转
01
JAX生态全景
定位差异 · 自动微分 · JIT编译 · AlphaFold/Gemini案例
02
JAX基础入门
jax.numpy · 随机数 · DeviceArray · @jit初体验
03
自动微分与梯度计算
grad/vmap/pmap · custom_vjp · 高阶梯度 · Hessian
04
JIT编译深度优化
静态/动态参数 · 三值规则 · 编译缓存 · lower/compile
05
函数式纯化与状态管理
纯函数 · lax.scan · Stateful类 · 随机数种子模拟
06
并行计算与多设备部署
pmap数据并行 · shard_map · 多主机 · NCCL all-reduce
07
模型构建与Flax框架
Module系统 · 参数初始化 · optimizer · 与原生JAX互操作
08
Haiku与Optax生态
模块化设计 · 优化器链式组合 · 梯度裁剪 · 自定义优化器
09
模型序列化与Checkpoint
orbax/flax.serialization · 分片检查点 · 模型互转
10
JAX Serving与模型导出
jax2tf · export导出XLA HLO · ONNX方案对比
11
TensorFlow Serving集成
TFServing部署 · gRPC/RESTful · 动态batching · 监控
12
NVIDIA Triton Inference Server
模型仓库 · 自定义后端 · 并发执行 · Triton Client
13
XLA编译与AOT编译
AOT原理 · jax.experimental.aot · PJRT架构 · 跨平台
14
量化与模型压缩
QAT · lax降精度 · 混合精度bf16/fp16 · TensorRT集成
15
模型蒸馏与知识迁移
蒸馏训练流程 · 教师-学生 · logits匹配 · 部署优化
16
推理引擎对比与选型
XLA:GPU/TPU vs PJRT · 原生vs ONNX · 基准测试 · 决策树
17
流水线并行与模型分片
device_put手动分片 · GSPMD · 微批次调度 · Megatron-LM
18
异步推理与流式处理
jit异步执行 · asyncio · Token-by-token · WebSocket
19
内存管理与显存优化
XLA Allocator · with_sharding_constraint · 梯度检查点 · 碎片整理
20
性能Profiling与调优
jax.profiler · TensorBoard/Chrome Trace · 计算通信重叠
21
分布式训练与推理一致性
pmap策略复用 · 参数同步 · 多副本负载均衡 · 容错
22
边缘设备部署
JAX转TFLite · ARM CPU推理 · 量化移动端适配 · MediaPipe
23
Web端推理（WASM）
编译到WebAssembly · 浏览器推理 · 内存精度权衡 · vs TF.js
24
安全与模型保护
加密推理 · lax同态计算HE · 模型水印 · SGX/TEE集成
25
CI/CD与自动化部署
Docker镜像 · Kubernetes配置 · A/B测试 · 版本管理回滚
26
监控与可观测性
延迟/吞吐量监控 · 模型漂移检测 · Prometheus · 告警
27
成本优化与资源调度
Spot实例 · 自动扩缩容HPA/VPA · GPU利用率 · 中断恢复
28
多框架混合部署
JAX+PyTorch共存 · 跨框架张量转换 · 统一推理网关 · 异构调度
29
大规模推理集群实战
千卡集群架构 · 全局负载均衡 · KV Cache共享 · 故障域隔离
30
前沿趋势与未来展望
MLIR/StableHLO · 稀疏计算MoE · 自动驾驶/医疗 · 社区资源