All-Gather Network Traffic

Problem

Compute the total bytes crossing the interconnect for a ring all-gather of a full tensor: each of num_gpus ranks sends num_gpus - 1 chunks of full_bytes / num_gpus, so the fabric carries (num_gpus - 1) * full_bytes in total.

The signature is all_gather_traffic(full_bytes, num_gpus). Return an integer number of bytes.

Hint

Every chunk travels across the ring on every link except its origin's outgoing edge.

Test cases

case 1

in: [4000,4]

expected: 12000

case 2

in: [1000,1]

expected: 0

case 3

in: [8000,8]

expected: 56000

case 4

in: [100,2]

expected: 100

case 5

in: [0,4]

expected: 0

Starter codePython

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

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