ZeRO-3 Parameter Shard

Problem

Compute the per-GPU parameter shard under ZeRO stage 3: ceil(num_params / num_gpus) parameters per rank at bytes_per_param bytes each. Parameters are gathered layer by layer just in time for compute and freed afterwards.

The signature is zero3_parameter_shard(num_params, num_gpus, bytes_per_param=2). Return an integer number of bytes.

Hint

At stage 3 even the parameters are partitioned, so per-GPU storage scales as 1 / num_gpus.

Test cases

case 1

in: [7000000000,64,2]

expected: 218750000

case 2

in: [1001,8,4]

expected: 504

case 3

in: [100,3,2]

expected: 68

case 4

in: [1,1,2]

expected: 2

Starter codePython

def zero3_parameter_shard(num_params, num_gpus, bytes_per_param=2):
    # Your code here
    pass

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