Tensor Parallel All-Reduce Count

Problem

Count the all-reduces in a transformer forward pass with tensor parallelism. Each block ends two row-parallel projections (attention output and MLP down) with a sum all-reduce. Set backward=True to include the backward pass, which repeats the same count.

The signature is tp_all_reduce_count(num_layers, backward=False). Return an integer.

Hint

Two row-parallel layers per block, and backward mirrors forward.

Test cases

case 1

in: [12]

expected: 24

case 2

in: [12,true]

expected: 48

case 3

in: [1,false]

expected: 2

case 4

in: [1,true]

expected: 4

case 5

in: [32]

expected: 64

Starter codePython

def tp_all_reduce_count(num_layers, backward=False):
    # Your code here
    pass

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