Tensor Parallelism’s Communication Boundary

Splitting a matrix across devices is easy to draw. The important question is what must cross the network before the next operation can run. A partition that looks balanced in isolation can create a large synchronization point when followed by a nonlinear activation.

Megatron-LM's feed-forward construction pairs a column-partitioned first weight matrix with a row-partitioned second one. Each worker can apply the elementwise activation to its own hidden coordinates, then produce a partial output. Those outputs are reduced together. Its attention construction likewise exploits separate heads before combining their outputs. Megatron-LM paper

Follow the coordinates

Use one input row x = [1, 2]. Let the first matrix have shape 2 × 4 and produce pre-activation values [1, 1, 0, 6]. For simple arithmetic, this illustration uses ReLU rather than the paper's GeLU; both operate coordinate by coordinate.

Worker 0 owns hidden coordinates 0 and 1. Worker 1 owns coordinates 2 and 3. Each worker has the full two-element input, but only its two columns of the first matrix. After the activation, their local vectors are [1, 1] and [0, 6].

The second matrix has shape 4 × 2. Give each worker the rows corresponding to its hidden coordinates. Worker 0 produces [1, 1]; worker 1 produces [6, −6]. Adding these partial outputs gives [7, −5], exactly the result of the unsplit block.

Neither worker needed the other's hidden vector. The two matrix partitions line up so that each worker can carry its own hidden coordinates through the activation and into the second product.

Count the object, then the collective

For a larger illustrative activation with 128 token rows, hidden width 1,024, expanded width 4,096, and two-byte elements, the expanded hidden tensor contains 1,048,576 bytes. The output-shaped reduction tensor contains 262,144 bytes.

These sizes describe logical tensors, not bytes sent by each device. An all-reduceCombine values from every worker and give the combined result back to every worker. may exchange multiple chunks, and traffic depends on the algorithm and number of workers. Treating tensor size as wire traffic would overstate the precision of this calculation.

Why not split the first matrix along its input dimension instead? Each worker would then calculate partial sums for every hidden coordinate. Those partial sums must be combined before applying the nonlinearity. In general ReLU(a+b) differs from ReLU(a)+ReLU(b); using a = 1 and b = −1 makes the mismatch immediate. The communication boundary has moved to the wider expanded tensor.

The accompanying CPU code checks the small numerical result and reports the two logical tensor sizes. It leaves backward communication, biases, dropout, normalization, and distributed collectives outside its scope. Bias handling in particular needs care: adding the same output bias independently before a reduction would duplicate it.

def row_times_matrix(row, matrix):
    return [sum(row[i] * matrix[i][j] for i in range(len(row)))
            for j in range(len(matrix[0]))]

x = [1, 2]
w1 = [[1, -1, 2, 0], [0, 1, -1, 3]]
w2 = [[1, 0], [0, 1], [2, 1], [1, -1]]
hidden = [max(0, v) for v in row_times_matrix(x, w1)]
dense = row_times_matrix(hidden, w2)
partials = []
for columns in ((0, 1), (2, 3)):
    local_w1 = [[row[j] for j in columns] for row in w1]
    local_hidden = [max(0, v) for v in row_times_matrix(x, local_w1)]
    local_w2 = [w2[j] for j in columns]
    partials.append(row_times_matrix(local_hidden, local_w2))
split = [sum(part[j] for part in partials) for j in range(2)]
assert split == dense == [7, -5]
print(f"dense={dense}; split={split}")
tokens, width, expanded, element_bytes = 128, 1024, 4096, 2
hidden_bytes = tokens * expanded * element_bytes
output_bytes = tokens * width * element_bytes
print(f"expanded_tensor_bytes={hidden_bytes}")
print(f"output_reduction_tensor_bytes={output_bytes}")
assert hidden_bytes == 4 * output_bytes

dense=[7, -5]; split=[7, -5]
expanded_tensor_bytes=1048576
output_reduction_tensor_bytes=262144