Gradient Accumulation Equivalence

Problem

Compute the effective learning rate when using gradient accumulation.

Return micro_lr * accumulation_steps, the learning rate that matches one large-batch step under this convention.

Hint

Accumulating k micro-batches with the same lr is equivalent to a k-times larger step.

case 1

in: [0.01,4]

expected: 0.04

case 2

in: [0.1,1]

expected: 0.1

case 3

in: [0.02,8]

expected: 0.16

case 4

in: [0,5]

expected: 0

Starter codePython

def gradient_accumulation_equivalence(micro_lr, accumulation_steps):
    # Your code here
    pass

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