ZeRO-2 Gradient Shard

Problem

Compute the per-GPU gradient shard under ZeRO stage 2: ceil(num_params / num_gpus) gradient elements per rank at bytes_per_grad bytes each. Stage 2 shards gradients on top of the stage 1 optimizer split.

The signature is zero2_gradient_shard(num_params, num_gpus, bytes_per_grad=4). Return an integer number of bytes.

Hint

Reduce-scatter leaves each rank holding one reduced slice of the gradient buffer.

Test cases

case 1

in: [1000000000,8,4]

expected: 500000000

case 2

in: [1001,8,4]

expected: 504

case 3

in: [100,3,2]

expected: 68

case 4

in: [1,1,4]

expected: 4

Starter codePython

def zero2_gradient_shard(num_params, num_gpus, bytes_per_grad=4):
    # Your code here
    pass

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