TPU 训练的真实效率往往取决于两个核心要素：Shape 的稳定性与算子的融合度。
很多时候，JAX 任务之所以出现严重的性能瓶颈，并非算法本身设计有问题，而是忽视了 XLA 编译器与底层硬件对“确定性”的极度偏好。基于大量实战调优经验，本文总结了八条能让 JAX 训练任务从“甚至跑不通”蜕变为“跑满 TPU 算力”的工程经验。
1、尽早锁定 Shape
TPU 喜欢静态 Shape，JAX 也是，所以动态 Shape 是性能杀手，它会触发重新编译（Recompile）。一旦发生重编译，Step time 和内存占用都会直接炸裂。所以解决方法也很简单，选定几个规范的尺寸，剩下的全填（Pad）满。
全局 Batch Size 要能被 TPU 核心数整除，然后就是对于变长序列，别指望它原本多长就多长，把它 Pad 到几个固定的“桶（Bucket）”里，比如 128、256 或 512，这步工作最好在输入（Input Pipeline）里就做完。
Python层面的条件判断尽量别依赖 Shape，真要分支逻辑，就老老实实让
lax.cond
或
lax.switch
来接管。
# Example: bucketing & padding (conceptual)
def pad_to_length(arr, L):
pad = L - arr.shape[0]
return jnp.pad(arr, ((0, pad), (0, 0)), mode='constant')
bucket_sizes = [128, 256, 512]
def bucket_len(n):
return next(b for b in bucket_sizes if n <= b)
def preprocess_batch(batch):
L = bucket_len(batch["tokens"].shape[1])
batch["tokens"] = pad_to_length(batch["tokens"], L)
batch["mask"] = pad_to_length(batch["mask"], L)
return batch
每个 Step 喂给 TPU 的 Shape 只要是固定的，XLA 编译器就不会找麻烦。
2、激活值默认用 bfloat16，主权重要 FP32
在 TPU 上
bfloat16
(bf16) 是个好东西，兼顾了速度、内存和数值稳定性。
工程上的常规操作是：激活（Activations）和梯度（Gradients）存成 bf16。但是，优化器状态里的权重必须保留一份 FP32 的“主副本”，不然跑久了数值就会漂移。所欲需要在模型边界做类型转换（Cast）的时候小心点。
class MLP(nn.Module):
features: int
@nn.compact
def __call__(self, x):
x = x.astype(jnp.bfloat16) # fast path on TPUs
x = nn.Dense(self.features, dtype=jnp.bfloat16)(x)
x = nn.gelu(x)
x = nn.Dense(self.features, dtype=jnp.bfloat16)(x)
return x
# Optimizer state stays in FP32 (conceptual)
params_fp32 = params.astype(jnp.float32)
grads_bf16 = compute_grads_bf16(...)
updates_fp32 = opt.update(grads_bf16.astype(jnp.float32), opt_state, params_fp32)
3、pjit和命名网格：切分要明确，别靠猜
JAX 在 TPU 上最强的一点就是通过
pjit
实现了 GSPMD。你通过 PartitionSpecs 告诉它想要什么切分方式，XLA 负责搞定如何在设备间搬运数据。
在 TPU 核心上建个命名网格（Mesh）。做数据并行（Data Parallelism）时，用
PartitionSpec('data', None)
这种模式。如果模型太大需要张量并行（Tensor Model Parallelism），就加个
'model'
轴。
import numpy as np
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, PartitionSpec as P
from jax.experimental import pjit
devices = np.array(jax.devices()).reshape(1, -1) # 1 x N mesh
mesh = Mesh(devices, ('data',))
def loss_fn(params, batch):
logits = model_apply(params, batch['x'])
return cross_entropy(logits, batch['y'])
@pjit.pjit(
in_shardings=(P(None), P('data')), # params replicated, batch sharded on 'data'
out_shardings=P(None), # scalar loss replicated
)
def step(params, batch):
grads = jax.grad(loss_fn)(params, batch)
# aggregate grads across cores
grads = jax.tree.map(lambda g: jax.lax.pmean(g, axis_name='data'), grads)
return grads
with mesh:
grads = step(params, sharded_batch)
切分（Sharding）这事必须显式。如果偷懒依赖自动推导，等到后期 debug 那些悄无声息的跨设备数据传输时，绝对会很痛苦。
4、jit, vmap, scan 三件套
TPU 喜欢大块头的 Kernel，讨厌成千上万个细碎的小算子。训练 Step 和任何中大型计算逻辑，必须用
jit
包起来。遇到 Python 循环，如果是时间步逻辑就换成
lax.scan
，如果是批次并行就用
vmap
。
把 Loss 计算、梯度计算和参数更新塞进同一个 jitted 函数里，这样编译器才有机会把它们融合成一个大算子。
import optax
import jax
optimizer = optax.adamw(3e-4)
def loss_and_grads(params, batch):
def _loss(p):
logits = model_apply(p, batch['x'])
return cross_entropy(logits, batch['y'])
loss, grads = jax.value_and_grad(_loss)(params)
return loss, grads
@jax.jit
def train_step(state, batch):
loss, grads = loss_and_grads(state.params, batch)
grads = jax.lax.pmean(grads, axis_name='data')
updates, new_opt_state = optimizer.update(grads, state.opt_state, state.params)
new_params = optax.apply_updates(state.params, updates)
return state.replace(params=new_params, opt_state=new_opt_state), loss
5、别让输入管道拖后腿
Host 到 Device 的数据传输一旦停顿，吞吐量就掉下来了，所以永远别让计算单元等数据。
用
tf.data
或者高效的 NumPy loader 配合 prefetch。数据预取到设备（Stage to device） 最好做双重缓冲。全局 Batch 尽量大（当然要能被核心数整除），数据增强这种脏活累活在 Host 上一次性做完。
# tf.data pipeline (conceptual)
ds = (tf.data.TFRecordDataset(files)
.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)
.batch(global_batch_size, drop_remainder=True)
.prefetch(tf.data.AUTOTUNE))
# Convert to NumPy and prefetch onto devices
from flax.jax_utils import prefetch_to_device
it = prefetch_to_device(map(npify, ds.as_numpy_iterator()), size=2)
with mesh:
for step_i in range(num_steps):
batch = next(it) # already sharded/prefetched
state, loss = train_step(state, batch)
6、PRNG要Fold 进 Step 和 Device ID
JAX 的 PRNG 是无状态的，这意味如果不小心，很容易在不同 Step 或者不同设备上用了一样的随机数 Key。
每个 Step 都要 Split 一次绝对别复用。所以说为了保证独立性必须把 Global Step 和 Device Index 都 Fold 进去。数据增强/Dropout 的 Key 和参数初始化的 Key 得分开管理。
def make_step_rng(rng, step):
step_key = jax.random.fold_in(rng, step)
dev_key = jax.random.fold_in(step_key, jax.lax.axis_index('data'))
return jax.random.split(dev_key, 1)[0]
@jax.jit
def train_step(state, batch, base_rng):
rng = make_step_rng(base_rng, state.step)
logits = model_apply(state.params, batch['x'], rngs={'dropout': rng})
...
7、Remat，智能 Checkpoint，梯度累积
TPU 内存看着大，模型一跑起来就不够用。深层网络可以直接用 Activation Checkpointing（
jax.checkpoint
或
nn.remat
），用计算换显存。想跑大 Batch 但显存不够，就用梯度累积（Gradient Accumulation） 把它切成小的 micro-step。
存盘的时候，推荐用 Orbax 做异步、分片（Sharded）的 Checkpoint，稳。
from flax import linen as nn
class DeepBlock(nn.Module):
@nn.compact
def __call__(self, x):
# recompute on backward to trim activation memory
f = nn.remat(lambda y: nn.gelu(nn.Dense(x.shape[-1])(y)))
return f(x)
# Gradient accumulation (conceptual)
@jax.jit
def accum_step(state, batch_slices):
def body(carry, micro):
state, grad_sum = carry
_, grads = loss_and_grads(state.params, micro)
return (state, jax.tree_util.tree_map(jnp.add, grad_sum, grads)), None
init_grads = jax.tree_util.tree_map(jnp.zeros_like, state.params)
(state, grad_sum), _ = jax.lax.scan(body, (state, init_grads), batch_slices)
grads = jax.tree_map(lambda g: g / len(batch_slices), grad_sum)
...
8、一定要跑 Profiler
把关键代码段用 Profiler Annotations 包起来，看 Step Timeline。重点找 Host Waits、Recompiles 和那些没融合好的细碎算子（Small op soup）。
稳态运行的时候，盯着 Tokens/sec 或者Images/sec，还有硬件利用率。
from jax.experimental import host_callback as hcb
from jax import profiler
def tagged(name, fn, *a, **k):
profiler.annotate_function(name=name)
return fn(*a, **k)
@jax.jit
def train_step(state, batch):
profiler.annotate_function(name="train_step")
# do work...
return state, loss
一定要在锁定 Shape 并且 JIT 完热点路径之后再做 Profile，不然全是噪音，根本看不到真正的瓶颈。
极简 TPU 训练示例
这基本包含了上面所有的内容
# Pseudo-skeleton (Flax + JAX + TPU)
mesh = Mesh(np.array(jax.devices()).reshape(1, -1), ('data',))
@pjit.pjit(in_shardings=(P(None), P('data'), P(None)), out_shardings=(P(None), P(None)))
def train_step(state, batch, base_rng):
rng = jax.random.fold_in(base_rng, state.step)
rng = jax.random.fold_in(rng, jax.lax.axis_index('data'))
def loss_fn(p):
logits = model_apply(p, batch['x'].astype(jnp.bfloat16),
rngs={'dropout': rng})
return cross_entropy(logits, batch['y'])
loss, grads = jax.value_and_grad(loss_fn)(state.params)
grads = jax.tree_map(lambda g: jax.lax.pmean(g, 'data'), grads)
updates, opt_state = optimizer.update(grads, state.opt_state, state.params)
params = optax.apply_updates(state.params, updates)
return state.replace(params=params, opt_state=opt_state, step=state.step+1), loss
with mesh:
for step_i, batch in enumerate(prefetched_iterator):
state, loss = train_step(state, batch, base_rng)
if step_i % log_every == 0:
# Pull back just tiny scalars; keep big tensors on device
host_loss = jax.device_get(loss)
print(f"[{step_i}] loss={host_loss:.4f}")
总结
TPU 需要的是 一致性：稳定的 Shape，融合的 Kernel，目的明确的切分，不掉链子的数据管道，把上面的这八件事做好，写 JAX 训练循环就非常顺畅了。
https://avoid.overfit.cn/post/16b582a493ba4eca8333314859665dd2
作者:Modexa
