ZeRO-3 Gather Bytes

Problem

Estimate ZeRO-3 communication for one layer. The transient buffer holds the full layer of per_layer_params * bytes, and the all-gather pushes num_layers * (num_gpus - 1) / num_gpus * per_layer_params * bytes in total. Return [transient_bytes, total_comm_bytes].

The signature is zero3_gather_bytes(per_layer_params, bytes_per_param, num_layers, num_gpus).

Hint

ZeRO-3 gathers parameters layer by layer so only one full layer exists at a time.

Test cases

case 1

in: [1000000,4,12,8]

expected: [4000000,42000000]

case 2

in: [1000,2,1,2]

expected: [2000,1000]

case 3

in: [100,4,2,4]

expected: [400,600]

case 4

in: [1000,2,4,1]

expected: [2000,0]

Starter codePython

def zero3_gather_bytes(per_layer_params, bytes_per_param, num_layers, num_gpus):
    # Your code here
    pass

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