Reduce-Scatter Bytes

Problem

Compute the communication volume of a reduce-scatter: each GPU sends (num_gpus - 1) / num_gpus of the full tensor and keeps one reduced shard of full_bytes / num_gpus.

The signature is reduce_scatter_bytes(full_bytes, num_gpus). Return [sent_bytes, result_bytes] as floats.

Hint

Reduce-scatter halves the traffic of an all-reduce when combined with all-gather.

Test cases

case 1

in: [4000,4]

expected: [3000,1000]

case 2

in: [1000,1]

expected: [0,1000]

case 3

in: [100,2]

expected: [50,50]

case 4

in: [8000,8]

expected: [7000,1000]

Starter codePython

def reduce_scatter_bytes(full_bytes, num_gpus):
    # Your code here
    pass

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