Data Parallel Epoch Traffic

Problem

Compute the total gradient all-reduce bytes moved by data parallelism over an epoch. Each of num_steps optimizer steps all-reduces the full gradient tensor at ring cost 2 * (num_gpus - 1) / num_gpus * num_params * bytes_per_param.

The signature is data_parallel_epoch_traffic(num_params, bytes_per_param, num_gpus, num_steps). Return a float number of bytes.

Hint

A single rank reduces nothing, so the traffic vanishes when num_gpus is 1.

Test cases

case 1

in: [1000000000,4,8,1000]

expected: 7000000000000

case 2

in: [100,4,1,10]

expected: 0

case 3

in: [1000,2,4,5]

expected: 15000

case 4

in: [7000000000,2,16,100]

expected: 2625000000000

Starter codePython

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

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