Sequence Parallel Split

Problem

Split a sequence of seq_len positions across num_gpus ranks as evenly as possible: the first seq_len % num_gpus ranks get one extra position. Return the per-rank chunk sizes.

The signature is sequence_parallel_split(seq_len, num_gpus).

Hint

Sequence parallelism splits along the token axis of LayerNorm activations.

Test cases

case 1

in: [8,4]

expected: [2,2,2,2]

case 2

in: [10,3]

expected: [4,3,3]

case 3

in: [5,8]

expected: [1,1,1,1,1,0,0,0]

case 4

in: [7,1]

expected: [7]

Starter codePython

def sequence_parallel_split(seq_len, num_gpus):
    # Your code here
    pass

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