Gradient Checkpointing Memory

Problem

Estimate gradient checkpointing memory. The baseline stores every layer activation: num_layers * activation_mb. With a checkpoint every k layers, the peak is (ceil(num_layers / k) + k) * activation_mb, since all checkpoints plus one recomputed segment are held.

The signature is checkpointing_memory(num_layers, activation_mb, checkpoint_every). Return [baseline_mb, checkpointed_mb, saved_mb].

Hint

Checkpointing trades compute for memory: activations are recomputed during backward.

Test cases

case 1

in: [12,100,3]

expected: [1200,700,500]

case 2

in: [24,50,4]

expected: [1200,500,700]

case 3

in: [10,20,5]

expected: [200,140,60]

case 4

in: [8,16,2]

expected: [128,96,32]

Starter codePython

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

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