ZeRO Stage Memory

Problem

Estimate per-GPU training memory for ZeRO stages. With P parameter bytes, G gradient bytes and O optimizer bytes (num_params * optimizer_states * bytes): stage 0 holds P + G + O; stage 1 shards the optimizer state; stage 2 also shards gradients; stage 3 also shards parameters. Sharded quantities use ceiling division by num_gpus.

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

Hint

Each increase in ZeRO stage shards one more training-state component across the data-parallel group.

Test cases

case 1

in: [1000000,4,8,0]

expected: 16000000

case 2

in: [1000000,4,8,1]

expected: 9000000

case 3

in: [1000000,4,8,2]

expected: 5500000

case 4

in: [1000000,4,8,3]

expected: 2000000

case 5

in: [1000000,4,1,3]

expected: 16000000

Starter codePython

def zero_stage_memory(num_params, bytes_per_param, num_gpus, stage, optimizer_states=2):
    # Your code here
    pass

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