Collective All-Reduce Bytes

Problem

Compute the communication volume of a ring all-reduce: 2 * (num_gpus - 1) / num_gpus * num_params * bytes_per_param.

The signature is all_reduce_bytes(num_params, bytes_per_param, num_gpus).

Hint

A single GPU has nothing to reduce; larger rings approach 2x the tensor size.

case 1

in: [1000000,4,8]

expected: 7000000

case 2

in: [100,4,2]

expected: 400

case 3

in: [1000,2,1]

expected: 0

case 4

in: [1000,2,4]

expected: 3000

Starter codePython

def all_reduce_bytes(num_params, bytes_per_param, num_gpus):
    # Your code here
    pass

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