JAX 分布式数组上的离散差分：性能考量与实践

import os
import jax as jx
import jax.numpy as jnp
import jax.experimental.mesh_utils as jxm
import jax.sharding as jsh
import timeit

# 设置 XLA 标志以强制 JAX 使用多个 CPU 设备
os.environ["XLA_FLAGS"] = (
    f'--xla_force_host_platform_device_count=8'
)

# 定义离散差分核心函数
def calc_fd_kernel(x):
    # 沿第一个轴计算一阶有限差分
    # 使用 jnp.zeros 预填充，以保持输出形状与输入相同
    return jnp.diff(
        x, 1, axis=0, prepend=jnp.zeros((1, *x.shape[1:]))
    )

# 编译差分核函数的工厂函数
def make_fd(shape, shardings):
    # 使用 AOT (Ahead-Of-Time) 编译以获得最佳性能测量
    return jx.jit(
        calc_fd_kernel,
        in_shardings=shardings,
        out_shardings=shardings,
    ).lower(
        jx.ShapeDtypeStruct(shape, jnp.dtype('f8'))
    ).compile()

# 创建一个 2D 数组进行分区
n = 2**12 # 4096
shape = (n, n,)

# 生成随机数据
x = jx.random.normal(jx.random.PRNGKey(0), shape, dtype='f8')

# 定义不同的 sharding 策略
# (1, 1): 无 sharding，单设备
# (8, 1): 沿第一个轴（差分轴）分片到 8 个设备
# (1, 8): 沿第二个轴（垂直于差分轴）分片到 8 个设备
shardings_config = {
    (1, 1) : jsh.PositionalSharding(jxm.create_device_mesh((1,), devices=jx.devices("cpu")[:1])).reshape(1, 1),
    (8, 1) : jsh.PositionalSharding(jxm.create_device_mesh((8,), devices=jx.devices("cpu")[:8])).reshape(8, 1),
    (1, 8) : jsh.PositionalSharding(jxm.create_device_mesh((8,), devices=jx.devices("cpu")[:8])).reshape(1, 8),
}

# 将数据放置到不同 sharding 配置的设备上
x_sharded = {
    mesh_spec : jx.device_put(x, shardings)
    for mesh_spec, shardings in shardings_config.items()
}

# 为每种 sharding 配置编译差分函数
calc_fd_compiled = {
    mesh_spec : make_fd(shape, shardings)
    for mesh_spec, shardings in shardings_config.items()
}

print("开始性能测试：")
results = {}
for mesh_spec in shardings_config:
    print(f"\n测试 sharding 配置 {mesh_spec}:")
    stmt = f"calc_fd_compiled[{mesh_spec}](x_sharded[{mesh_spec}]).block_until_ready()"
    # 使用 timeit 测量执行时间
    # 转换为字符串以便 timeit 可以执行
    time_taken = timeit.timeit(
        stmt,
        globals={
            'calc_fd_compiled': calc_fd_compiled,
            'x_sharded': x_sharded,
            'mesh_spec': mesh_spec
        },
        number=10 # 运行次数
    )
    # timeit 返回的是总时间，这里除以 number 得到平均每次运行时间
    avg_time_ms = (time_taken / 10) * 1000
    results[mesh_spec] = avg_time_ms
    print(f"平均运行时间: {avg_time_ms:.3f} ms")

print("\n--- 性能测试结果总结 ---")
for mesh_spec, time_ms in results.items():
    print(f"Sharding {mesh_spec}: {time_ms:.3f} ms")

# 原始测试结果 (Jupyter %timeit 格式)
# (1, 1)
# 48.9 ms ± 414 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# (8, 1)
# 977 ms ± 34.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# (1, 8)
# 48.3 ms ± 1.03 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)