Checkpointing Recompute Factor

Problem

Estimate the compute factor from activation checkpointing every checkpoint_every layers. The baseline is 3 * num_layers forward-equivalents (one forward plus a two-forward backward); checkpointing adds one forward-equivalent per checkpointed segment, giving (3 * num_layers + ceil(num_layers / checkpoint_every)) / (3 * num_layers).

The signature is checkpointing_recompute_factor(num_layers, checkpoint_every).

Hint

Checkpointing every layer recomputes the whole forward pass once, so the factor approaches 4/3.

Test cases

case 1

in: [24,4]

expected: 1.0833333333

case 2

in: [24,1]

expected: 1.3333333333

case 3

in: [24,24]

expected: 1.0138888889

case 4

in: [12,3]

expected: 1.1111111111

case 5

in: [6,4]

expected: 1.1111111111

Starter codePython

def checkpointing_recompute_factor(num_layers, checkpoint_every):
    # Your code here
    pass

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