All-Gather Bytes

Problem

Compute the communication volume of an all-gather: each GPU sends its shard to the other num_gpus - 1 ranks, so sent bytes = (num_gpus - 1) * shard_bytes and every GPU ends with num_gpus * shard_bytes.

The signature is all_gather_bytes(shard_bytes, num_gpus). Return [sent_bytes, gathered_bytes].

Hint

ZeRO-3 gathers parameters with all-gather before each layer's forward.

Test cases

case 1

in: [1000,4]

expected: [3000,4000]

case 2

in: [1,1]

expected: [0,1]

case 3

in: [500,2]

expected: [500,1000]

case 4

in: [100,8]

expected: [700,800]

Starter codePython

def all_gather_bytes(shard_bytes, num_gpus):
    # Your code here
    pass

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