Tensor Parallel Column Split

Problem

Split a weight matrix along its input (column) dimension for tensor parallelism. Each partition keeps all output rows but only in_dim / num_partitions consecutive columns.

The signature is column_split(weight, num_partitions). Return the list of partitioned matrices.

Hint

Column parallelism splits each projection so every GPU computes a slice of the output.

Test cases

case 1

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

expected: [[[1,2],[5,6]],[[3,4],[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],[4]],[[2],[5]],[[3],[6]]]

Starter codePython

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

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