Tensor Parallelism’s Communication Boundary notebook

Python starts when you choose Run.

Tensor Parallelism’s Communication Boundary

A paired matrix partition shows where tensor parallelism needs communication, and why the nonlinearity determines that boundary.

Original CPU forward-pass illustration using ReLU; not a distributed Megatron-LM implementation or GPU scaling test.

The code uses Python 3 and the standard library. Saved output is from the checked CPU run.

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

Sources

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism: https://arxiv.org/html/1909.08053v4

Python runs locally in this page. Saved results and figures are visible before execution.

Download full experiment notebook