Tensor Parallel Row Split

Problem

Split a weight matrix along its output (row) dimension for tensor parallelism: each partition keeps a contiguous block of rows and all input columns. Row-parallel layers produce partial outputs that are summed with an all-reduce.

The signature is row_split(weight, num_partitions). Return the list of row blocks.

Hint

Row parallelism splits the output features so each rank computes a partial result of the same shape.

Test cases

case 1

in: [[[1,2],[3,4],[5,6],[7,8]],2]

expected: [[[1,2],[3,4]],[[5,6],[7,8]]]

case 2

in: [[[1,2],[3,4]],1]

expected: [[[1,2],[3,4]]]

case 3

in: [[[1],[2],[3],[4]],4]

expected: [[[1]],[[2]],[[3]],[[4]]]

case 4

in: [[[1,2],[3,4],[5,6]],3]

expected: [[[1,2]],[[3,4]],[[5,6]]]

Starter codePython

def row_split(weight, num_partitions):
    # Your code here
    pass

4 test cases · run your code against them in the editor