1 / 1

Tensor-Parallel MLP

JAX

Medium

Run a two-layer MLP with its weights split across 8 devices, so that no device ever holds a full weight matrix.

h = \operatorname{relu}(x W_1 + b_1), \qquad \text{out} = h W_2 + b_2

Here x has shape (B, D), $W_1$ is w1 with shape (D, H), $b_1$ is b1 with shape (H,), $W_2$ is w2 with shape (H, O), $b_2$ is b2 with shape (O,), and relu replaces every negative entry with 0. Build a one-axis mesh named model over all 8 devices and use jax.shard_map. Split w1 and b1 along the hidden dimension, so device d holds hidden units dH/8 up to (d + 1)H/8 minus 1, and split w2 along its rows in the same way. Every device receives the full x. With this split each device computes its own H/8 columns of $h$ without talking to any other device, multiplies them by its rows of $W_2$ , and obtains a partial output of shape (B, O). The 8 partial outputs add up to $h W_2$ , and jax.lax.psum performs that addition across devices.

Return a dict with two keys. out is a float32 JAX array of shape (B, O), replicated so every device holds all of it. hidden is the float32 JAX array $h$ of shape (B, H), left split by columns so each device holds its own (B, H/8) block.

Examples

Input: x = [[0.0,0.5]], w1 = [[-0.8,0.2,-1.2,-0.7,-0.1,0.9,-1.4,1.0],[1.3,-0.4,0.7,1.4,1.7,0.6,-0.6,-0.9]], b1 = [-0.1,-0.1,0.0,-0.1,0.3,0.4,-0.3,0.1], w2 = [[1.1],[0.0],[1.4],[0.2],[0.5],[0.0],[0.2],[-1.0]], b2 = [0.6]

Output: {"out":[[2.39]],"hidden":[[0.55,0.0,0.35,0.6,1.15,0.7,0.0,0.0]]}

Explanation: H is 8, so device d holds hidden unit d. The pre-activations x W1 + b1 are [0.55, -0.3, 0.35, 0.6, 1.15, 0.7, -0.6, -0.35], and relu turns the three negative ones into 0. Each device multiplies its single hidden value by its row of w2, giving the partial outputs 0.605, 0, 0.49, 0.12, 0.575, 0, 0 and 0. Their psum is 1.79, and adding b2 = 0.6 once gives 2.39.

Input: x = [[0.0,0.8,0.7],[-0.1,0.4,-0.3]], w1 = [[-0.8,0.7,0.5,0.2,0.0,-0.9,0.8,0.1],[0.0,0.2,-0.6,0.2,0.4,1.5,0.1,-0.7],[1.3,0.2,-0.1,0.2,-1.3,0.0,0.0,-1.2]], b1 = [0.6,-0.4,0.3,0.2,-0.3,0.4,0.0,0.3], w2 = [[0.0,-0.4],[0.8,-0.5],[-0.6,0.5],[0.0,-0.4],[0.0,-0.4],[0.9,-0.2],[0.9,-0.1],[0.6,0.2]], b2 = [-0.4,0.1]

Output: {"out":[[1.112,-1.032],[0.779,-0.32]],"hidden":[[1.51,0.0,0.0,0.5,0.0,1.6,0.08,0.0],[0.29,0.0,0.04,0.2,0.25,1.09,0.0,0.37]]}

Input: x = [[-1.6,1.8],[-0.3,-0.9],[-0.5,-0.3]], w1 = [[0.5,-0.6,-0.8,-0.7,0.3,0.2,0.9,1.2,0.2,0.3,-0.8,0.8,1.5,-1.2,-1.2,-0.9],[0.8,1.0,0.1,0.0,-1.4,-1.0,-0.3,1.2,-0.5,-0.6,0.5,0.5,1.5,0.0,0.6,-0.2]], b1 = [0.3,0.1,0.0,0.0,0.4,-0.1,0.7,0.0,0.5,-0.1,0.6,-0.3,-0.2,-0.2,-0.2,-0.2], w2 = [[0.2,-0.4],[-0.5,0.4],[0.3,-0.4],[0.0,0.4],[-0.1,0.0],[0.3,-0.4],[0.3,0.4],[0.2,-0.7],[0.2,0.1],[0.0,-0.2],[-0.7,0.8],[-0.1,-0.1],[0.6,-0.7],[-0.9,0.3],[0.4,0.4],[-0.3,-0.3]], b2 = [-0.2,0.5]

