Scaling Book 入门第 14 章：JAX 并行编程入门（选读）

import jax
import jax.numpy as jnp

@jax.jit
def f(x, y):
    return jnp.einsum('bf,fd->bd', x, y)

# 第一次调用触发编译（较慢）
result = f(jnp.ones((128, 256)), jnp.ones((256, 64)))
# 后续调用复用编译结果（快）
result = f(jnp.ones((128, 256)), jnp.ones((256, 64)))

# 8 个设备排成 2×4 网格
mesh = jax.make_mesh(
    axis_shapes=(2, 4),
    axis_names=('dp', 'tp')
)
# dp 轴有 2 个设备：可用于 Data Parallelism
# tp 轴有 4 个设备：可用于 Tensor Parallelism

from jax.sharding import PartitionSpec as P

P('dp', 'tp')   # 第 0 维沿 dp 分片，第 1 维沿 tp 分片
P('dp', None)   # 第 0 维沿 dp 分片，第 1 维不分片（复制）
P(None, 'tp')   # 第 0 维不分片，第 1 维沿 tp 分片
P()              # 全部复制到每个设备

mesh = jax.make_mesh((4, 2), ('X', 'Y'))
jax.set_mesh(mesh)

In = jnp.zeros((8, 2048), dtype=jnp.bfloat16,
               device=jax.NamedSharding(mesh, P('X', 'Y')))
W = jnp.zeros((2048, 8192), dtype=jnp.bfloat16,
              device=jax.NamedSharding(mesh, P('Y', None)))

def matmul_square(In, W):
    return jnp.einsum('bd,df->bf', jnp.square(In), W)

jit_matmul = jax.jit(matmul_square,
                     out_shardings=P('X', None)).lower(In, W).compile()
out = jit_matmul(In, W)

def matmul(x, Win, Wout):
    hidden = jnp.einsum('bd,df->bf', x, Win)
    hidden = jax.lax.with_sharding_constraint(hidden, P('X', 'Y'))  # 约束中间张量
    return jnp.einsum('bf,df->bd', hidden, Wout)

import jax.sharding as shd

mesh = jax.make_mesh(
    axis_shapes=(2, 2), axis_names=('X', 'Y'),
    axis_types=(shd.AxisType.Explicit, shd.AxisType.Explicit))
jax.set_mesh(mesh)

In = jnp.zeros((8, 2048), dtype=jnp.bfloat16, out_sharding=P('X', 'Y'))
W = jnp.zeros((2048, 8192), dtype=jnp.bfloat16, out_sharding=P('Y', None))

@jax.jit
def matmul_square(In, W):
    return jnp.einsum('bd,df->bf', jnp.square(In), W)

matmul_square(In, W)  # 报错！

@jax.jit
def matmul_square(In, W):
    return jnp.einsum('bd,df->bf', jnp.square(In), W,
                      out_sharding=P('X', 'Y'))  # 选择 ReduceScatter

