Ring All-Reduce Time

Problem

Estimate ring all-reduce wall time: reduce-scatter and all-gather each run num_gpus - 1 steps that send full_bytes / num_gpus over the link. Each step costs chunk / (bandwidth_gbps * 1e9) seconds plus latency_us microseconds.

The signature is ring_all_reduce_time(full_bytes, num_gpus, bandwidth_gbps, latency_us=0.0). Return seconds as a float.

Hint

Bandwidth time shrinks as 1 / num_gpus while latency time grows with num_gpus - 1.

Test cases

case 1

in: [1000000000,8,100,0]

expected: 0.0175

case 2

in: [1000000000,8,100,10]

expected: 0.01764

case 3

in: [4000000000,4,200,5]

expected: 0.03003

case 4

in: [1000000,2,10,0]

expected: 0.0001

Starter codePython

def ring_all_reduce_time(full_bytes, num_gpus, bandwidth_gbps, latency_us=0.0):
    # Your code here
    pass

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