基于Sharding API的JAX多节点TPU训练实现问询

前置操作 #

每个节点启动脚本时，需先完成集群初始化：

import jax
from jax.sharding import Mesh, PartitionSpec as P
from jax.experimental import mesh_utils

# 主节点运行（假设主节点IP为10.0.0.1，端口8080）
jax.distributed.initialize(coordinator_address="10.0.0.1:8080", num_processes=2, process_id=0)

# 从节点运行
# jax.distributed.initialize(coordinator_address="10.0.0.1:8080", num_processes=2, process_id=1)

1. 创建多节点设备Mesh #

针对2台v3-8节点（共16个TPU设备），构建3D设备Mesh，轴分别对应节点、本地设备、模型并行预留轴：

# 获取全局所有设备，按节点分组
devices = mesh_utils.create_device_mesh((2, 8, 1))
mesh = Mesh(devices, axis_names=("node", "local_dev", "model"))

2. DDP风格数据并行示例 #

将输入数据按node+local_dev轴分片，实现16路数据并行：

import jax.numpy as jnp

# 模拟一批输入数据（64条样本，每条128维特征）
batch_data = jnp.random.normal(size=(64, 128))

# 指定分片策略：数据的样本轴分片到node和local_dev轴
data_sharding = P(("node", "local_dev"), None)

# 将数据分发到多节点设备
with mesh:
  sharded_data = jax.device_put(batch_data, data_sharding)

# 验证分片结果：每个设备拿到64/16=4条样本
print(f"单设备数据形状: {sharded_data.shape}")  # 输出 (4, 128)

3. 模型并行示例 #

以线性层参数为例，将权重的隐藏层维度分片到node轴，实现跨节点模型并行：

from jax.nn import linear

# 模拟线性层权重（输入128维，输出256维）
params = {"w": jnp.random.normal(size=(128, 256)), "b": jnp.random.normal(size=(256,))}

# 指定权重分片策略：输出维度分片到node轴，输入维度不分片
param_sharding = {
    "w": P(None, "node"),
    "b": P("node")
}

# 将参数分发到多节点设备
with mesh:
  sharded_params = jax.device_put(params, param_sharding)

# 验证权重分片：每个节点的8个设备共享同一份权重分片，输出维度拆分为256/2=128
print(f"单设备权重形状: {sharded_params['w'].shape}")  # 输出 (128, 128)

# 执行模型前向（XLA自动处理跨节点通信）
with mesh:
  output = linear(sharded_data, sharded_params["w"], sharded_params["b"])
print(f"输出形状: {output.shape}")  # 输出 (4, 128)

关键说明 #

• Mesh的轴命名可自定义，核心是匹配分片策略（PartitionSpec）的轴名
• 切换并行模式只需调整PartitionSpec：数据并行对应样本轴分片到所有设备轴；模型并行对应模型参数的某维度分片到指定设备轴
• 所有跨节点通信由XLA自动完成，无需手动调用集合通信API

内容的提问来源于stack exchange，提问作者neel g