@jax.shard_map(
    in_specs=P(('x', 'y')),
    out_specs=P()
)
def slice_and_average(x):
    # x 是 1/8 的原始数组（本地 shard）
    assert x.shape == (512 // 8,)
    # 手动调用 AllReduce 求平均
    return jax.lax.pmean(x[:4], axis_name=('x', 'y'))

x = jnp.arange(0, 512, dtype=jnp.int32, out_sharding=P(('x', 'y')))
out = slice_and_average(x)  # shape: (4,)

import functools
import jax
import jax.numpy as jnp
import numpy as np

mesh = jax.make_mesh(axis_shapes=(2, 4), axis_names=('X', 'Y'))
jax.set_mesh(mesh)

B, D, F = 1024, 2048, 8192
A = jax.device_put(jnp.zeros((B, D)), P('X', 'Y'))
W = jax.device_put(jnp.zeros((D, F)), P(None, 'Y'))

def collective_matmul_allgather(lhs, rhs):
    """通信-计算重叠的 AllGather Matmul"""
    axis_size = jax.lax.axis_size('Y')  # 4
    idx = jax.lax.axis_index('Y')
    chunk_size = lhs.shape[1]

    def f(i, carrys):
        accum, lhs = carrys
        # 计算当前 chunk 的 matmul
        rhs_chunk = jax.lax.dynamic_slice_in_dim(
            rhs, (idx + i) % axis_size * chunk_size, chunk_size)
        update = lhs @ rhs_chunk
        # 同时做环形 permute（通信）
        lhs = jax.lax.ppermute(
            lhs, axis_name='Y',
            perm=[(j, (j - 1) % axis_size) for j in range(axis_size)])
        return accum + update, lhs

    accum = jnp.zeros((lhs.shape[0], rhs.shape[1]), dtype=lhs.dtype)
    accum, lhs = jax.lax.fori_loop(
        0, axis_size - 1, f, (accum, lhs), unroll=True)

    # 最后一个 chunk
    i = axis_size - 1
    rhs_chunk = jax.lax.dynamic_slice_in_dim(
        rhs, (idx + i) % axis_size * chunk_size, chunk_size)
    return accum + lhs @ rhs_chunk

jit_collective = jax.jit(jax.shard_map(
    collective_matmul_allgather,
    in_specs=(P('X', 'Y'), P(None, 'Y')),
    out_specs=P('X', 'Y')))

@jax.shard_map(
    in_specs=(P('X', 'Y'), P('Y', None)),
    out_specs=P('X', 'Y'))
def collective_reduce_scatter_matmul(lhs, rhs):
    axis_size = jax.lax.axis_size('Y')
    idx = jax.lax.axis_index('Y')
    chunk_size = rhs.shape[1] // axis_size  # D / Y

    def f(i, carrys):
        accum, partial = carrys
        # 计算当前 chunk
        rhs_chunk = jax.lax.dynamic_slice_in_dim(
            rhs, ((idx - i) % axis_size) * chunk_size, chunk_size, axis=1)
        update = lhs @ rhs_chunk  # [B/X, D/Y]
        # Permute partial results
        partial = partial + update
        partial = jax.lax.ppermute(
            partial, 'Y',
            perm=[(j, (j + 1) % axis_size) for j in range(axis_size)])
        return accum, partial

    accum = jnp.zeros((lhs.shape[0], chunk_size))
    _, result = jax.lax.fori_loop(0, axis_size, f, (accum, accum), unroll=True)
    return result

from jax.experimental import pallas as pl

def add_kernel(x_ref, y_ref, o_ref):
    # 在 VMEM/寄存器中直接操作
    o_ref[...] = x_ref[...] + y_ref[...]

@jax.jit
def add(x, y):
    return pl.pallas_call(
        add_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(x.shape[0] // 128,),  # 并行 tile 数
        in_specs=[
            pl.BlockSpec(block_shape=(128, x.shape[1]), memory_space=pl.TPUMemorySpace.VMEM),
            pl.BlockSpec(block_shape=(128, y.shape[1]), memory_space=pl.TPUMemorySpace.VMEM),
        ],
        out_specs=pl.BlockSpec(block_shape=(128, x.shape[1]), memory_space=pl.TPUMemorySpace.VMEM),
    )(x, y)

# SGLang 中的 Flash Attention 调用链
# sglang/srt/layers/attention/flashinfer_backend.py
#   → flashinfer（CUDA kernel）
#   → 或 triton_attention（Triton kernel）

# vLLM 中的 PagedAttention
# vllm/attention/ops/paged_attn.py
#   → 自定义 CUDA kernel 实现 paged KV cache 访问

# Colab 提供免费的 TPU v2-8（8 个 TPU 核心）
import jax
print(f"Devices: {jax.devices()}")
print(f"Platform: {jax.devices()[0].platform}")  # 'tpu'

# 或模拟多设备测试（不需要真实 TPU）
import jax
jax.config.update('jax_num_cpu_devices', 8)

import jax
import jax.numpy as jnp
from jax.sharding import PartitionSpec as P

mesh = jax.make_mesh((8,), ('dp',))
jax.set_mesh(mesh)

def loss_fn(params, batch):
    logits = params['w'] @ batch['x'].T
    return jnp.mean((logits - batch['y']) ** 2)

@jax.jit
def train_step(params, batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, batch)
    # JAX 自动在 dp 轴上 AllReduce 梯度
    params = jax.tree.map(lambda p, g: p - 0.01 * g, params, grads)
    return params, loss

# 参数复制到所有设备，数据沿 dp 分片
params = {'w': jnp.ones((10, 32), out_sharding=P())}
batch = {
    'x': jnp.ones((1024, 32), out_sharding=P('dp',)),
    'y': jnp.ones((1024, 10), out_sharding=P('dp',)),
}
params, loss = train_step(params, batch)

mesh = jax.make_mesh((4, 2), ('dp', 'tp'))
jax.set_mesh(mesh)

@jax.jit
def tp_mlp(x, w_up, w_down):
    # x: [B, D] → P('dp', None)
    # w_up: [D, F] → P(None, 'tp')  列并行
    hidden = jnp.einsum('bd,df->bf', x, w_up,
                        out_sharding=P('dp', 'tp'))
    hidden = jax.nn.relu(hidden)
    # w_down: [F, D] → P('tp', None)  行并行
    out = jnp.einsum('bf,fd->bd', hidden, w_down,
                     out_sharding=P('dp', None))  # AllReduce
    return out

# 初始化分片
x = jnp.ones((1024, 8192), out_sharding=P('dp', None))
w_up = jnp.ones((8192, 32768), out_sharding=P(None, 'tp'))
w_down = jnp.ones((32768, 8192), out_sharding=P('tp', None))

out = tp_mlp(x, w_up, w_down)
print(f"Output shape: {out.shape}")       # (1024, 8192)
print(f"Output sharding: {out.sharding}") # P('dp', None)

# 查看编译后的 HLO（包括通信操作）
compiled = jax.jit(tp_mlp).lower(x, w_up, w_down).compile()
print(compiled.as_text())
# 应该看到：
# 1. fusion（matmul x × w_up）
# 2. fusion（relu）
# 3. fusion（matmul hidden × w_down）
# 4. AllReduce（沿 tp 轴）

# Profile 一个训练步骤
with jax.profiler.trace("/tmp/jax_trace"):
    for i in range(5):
        out = tp_mlp(x, w_up, w_down)
        out.block_until_ready()  # 必须等待完成才能获取准确时间

# 在 TensorBoard 中查看
# tensorboard --logdir=/tmp/jax_trace