ZeRO-1 Optimizer Shard

Problem

Compute the per-GPU optimizer-state shard under ZeRO stage 1: ceil(num_params / num_gpus) parameters per rank, times optimizer_states tensors at bytes_per_value bytes each.

The signature is zero1_optimizer_shard(num_params, num_gpus, optimizer_states=2, bytes_per_value=4). Return an integer number of bytes.

Hint

Only the optimizer state is partitioned at stage 1; parameters and gradients stay replicated.

Test cases

case 1

in: [1000000000,8,2,4]

expected: 1000000000

case 2

in: [1001,8,2,4]

expected: 1008

case 3

in: [100,3,2,4]

expected: 272

case 4

in: [1,1,2,2]

expected: 4

Starter codePython

def zero1_optimizer_shard(num_params, num_gpus, optimizer_states=2, bytes_per_value=4):
    # Your code here
    pass

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