📅 创建时间：2026-06-03 🏷️ 标签：#XLA #HLO #JAX #Google #SPMD #CollectiveOps #Fusion #CompilerBackend 📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络） 📚 相关知识：[[06-frontend-formats]]（前端格式） [[12-scheduling]]（调度）
XLA 内部机制：HLO、融合与 SPMD / XLA Internals, HLO, Fusion, and SPMD
┌─────────────────────────────────────────────────────────────────────────────┐
│ 场景：分布式训练中的 SPMD 困惑 │
├─────────────────────────────────────────────────────────────────────────────┤
│ 你用 JAX 写了一个分布式训练代码： │
│ │
│ @jax.jit │
│ def train_step(params, batch): │
│ def loss_fn(params): │
│ return loss(params, batch) │
│ grads = jax.grad(loss_fn)(params) │
│ return optax.apply_updates(params, grads) │
│ │
│ # 运行在 8 个 TPU 核心上 │
│ train_step = pmap(train_step, axis_name='devices') │
│ │
│ 每个 TPU 上跑同一个 jitted 函数，但 XLA 如何知道如何切分数据和计算？ │
│ 这套 SPMD 自动并行的魔法到底是怎么实现的？ │
└─────────────────────────────────────────────────────────────────────────────┘第1节 XLA 架构总览
1.1 为什么需要 XLA
XLA（Accelerated Linear Algebra） 是 Google 开源的深度学习编译器，最初为 TensorFlow 设计，后来成为 JAX 的核心后端。它的核心价值在于：
- 硬件抽象：同一套代码可以跑在 CPU、GPU、TPU 上
- 自动优化：算子融合、常量折叠、代数化简
- 内存优化：减少 HBM 访问，隐藏延迟
- 自动并行：SPMD 自动分区，无需手动管理通信
1.2 XLA 编译流水线架构
┌─────────────────────────────────────────────────────────────────────────────┐
│ XLA Compilation Pipeline │
└─────────────────────────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────────────────┐
│ Client Layer（Python API） │
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
│ │ JAX │ │ TensorFlow │ │ PyTorch │ │
│ │ (XLA API) │ │ (XLA API) │ │ (XLA via │ │
│ │ │ │ │ │ LazyTensor) │ │
│ └──────┬──────┘ └──────┬──────┘ └──────┬──────┘ │
└─────────┼────────────────┼────────────────┼────────────────────────┘
│ │ │
▼ ▼ ▼
┌──────────────────────────────────────────────────────────────────────┐
│ XLA Client → Server (gRPC) │
│ ┌────────────────────────────────────────────────────────────────┐ │
│ │ HLO IR (High-Level Operation Intermediate Representation) │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │CustomCall│ │ Fusion │ │ AllReduce│ │ Dot │ │ │
│ │ │ │ │ │ │ │ │ │ │ │
│ │ └─────────┘ └─────────┘ └─────────┘ └─────────┘ │ │
│ └────────────────────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────────┐
│ XLA Compiler Server (可分离部署) │
│ ┌────────────────────────────────────────────────────────────────┐ │
│ │ HLO Optimization Pipeline │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │ Algebra │ │ Fusion │ │ SPMD │ │ Layout │ │ │
│ │ │Simplify │ │ Pass │ │ Partitioner│ │ Assignment│ │ │
│ │ └─────────┘ └─────────┘ └─────────┘ └─────────┘ │ │
│ └────────────────────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────────┐
│ Backend CodeGen │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ CPU │ │ GPU │ │ TPU │ │
│ │ (LLVM) │ │ (CUDA/NVVM) │ │ (XTPUwre) │ │
│ │ │ │ │ │ │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
└──────────────────────────────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────────────────┐
│ Executable (HLO Module) │
│ 编译产物可以是： │
│ - PTX (NVIDIA GPU) │
│ - NVVM (NVIDIA compute dialect) │
│ - LLVM IR → 机器码 (CPU) │
│ - XTPUwre (TPU) │
└──────────────────────────────────────────────────────────────────────┘1.3 Client-Server 部署模式
XLA 支持两种部署模式：
python
# 模式1：In-Process（默认，JAX 常用）
# XLA 编译器在同一个进程内，加载到内存
import jax
import jax.numpy as jnp
@jax.jit
def forward(x):
return jnp.dot(x, x.T)
# JIT 编译在首次调用时发生，编译产物缓存
result = forward(jnp.ones((1024, 1024)))
# 模式2：Out-of-Process（XLA Server，生产环境）
# 编译器进程独立，通过 gRPC 与客户端通信
# 适用于：多客户端共享编译器、多版本共存
#
# 启动 XLA Server：
# $ xla_server --num_replicas=8 --topology=2x2x2第2节 HLO 指令集详解
2.1 什么是 HLO
HLO（High-Level Operation） 是 XLA 的中间表示，每条 HLO 指令代表一个抽象的代数操作，不绑定具体硬件。
python
# HLO 的文本表示（HLO Text Format）
# 可以在 XLA 调试时打印出来
"""
ENTRY main.6 {
# 参数：形状为 f32[128, 256] 的输入
param.0 = f32[128, 256] parameter(0)
param.1 = f32[256, 512] parameter(1)
# 矩阵乘法：f32[128, 512]
dot.2 = f32[128, 512] dot(param.0, param.1),
lhs_contracting_dims={1},
rhs_contracting_dims={0}
# 加偏置：f32[128, 512]
add.3 = f32[128, 512] add(dot.2, param.2)
# ReLU 激活
ROOT tuple.4 = (f32[128, 512]) tuple(add.3)
}
"""2.2 常用 HLO 操作分类
| 类别 | HLO Op | 语义 | 示例 |
|---|---|---|---|
| Element-wise | Add, Sub, Mul, Relu, Sigmoid | 逐元素操作 | add(x, y) |
| Reshape | Reshape, Transpose, Broadcast | 形状变换 | reshape(x, [a, b]) |
| Reduction | Reduce, AllReduce | 聚合操作 | reduce(x, init, sum) |
| Dot | Dot, Convolution | 矩阵乘/卷积 | dot(a, b) |
| Collective | AllReduce, AllGather, CollectivePermute | 分布式通信 | all-reduce(x) |
| Fusion | Fusion, FusedMha | 融合操作 | fusion(...) |
| Control Flow | While, Conditional | 循环/分支 | while(cond, body) |
2.3 HLO 计算图可视化
python
import jax
import os
# 方法1：设置环境变量导出 HLO
os.environ['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_hlo/'
@jax.jit
def matmul_relu(x, w, b):
y = jnp.dot(x, w) # Dot operation
return jax.nn.relu(y + b) # Add + Relu
import jax.numpy as jnp
x = jnp.ones((64, 128))
w = jnp.ones((128, 256))
b = jnp.zeros((256,))
out = matmul_relu(x, w, b)
# 会在 /tmp/xla_hlo/ 下生成：
# - .optimized.txt (优化后的 HLO)
# - .hlo (原始 HLO)
# - .dot (GraphViz 文件)
# 方法2：用 JAX 的 debug 工具
from jax import make_jaxpr
# 打印 JIT 函数的 JAXPR（HLO 的前端表示）
jaxpr = make_jaxpr(matmul_relu)(x, w, b)
print(jaxpr)
"""
{ lambda ; a:f32[64,128] b:f32[128,256] c:f32[256].
let d:f32[64,256] = dot a, b
e:f32[64,256] = add d, c
f:f32[64,256] = relu e
in (f,) }
"""2.4 HLO 优化 Pass 顺序
XLA 的优化流水线按顺序执行多个 Pass，每个 Pass 负责一类优化：
| Pass 顺序 | Pass 名称 | 作用 | 示例 |
|---|---|---|---|
| 1 | Simplifier | 代数化简、常量折叠 | x + 0 → x, x * 1 → x |
| 2 | AlgebraicSimplifier | 代数恒等变换 | (a+b)^2 → a^2 + 2ab + b^2 |
| 3 | TupleSimplifier | 删除无用 Tuple 提取 | 移除死代码 |
| 4 | WhileLoopSimplifier | While 循环简化 | 循环不变式外提 |
| 5 | ScatterSimplifier | Scatter 索引优化 | 合并重复索引 |
| 6 | ConvolutionSimplifier | 卷积维度化简 | 1x1 卷积 → 矩阵乘 |
| 7 | GemmSimplifier | 矩阵乘法优化 | Reorder operands |
| 8 | DepthwiseConvolutionSimplifier | DW 卷积优化 | 特殊 case 加速 |
| 9 | CholeskyExpander | Cholesky 分解展开 | 展开为三角解 |
| 10 | ReshapeMover | Reshape 移动优化 | 减少不必要的 reshape |
| 11 | TransposeFolding | 转置折叠 | 合并相邻转置 |
| 12 | DotOpEmitter | Dot 操作特殊化 | 根据 shape 选择算法 |
| 13 | CallInliner | 内联函数调用 | 减少函数调用开销 |
| 14 | AllReduceFolder | AllReduce 折叠 | 合并相邻 AllReduce |
| 15 | MultiplicationFinder | 寻找矩阵乘法模式 | 识别 GEMM |
| 16 | Fusion | 算子融合 | 见第3节 |
| 17 | SPMDPartitioner | SPMD 分区 | 见第4节 |
| 18 | LayoutAssignment | 布局分配 | HBM 访问优化 |
第3节 Fusion（算子融合）
3.1 Fusion 的本质
Fusion（融合） 是深度学习编译器最重要的优化：把多个独立的 HLO 操作合并为一个 FusionOp，减少内存带宽压力。
融合前（3次 HBM 访问）：
┌─────────┐ ┌─────────┐ ┌─────────┐
│ Relu │────▶│ Matmul │────▶│ Add │
└─────────┘ └─────────┘ └─────────┘
HBM HBM HBM
融合后（1次 HBM 访问）：
┌─────────────────────────────────────┐
│ FusionOp (ReLU+Matmul+Add) │
│ │
│ Read once from HBM → Compute → Write│
└─────────────────────────────────────┘3.2 Fusion 的种类
python
# XLA 支持多种 Fusion 类型
# 1. Input Fusion（输入融合）
# 多个操作共享相同输入时融合
# 例：同一个输入的 Relu 和 Sigmoid
"""
before:
input ──▶ Relu
input ──▶ Sigmoid
after (input fusion):
input ──▶ FusionOp { Relu, Sigmoid }
"""
# 2. Output Fusion（输出融合）
# 多个操作产生相同输出时融合
# 例：同时需要 loss 和 accuracy 时
"""
before:
output ──▶ Accuracy
output ──▶ Loss
after (output fusion):
input ──▶ MainOp ──▶ FusionOp { Accuracy, Loss }
"""
# 3. Loop Fusion（循环融合）
# 嵌套循环中的操作融合
# 例：Conv → BatchNorm → ReLU
"""
before:
Loop(i) {
conv(i) ──▶ bn(i) ──▶ relu(i)
}
after (loop fusion):
Loop(i) {
FusionOp { conv(i), bn(i), relu(i) }
}
"""
# 4. Diamond Fusion（钻石融合）
# 多个分支汇聚的操作融合
"""
before:
──▶ Op1 ──┐
input ├──▶ Op4
──▶ Op2 ──┤
──▶ Op3 ──┘
after (diamond fusion):
input ──▶ FusionOp { Op1, Op2, Op3 } ──▶ Op4
"""3.3 Fusion Heuristics（融合决策）
不是所有操作都值得融合，XLA 根据启发式规则决定：
python
# XLA fusion heuristics 伪代码
def should_fuse(producer, consumer):
# 融合收益分析
producer_reads = count_hbm_reads(producer)
consumer_writes = count_hbm_writes(consumer)
# 计算融合后的 HBM 访问减少量
hbm_savings = producer_reads + consumer_writes
# 计算融合成本
fusion_cost = estimate_code_size_increase() # 代码膨胀
fusion_cost += estimate_compile_time_increase()
# 决定是否融合
# 融合条件：HBM 节省 > 成本，且不违反约束
return (hbm_savings > fusion_cost and
not exceeds_code_size_limit() and
not exceeds_register_pressure())| 场景 | 是否融合 | 原因 |
|---|---|---|
| Matmul + Element-wise | ✅ 融合 | 消除中间 HBM 访问 |
| 小 Shape Element-wise | ❌ 不融合 | 融合收益 < 代码膨胀 |
| 跨 Batch 维度融合 | ❌ 不融合 | 寄存器压力过大 |
| Conv + ReLU | ✅ 融合 | 经典融合模式 |
| 多个 Matmul 链式 | ⚠️ 部分融合 | 受限于代码大小 |
3.4 GEMM Fusion 详解
矩阵乘法是深度学习最核心的操作，XLA 有专门的 GEMM Fusion：
python
# JAX 代码
@jax.jit
def transformer_block(x, w_q, w_k, w_v, w_o):
# QKV 投影：3 个独立的 Matmul
q = jnp.dot(x, w_q) # Matmul 1
k = jnp.dot(x, w_k) # Matmul 2
v = jnp.dot(x, w_v) # Matmul 3
# 融合为 MultiMatmulFusion
# XLA 会生成一个 fused kernel，同时执行 3 个 matmul
# Attention 后的输出投影
attn = dot_product_attention(q, k, v)
out = jnp.dot(attn, w_o) # Matmul 4
return out
# 生成的 HLO（简化）
"""
ENTRY %computation {
param.0 = f32[32, 512, 768] parameter(0)
param.1 = f32[768, 2304] parameter(1) # 合并的 QKV weight
# MultiMatmulFusion：一次性完成 QKV 三个投影
ROOT %fusion.1 = f32[32, 512, 768], f32[32, 512, 768], f32[32, 512, 768]
fusion(param.0, param.1),
kind=kCustom,
calls={%qkv_fusion}
}
"""第4节 SPMD 自动并行
4.1 什么是 SPMD
SPMD（Single Program Multiple Data） 是一种并行编程模型：所有进程/设备运行相同的程序，但操作不同的数据分片。
python
import jax
import jax.numpy as jnp
from jax import pmap
# JAX SPMD 示例
@jax.jit
def model_step(params, batch):
def loss_fn(params):
return compute_loss(params, batch)
grads = jax.grad(loss_fn)(params)
return optax.apply_updates(params, grads)
# pmap：自动 SPMD 并行
# devices = [TPU0, TPU1, TPU2, TPU3, TPU4, TPU5, TPU6, TPU7]
parallel_step = pmap(model_step, axis_name='devices')
# 数据自动分片
# batch shape: [8, batch_size, seq_len, hidden]
# axis_name='devices' 表示第一个维度分片到 8 个设备
params = ... # 参数会广播到所有设备
batch = jnp.ones((8, 32, 512, 768)) # 每个设备处理 batch_size=32
output = parallel_step(params, batch) # 8x 并行4.2 pmap 如何自动分区计算图
python
# pmap 的 SPMD 分区过程
"""
原始计算图（单设备）：
┌────────────────────────────────────────┐
│ input: [B, H] │
│ │ │
│ ▼ │
│ LayerNorm ──────────────────────┐ │
│ │ │ │
│ ▼ │ │
│ Matmul(W1) ──▶ ReLU ──▶ Matmul(W2) │
│ │ │ │
│ ▼ ▼ │
│ residual + ──▶ LayerNorm │ │
│ │ │ │
│ ▼ ▼ │
│ output: [B, H] output │
└────────────────────────────────────────┘
SPMD 分区后（8 设备）：
┌────────────────────────────────────────────────────────────┐
│ Device 0 │ Device 1 │ ... │ Device 7 │
│ input[0] │ input[1] │ │ input[7] │
│ │ │ │ │ │ │ │
│ LayerNorm│ LayerNorm│ │ LayerNorm │
│ │ │ │ │ │ │ │
│ Matmul │ Matmul │ │ Matmul │
│ │ │ │ │ │ │ │
│ ├─────┴──────────┴─────┘ │ ← AllReduce (梯度同步) │
│ ▼ ▼ │
│ residual + ──▶ AllReduce ──▶ residual + │
│ │ │ │
│ ▼ ▼ │
│ output[0]│output[1]│ ... │output[7] │
└────────────────────────────────────────────────────────────┘
"""4.3 Collective 操作详解
SPMD 并行的核心是 Collective（集合）操作，用于设备间通信：
| 操作 | 语义 | XLA HLO | 使用场景 |
|---|---|---|---|
| AllReduce | 所有设备做同一 Reduce，结果广播 | all-reduce(add, arrays) | 梯度同步 |
| AllGather | 收集所有设备的数据 | all-gather(data, dim) | 聚合特征 |
| AllToAll | 设备间全交换 | all-to-all(data) | 张量重分区 |
| CollectivePermute | 点对点循环移位 | collective-permute(data) | 流水线并行 |
| ReduceScatter | Reduce 后分片 | reduce-scatter(data) | 梯度分片 |
| AllBatchNorm | 跨设备 BatchNorm | collective-batch-norm | 大 Batch 训练 |
python
# JAX 中使用 Collective 操作
import jax
from jax.lax import all_gather, all_reduce, pmean
# 方法1：通过 pmap 自动插入
@pmap
def data_parallel_training(params, batch):
def loss_fn(params):
return compute_loss(params, batch)
# jax.grad 会自动插入 AllReduce
grads = jax.grad(loss_fn)(params)
# 手动同步（等效于上面的自动插入）
grads = pmean(grads, axis_name='devices')
return optax.apply_updates(params, grads)
# 方法2：手动调用 collective
@jax.jit
def manual_collective_example(x):
# AllReduce：所有设备求和
sum_x = all_reduce(x, op=jax.lax.Add, axis_name='devices')
# AllGather：收集所有设备的数据
all_x = all_gather(x, axis_name='devices')
return sum_x, all_x4.4 通信和计算 Overlap
XLA 的核心优化之一是把通信和计算重叠执行：
python
# XLA 的 Pipelining（流水线）策略
"""
时间线：
Device 0: [Compute]────[Comm]────────[Compute]────[Comm]────────
Device 1: [Wait]──[Compute]────[Comm]────────[Compute]────[Comm]
Device 2: [Wait]──[Wait]──[Compute]────[Comm]────────[Compute]
Device 3: [Wait]──[Wait]──[Wait]──[Compute]────[Comm]─────
◄── Computation ──►◄── Communication ──►
XLA 的调度策略：
1. 在上一个 collective 完成前，提前调度下一个 compute
2. 使用 Double Buffering：计算当前 micro-batch 时，通信下一个 micro-batch
3. 使用待执行队列隐藏启动延迟
"""
# JAX 中启用 pipeline 调度
from jax.experimental import maps
# 设置启用异步调度
maps.thread_resources.env = maps.ThreadingEnvironment()第5节 延迟隐藏机制
5.1 为什么需要延迟隐藏
在分布式训练中，AllReduce 等通信操作的延迟可能高达毫秒级，而 GPU/TPU 的计算延迟仅为微秒级。如果串行执行，总时间 = 计算时间 + 通信时间，效率极低。
5.2 XLA 的调度策略
python
# XLA 延迟隐藏的调度原则
"""
XLA Delay Hiding 调度器工作原理：
1. 依赖分析
- 构建 HLO 指令的依赖图
- 识别可以并行执行的指令集合
2. 资源感知调度
- 每个设备有计算单元和通信单元
- 计算单元和通信单元可以同时工作
3. 调度决策
- 当通信单元空闲时，尽早发起通信
- 通信进行时，计算单元执行其他独立计算
- 通信完成后，合并结果继续计算
示例：
AllReduce(reduce_scatter(gemm(input)))
│ │
▼ ▼
发起 AllReduce 计算下个 micro-batch
│ │
▼ ▼
等待通信完成 与 AllReduce 结果合并
"""5.3 延迟预算配置
python
# JAX 中配置延迟隐藏参数
import jax
# 设置 SPMD 的通信延迟预算（单位：微秒）
jax.config.update('jax_spmd_mode', 'pjit') # 使用 pjit API
# 在 pjit 中显式指定资源约束
from jax.experimental import pjit
# pjit 的并行策略
mesh = jax.sharding.Mesh(devices, ('devices',))
sharding = jax.sharding.NamedSharding(mesh, P('devices'))
# 编译时优化通信
compiled = pjit.pjit(
model_step,
in_shardings=sharding,
out_shardings=sharding,
# 告诉 XLA 优化这个 collective
static_argnums=()
)(params, batch)第6节 XLA 与 JAX 的关系
6.1 分层架构
┌─────────────────────────────────────────────────────────────┐
│ User Code (Python) │
│ @jax.jit / @pmap / @pjit │
└─────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────┐
│ JAX（Python 库） │
│ ┌───────────┐ ┌───────────┐ ┌───────────┐ │
│ │ autodiff │ │ sharding │ │ vmap │ │
│ │ (grad) │ │ (pjit) │ │ │ │
│ └─────┬─────┘ └─────┬─────┘ └───────────┘ │
│ │ │ │
│ └──────┬───────┘ │
│ ▼ │
│ ┌────────────────────────────────────────┐ │
│ │ XLA Translation │ │
│ │ JAXPR → HLO Conversion │ │
│ └─────────────────┬──────────────────────┘ │
└────────────────────┼────────────────────────────────────────┘
▼
┌─────────────────────────────────────────────────────────────┐
│ XLA Compiler（可移植层） │
│ ┌───────────┐ ┌───────────┐ ┌───────────┐ │
│ │ HLO IR │ │ Optimizer │ │ Backend │ │
│ │ │ │ │ │ (GPU/CPU)│ │
│ └───────────┘ └───────────┘ └───────────┘ │
└─────────────────────────────────────────────────────────────┘
│
┌───────────┼───────────┐
▼ ▼ ▼
┌─────────┐ ┌─────────┐ ┌─────────┐
│ GPU │ │ TPU │ │ CPU │
│ (NVVM) │ │ (XTPU) │ │ (LLVM) │
└─────────┘ └─────────┘ └─────────┘6.2 JAX 到 XLA 的转换流程
python
# JAX → XLA 转换的完整流程
"""
1. Python 函数定义
def forward(params, x):
return jnp.dot(x, params['w']) + params['b']
2. JAXPR 生成（make_jaxpr）
{ lambda ; a:f32[784, 256] b:f32[256] c:f32[784].
let d:f32[256] = b
e:f32[784, 256] = a
f:f32[784, 256] = e
g:f32[256] = dot(f, a)
h:f32[256] = add(g, d)
in (h,) }
3. HLO 生成（XLA Client）
ENTRY forward {
param.0 = f32[784, 256] parameter(0)
param.1 = f32[256] parameter(1)
param.2 = f32[784] parameter(2)
dot.3 = f32[784, 256] dot(param.2, param.0)
ROOT add.4 = f32[784, 256] add(dot.3, param.1)
}
4. HLO 优化（XLA Compiler）
- Fusion
- AlgebraicSimplify
- SPMD Partitioning
- Layout Assignment
5. 目标代码生成
- GPU: PTX/NVVM
- CPU: LLVM IR → 机器码
- TPU: XTPUwre
"""第7节 XLA 调试技术
7.1 HLO Dump 和分析
bash
# 环境变量配置
export XLA_FLAGS="
--xla_dump_to=/tmp/xla_hlo_dump # 输出目录
--xla_dump_hlo_as_text=true # 文本格式
--xla_dump_hlo_as_dot=false # 不生成 DOT 图
--xla_dump_hlo_as_proto=false # 不生成 protobuf
--xla_dump_fusion_visualization=true # 可视化 fusion
--xla_enable_hlo_profiling=true # 启用性能分析
"
# 运行 JAX 代码
python train.py
# 查看生成的 HLO 文件
ls /tmp/xla_hlo_dump/
# ├── forward.0.txt # 原始 HLO
# ├── forward.1.txt # 优化后 HLO
# ├── forward.2_after_mlir.txt # MLIR 格式
# └── forward.3_after_gpu.txt # GPU 代码生成后7.2 使用 XLA Profile
python
import jax
import jax.profiler as profiler
# 启动 profiler
profiler.start_trace('/tmp/jax_profile')
# 运行训练
for step in range(1000):
params = train_step(params, batch)
profiler.stop_trace()
# 用 tensorboard 查看
# tensorboard --logdir=/tmp/jax_profile
# 或者用 Python API 分析
from jax.profiler import device_memory_profile
# 内存分析
memory_profile = device_memory_profile()
print(f"Peak memory: {memory_profile.bytes_accessed / 1e9:.2f} GB")7.3 常见问题诊断
| 问题 | 诊断方法 | 解决方案 |
|---|---|---|
| 编译太慢 | XLA_FLAGS="--xla_dump_hlo_as_text" | 简化模型、减少 compilation cache |
| Fusion 失败 | 查看 .1.txt 优化后 HLO | 拆分大的 fusion |
| 内存溢出 | device_memory_profile() | 启用 staged computation |
| 通信瓶颈 | 分析 timeline 中的 collective 时长 | 调整 batch size、启用 overlap |
| 数值错误 | 比较 eager vs jit 结果 | 检查 dtype、精度设置 |
升华
┌─────────────────────────────────────────────────────────────────────────────┐
│ XLA 编译器核心原则 │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ 1. 层次化抽象 │
│ HLO 提供硬件无关的抽象，Backend 负责硬件相关优化 │
│ │
│ 2. 融合优于分离 │
│ 减少 HBM 访问是深度学习编译器的第一优先级 │
│ │
│ 3. 自动并行优于手动并行 │
│ SPMD 自动分区让用户专注于算法，而非通信 │
│ │
│ 4. 延迟隐藏是性能关键 │
│ 计算和通信 overlap 是分布式训练效率的核心 │
│ │
└─────────────────────────────────────────────────────────────────────────────┘"AI 可查 vs 必须理解"清单
必须理解（不理解就等于不会）：
- 🔴 HLO 是什么：HLO 是 XLA 的中间表示，每条 HLO 指令代表一个抽象代数操作
- 🔴 Fusion 的本质：把多个 HLO 操作合并为一个 FusionOp，减少 HBM 访问
- 🔴 SPMD 如何分区：pmap 自动把计算图按 axis_name 分区，插入 Collective 操作
- 🔴 AllReduce 的作用：分布式训练中同步梯度的核心机制，所有设备求和后广播
- 🔴 XLA 与 JAX 的关系：JAX 是 XLA 的 Python 前端，负责生成 HLO
AI 可查（知道去哪查就行）：
- ✅ XLA 具体 Pass 顺序：可以通过
--xla_dump_to导出查看具体 Pass 执行 - ✅ 特定 HLO 操作语义：XLA 官方文档列出所有 HLO op 的完整语义
- ✅ 融合 Heuristics 具体阈值：XLA 源码中的 hard-coded 常量
- ✅ TPU/GPU 特定后端细节：各后端文档详细说明指令集差异
- ✅ XLA Profiler 使用方法：JAX/TensorFlow 官方 profiling 教程
学习状态：🟡 开始学习