Output: {"out":[[-3.534,4.49],[-0.194,0.812],[-0.924,1.478]],"hidden":[[0.94,2.86,1.46,1.12,0.0,0.0,0.0,0.24,0.0,0.0,2.78,0.0,0.1,1.72,2.8,0.88],[0.0,0.0,0.15,0.21,1.57,0.74,0.7,0.0,0.89,0.35,0.39,0.0,0.0,0.16,0.0,0.25],[0.0,0.1,0.37,0.35,0.67,0.1,0.34,0.0,0.55,0.0,0.85,0.0,0.0,0.4,0.22,0.31]]}

Hint 1

P(None, "model") splits columns and P("model", None) splits rows.

Hint 2

jax.lax.psum(partial, "model") adds the partial outputs across devices.

Hint 3

Pass axis_types=(AxisType.Auto,) to jax.make_mesh so shard_map accepts plain inputs.

Requirements

The function is traced with jax.make_jaxpr, and the shard_map body must receive weight blocks of shape (D, H/8) and (H/8, O) through its in_specs and must contain a psum.
hidden must come straight out of shard_map still split by columns, so gathering it into one replicated array is rejected.
Add b2 exactly once to the combined output.
The function is called inside jax.jit with every array traced, so it must not convert traced values to Python numbers.

Constraints

B is between 1 and 4, D is between 1 and 5, H is 8, 16, 24 or 32, and O is between 1 and 4.
x, w1, b1, w2 and b2 are float32 JAX arrays with the shapes given above, and every value is finite.

Try Similar Problems

Data-Parallel Step with psumMedium Device Mesh and Named ShardingEasy Ring All-Reduce with ppermuteHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium

Tensor-Parallel MLP

JAX

Medium

Run a two-layer MLP with its weights split across 8 devices, so that no device ever holds a full weight matrix.

h = \operatorname{relu}(x W_1 + b_1), \qquad \text{out} = h W_2 + b_2

Examples

Output: {"out":[[2.39]],"hidden":[[0.55,0.0,0.35,0.6,1.15,0.7,0.0,0.0]]}

Output: {"out":[[1.112,-1.032],[0.779,-0.32]],"hidden":[[1.51,0.0,0.0,0.5,0.0,1.6,0.08,0.0],[0.29,0.0,0.04,0.2,0.25,1.09,0.0,0.37]]}

Hint 1

P(None, "model") splits columns and P("model", None) splits rows.

Hint 2

jax.lax.psum(partial, "model") adds the partial outputs across devices.

Hint 3

Pass axis_types=(AxisType.Auto,) to jax.make_mesh so shard_map accepts plain inputs.

Requirements

The function is traced with jax.make_jaxpr, and the shard_map body must receive weight blocks of shape (D, H/8) and (H/8, O) through its in_specs and must contain a psum.
hidden must come straight out of shard_map still split by columns, so gathering it into one replicated array is rejected.
Add b2 exactly once to the combined output.
The function is called inside jax.jit with every array traced, so it must not convert traced values to Python numbers.

Constraints

B is between 1 and 4, D is between 1 and 5, H is 8, 16, 24 or 32, and O is between 1 and 4.
x, w1, b1, w2 and b2 are float32 JAX arrays with the shapes given above, and every value is finite.

Try Similar Problems

x =

Accepts: array

w1 =

Accepts: array

b1 =

Accepts: array

w2 =

Accepts: array

b2 =

Accepts: array