Training Checkpoint Size

Problem

Compute the size in bytes of a training checkpoint holding weights and optimizer state (gradients are typically not saved): num_params * bytes_per_param + num_params * optimizer_states * optimizer_bytes.

The signature is training_checkpoint_size(num_params, bytes_per_param=2, optimizer_states=2, optimizer_bytes=4). Return an integer number of bytes.

Hint

Adam-style state usually dwarfs the fp16 weights.

Test cases

case 1

in: [1000000000,2,2,4]

expected: 10000000000

case 2

in: [7000000000,2,2,4]

expected: 70000000000

case 3

in: [1000,4,0,4]

expected: 4000

case 4

in: [1,2,2,4]

expected: 10

Starter codePython

def training_checkpoint_size(num_params, bytes_per_param=2, optimizer_states=2, optimizer_bytes=4):
    # Your code here
    pass

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