Gradient Accumulation Micro-Steps

Problem

Compute the number of micro-steps each data-parallel rank runs per optimizer step: ceil(global_batch_size / (micro_batch_size * data_parallel)). Each rank processes micro_batch_size samples at a time and gradients are accumulated before the optimizer updates.

The signature is gradient_accumulation_micro_steps(global_batch_size, micro_batch_size, data_parallel=1). Return an integer.

Hint

Data-parallel ranks each contribute micro_batch_size samples to every micro-step.

Test cases

case 1

in: [1024,4,8]

expected: 32

case 2

in: [1000,4,8]

expected: 32

case 3

in: [8,4,2]

expected: 1

case 4

in: [3,4,1]

expected: 1

case 5

in: [64,8,1]

expected: 8

Starter codePython

def gradient_accumulation_micro_steps(global_batch_size, micro_batch_size, data_parallel=1):
    # Your code here
    pass

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