NNX Tensor-Parallel Linear

Why this matters

A linear layer computes x @ kernel + bias. Each output column depends on the input and one column of the kernel, independently of the other output columns. Column-parallel computation uses this fact to divide a layer's output features into groups.

This exercise simulates those groups with a loop over array slices. It creates the full kernel and returns the full output. It does not allocate a device mesh, communicate between devices, or demonstrate reduced memory use or faster execution.

Split output columns

For input x of shape (N, in_dim), an NNX Linear kernel has shape (in_dim, features) and its bias has shape (features,). Let chunk = features // tp_size, assuming exact division.

Group i uses:

• kernel[:, i * chunk : (i + 1) * chunk], shape (in_dim, chunk).
• bias[i * chunk : (i + 1) * chunk], shape (chunk,).

Apply the full x to each group: x @ kernel_slice + bias_slice. Each result has shape (N, chunk). Concatenate results in group order along the last axis to recover (N, features).

The sum over input features still occurs inside each matrix multiplication. Do not sum the groups' outputs: they represent different output columns, not partial sums of the same columns.

A complete example

The following two groups reproduce one affine layer, including a nonzero bias:

import jax.numpy as jnp
import numpy as np

x = jnp.array([[1., 2.]])
kernel = jnp.array([[1., 0., 2., -1.],
                    [0., 1., 1.,  3.]])
bias = jnp.array([0.5, -0.5, 1., 2.])
pieces = [x @ kernel[:, start:start + 2] + bias[start:start + 2]
          for start in (0, 2)]
output = jnp.concatenate(pieces, axis=-1)
np.testing.assert_allclose(output, [[1.5, 1.5, 5., 7.]])
np.testing.assert_allclose(output, x @ kernel + bias)

The equality follows from matrix multiplication over separate output columns. In floating-point implementations, different compiled kernels may give small rounding differences, so compare with a suitable tolerance rather than requiring bitwise identity.

How this relates to real tensor parallelism

On a model-parallel mesh, column layouts are commonly expressed as P(None, "model") for the kernel and P("model") for the bias. In the simple scheme shown here, each model-axis device needs the full input-feature vector and computes different output columns. Inputs can still be split along the batch axis on a separate mesh axis.

The global output can remain sharded. Materializing every output column on every device would require communication such as an all-gather; a distributed array's global shape does not imply each device holds that entire array. This exercise's local concatenation does not simulate that communication.

Row-parallel computation instead splits the kernel's input rows and the matching input features. Those products contribute partial sums to the same outputs, so they must be combined by addition. A replicated full output can use an all-reduce; other output layouts can use other collective patterns. Add the bias once to the combined result, rather than once per partial sum.

Read the initialized NNX parameters

Construct one nnx.Linear(in_dim, features, rngs=nnx.Rngs(seed)) with its default kernel initialization and enabled zero-initialized bias. Read its arrays with model.kernel[...] and model.bias[...]. These parameter wrappers let NNX track model state; the array slices can be passed directly to JAX operations.

Do not create a separately initialized Linear for each group. That would generally draw different random weights from the slices of the one full layer required here. There is no training step.